Problem

For Ray Data pipelines with an explicit download stage, NeMo Curator currently lets the download stage scale like a normal CPU stage. This is a poor fit for internet / object-store downloads because download throughput is usually bounded by node-level network / disk capacity, not by total cluster CPU.

In a large Nemotron-CC Ray Data run, the download stage was a task operator that emitted small path-like outputs. Because it had low object-store pressure and was immediately dispatchable, Ray Data allowed it to run far ahead of the downstream iterate/extract stage:

• Cluster shown by Ray Data progress: 900 CPUs.
• DocumentDownloadStageTask reached 591 tasks by 35s.
• It later peaked at 1198 tasks / 599 CPU around 1:40:40.
• At that same 1:40:40 point, the downstream fused DocumentIterateExtractStageTask -> ScoreFilterActor operator had only 112 actors.
• By the final snapshot, download's object-store footprint was tiny relative to downstream materialized data: download ~52.6 MiB vs fused iterate/extract/score ~57.2 GiB.

This is suboptimal for download-heavy curation pipelines: a cheap producer can overfill the pipeline with downloaded paths/files, while the expensive downstream stages create the large dataframes/text blocks and control end-to-end throughput.

Why node-count scaling is the right model

For these download stages, more concurrent workers on the same node does not necessarily improve throughput linearly. Each node has finite network, disk, socket, and remote service capacity. A better default is:

download_workers ~= download_workers_per_node * num_ray_worker_nodes

For example, use 2 download workers per worker node rather than letting download fan out to hundreds of tasks just because CPU is available.

Ray limitation

Ray Data has native global concurrency controls:

ray.data.TaskPoolStrategy(size=n)
ray.data.ActorPoolStrategy(size=n)
ray.data.ActorPoolStrategy(min_size=m, max_size=n)

and Ray Core supports best-effort spread:

scheduling_strategy="SPREAD"

Note: Ray Data already uses SPREAD by default for normal map_batches tasks and actor creation. For task inputs larger than Ray Data's large-args threshold, Ray Data may use DEFAULT to preserve locality, and user-provided remote args can override the default. We should still set or document scheduling_strategy="SPREAD" explicitly for this policy so the placement intent is visible; the main behavior change is the num_workers_per_node * num_nodes global cap.

But Ray Data does not have a hard native per_node_concurrency=2 operator knob. A hard per-node cap would require custom resources or explicit node affinity. We do not want to require custom resources for users.

Proposed Curator fix

Add a Curator-level num_workers_per_node execution policy for download-like
stages. This is similar to the Xenna-style configuration where I/O stage
parallelism is expressed per node instead of as a raw global worker count.

1. Add stage spec keys such as:

   RayStageSpecKeys.NUM_WORKERS_PER_NODE = "num_workers_per_node"

2. Allow a stage to declare:

   {
       RayStageSpecKeys.NUM_WORKERS_PER_NODE: 2,
   }

3. In the Ray Data adapter, compute a global pool size from the current Ray
   cluster:

   num_nodes = count_alive_ray_nodes(exclude_head_node=...)
   size = num_workers_per_node * num_nodes

4. Then translate that into the appropriate Ray Data compute strategy:

   if stage_is_actor:
       map_batches_kwargs["compute"] = ray.data.ActorPoolStrategy(
           size=size,
       )
   else:
       map_batches_kwargs["compute"] = ray.data.TaskPoolStrategy(size=size)

5. Keep the placement intent explicit. Ray Data's default scheduling strategy is already SPREAD, but when num_workers_per_node is used the adapter should either leave that default visible in logs or set it explicitly unless the user overrides it:

   map_batches_kwargs.setdefault("scheduling_strategy", "SPREAD")

6. For DocumentDownloadStage / Common Crawl WARC download, use this policy by
   default or expose it as an option.

dataset.map_batches(
    DownloadActor,
    compute=ray.data.ActorPoolStrategy(
        size=num_workers_per_node * num_download_nodes,
        max_tasks_in_flight_per_actor=1,
    ),
    num_cpus=<small value>,
    scheduling_strategy="SPREAD",
)

dataset.map_batches(
    download_fn,
    compute=ray.data.TaskPoolStrategy(size=num_workers_per_node * num_download_nodes),
    num_cpus=<small value>,
    scheduling_strategy="SPREAD",
)

Suggested API shape

This could be generalized through ray_stage_spec, not hardcoded only for Common Crawl:

{
    "is_actor_stage": True,
    "num_workers_per_node": 2,
}

Then the Ray Data adapter can translate that into either:

• ActorPoolStrategy(size=num_workers_per_node * num_nodes, max_tasks_in_flight_per_actor=1) plus scheduling_strategy="SPREAD", or
• TaskPoolStrategy(size=num_workers_per_node * num_nodes) plus scheduling_strategy="SPREAD".