Skip to content

Add num_workers_per_node() as a method for stages + Per-node worker sizing for download stages#2198

Draft
praateekmahajan wants to merge 5 commits into
NVIDIA-NeMo:mainfrom
praateekmahajan:praateek/ray_data_max_workers_per_node
Draft

Add num_workers_per_node() as a method for stages + Per-node worker sizing for download stages#2198
praateekmahajan wants to merge 5 commits into
NVIDIA-NeMo:mainfrom
praateekmahajan:praateek/ray_data_max_workers_per_node

Conversation

@praateekmahajan

@praateekmahajan praateekmahajan commented Jul 11, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Add ProcessingStage.num_workers_per_node() (method-only) and with_(num_workers_per_node=...).
  • Enforce at most one worker-sizing option (num_workers / num_workers_per_node / actor-pool min/max/initial_workers) once at stage construction, and drop the duplicated checks from the Ray Data, Ray Actor Pool, and Xenna backends.
  • Translate per-node sizing across all three backends using shared alive-Ray-node helpers, and apply default per-node limits to download-like stages.
  • Add matching backend/stage tests.

Behavior/API changes

  • Removed the legacy xenna_stage_spec()["num_workers_per_node"] path and the ImageDuplicatesRemovalStage(num_workers_per_node=...) constructor arg; use .with_(num_workers_per_node=...) instead (migrated ImageDuplicatesRemovalStage and the CallHome tutorial).

Fixes #2197

Signed-off-by: Praateek <praateekm@gmail.com>
@copy-pr-bot

copy-pr-bot Bot commented Jul 11, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@praateekmahajan

Copy link
Copy Markdown
Contributor Author

/ok to test 82603dd

Fold _compute_kwargs and _build_resource_kwargs into a single
_managed_map_batches_kwargs that builds the full Curator-managed set
(compute strategy + CPU/GPU reservations + max_calls) in one place.

Signed-off-by: Praateek <praateekm@gmail.com>
@praateekmahajan

Copy link
Copy Markdown
Contributor Author

/ok to test 59e866e

The value is num_workers_per_node * node_count, i.e. the total cluster-wide
Ray Data pool size (ActorPoolStrategy/TaskPoolStrategy size is a total replica
count), not a per-node count. SPREAD scheduling then distributes it so each node
gets ~num_workers_per_node replicas. Rename for clarity and document the scaling.

Signed-off-by: Praateek <praateekm@gmail.com>
@praateekmahajan

Copy link
Copy Markdown
Contributor Author

/ok to test fc50118

…ax_workers_per_node

Signed-off-by: Praateek <praateekm@gmail.com>
@praateekmahajan

Copy link
Copy Markdown
Contributor Author

/ok to test a90f090

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Ray Data backend should cap download-stage concurrency by node count and spread workers

1 participant