> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/curator/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/curator/_mcp/server.

# nemo_curator.stages.base

## Module Contents

### Classes

| Name                                                           | Description                                                       |
| -------------------------------------------------------------- | ----------------------------------------------------------------- |
| [`CompositeStage`](#nemo_curator-stages-base-CompositeStage)   | Base class for high-level composite stages.                       |
| [`ProcessingStage`](#nemo_curator-stages-base-ProcessingStage) | Base class for all processing stages.                             |
| [`StageMeta`](#nemo_curator-stages-base-StageMeta)             | Metaclass that automatically registers concrete Stage subclasses. |
| [`_UnsetType`](#nemo_curator-stages-base-_UnsetType)           | -                                                                 |

### Functions

| Name                                                                   | Description                                            |
| ---------------------------------------------------------------------- | ------------------------------------------------------ |
| [`_num_workers_method`](#nemo_curator-stages-base-_num_workers_method) | -                                                      |
| [`_stage_spec_method`](#nemo_curator-stages-base-_stage_spec_method)   | -                                                      |
| [`get_stage_class`](#nemo_curator-stages-base-get_stage_class)         | Retrieve a registered stage class by its *class name*. |

### Data

[`X`](#nemo_curator-stages-base-X)

[`Y`](#nemo_curator-stages-base-Y)

[`_STAGE_REGISTRY`](#nemo_curator-stages-base-_STAGE_REGISTRY)

[`_UNSET`](#nemo_curator-stages-base-_UNSET)

### API

```python
class nemo_curator.stages.base.CompositeStage()
```

Abstract

**Bases:** [ProcessingStage\[X, Y\]](#nemo_curator-stages-base-ProcessingStage)

Base class for high-level composite stages.

Composite stages are user-facing stages that decompose into multiple
low-level execution stages during pipeline planning. They provide a
simplified API while maintaining fine-grained control at execution time.

Composite stages never actually execute - they only exist to be decomposed
into their constituent execution stages.

**`_with_operations`** `= []`

---

```python
nemo_curator.stages.base.CompositeStage._apply_with_(
    stages: list[nemo_curator.stages.base.ProcessingStage]
) -> list[nemo_curator.stages.base.ProcessingStage]
```

Apply configuration changes to this stage.

```python
nemo_curator.stages.base.CompositeStage.decompose() -> list[nemo_curator.stages.base.ProcessingStage]
```

abstract

Decompose into execution stages.

This method must be implemented by composite stages to define
what low-level stages they represent.

Returns (list\[ProcessingStage]):
List of execution stages that will actually run

```python
nemo_curator.stages.base.CompositeStage.decompose_and_apply_with() -> list[nemo_curator.stages.base.ProcessingStage]
```

Decompose and apply configuration changes to this stage.

```python
nemo_curator.stages.base.CompositeStage.get_description() -> str
```

Get a description of what this composite stage does.

Override this to provide user-friendly documentation.

```python
nemo_curator.stages.base.CompositeStage.inputs() -> tuple[list[str], list[str]]
```

Get the inputs for this stage.

```python
nemo_curator.stages.base.CompositeStage.outputs() -> tuple[list[str], list[str]]
```

Get the outputs for this stage.

```python
nemo_curator.stages.base.CompositeStage.process(
    task: nemo_curator.stages.base.X
) -> nemo_curator.stages.base.Y | list[nemo_curator.stages.base.Y]
```

Composite stages should never be executed directly.

```python
nemo_curator.stages.base.CompositeStage.with_(
    stage_with_dict: dict[str, typing.Any]
) -> nemo_curator.stages.base.CompositeStage
```

Apply configuration changes to this stage.

```python
class nemo_curator.stages.base.ProcessingStage()
```

Abstract

**Bases:** `Generic[X, Y]`

Base class for all processing stages.
Processing stages operate on Task objects (or subclasses like DocumentBatch).
Each stage type can declare what type of Task it processes as input (X)
and what type it produces as output (Y).
Stages can return either:

* A single task (typical for transformations)
* A list of tasks (for stages that split work, like readers)
* None (for filtered out tasks)

**`_batch_size`** `int | None`

Number of tasks to process in a batch.

---

**`_name`** `str`

---

**`_resources`** `Resources`

---

**`batch_size`** `= 1`

---

**`is_resumable`** `bool = True`

---

**`is_sink_stage`** `bool = False`

---

**`is_source_stage`** `bool = False`

---

**`name`** `= 'ProcessingStage'`

---

**`resources`** `= Resources(cpus=1.0)`

---

**`runtime_env`** `dict[str, Any] | None = None`

---

```python
nemo_curator.stages.base.ProcessingStage.__init_subclass__(
    kwargs = {}
)
```

```python
nemo_curator.stages.base.ProcessingStage.__repr__() -> str
```

String representation of the stage.

```python
nemo_curator.stages.base.ProcessingStage._annotation_contains_list(
    annotation: object
) -> bool
```

classmethod

```python
nemo_curator.stages.base.ProcessingStage._consume_custom_metrics() -> dict[str, float]
```

Return and clear metrics recorded during the last process call.

```python
nemo_curator.stages.base.ProcessingStage._log_metric(
    name: str,
    value: float
) -> None
```

```python
nemo_curator.stages.base.ProcessingStage._log_metrics(
    metrics: dict[str, float]
) -> None
```

Record custom metrics for this stage (e.g., sub-stage timings).

```python
nemo_curator.stages.base.ProcessingStage._time_metric(
    name: str
) -> contextlib.AbstractContextManager[None]
```

Record elapsed time for a code block as a custom stage metric.

```python
nemo_curator.stages.base.ProcessingStage.get_config() -> dict[str, typing.Any]
```

Get configuration for this stage.
Returns (dict\[str, Any]):
Dictionary containing configuration for this stage

```python
nemo_curator.stages.base.ProcessingStage.inputs() -> tuple[list[str], list[str]]
```

Define stage input requirements.

Returns (tuple\[list\[str], list\[str]]):
Tuple of (required\_attributes, required\_columns) where:

* required\_top\_level\_attributes: List of task attributes that must be present
* required\_data\_attributes: List of attributes within the data that must be present

```python
nemo_curator.stages.base.ProcessingStage.is_fanout_stage() -> bool
```

Infer whether `process()` can fan out one input task into many outputs.

```python
nemo_curator.stages.base.ProcessingStage.num_workers() -> int | None
```

Number of workers required. If None, then executor will determine the number of workers.

```python
nemo_curator.stages.base.ProcessingStage.outputs() -> tuple[list[str], list[str]]
```

Define stage output specification.

Returns (tuple\[list\[str], list\[str]]):
Tuple of (output\_attributes, output\_columns) where:

* output\_top\_level\_attributes: List of task attributes this stage adds/modifies
* output\_data\_attributes: List of attributes within the data that this stage adds/modifies

```python
nemo_curator.stages.base.ProcessingStage.process(
    task: nemo_curator.stages.base.X
) -> nemo_curator.stages.base.Y | list[nemo_curator.stages.base.Y]
```

abstract

Process a task and return the result.
Args:
task (X): Input task to process
Returns (Y | list\[Y]):

* Single task: For 1-to-1 transformations
* List of tasks: For 1-to-many transformations (e.g., readers)
* None: If the task should be filtered out

```python
nemo_curator.stages.base.ProcessingStage.process_batch(
    tasks: list[nemo_curator.stages.base.X]
) -> list[nemo_curator.stages.base.Y]
```

Process a batch of tasks and return results.
Override this method to enable batch processing for your stage.
If not overridden, the stage will only support single-task processing.
Args:
tasks (list\[X]): List of input tasks to process
Returns (list\[Y]):
List of results, where each result can be:

* Single task: For 1-to-1 transformations
* List of tasks: For 1-to-many transformations
* None: If the task should be filtered out
  Note: The returned list should have the same length as the input list,
  with each element corresponding to the result of processing the task
  at the same index.

`task_id` is framework-owned: stages must NOT set it. The executor
adapter (`BaseStageAdapter._post_process_task_ids`) assigns a
deterministic id to every emitted task — regardless of whether
a stage uses this default or overrides `process_batch`. Where the
input→output mapping is ambiguous (e.g. a batch aggregation), the
adapter falls back to a random `"r"`-prefixed id (see
`Task.task_id`); there is no way for a stage to supply its own.

```python
nemo_curator.stages.base.ProcessingStage.ray_stage_spec() -> dict[str, typing.Any]
```

Get Ray configuration for this stage.
Note : This is only used for Ray Data backend.
The keys are defined in RayStageSpecKeys in backends/ray\_data/utils.py

Returns (dict\[str, Any]):
Dictionary containing Ray-specific configuration

```python
nemo_curator.stages.base.ProcessingStage.setup(
    worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None
) -> None
```

Setup method called once before processing begins.
Override this method to perform any initialization that should
happen once per worker.
Args:
worker\_metadata (WorkerMetadata, optional): Information about the worker (provided by some backends)

```python
nemo_curator.stages.base.ProcessingStage.setup_on_node(
    node_info: nemo_curator.backends.base.NodeInfo | None = None,
    worker_metadata: nemo_curator.backends.base.WorkerMetadata | None = None
) -> None
```

Setup method called once per node in distributed settings.
Override this method to perform node-level initialization.
Args:
node\_info (NodeInfo, optional): Information about the node (provided by some backends)
worker\_metadata (WorkerMetadata, optional): Information about the worker (provided by some backends)

```python
nemo_curator.stages.base.ProcessingStage.supports_batch_processing() -> bool
```

Whether this stage supports vectorized batch processing.
This is automatically determined by checking if the stage has
overridden the process\_batch method from the base class.

```python
nemo_curator.stages.base.ProcessingStage.teardown() -> None
```

Teardown method called once after processing ends.
Override this method to perform any cleanup.

```python
nemo_curator.stages.base.ProcessingStage.validate_input(
    task: nemo_curator.tasks.Task
) -> bool
```

Validate input task meets requirements.
Args:
task: Task to validate
Returns:
True if valid, False otherwise

```python
nemo_curator.stages.base.ProcessingStage.with_(
    name: str | None = None,
    resources: nemo_curator.stages.resources.Resources | None = None,
    batch_size: int | None = None,
    runtime_env: dict[str, typing.Any] | None = None,
    ray_stage_spec: dict[str, typing.Any] | None = None,
    xenna_stage_spec: dict[str, typing.Any] | None = None,
    num_workers: int | None | nemo_curator.stages.base._UnsetType = _UNSET
) -> nemo_curator.stages.base.ProcessingStage
```

Apply configuration changes to this stage with overridden properties.

Note: This method uses class-level attributes and instance attributes interchangeably which can sometimes
lead to unexpected behavior. Please see [https://github.com/NVIDIA-NeMo/Curator/pull/764](https://github.com/NVIDIA-NeMo/Curator/pull/764) for more details.
Args:
name: Override the name property
resources: Override the resources property
batch\_size: Override the batch\_size property
runtime\_env: Override the runtime\_env (Ray runtime environment dict)
ray\_stage\_spec: Merge overrides into the Ray stage spec. User-provided keys win.
xenna\_stage\_spec: Merge overrides into the Xenna stage spec. User-provided keys win.
Use num\_workers instead of setting num\_workers in xenna\_stage\_spec.
num\_workers: Override the num\_workers() result. Passing None explicitly resets to executor default behavior.

```python
nemo_curator.stages.base.ProcessingStage.xenna_stage_spec() -> dict[str, typing.Any]
```

Get Xenna configuration for this stage.

Returns (dict\[str, Any]):
Dictionary containing Xenna-specific configuration

```python
class nemo_curator.stages.base.StageMeta()
```

Abstract

**Bases:** `ABCMeta`

Metaclass that automatically registers concrete Stage subclasses.
A class is considered *concrete* if it directly inherits from
:class:`ProcessingStage` **and** implements a `name` property.  Abstract
helper classes (e.g. *ProcessingStage* itself) will not be added to the
registry because they have the `_is_abstract` attribute set.

```python
nemo_curator.stages.base.StageMeta.__new__(
    mcls,
    name,
    bases,
    namespace,
    kwargs = {}
)
```

```python
class nemo_curator.stages.base._UnsetType()
```

**`__slots__`** `= ()`

---

```python
nemo_curator.stages.base._num_workers_method(
    num_workers: int | None
) -> collections.abc.Callable[[], int | None]
```

```python
nemo_curator.stages.base._stage_spec_method(
    stage_spec: dict[str, typing.Any]
) -> collections.abc.Callable[[], dict[str, typing.Any]]
```

```python
nemo_curator.stages.base.get_stage_class(
    name: str
) -> type[nemo_curator.stages.base.ProcessingStage]
```

Retrieve a registered stage class by its *class name*.
Raises
------

KeyError
If no stage with that name is registered.

```python
nemo_curator.stages.base.X = TypeVar('X', bound=Task)
```

```python
nemo_curator.stages.base.Y = TypeVar('Y', bound=Task)
```

```python
nemo_curator.stages.base._STAGE_REGISTRY: dict[str, type[ProcessingStage]] = {}
```

```python
nemo_curator.stages.base._UNSET = _UnsetType()
```