Implementation:NVIDIA NeMo Curator ShuffleStage

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

Overview

Defines ShuffleStage, a generic processing stage for GPU-based distributed data shuffling on specified columns, used in deduplication pipelines.

Description

ShuffleStage extends ProcessingStage[FileGroupTask, FileGroupTask] and wraps BulkRapidsMPFShuffler as its actor class. It separates shuffle coordination logic (this stage) from shuffle execution logic (the actor), enabling the executor to manage the distributed shuffle lifecycle. Key characteristics:

• Actor-based execution: The stage's process() method raises NotImplementedError because it is designed to be used with an actor-based executor, not direct invocation.

• Shuffle lifecycle methods:
  • read_and_insert(task) -- Reads files from a FileGroupTask and inserts data into the shuffler actor.
  • insert_finished() -- Signals to the shuffler that all data has been inserted.
  • extract_and_write() -- Extracts shuffled partitions and writes them to Parquet, returning a list of FileGroupTask objects with partition metadata.
  • teardown() -- Cleans up the shuffler actor resources.

• Configuration: Supports configuring RMM GPU memory pool size (rmm_pool_size), device memory spill limits (spill_memory_limit), output partition count (total_nparts), and read/write keyword arguments for cuDF Parquet I/O.

• Ray integration: The ray_stage_spec() method returns IS_SHUFFLE_STAGE: True, signaling to the Ray executor that this stage requires special shuffle handling.

Usage

Use ShuffleStage within deduplication pipelines that require distributed data shuffling on specific columns (e.g., shuffling by document hash for dedup bucket creation). It requires a Ray-based actor executor and GPU resources.

Code Reference

Source Location

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

Signature

class ShuffleStage(ProcessingStage[FileGroupTask, FileGroupTask]):
    name = "ShuffleStage"
    resources = Resources(gpus=1.0)
    actor_class = BulkRapidsMPFShuffler

    def __init__(
        self,
        shuffle_on: list[str],
        total_nparts: int | None = None,
        output_path: str = "./",
        read_kwargs: dict[str, Any] | None = None,
        write_kwargs: dict[str, Any] | None = None,
        rmm_pool_size: int | Literal["auto"] | None = "auto",
        spill_memory_limit: int | Literal["auto"] | None = "auto",
        enable_statistics: bool = False,
    ): ...

    def process(self, task: FileGroupTask) -> FileGroupTask: ...
    def ray_stage_spec(self) -> dict[str, Any]: ...
    def read_and_insert(self, task: FileGroupTask) -> FileGroupTask: ...
    def insert_finished(self) -> None: ...
    def extract_and_write(self) -> list[FileGroupTask]: ...
    def teardown(self) -> None: ...

Import

from nemo_curator.stages.deduplication.shuffle_utils.stage import ShuffleStage

I/O Contract

Constructor Inputs

Name	Type	Required	Description
shuffle_on	list[str]	Yes	List of column names to shuffle on
total_nparts	int or None	No	Total number of output partitions. None lets the executor decide (default: None)
output_path	str	No	Path to write output Parquet files (default: "./")
read_kwargs	dict[str, Any]	No	Keyword arguments for cudf.read_parquet() (default: {})
write_kwargs	dict[str, Any]	No	Keyword arguments for cudf.to_parquet() (default: {})
rmm_pool_size	int, "auto", or None	No	RMM GPU memory pool size in bytes. "auto" = 90% of free GPU memory. None = 50% expandable (default: "auto")
spill_memory_limit	int, "auto", or None	No	Device memory limit for spilling to host. "auto" = 80% of pool. None = disabled (default: "auto")
enable_statistics	bool	No	Whether to collect shuffle statistics (default: False)

Process Input

Name	Type	Required	Description
task	FileGroupTask	Yes	Task containing file paths to read and shuffle

Outputs

Name	Type	Description
result	list[FileGroupTask]	List of FileGroupTask objects, one per output partition, each containing the path to a shuffled Parquet file and partition metadata

Usage Examples

Basic Usage

from nemo_curator.stages.deduplication.shuffle_utils.stage import ShuffleStage

shuffle_stage = ShuffleStage(
    shuffle_on=["_bucket_id"],
    total_nparts=128,
    output_path="/data/dedup/shuffled/",
    rmm_pool_size="auto",
    spill_memory_limit="auto",
)

With Custom I/O Options

from nemo_curator.stages.deduplication.shuffle_utils.stage import ShuffleStage

shuffle_stage = ShuffleStage(
    shuffle_on=["hash_column"],
    output_path="s3://my-bucket/shuffled/",
    write_kwargs={"storage_options": {"key": "...", "secret": "..."}},
    enable_statistics=True,
)

Related Pages

• Environment:NVIDIA_NeMo_Curator_Python_Linux_Base