Implementation:NVIDIA NeMo Curator RapidsMPFShuffler

Knowledge Sources	NVIDIA NeMo Curator
Domains	GPU Computing, Distributed Systems, Deduplication
Last Updated	2026-02-14 00:00 GMT

Overview

The BulkRapidsMPFShuffler class implements a GPU-based distributed data shuffler using the RapidsMPF library for efficient cross-worker data redistribution in deduplication pipelines.

Description

BulkRapidsMPFShuffler extends BaseShufflingActor (Ray-compatible) to perform bulk shuffle operations where data is redistributed across workers based on hash-partitioned columns. This is core infrastructure for GPU-accelerated distributed shuffling used by LSH and other deduplication stages to co-locate related records on the same worker for processing.

The class manages the complete lifecycle of a shuffle operation:

Memory management: On worker setup, it initializes an RMM (RAPIDS Memory Manager) GPU memory pool. The pool size defaults to 90% of free GPU memory when set to "auto", or can be explicitly specified. It also supports optional spill-to-host memory with a configurable limit (defaulting to 80% of the RMM pool size) via LimitAvailableMemory.

Communication: Uses UCXX for inter-worker communication, initialized via setup_worker with root address bytes for the communication group.

Data flow:

1. Read: read_batch reads Parquet files into cuDF DataFrames.
2. Partition and insert: insert_chunk hash-partitions data on specified columns using partition_and_pack and inserts packed chunks into the RapidsMPF shuffler.
3. Signal completion: insert_finished signals the shuffler that all data has been inserted for all partition IDs.
4. Extract: extract yields shuffled partitions as they become ready via unpack_and_concat.
5. Write: write_table / extract_and_write writes pylibcudf tables to Parquet files.

The convenience method read_and_insert combines reading and insertion with configurable batch sizes, while extract_and_write iterates over ready partitions and writes them to disk.

Usage

Use BulkRapidsMPFShuffler when you need to redistribute large datasets across GPU workers based on hash columns during deduplication stages. It is typically used as a Ray actor within the NeMo Curator deduplication pipeline, handling the data movement needed for LSH-based and other hash-partitioned deduplication strategies.

Code Reference

Source Location

• Repository: NeMo-Curator
• File: nemo_curator/stages/deduplication/shuffle_utils/rapidsmpf_shuffler.py
• Lines: 1-302

Signature

class BulkRapidsMPFShuffler(BaseShufflingActor):
    def __init__(
        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",
        *,
        enable_statistics: bool = False,
        read_kwargs: dict[str, Any] | None = None,
        write_kwargs: dict[str, Any] | None = None,
    ): ...

    def setup_worker(self, root_address_bytes: bytes) -> None: ...
    def cleanup(self) -> None: ...
    def read_batch(self, paths: list[str]) -> tuple[cudf.DataFrame | None, list[str]]: ...
    def write_table(self, table: plc.Table, output_path: str, partition_id: int | str, column_names: list[str]) -> str: ...
    def insert_chunk(self, table: plc.Table | cudf.DataFrame, column_names: list[str]) -> None: ...
    def read_and_insert(self, paths: list[str], batchsize: int | None = None) -> list[str]: ...
    def insert_finished(self) -> None: ...
    def extract(self) -> Iterator[tuple[int, plc.Table]]: ...
    def extract_and_write(self, column_names: list[str]) -> list[tuple[int, str]]: ...

Import

from nemo_curator.stages.deduplication.shuffle_utils.rapidsmpf_shuffler import BulkRapidsMPFShuffler

I/O Contract

Inputs

Name	Type	Required	Description
nranks	int	Yes	Number of ranks (workers) in the communication group
total_nparts	int	Yes	Total number of output partitions for the shuffle
shuffle_on	list[str]	Yes	Column names to hash-partition (shuffle) data on
output_path	str	No (default: "./")	Path to write output Parquet files
rmm_pool_size	Literal["auto"] | None	No (default: "auto")	Size of the RMM GPU memory pool in bytes; "auto" sets to 90% of free GPU memory; None sets to 50% expandable
spill_memory_limit	Literal["auto"] | None	No (default: "auto")	Device memory limit in bytes for spilling to host; "auto" sets to 80% of RMM pool; None disables spilling
enable_statistics	bool	No (default: False)	Whether to collect shuffle performance statistics
read_kwargs	None	No	Additional keyword arguments for cudf.read_parquet
write_kwargs	None	No	Additional keyword arguments for cudf.DataFrame.to_parquet

Outputs

Name	Type	Description
read_batch	None, list[str]]	Tuple of the DataFrame read from Parquet and its column names
read_and_insert	list[str]	Column names from the read data
extract	Iterator[tuple[int, plc.Table]]	Iterator yielding (partition_id, pylibcudf_table) for each shuffled partition
extract_and_write	list[tuple[int, str]]	List of (partition_id, file_path) for each written partition
write_table	str	Path of the written Parquet file

Memory Management

The shuffler manages GPU memory through a layered approach:

Layer	Component	Purpose
Base	rmm.mr.CudaMemoryResource	Raw CUDA device memory allocation
Pool	rmm.mr.PoolMemoryResource	Pre-allocated memory pool for fast allocation/deallocation
Adaptor	RmmResourceAdaptor	Bridges RMM with RapidsMPF buffer management
Spill control	LimitAvailableMemory	Limits device memory usage, enabling spill-to-host when limit is reached
Buffer	BufferResource	Manages device and host memory for the shuffler

Memory sizes are aligned down to 256-byte boundaries via align_down_to_256.

Usage Examples

Basic Usage

from nemo_curator.stages.deduplication.shuffle_utils.rapidsmpf_shuffler import BulkRapidsMPFShuffler

# Create a shuffler actor
shuffler = BulkRapidsMPFShuffler(
    nranks=4,
    total_nparts=128,
    shuffle_on=["hash_column"],
    output_path="/output/shuffled",
    rmm_pool_size="auto",
    spill_memory_limit="auto",
    enable_statistics=True,
)

# Setup worker communication
shuffler.setup_worker(root_address_bytes)

# Read, partition, and insert data
column_names = shuffler.read_and_insert(
    paths=["/data/part.0.parquet", "/data/part.1.parquet"],
    batchsize=1,
)

# Signal insertion complete
shuffler.insert_finished()

# Extract and write shuffled partitions
partition_paths = shuffler.extract_and_write(column_names)

# Cleanup
shuffler.cleanup()

Related Pages

• Environment:NVIDIA_NeMo_Curator_Python_Linux_Base