chore: import upstream snapshot with attribution

python/ray/data/_internal/gpu_shuffle/rapidsmpf_backend.py

@@ -0,0 +1,270 @@
+# 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}")