wehub-resource-sync
271 lines
10 KiB
Python
271 lines
10 KiB
Python
# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
# Adapted from: https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/deduplication/shuffle_utils/rapidsmpf_shuffler.py
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Any, Literal
|
|
|
|
from packaging.version import parse as parse_version
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Iterator
|
|
|
|
import pylibcudf as plc
|
|
|
|
|
|
def align_down_to_256(value: int) -> int:
|
|
return (value >> 8) << 8
|
|
|
|
|
|
def get_device_free_memory() -> int | None:
|
|
try:
|
|
import pynvml
|
|
|
|
import ray
|
|
|
|
pynvml.nvmlInit()
|
|
index = int(ray.get_gpu_ids()[0]) if ray.is_initialized() else 0
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(index)
|
|
info = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
|
return int(info.free)
|
|
except Exception:
|
|
return None
|
|
finally:
|
|
pynvml.nvmlShutdown()
|
|
|
|
|
|
def lazy_load() -> type[Any]:
|
|
"""Lazy load the BulkRapidsMPFShuffler class to allow for CPU-only environments.
|
|
|
|
This is necessary because the BulkRapidsMPFShuffler inherits from BaseShufflingActor.
|
|
"""
|
|
import cudf
|
|
import rapidsmpf
|
|
import rmm.mr
|
|
from rapidsmpf.integrations.cudf.partition import (
|
|
partition_and_pack,
|
|
unpack_and_concat,
|
|
unspill_partitions,
|
|
)
|
|
from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor
|
|
from rapidsmpf.shuffler import Shuffler
|
|
from rapidsmpf.utils.cudf import cudf_to_pylibcudf_table
|
|
|
|
if parse_version(rapidsmpf.__version__) >= parse_version("26.4.0"):
|
|
from rapidsmpf.integrations.ray import RapidsMPFActor as RapidsMPFActorBase
|
|
else:
|
|
from rapidsmpf.progress_thread import ProgressThread
|
|
from rapidsmpf.utils.ray_utils import BaseShufflingActor as RapidsMPFActorBase
|
|
|
|
if parse_version(rapidsmpf.__version__) >= parse_version("26.2.0"):
|
|
from rapidsmpf.memory.buffer import MemoryType
|
|
from rapidsmpf.memory.buffer_resource import (
|
|
BufferResource,
|
|
LimitAvailableMemory,
|
|
)
|
|
else:
|
|
from rapidsmpf.buffer.buffer import MemoryType
|
|
from rapidsmpf.buffer.resource import BufferResource, LimitAvailableMemory
|
|
|
|
# Exempt this class from coverage is it's indirectly tested by the ShuffleStage which coverage tools don't pick up.
|
|
class BulkRapidsMPFShuffler(RapidsMPFActorBase): # pragma: no cover
|
|
"""Class that performs a bulk shuffle operation.
|
|
|
|
This class is compatible with Ray Actors communicating with each
|
|
other using UCXX communication.
|
|
|
|
Args:
|
|
nranks: Number of ranks in the communication group.
|
|
total_nparts: Total number of output partitions.
|
|
shuffle_on: List of column names to shuffle on.
|
|
output_path: Path to write output files.
|
|
rmm_pool_size: Size of the RMM GPU memory pool in bytes.
|
|
If "auto", the memory pool is set to 90% of the free GPU memory.
|
|
If None, the memory pool is set to 50% of the free GPU memory
|
|
that can expand if needed.
|
|
spill_memory_limit: Device memory limit in bytes for spilling to host.
|
|
If "auto", the limit is set to 80% of the RMM pool size.
|
|
If None spilling is disabled.
|
|
read_kwargs: Keyword arguments for cudf.read_parquet method.
|
|
write_kwargs: Keyword arguments for cudf.to_parquet method.
|
|
"""
|
|
|
|
def __init__( # noqa: PLR0913
|
|
self,
|
|
nranks: int,
|
|
total_nparts: int,
|
|
shuffle_on: list[str],
|
|
output_path: str = "./",
|
|
rmm_pool_size: int | Literal["auto"] | None = "auto",
|
|
spill_memory_limit: int | Literal["auto"] | None = "auto",
|
|
*,
|
|
read_kwargs: dict[str, Any] | None = None,
|
|
write_kwargs: dict[str, Any] | None = None,
|
|
):
|
|
super().__init__(nranks)
|
|
self.shuffle_on = shuffle_on
|
|
self.output_path = output_path
|
|
self.total_nparts = total_nparts
|
|
|
|
if isinstance(rmm_pool_size, int):
|
|
self.rmm_pool_size = align_down_to_256(rmm_pool_size)
|
|
elif rmm_pool_size == "auto":
|
|
free_memory = get_device_free_memory()
|
|
self.rmm_pool_size = (
|
|
align_down_to_256(int(free_memory * 0.9))
|
|
if free_memory is not None
|
|
else None
|
|
)
|
|
elif rmm_pool_size is None:
|
|
self.rmm_pool_size = None
|
|
else:
|
|
err_msg = f"Invalid rmm_pool_size: {rmm_pool_size}"
|
|
raise ValueError(err_msg)
|
|
if isinstance(spill_memory_limit, int):
|
|
self.spill_memory_limit = align_down_to_256(spill_memory_limit)
|
|
elif spill_memory_limit == "auto":
|
|
self.spill_memory_limit = (
|
|
align_down_to_256(int(0.8 * self.rmm_pool_size))
|
|
if self.rmm_pool_size is not None
|
|
else None
|
|
)
|
|
elif spill_memory_limit is None:
|
|
self.spill_memory_limit = None
|
|
else:
|
|
err_msg = f"Invalid spill_memory_limit: {spill_memory_limit}"
|
|
raise ValueError(err_msg)
|
|
|
|
self.read_kwargs = read_kwargs if read_kwargs is not None else {}
|
|
self.write_kwargs = write_kwargs if write_kwargs is not None else {}
|
|
|
|
def setup_worker(self, root_address_bytes: bytes) -> None:
|
|
"""Setup the UCXX communication and a shuffle operation.
|
|
|
|
Args:
|
|
root_address_bytes: Address of the root worker for UCXX
|
|
initialization.
|
|
"""
|
|
super().setup_worker(root_address_bytes)
|
|
|
|
# Initialize the RMM memory resource
|
|
mr = RmmResourceAdaptor(
|
|
rmm.mr.PoolMemoryResource(
|
|
rmm.mr.CudaMemoryResource(),
|
|
initial_pool_size=self.rmm_pool_size,
|
|
maximum_pool_size=None,
|
|
)
|
|
)
|
|
rmm.mr.set_current_device_resource(mr)
|
|
# Create a buffer resource that limits device memory if spill_memory_limit is set
|
|
memory_available = (
|
|
None
|
|
if self.spill_memory_limit is None
|
|
else {
|
|
MemoryType.DEVICE: LimitAvailableMemory(
|
|
mr, limit=self.spill_memory_limit
|
|
)
|
|
}
|
|
)
|
|
self.br = BufferResource(device_mr=mr, memory_available=memory_available)
|
|
# Create a shuffler
|
|
if parse_version(rapidsmpf.__version__) >= parse_version("26.4.0"):
|
|
self.shuffler = Shuffler(
|
|
self.comm, 0, total_num_partitions=self.total_nparts, br=self.br
|
|
)
|
|
else:
|
|
self.shuffler = Shuffler(
|
|
self.comm,
|
|
ProgressThread(self.comm),
|
|
0,
|
|
self.total_nparts,
|
|
self.br,
|
|
)
|
|
|
|
def cleanup(self) -> None:
|
|
"""Cleanup the UCXX communication and the shuffle operation."""
|
|
if self.shuffler is not None:
|
|
self.shuffler.shutdown()
|
|
|
|
def insert_chunk(
|
|
self, table: plc.Table | cudf.DataFrame, column_names: list[str]
|
|
) -> None:
|
|
"""Insert a pylibcudf Table or cuDF DataFrame into the shuffler.
|
|
|
|
Args:
|
|
table: The table or DataFrame to insert.
|
|
column_names: The column names of the table.
|
|
"""
|
|
from rmm.pylibrmm.stream import DEFAULT_STREAM
|
|
|
|
if isinstance(table, cudf.DataFrame):
|
|
table = cudf_to_pylibcudf_table(table)
|
|
columns_to_hash = tuple(column_names.index(val) for val in self.shuffle_on)
|
|
packed_inputs = partition_and_pack(
|
|
table=table,
|
|
br=self.br,
|
|
columns_to_hash=columns_to_hash,
|
|
num_partitions=self.total_nparts,
|
|
stream=DEFAULT_STREAM,
|
|
)
|
|
self.shuffler.insert_chunks(packed_inputs)
|
|
|
|
def insert_finished(self) -> None:
|
|
"""Tell the shuffler that we are done inserting data."""
|
|
for pid in range(self.total_nparts):
|
|
self.shuffler.insert_finished(pid)
|
|
self.comm.logger.info("Insert finished")
|
|
|
|
def extract(self) -> Iterator[tuple[int, plc.Table]]:
|
|
"""Extract shuffled partitions as they become ready.
|
|
|
|
Partitions are yielded in completion order (via ``wait_any()``),
|
|
not partition order. Callers are responsible for reordering.
|
|
|
|
Yields:
|
|
tuple[int, plc.Table]: ``(partition_id, partition)`` tuples.
|
|
"""
|
|
from rmm.pylibrmm.stream import DEFAULT_STREAM
|
|
|
|
first_time = True
|
|
while not self.shuffler.finished() or first_time:
|
|
first_time = False
|
|
partition_id = self.shuffler.wait_any()
|
|
packed_chunks = self.shuffler.extract(partition_id)
|
|
partition = unpack_and_concat(
|
|
unspill_partitions(
|
|
packed_chunks,
|
|
br=self.br,
|
|
allow_overbooking=True,
|
|
),
|
|
br=self.br,
|
|
stream=DEFAULT_STREAM,
|
|
)
|
|
yield partition_id, partition
|
|
|
|
return BulkRapidsMPFShuffler
|
|
|
|
|
|
_BulkRapidsMPFShuffler = None
|
|
|
|
|
|
def __getattr__(name: str) -> type:
|
|
global _BulkRapidsMPFShuffler
|
|
if name == "BulkRapidsMPFShuffler":
|
|
if _BulkRapidsMPFShuffler is None:
|
|
_BulkRapidsMPFShuffler = lazy_load()
|
|
return _BulkRapidsMPFShuffler
|
|
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|