fix: honor NEMO_RL_PY_EXECUTABLES_SYSTEM for all actor environments - #4020
Conversation
github-actions
Bot
added
the
Documentation
e5db960 to
46e5cdb
Compare
yuki-97
left a comment
yuki-97
left a comment
There was a problem hiding this comment.
Reviewed by a team of specialized agents (RL/codebase, bug-finder, design, tests, devil's advocate). 3 inline comments below — one test gap worth acting on, plus two docs nits. No blockers.
The design call here is right. Resolving the flag at the PY_EXECUTABLES layer rather than in the registry is what makes "a plain PY_EXECUTABLES.X silently opts out" structurally impossible. Concretely, single_controller_utils/setup.py:1028 and research/template_project/single_update.py:60 write named constants into ACTOR_ENVIRONMENT_REGISTRY after it is built — no registry-layer check can reach those, which is why this isn't redundant with an alias-based approach. It also incidentally fixes trtllm_worker_async.py:68, where os.path.dirname(PY_EXECUTABLES.TRTLLM) was prepending the literal string uv run --locked --extra trtllm --directory /opt to PATH under the flag.
Verified clean: no remaining reference to the deleted nemo_rl/modelopt/registry.py or the *_EXECUTABLE aliases; all five ModelOpt FQNs byte-identical to the deleted file and resolving to real classes; 27 registry keys with no duplicates; the --extra modelopt --extra vllm|automodel|mcore combinations absent from the [tool.uv] conflicts table; pyrefly.toml removal correct; the new test's finally restore complete across both parametrizations.
Scope caveat: no tests were executed for this review (no uv/GPU available), and CI has not run on this PR yet — copy-pr-bot is still gating it. Conclusions come from reading the code plus standalone semantic models.
Generated by Claude Code
yuki-97
added
the
CI:Lfast
|
/ok to test 6e1e16d |
terrykong
left a comment
terrykong
left a comment
There was a problem hiding this comment.
Nice change — two things I liked, one thing to watch on the rebase, and a scheduling question for whoever merges these.
This needs a rebase, and the rebase is the risky part. Main has rewritten the same block since you branched. Git will show you a conflict only in the alias block you delete, and will quietly auto-merge the dict body to your version — which drops the nemo_gym extra from the two vLLM workers. The comment on ray_actor_environment_registry.py has the details and a one-command check to run after rebasing.
Resolving the flag inside PY_EXECUTABLES is the right place for it, not just the shorter one. Some code reads PY_EXECUTABLES.X without going through the registry at all — trtllm_worker_async.py reads PY_EXECUTABLES.TRTLLM directly — so a fix inside get_actor_python_env() would have missed those.
Folding nemo_rl/modelopt/registry.py into the one registry also removes a backwards import. The deleted file had to import git_root back out of virtual_cluster, and the registry imported the modelopt module at the bottom of the file. After this PR nemo_rl/distributed/ has no import edge into nemo_rl.modelopt, because the entries are plain strings. Central registration is what dynamo, nemo_gym, trtllm and sglang already do; modelopt was the only one off on its own.
Overlapping PRs — a scheduling question, not a complaint. What this PR does that the others do not is cover every actor environment in one place: jepio's PR (#3982) is the same fix by a different author, but its own description says ModelOpt keeps its separate handling of the flag, so #4020 is a superset. tdene's PR (#3947) and #4002 (mine) touch the same mapping from the dependency side. Two open PRs rewriting the same block will conflict, so someone has to pick which lands first.
Lint note: the pinned ruff 0.9.9 reports the four changed Python files clean, formatting included. That is a single-machine result rather than a CI claim — pre-commit needs the Automodel submodule, which was not checked out here.
Generated by Claude Code
Signed-off-by: Yuki Huang <yukih@nvidia.com> Co-authored-by: Jeremi Piotrowski <jpiotrowski@nvidia.com>
…T_REGISTRY, fix stale SGLANG_EXECUTABLE and registry line references Signed-off-by: Yuki Huang <yukih@nvidia.com>
…ert PY_EXECUTABLES directly in the system-flag probe, clarify NEMO_RL_PY_EXECUTABLES_SYSTEM scope in uv.md Signed-off-by: Yuki Huang <yukih@nvidia.com>
6e1e16d to
4926437
Compare
|
/ok to test 4926437 |
#4020 moved the flag into PY_EXECUTABLES._resolve_system_overrides and says callers should not check it themselves. That rewrite only touches the class attributes, so uv_py_executable -- which builds a fresh string -- was not covered, and this branch had to keep its own USE_SYSTEM_EXECUTABLE check. Put the check in uv_py_executable instead and drop the one in the registry, so the flag is handled in one place. Verified: with the flag set, all 28 actors resolve to the driver interpreter. Signed-off-by: Terry Kong <terryk@nvidia.com>
2 participants
What does this PR do ?
Resolves
NEMO_RL_PY_EXECUTABLES_SYSTEMinsidePY_EXECUTABLESitself, so every consumer of aPY_EXECUTABLES.*constant honors the flag instead of each call site re-deriving it.Before, the check was duplicated as seven ternaries in
ray_actor_environment_registry.pyand three more innemo_rl/modelopt/registry.py, and any entry written as a plainPY_EXECUTABLES.Xsilently opted out. Seven actors were in that state:DTensorPolicyWorker,DTensorPolicyWorkerV2,DTensorValueWorkerV2,AsyncTrajectoryCollector,ReplayBuffer,SyncRolloutActorandNemoGym.Changes:
PY_EXECUTABLES._resolve_system_overrides()rewrites every uv command constant toSYSTEMwhen the flag is set, and runs once at import.*_EXECUTABLEaliases inray_actor_environment_registry.pyare gone; every entry is now aPY_EXECUTABLES.*constant.nemo_rl/modelopt/registry.pyis deleted. Its three executables becomePY_EXECUTABLES.MODELOPT_VLLM/MODELOPT_AUTOMODEL/MODELOPT_MCOREand its five actors move intoACTOR_ENVIRONMENT_REGISTRY.docs/design-docs/uv.mddocuments the flag;dependency-management.mddrops a snippet referencing the removed aliases.Resolving at the
PY_EXECUTABLESlayer rather than at the registry also covers the runtime overrides that write into the registry after it is built, which a registry-level check cannot reach —single_controller_utils/setup.py:1030(PY_EXECUTABLES.VLLM_GYMundertoken_capture.enabled) andresearch/template_project/single_update.py:60.Issues
N/A
Usage
Set
NEMO_RL_PY_EXECUTABLES_SYSTEM=1when every actor's dependencies are already installed in the driver's interpreter, such as in a single-environment container image.Before your PR is "Ready for review"
Pre checks:
Additional Information
Overlaps with three open PRs that carry the same seven-actor behavior change: #3982 (alias-based), and #3947 / #4002 (which move the mapping into a separate actor-environments module). Worth deciding which lands first.
test_actor_registry_honors_system_flagintests/unit/distributed/test_virtual_cluster.pycovers both flag states throughACTOR_ENVIRONMENT_REGISTRY, in a subprocess so the flag is resolved at import time like it is in production.