Skip to content

Move ray actor pool out of experimental - #1748

Closed
oyilmaz-nvidia wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
oyilmaz-nvidia:onur/move-ray-actor-pool-from-experimental
Closed

Move ray actor pool out of experimental#1748
oyilmaz-nvidia wants to merge 1 commit into
NVIDIA-NeMo:mainfrom
oyilmaz-nvidia:onur/move-ray-actor-pool-from-experimental

Conversation

@oyilmaz-nvidia

@oyilmaz-nvidia oyilmaz-nvidia commented Apr 6, 2026

Copy link
Copy Markdown
Contributor

Description

Summary

Moves RayActorPoolExecutor from nemo_curator.backends.experimental.ray_actor_pool to nemo_curator.backends.ray_actor_pool, following the same pattern as #1619 which promoted RayDataExecutor
Merges nemo_curator.backends.experimental.utils into nemo_curator.backends.utils, since those utilities were already being used by non-experimental code (RayDataExecutor, stages, tutorials)
Removes the now-empty nemo_curator/backends/experimental/ package entirely
Updates all import paths across stages, tests, benchmarking scripts, and documentation

Changes

nemo_curator/backends/ray_actor_pool/ — moved from experimental/ray_actor_pool/
nemo_curator/backends/utils.py — merged in RayStageSpecKeys, execute_setup_on_node, get_available_cpu_gpu_resources, and related helpers from the deleted experimental/utils.py
tests/backends/ray_actor_pool/ — moved from tests/backends/experimental/ray_actor_pool/
tests/backends/test_utils.py — moved from tests/backends/experimental/test_utils.py
Import updates across 30+ files in stages/, tests/, benchmarking/, and fern/ docs
nemo_curator/backends/experimental/ — deleted (no longer needed)

NOTE:

Ran ruff format and check on the repo. So there are some files with just format update. FYI.

Checklist

  • I am familiar with the Contributing Guide.
  • New or Existing tests cover these changes.
  • The documentation is up to date with these changes.

@copy-pr-bot

copy-pr-bot Bot commented Apr 6, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@oyilmaz-nvidia
oyilmaz-nvidia marked this pull request as ready for review April 6, 2026 22:14
@oyilmaz-nvidia

Copy link
Copy Markdown
Contributor Author

/ok to test b25632f

@greptile-apps

greptile-apps Bot commented Apr 6, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR promotes RayActorPoolExecutor from nemo_curator.backends.experimental to nemo_curator.backends, deletes the experimental/ sub-package, and updates ~30+ import sites. The refactor is mostly clean, but two breaking bugs were introduced during the merge of experimental/utils.py into backends/utils.py.

  • RAY_REMOTE_ARGS omitted from RayStageSpecKeysnemo_curator/backends/utils.py is missing this enum entry that was present in experimental/utils.py. RayDataStageAdapter accesses RayStageSpecKeys.RAY_REMOTE_ARGS at runtime, causing an immediate AttributeError.
  • Broken test importtests/backends/ray_data/test_utils.py still imports reset_head_node_cache from the now-deleted tests.backends.experimental.test_utils, causing ModuleNotFoundError on collection."

Confidence Score: 4/5

Not safe to merge as-is — two P1 regressions exist: a missing enum key that breaks RayData execution and a broken test import.

Score of 4 reflects that the overall restructuring is correct and well-executed, but two concrete defects in the changed code need to be fixed before merging.

nemo_curator/backends/utils.py (missing RAY_REMOTE_ARGS) and tests/backends/ray_data/test_utils.py (stale import path)

Important Files Changed

