Move ray actor pool out of experimental - #1748
Conversation
|
/ok to test b25632f |
Greptile SummaryThis PR promotes
Confidence Score: 4/5Not 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
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
|
sarahyurick
left a comment
There was a problem hiding this comment.
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: |
There was a problem hiding this comment.
Is this test class being re-added somewhere?
praateekmahajan
left a comment
There was a problem hiding this comment.
@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>
c8ec70c to
b65410d
Compare
| for node in ray.nodes(): | ||
| if is_head_node(node): | ||
| _HEAD_NODE_ID_CACHE = node["NodeID"] | ||
| return _HEAD_NODE_ID_CACHE | ||
|
|
||
| return None | ||
|
|
||
|
|
There was a problem hiding this comment.
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.
| 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" |
|
Will create a new PR. It's getting harder to deal with ruff issues. And won't run ruff on all files. |
|
New PR: #1754. Closing this one. |
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