Stage Worker Sizing

Use ProcessingStage.num_workers() or stage.with_(num_workers=...) to set a backend-neutral worker request. Use ray_stage_spec and xenna_stage_spec for backend-specific execution behavior.

Worker count and per-worker resources are separate controls:

• num_workers() controls how many workers the executor requests.
• Resources controls CPU and GPU resources reserved for each worker.
• ray_stage_spec() and xenna_stage_spec() provide backend-specific scheduling options.

Common Stage for the Examples

The examples on this page configure this minimal stage:

1
from nemo_curator.stages.base import ProcessingStage
2
from nemo_curator.stages.resources import Resources
3
from nemo_curator.tasks import EmptyTask
4
5
6
class PassthroughStage(ProcessingStage[EmptyTask, EmptyTask]):
7
    name = "passthrough"
8
    resources = Resources(cpus=1.0)
9
10
    def process(self, task: EmptyTask) -> EmptyTask:
11
        return task

Worker Controls by Backend

Ray Data and Ray Actor Pool also treat 0 and negative values as no fixed-size request. Xenna passes every non-None value to its stage specification, so use a positive count or None; do not use zero or negative Xenna worker counts.

Ray Data Worker Pools

Ray Data uses actor stages for stateful setup or GPU-backed work and task stages for stateless work. A stage is an actor stage when either condition is true:

• ray_stage_spec[RayStageSpecKeys.IS_ACTOR_STAGE] is explicitly True.
• The stage overrides setup(), or requests both CPU and GPU resources.

Set IS_ACTOR_STAGE to False to force task execution. Do this only when the stage does not depend on worker-local setup state.

Fixed actor pool

Set a positive num_workers and mark the stage as an actor stage:

1
from nemo_curator.backends.ray_data import RayDataExecutor
2
from nemo_curator.backends.utils import RayStageSpecKeys
3
4
stage = PassthroughStage().with_(
5
    num_workers=4,
6
    ray_stage_spec={
7
        RayStageSpecKeys.IS_ACTOR_STAGE: True,
8
    },
9
)
10
11
executor = RayDataExecutor()

Ray Data receives ActorPoolStrategy(size=4). Cluster resources must be able to place four actors with the stage’s per-worker Resources.

Autoscaling actor pool

Use the exact RayStageSpecKeys members MIN_WORKERS, MAX_WORKERS, and INITIAL_WORKERS:

1
from nemo_curator.backends.ray_data import RayDataExecutor
2
from nemo_curator.backends.utils import RayStageSpecKeys
3
4
stage = PassthroughStage().with_(
5
    # Explicitly clear an inherited fixed worker count, if the stage has one.
6
    num_workers=None,
7
    ray_stage_spec={
8
        RayStageSpecKeys.IS_ACTOR_STAGE: True,
9
        RayStageSpecKeys.MIN_WORKERS: 2,
10
        RayStageSpecKeys.MAX_WORKERS: 8,
11
        RayStageSpecKeys.INITIAL_WORKERS: 4,
12
    },
13
)
14
15
executor = RayDataExecutor()

Ray Data receives ActorPoolStrategy(min_size=2, max_size=8, initial_size=4).

Ray validates the relationship among the three values. For example, an initial size greater than the maximum raises ValueError with the stage name and the underlying actor-pool validation message.

Fixed task pool

For a stateless task stage, a positive worker count selects TaskPoolStrategy(size=N):

1
from nemo_curator.backends.ray_data import RayDataExecutor
2
from nemo_curator.backends.utils import RayStageSpecKeys
3
4
stage = PassthroughStage().with_(
5
    num_workers=6,
6
    ray_stage_spec={
7
        RayStageSpecKeys.IS_ACTOR_STAGE: False,
8
    },
9
)
10
11
executor = RayDataExecutor()

Without a positive num_workers, a task stage does not set compute; Ray Data uses its default task scheduler.

Ray Data precedence and errors

Xenna Worker Counts

Xenna supports either a cluster-wide fixed count or a per-node count, but not both on one stage.

1
from nemo_curator.backends.xenna import XennaExecutor
2
3
stage = PassthroughStage().with_(num_workers=12)
4
executor = XennaExecutor()

The following combinations are invalid:

When neither setting is present, Xenna controls worker allocation using stage resources and pipeline demand.

Ray Actor Pool Resource Capping

RayActorPoolExecutor calculates how many actors fit from available CPU and GPU resources. A positive num_workers() is honored when it fits. If it exceeds capacity, Curator warns and caps the actor count instead of waiting for impossible placements.

For example, a stage that requests Resources(cpus=2.0) and num_workers=8 can place at most four actors when eight CPUs are available. Work batch count can further reduce the automatically chosen pool size when num_workers is unset.

Configure a Stage with with_()

with_() deep-copies the stage and returns the configured copy. It can override portable properties and merge backend stage specs in one call:

1
from nemo_curator.backends.utils import RayStageSpecKeys
2
from nemo_curator.stages.resources import Resources
3
4
configured = PassthroughStage().with_(
5
    name="gpu_passthrough",
6
    resources=Resources(cpus=2.0, gpus=1.0),
7
    batch_size=8,
8
    runtime_env={"pip": ["transformers==4.45.0"]},
9
    num_workers=4,
10
    ray_stage_spec={
11
        RayStageSpecKeys.RAY_NUM_CPUS: 1.0,
12
        RayStageSpecKeys.RAY_REMOTE_ARGS: {
13
            "scheduling_strategy": "SPREAD",
14
        },
15
    },
16
    xenna_stage_spec={
17
        "worker_max_lifetime_m": 45,
18
    },
19
)

The original stage remains unchanged. resources, batch_size, runtime_env, and num_workers apply across supported executors. ray_stage_spec configures Ray Data and the Ray-specific flags used by Ray Actor Pool; xenna_stage_spec configures Xenna.

Shallow-merge behavior

ray_stage_spec and xenna_stage_spec are shallow-merged with the stage’s existing dictionary. New top-level keys are preserved, and user values replace existing values at the same top-level key. Nested dictionaries are replaced, not recursively merged.

For example, if a stage returns:

1
{
2
    RayStageSpecKeys.MAX_WORKERS: 8,
3
    RayStageSpecKeys.RAY_REMOTE_ARGS: {
4
        "scheduling_strategy": "SPREAD",
5
        "max_retries": 2,
6
    },
7
}

then this override:

1
stage.with_(
2
    ray_stage_spec={
3
        RayStageSpecKeys.MAX_WORKERS: None,
4
        RayStageSpecKeys.RAY_REMOTE_ARGS: {"max_retries": 0},
5
    }
6
)

produces a spec where max_workers is explicitly None and ray_remote_args contains only {"max_retries": 0}. The original scheduling_strategy nested key is not retained.

Omitted values and explicit None

Chained with_() calls merge against the already configured copy, so later values win while unrelated top-level keys from earlier calls remain.

Reserved num_workers Hook

num_workers is reserved as a method on ProcessingStage. Do not define it as a dataclass field or ordinary stage attribute:

1
class FixedStage(PassthroughStage):
2
    def num_workers(self) -> int:
3
        return 4

Alternatively, keep the class executor-neutral and use stage.with_(num_workers=4) at pipeline construction time. Defining num_workers: int = 4 as a class or dataclass field raises TypeError when the subclass is created.

Some built-in single-input and fanout stages override num_workers() to return 1, preventing duplicate source enumeration. Use with_(num_workers=None) only when executor-controlled parallelism is safe for that stage’s input contract.

For ASR stages, the stage-local NeMo DataLoader setting is named dataloader_num_workers. When migrating an ASR constructor or configuration, rename the previous stage field num_workers to dataloader_num_workers; the common num_workers() hook is now reserved for Curator backend workers:

1
from nemo_curator.stages.audio.tagging.inference.nemo_asr_align import NeMoASRAlignerStage
2
3
asr_stage = NeMoASRAlignerStage(
4
    dataloader_num_workers=10,
5
).with_(num_workers=2)

Resource Matching and Ray Data Fusion

Ray Data may fuse compatible adjacent operators to avoid materializing and serializing intermediate blocks. Fusion is an optimizer decision, not a guaranteed stage behavior.

Ray Data determines the per-worker reservation in this order:

Use RAY_NUM_CPUS when Ray Data needs a different scheduling reservation from other executors:

1
stage = PassthroughStage().with_(
2
    resources=Resources(cpus=4.0),
3
    ray_stage_spec={
4
        RayStageSpecKeys.RAY_NUM_CPUS: 1.0,
5
    },
6
)

The CPU paths of VideoFrameExtractionStage and ClipTranscodingStage use a Ray Data reservation of 1.0 CPU by default while preserving their larger Resources.cpus values for Xenna. Matching the adjacent video operators’ Ray Data reservations makes fusion possible and can remove inter-stage serialization overhead. Different resources, execution constraints, or optimizer decisions can still keep stages separate; inspect the Ray Data execution plan or dashboard instead of assuming fusion occurred.

Related References

• ProcessingStage API
• Pipeline Execution Backends
• Per-Stage Runtime Environments
• Resources API