nemo_curator.backends.ray_data.adapter

Module Contents

Classes

Name	Description
`RayDataStageAdapter`	Adapts ProcessingStage to Ray Data operations.

Functions

Name	Description
`create_actor_from_stage`	Create a StageProcessor class with the proper stage name for display.
`create_task_from_stage`	Create a named Ray Data stage adapter function.

Data

CURATOR_MANAGED_MAP_BATCHES_KWARGS

API

class nemo_curator.backends.ray_data.adapter.RayDataStageAdapter(
    stage: nemo_curator.stages.base.ProcessingStage
)

Bases: BaseStageAdapter

Adapts ProcessingStage to Ray Data operations.

This adapter converts stages to work with Ray Data datasets by:

Working directly with Task objects (no dictionary conversion)
Using Ray Data’s map_batches for parallel processing a. If stage has both gpus and cpus specified, then we use actors b. If stage.setup is overridden, then we use actors c. Else we use tasks

_batch_size

= self.stage.batch_size

batch_size

int | None

Get the batch size for this stage.

nemo_curator.backends.ray_data.adapter.RayDataStageAdapter._build_resource_kwargs(
    ray_stage_spec: dict
) -> dict[str, float]

Build num_cpus/num_gpus kwargs for map_batches.

Checks ray_stage_spec for RAY_NUM_CPUS first so stages can request a different CPU reservation for Ray Data (e.g. cpus=1.0 to enable stage fusion) without changing resources.cpus used by other executors.

nemo_curator.backends.ray_data.adapter.RayDataStageAdapter._process_batch_internal(
    batch: dict[str, typing.Any]
) -> dict[str, typing.Any]

Internal method that handles the actual batch processing logic.

Parameters:

batch

dict[str, Any]

Dictionary with arrays/lists representing a batch of Task objects

Returns: dict[str, Any]

Dictionary with arrays/lists representing processed Task objects

nemo_curator.backends.ray_data.adapter.RayDataStageAdapter.process_dataset(
    dataset: ray.data.Dataset
) -> ray.data.Dataset

Process a Ray Data dataset through this stage.

Parameters:

dataset

Dataset

Ray Data dataset containing Task objects

Returns: Dataset

Processed Ray Data dataset

nemo_curator.backends.ray_data.adapter.create_actor_from_stage(
    stage: nemo_curator.stages.base.ProcessingStage
) -> type[nemo_curator.backends.ray_data.adapter.RayDataStageAdapter]

Create a StageProcessor class with the proper stage name for display.

nemo_curator.backends.ray_data.adapter.create_task_from_stage(
    stage: nemo_curator.stages.base.ProcessingStage
) -> collections.abc.Callable[[dict[str, Any]], dict[str, typing.Any]]

Create a named Ray Data stage adapter function.

This creates a standalone function that wraps the stage processing logic with a clean name that doesn’t include the class qualification.

Parameters:

stage

ProcessingStage

Processing stage to adapt

Returns: Callable[[dict[str, Any]], dict[str, Any]]

A function that can be used directly with Ray Data’s map_batches

nemo_curator.backends.ray_data.adapter.CURATOR_MANAGED_MAP_BATCHES_KWARGS = {'compute', 'max_calls', 'num_cpus', 'num_gpus'}