Filename Overview
nemo_curator/backends/utils.py Merged experimental utilities into this file but omitted RAY_REMOTE_ARGS from RayStageSpecKeys, breaking RayDataStageAdapter at runtime
nemo_curator/backends/ray_data/adapter.py Import updated from experimental to backends/utils; accesses RayStageSpecKeys.RAY_REMOTE_ARGS which is now missing from the enum
tests/backends/ray_data/test_utils.py Still imports reset_head_node_cache from deleted tests.backends.experimental.test_utils module, causing ModuleNotFoundError
nemo_curator/backends/experimental/utils.py Stray file — no init.py sibling, so this directory/file was incompletely removed
nemo_curator/backends/ray_actor_pool/init.py New location for RayActorPoolExecutor package; correctly exports executor
nemo_curator/backends/ray_actor_pool/executor.py Moved from experimental with correct import updates to nemo_curator.backends.utils
tests/backends/test_utils.py Merged experimental test_utils into this file; correct imports from nemo_curator.backends.utils
nemo_curator/backends/ray_data/executor.py Import of execute_setup_on_node correctly consolidated into nemo_curator.backends.utils
nemo_curator/backends/ray_data/utils.py Import of get_available_cpu_gpu_resources correctly updated to nemo_curator.backends.utils

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    subgraph Before
        E1[nemo_curator.backends.experimental.ray_actor_pool] -->|executor| RAP[RayActorPoolExecutor]
        E2[nemo_curator.backends.experimental.utils] -->|RayStageSpecKeys\nexecute_setup_on_node\nget_available_cpu_gpu_resources\nRAY_REMOTE_ARGS ✓| EU[experimental utils]
        E3[nemo_curator.backends.experimental.ray_data] -->|executor| RDE[RayDataExecutor]
    end
    subgraph After
        N1[nemo_curator.backends.ray_actor_pool] -->|executor| RAP2[RayActorPoolExecutor]
        N2[nemo_curator.backends.utils] -->|RayStageSpecKeys\nexecute_setup_on_node\nget_available_cpu_gpu_resources\nRAY_REMOTE_ARGS ✗ MISSING| NU[backends utils]
        N3[nemo_curator.backends.ray_data] -->|executor| RDE2[RayDataExecutor]
        STRAY[nemo_curator/backends/experimental/utils.py\nstray file - no __init__.py] -.->|unreachable| N2
    end
    N2 -->|imported by| adapter[ray_data/adapter.py\nuses RAY_REMOTE_ARGS → AttributeError]
    style N2 fill:#f88,stroke:#f00
    style adapter fill:#f88,stroke:#f00
    style STRAY fill:#ffa,stroke:#aa0
Loading

Comments Outside Diff (1)

  1. tests/backends/ray_data/test_utils.py, line 25 (link)

    P1 Broken import — module deleted in this PR

    This line imports reset_head_node_cache from tests.backends.experimental.test_utils, but tests/backends/experimental/ was deleted in this PR. The test file will fail to load with ModuleNotFoundError. The equivalent fixture now lives in tests/backends/test_utils.py.

Reviews (3): Last reviewed commit: "Fix merge conflicts" | Re-trigger Greptile

@sarahyurick sarahyurick left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm fine with the style changes but why aren't they being flagged by our Ruff job? Is there some version mismatch?

from nemo_curator.backends.utils import merge_executor_configs


class TestMergeExecutorConfig:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this test class being re-added somewhere?

@praateekmahajan praateekmahajan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@oyilmaz-nvidia can you please double check if your ruff version is 0.14.10 and looks like we have some failing tests?

Signed-off-by: Onur Yilmaz <oyilmaz@nvidia.com>
@oyilmaz-nvidia
oyilmaz-nvidia force-pushed the onur/move-ray-actor-pool-from-experimental branch from c8ec70c to b65410d Compare April 7, 2026 16:35
Comment on lines +144 to +151
for node in ray.nodes():
if is_head_node(node):
_HEAD_NODE_ID_CACHE = node["NodeID"]
return _HEAD_NODE_ID_CACHE

return None


Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Missing RAY_REMOTE_ARGS key in RayStageSpecKeys

The old experimental/utils.py contained RAY_REMOTE_ARGS = "ray_remote_args" in RayStageSpecKeys, but it was not included when merging into nemo_curator/backends/utils.py. nemo_curator/backends/ray_data/adapter.py (line 106) accesses RayStageSpecKeys.RAY_REMOTE_ARGS, which will raise AttributeError: 'RAY_REMOTE_ARGS' is not a valid RayStageSpecKeys at runtime. Additionally, the key-validation loop in RayDataStageAdapter.__init__ will reject any stage spec containing ray_remote_args.

Suggested change
for node in ray.nodes():
if is_head_node(node):
_HEAD_NODE_ID_CACHE = node["NodeID"]
return _HEAD_NODE_ID_CACHE
return None
class RayStageSpecKeys(str, Enum):
"""String enum of different flags that define keys inside ray_stage_spec."""
IS_ACTOR_STAGE = "is_actor_stage"
IS_FANOUT_STAGE = "is_fanout_stage"
IS_RAFT_ACTOR = "is_raft_actor"
IS_LSH_STAGE = "is_lsh_stage"
IS_SHUFFLE_STAGE = "is_shuffle_stage"
MAX_CALLS_PER_WORKER = "max_calls_per_worker"
RAY_REMOTE_ARGS = "ray_remote_args"

@oyilmaz-nvidia

Copy link
Copy Markdown
Contributor Author

Will create a new PR. It's getting harder to deal with ruff issues. And won't run ruff on all files.

@oyilmaz-nvidia

Copy link
Copy Markdown
Contributor Author

New PR: #1754. Closing this one.

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.

3 participants