Implementation:NVIDIA NeMo Curator ShuffleAdapter

Knowledge Sources	NVIDIA NeMo Curator
Domains	Backend Architecture, Ray Integration, Data Shuffling, UCXX
Last Updated	2026-02-14 00:00 GMT

Overview

Wraps shuffle and LSH stages as Ray remote actors with UCXX-based inter-actor communication, managing the three-phase shuffle lifecycle of insert, signal completion, and extract results.

Description

ShuffleStageAdapter is a specialized Ray remote actor (decorated with @ray.remote) that adapts shuffle-type stages (such as ShuffleStage and LSHStage) for distributed data redistribution. It is decorated with a runtime environment that sets UCX_RNDV_FRAG_MEM_TYPES=host to work around a known UCX memory issue with GPU staging buffers.

The adapter performs the following during initialization:

Retrieves Ray runtime context (node ID, worker ID) via get_worker_metadata_and_node_id().
Reads the stage's batch size, defaulting to 1 if not set.
Computes output partition count: for LSH stages it rounds down to the nearest power of 2 (empirically improving shuffle performance), otherwise it matches the input task count. If total_nparts is explicitly set on the stage, that value is used.
Initializes the underlying shuffle actor object with the computed parameters.

The shuffle lifecycle is driven by the executor through three phases:

read_and_insert() -- Feeds tasks (and optional band ranges for LSH) into the shuffler.
insert_finished() -- Signals that all data has been inserted, triggering the actual shuffle operation.
extract_and_write() -- Extracts the shuffled data and writes output files, returning the resulting FileGroupTask objects.

Communication setup uses UCXX: setup_root() initializes the root endpoint and returns its address, while setup() (or setup_worker()) connects all workers to the root.

Usage

This adapter is used internally by the Ray ActorPool executor for stages that require distributed data shuffling. The executor creates multiple ShuffleStageAdapter actors, sets up UCXX communication between them, and orchestrates the insert-shuffle-extract lifecycle. Users do not instantiate this class directly.

Code Reference

Source Location

Repository: NeMo-Curator
File: nemo_curator/backends/experimental/ray_actor_pool/shuffle_adapter.py
Lines: 1-147

Signature

@ray.remote(runtime_env={"env_vars": {"UCX_RNDV_FRAG_MEM_TYPES": "host"}})
class ShuffleStageAdapter(BaseStageAdapter):
    def __init__(
        self,
        stage: "ShuffleStage | LSHStage",
        rank: int,
        nranks: int,
        num_input_tasks: int | None = None,
    ): ...
    def get_batch_size(self) -> int: ...
    def setup_on_node(self) -> None: ...
    def setup(self, root_address: bytes, worker_metadata: WorkerMetadata | None = None) -> None: ...
    def setup_root(self) -> bytes: ...
    def setup_worker(self, root_address: bytes) -> None: ...
    def read_and_insert(
        self, tasks: list[FileGroupTask], band_range: tuple[int, int] | None = None
    ) -> list[FileGroupTask]: ...
    def insert_finished(self) -> None: ...
    def extract_and_write(self) -> list[FileGroupTask]: ...
    def teardown(self) -> None: ...

Import

from nemo_curator.backends.experimental.ray_actor_pool.shuffle_adapter import ShuffleStageAdapter

I/O Contract

Inputs

Name	Type	Required	Description
stage	ShuffleStage or LSHStage	Yes	The shuffle stage to wrap
rank	int	Yes	This actor's rank in the group
nranks	int	Yes	Total number of actors in the group
num_input_tasks	int or None	No	Total number of input tasks (used to auto-compute output partitions)

Outputs

Name	Type	Description
setup_root()	bytes	The UCXX root address for other workers to connect to
read_and_insert()	list[FileGroupTask]	Intermediate results from the insert phase
extract_and_write()	list[FileGroupTask]	Output tasks containing shuffled data written to files

Shuffle Lifecycle

The executor drives the following three-phase workflow:

Phase 1 - Insert: The executor calls read_and_insert() on each actor with its assigned tasks. For LSH stages, a band_range tuple is also provided.
Phase 2 - Shuffle: The executor calls insert_finished() on all actors, triggering the underlying UCXX-based data redistribution.
Phase 3 - Extract: The executor calls extract_and_write() on all actors to produce the final output FileGroupTask objects.

Usage Examples

Shuffle Adapter Lifecycle (Internal)

import ray

# 1. Create shuffle actors
actors = [
    ShuffleStageAdapter.remote(stage=my_shuffle_stage, rank=i, nranks=n, num_input_tasks=total)
    for i in range(n)
]

# 2. Setup UCXX communication
root_address = ray.get(actors[0].setup_root.remote())
ray.get([a.setup.remote(root_address) for a in actors])

# 3. Insert data
ray.get([a.read_and_insert.remote(task_batch) for a, task_batch in zip(actors, task_batches)])

# 4. Signal insertion complete
ray.get([a.insert_finished.remote() for a in actors])

# 5. Extract shuffled results
results = ray.get([a.extract_and_write.remote() for a in actors])

# 6. Teardown
ray.get([a.teardown.remote() for a in actors])

Related Pages

Environment:NVIDIA_NeMo_Curator_Python_Linux_Base
NVIDIA_NeMo_Curator_Backend_Base_Classes -- Parent base classes
NVIDIA_NeMo_Curator_RayActorPoolAdapter -- Base actor pool adapter
NVIDIA_NeMo_Curator_LSHStage -- LSH stage that uses this adapter for band shuffling

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment