Skip to content

Add per-node worker sizing for download stages#2198

Draft
praateekmahajan wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
praateekmahajan:praateek/ray_data_max_workers_per_node
Draft

Add per-node worker sizing for download stages#2198
praateekmahajan wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
praateekmahajan:praateek/ray_data_max_workers_per_node

Conversation

@praateekmahajan

Copy link
Copy Markdown
Contributor

Summary

  • Add ProcessingStage.num_workers_per_node() and with_(num_workers_per_node=...), with base validation rejecting simultaneous num_workers and per-node sizing.
  • Translate per-node sizing in Ray Data, Ray Actor Pool, and Xenna while reusing shared alive-Ray-node helpers.
  • Apply default per-node limits to download-like stages and add matching backend/stage tests.

Fixes #2197

Test Plan

  • ruff check nemo_curator/stages/base.py nemo_curator/backends/utils.py nemo_curator/backends/xenna/executor.py nemo_curator/backends/ray_data/adapter.py nemo_curator/backends/ray_data/executor.py nemo_curator/backends/ray_actor_pool/utils.py nemo_curator/utils/ray_utils.py nemo_curator/stages/text/download/base/download.py nemo_curator/stages/text/download/common_crawl/download.py nemo_curator/stages/text/download/arxiv/download.py nemo_curator/stages/text/download/wikipedia/download.py nemo_curator/stages/video/io/video_reader.py tests/stages/common/test_base.py tests/backends/test_xenna_executor.py tests/backends/ray_data/test_adapter.py tests/backends/ray_actor_pool/test_executor.py tests/utils/test_ray_utils.py tests/stages/text/download/base/test_download.py tests/stages/text/download/common_crawl/test_download.py tests/stages/text/download/arxiv/test_download.py tests/stages/text/download/wikipedia/test_download.py tests/stages/video/io/test_video_reader.py
  • ruff format --check nemo_curator/stages/base.py nemo_curator/backends/utils.py nemo_curator/backends/xenna/executor.py nemo_curator/backends/ray_data/adapter.py nemo_curator/backends/ray_data/executor.py nemo_curator/backends/ray_actor_pool/utils.py nemo_curator/utils/ray_utils.py nemo_curator/stages/text/download/base/download.py nemo_curator/stages/text/download/common_crawl/download.py nemo_curator/stages/text/download/arxiv/download.py nemo_curator/stages/text/download/wikipedia/download.py nemo_curator/stages/video/io/video_reader.py tests/stages/common/test_base.py tests/backends/test_xenna_executor.py tests/backends/ray_data/test_adapter.py tests/backends/ray_actor_pool/test_executor.py tests/utils/test_ray_utils.py tests/stages/text/download/base/test_download.py tests/stages/text/download/common_crawl/test_download.py tests/stages/text/download/arxiv/test_download.py tests/stages/text/download/wikipedia/test_download.py tests/stages/video/io/test_video_reader.py
  • pytest tests/stages/common/test_base.py tests/backends/test_xenna_executor.py tests/backends/ray_data/test_adapter.py tests/backends/ray_actor_pool/test_executor.py tests/utils/test_ray_utils.py tests/stages/text/download/base/test_download.py tests/stages/text/download/common_crawl/test_download.py tests/stages/text/download/arxiv/test_download.py tests/stages/text/download/wikipedia/test_download.py tests/stages/text/download/base/test_url_generation.py tests/stages/video/io/test_video_reader.py -q

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

@copy-pr-bot copy-pr-bot Bot deployed to nemo-ci July 11, 2026 01:01 Active
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