diff --git a/docs/design-docs/dependency-management.md b/docs/design-docs/dependency-management.md index 39d87aebc24..9cf3d9c3be5 100644 --- a/docs/design-docs/dependency-management.md +++ b/docs/design-docs/dependency-management.md @@ -94,12 +94,12 @@ Within the driver script, NeMo RL starts multiple [`RayWorkerGroup`](https://git - **Generation workers** (e.g., vLLM): Require `vllm` dependencies - **Environment workers** (e.g., math evaluation): Use system/base dependencies -Each worker type is mapped to a specific Python executable configuration in the [`ACTOR_ENVIRONMENT_REGISTRY`](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/distributed/ray_actor_environment_registry.py#L27-L46). This registry defines which virtual environment should be used for each actor type: +Each worker type is mapped to a specific Python executable configuration in the [`ACTOR_ENVIRONMENT_REGISTRY`](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/distributed/ray_actor_environment_registry.py#L17-L55). This registry defines which virtual environment should be used for each actor type: ```python ACTOR_ENVIRONMENT_REGISTRY: dict[str, str] = { - "nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker": VLLM_EXECUTABLE, - "nemo_rl.models.policy.megatron_policy_worker.MegatronPolicyWorker": MCORE_EXECUTABLE, + "nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker": PY_EXECUTABLES.VLLM, + "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker": PY_EXECUTABLES.MCORE, "nemo_rl.environments.math_environment.MathEnvironment": PY_EXECUTABLES.SYSTEM, # ... more mappings } diff --git a/docs/design-docs/uv.md b/docs/design-docs/uv.md index 368774e0ffb..62b2bcb6973 100644 --- a/docs/design-docs/uv.md +++ b/docs/design-docs/uv.md @@ -59,6 +59,8 @@ class PY_EXECUTABLES: To ensure consistent dependencies between actors, we run with `--locked` to make sure the dependencies are consistent with the contents of `uv.lock`. +Setting the `NEMO_RL_PY_EXECUTABLES_SYSTEM=1` environment variable rewrites every `uv run` command in `PY_EXECUTABLES` — not only the ones shown here — to `SYSTEM`, so all actors launch on the driver's interpreter and no per-actor venv is created. Use it only in an environment where every actor's dependencies are already installed, such as a single-environment container image. + ### Customization If you need a different Python executable configuration, you can override the default one by passing your own in {py:class}`RayWorkerBuilder.__call__ `. This provides flexibility for special use cases without modifying the core configurations. diff --git a/nemo_rl/distributed/ray_actor_environment_registry.py b/nemo_rl/distributed/ray_actor_environment_registry.py index a2788019ab4..dbad277abba 100644 --- a/nemo_rl/distributed/ray_actor_environment_registry.py +++ b/nemo_rl/distributed/ray_actor_environment_registry.py @@ -12,40 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -import os - from nemo_rl.distributed.virtual_cluster import PY_EXECUTABLES -USE_SYSTEM_EXECUTABLE = os.environ.get("NEMO_RL_PY_EXECUTABLES_SYSTEM", "0") == "1" -# vLLM workers always get the vllm + nemo_gym extras. Token capture -# (token_capture.enabled) needs nemo_gym inside the worker, and worker venvs -# are cached by actor class name, so the extras must be fixed here rather than -# swapped in at runtime (a venv prebuilt with plain `--extra vllm` would be -# reused as-is and the nemo_gym import would fail). -VLLM_EXECUTABLE = ( - PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else PY_EXECUTABLES.VLLM_GYM -) -SGLANG_EXECUTABLE = ( - PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else PY_EXECUTABLES.SGLANG -) -MCORE_EXECUTABLE = ( - PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else PY_EXECUTABLES.MCORE -) -TRTLLM_EXECUTABLE = ( - PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else PY_EXECUTABLES.TRTLLM -) ACTOR_ENVIRONMENT_REGISTRY: dict[str, str] = { - "nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker": VLLM_EXECUTABLE, - "nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker": VLLM_EXECUTABLE, - "nemo_rl.models.generation.sglang.sglang_worker.SGLangGenerationWorker": SGLANG_EXECUTABLE, + "nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker": PY_EXECUTABLES.VLLM_GYM, + "nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker": PY_EXECUTABLES.VLLM_GYM, + "nemo_rl.models.generation.sglang.sglang_worker.SGLangGenerationWorker": PY_EXECUTABLES.SGLANG, + "nemo_rl.models.generation.trtllm.trtllm_worker_async.TrtllmAsyncGenerationWorker": PY_EXECUTABLES.TRTLLM, "nemo_rl.models.generation.dynamo.dynamo_worker.DynamoVllmWorker": PY_EXECUTABLES.SYSTEM, "nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker": PY_EXECUTABLES.FSDP, "nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2": PY_EXECUTABLES.AUTOMODEL, "nemo_rl.models.value.workers.dtensor_value_worker_v2.DTensorValueWorkerV2": PY_EXECUTABLES.AUTOMODEL, - "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker": MCORE_EXECUTABLE, - "nemo_rl.data.energon.sft_worker.SFTMegatronPolicyWorker": MCORE_EXECUTABLE, - "nemo_rl.models.value.workers.megatron_value_worker.MegatronValueWorker": MCORE_EXECUTABLE, - "nemo_rl.models.generation.trtllm.trtllm_worker_async.TrtllmAsyncGenerationWorker": TRTLLM_EXECUTABLE, + "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker": PY_EXECUTABLES.MCORE, + "nemo_rl.models.value.workers.megatron_value_worker.MegatronValueWorker": PY_EXECUTABLES.MCORE, + "nemo_rl.data.energon.sft_worker.SFTMegatronPolicyWorker": PY_EXECUTABLES.MCORE, "nemo_rl.environments.math_environment.MathEnvironment": PY_EXECUTABLES.SYSTEM, "nemo_rl.environments.math_environment.MathMultiRewardEnvironment": PY_EXECUTABLES.SYSTEM, "nemo_rl.environments.vlm_environment.VLMEnvironment": PY_EXECUTABLES.SYSTEM, @@ -66,12 +46,14 @@ "nemo_rl.experience.sync_rollout_actor.SyncRolloutActor": PY_EXECUTABLES.VLLM, "nemo_rl.environments.tools.retriever.RAGEnvironment": PY_EXECUTABLES.SYSTEM, "nemo_rl.environments.nemo_gym.NemoGym": PY_EXECUTABLES.NEMO_GYM, + # ModelOpt actors need the modelopt extra on top of their backend extra. + "nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantGenerationWorker": PY_EXECUTABLES.MODELOPT_VLLM, + "nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantAsyncGenerationWorker": PY_EXECUTABLES.MODELOPT_VLLM, + "nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker.DTensorQuantPolicyWorker": PY_EXECUTABLES.MODELOPT_AUTOMODEL, + "nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker_v2.DTensorQuantPolicyWorkerV2": PY_EXECUTABLES.MODELOPT_AUTOMODEL, + "nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker.MegatronQuantPolicyWorker": PY_EXECUTABLES.MODELOPT_MCORE, } -from nemo_rl.modelopt.registry import MODELOPT_ACTOR_REGISTRY - -ACTOR_ENVIRONMENT_REGISTRY.update(MODELOPT_ACTOR_REGISTRY) - def get_actor_python_env(actor_class_fqn: str) -> str: if actor_class_fqn in ACTOR_ENVIRONMENT_REGISTRY: diff --git a/nemo_rl/distributed/virtual_cluster.py b/nemo_rl/distributed/virtual_cluster.py index 3f8ae284a62..75f9ea1051f 100644 --- a/nemo_rl/distributed/virtual_cluster.py +++ b/nemo_rl/distributed/virtual_cluster.py @@ -55,6 +55,12 @@ class ClusterConfig(TypedDict): class PY_EXECUTABLES: + """Command each Ray actor launches under, one entry per uv extra combination. + + Every uv command below is rewritten to SYSTEM when NEMO_RL_PY_EXECUTABLES_SYSTEM + is set to 1, so callers never apply that check themselves. + """ + SYSTEM = sys.executable # Use NeMo-RL direct dependencies. @@ -89,6 +95,28 @@ class PY_EXECUTABLES: # Use NeMo-RL direct dependencies and TRT-LLM. TRTLLM = f"uv run --locked --extra trtllm --directory {git_root}" + # Use NeMo-RL direct dependencies and ModelOpt. + MODELOPT_VLLM = ( + f"uv run --locked --extra modelopt --extra vllm --directory {git_root}" + ) + MODELOPT_AUTOMODEL = ( + f"uv run --locked --extra modelopt --extra automodel --directory {git_root}" + ) + MODELOPT_MCORE = ( + f"uv run --locked --extra modelopt --extra mcore --directory {git_root}" + ) + + @classmethod + def _resolve_system_overrides(cls) -> None: + """Rewrite every uv command constant to the system executable when the flag is set.""" + if os.environ.get("NEMO_RL_PY_EXECUTABLES_SYSTEM", "0") != "1": + return + for name in [n for n in vars(cls) if n.isupper()]: + setattr(cls, name, cls.SYSTEM) + + +PY_EXECUTABLES._resolve_system_overrides() + # Default port ranges — kept below the OS ephemeral range. On some DGX/GB200 # nodes the ephemeral floor is as low as 9000 (32768 on stock Linux), so every diff --git a/nemo_rl/modelopt/registry.py b/nemo_rl/modelopt/registry.py deleted file mode 100644 index 31f6e4f07fc..00000000000 --- a/nemo_rl/modelopt/registry.py +++ /dev/null @@ -1,43 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import os - -from nemo_rl.distributed.virtual_cluster import PY_EXECUTABLES, git_root - -USE_SYSTEM_EXECUTABLE = os.environ.get("NEMO_RL_PY_EXECUTABLES_SYSTEM", "0") == "1" - -MODELOPT_VLLM_EXECUTABLE = ( - PY_EXECUTABLES.SYSTEM - if USE_SYSTEM_EXECUTABLE - else f"uv run --locked --extra modelopt --extra vllm --directory {git_root}" -) -MODELOPT_AUTOMODEL_EXECUTABLE = ( - PY_EXECUTABLES.SYSTEM - if USE_SYSTEM_EXECUTABLE - else f"uv run --locked --extra modelopt --extra automodel --directory {git_root}" -) -MODELOPT_MCORE_EXECUTABLE = ( - PY_EXECUTABLES.SYSTEM - if USE_SYSTEM_EXECUTABLE - else f"uv run --locked --extra modelopt --extra mcore --directory {git_root}" -) - -MODELOPT_ACTOR_REGISTRY: dict[str, str] = { - "nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantGenerationWorker": MODELOPT_VLLM_EXECUTABLE, - "nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantAsyncGenerationWorker": MODELOPT_VLLM_EXECUTABLE, - "nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker.DTensorQuantPolicyWorker": MODELOPT_AUTOMODEL_EXECUTABLE, - "nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker_v2.DTensorQuantPolicyWorkerV2": MODELOPT_AUTOMODEL_EXECUTABLE, - "nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker.MegatronQuantPolicyWorker": MODELOPT_MCORE_EXECUTABLE, -} diff --git a/nemo_rl/models/generation/sglang/sglang_router.py b/nemo_rl/models/generation/sglang/sglang_router.py index d523e99d684..6e84f705215 100644 --- a/nemo_rl/models/generation/sglang/sglang_router.py +++ b/nemo_rl/models/generation/sglang/sglang_router.py @@ -43,7 +43,7 @@ def run_router(args): class RouterActor: """Starts and owns the sglang router subprocess. - Runs under SGLANG_EXECUTABLE so it can import sglang_router. + Runs under PY_EXECUTABLES.SGLANG so it can import sglang_router. The driver (SYSTEM env) holds a handle to this actor and retrieves (router_ip, router_port) without ever importing sglang_router itself. """ diff --git a/nemo_rl/utils/venvs.py b/nemo_rl/utils/venvs.py index 65f6d8b0930..6e632b8d9fb 100644 --- a/nemo_rl/utils/venvs.py +++ b/nemo_rl/utils/venvs.py @@ -69,6 +69,17 @@ def create_local_venv( Returns: str: Path to the python executable in the created virtual environment """ + # A single token that is an executable file is already an interpreter, not a + # command line to run under uv -- there is nothing to build. + parts = shlex.split(py_executable) + if len(parts) == 1 and os.path.isfile(parts[0]) and os.access(parts[0], os.X_OK): + logger.warning( + f"{py_executable} is an interpreter, not a uv command, so no venv was built " + f"for {venv_name}; using it as-is (NEMO_RL_PY_EXECUTABLES_SYSTEM=1 sets every " + "PY_EXECUTABLES entry to sys.executable)." + ) + return py_executable + # This directory is where virtual environments will be installed # It is local to the driver process but should be visible to all worker nodes # If this directory is not accessible from worker nodes (e.g., on a distributed diff --git a/pyrefly.toml b/pyrefly.toml index 6a40aea185c..1b54fdaab50 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -194,7 +194,6 @@ project-includes = [ "nemo_rl/modelopt/models/policy/workers/__init__.py", "nemo_rl/modelopt/models/policy/workers/dtensor_quant_policy_worker.py", "nemo_rl/modelopt/models/policy/workers/dtensor_quant_policy_worker_v2.py", - "nemo_rl/modelopt/registry.py", "nemo_rl/models/__init__.py", "nemo_rl/models/automodel/__init__.py", "nemo_rl/models/dtensor/__init__.py", diff --git a/tests/unit/distributed/test_virtual_cluster.py b/tests/unit/distributed/test_virtual_cluster.py index a832b69e5ac..93a39d70099 100644 --- a/tests/unit/distributed/test_virtual_cluster.py +++ b/tests/unit/distributed/test_virtual_cluster.py @@ -11,10 +11,12 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import json import os import re import socket import subprocess +import sys from pathlib import Path from tempfile import TemporaryDirectory from unittest.mock import MagicMock, patch @@ -328,6 +330,10 @@ def test_init_ray_alone_has_no_data_plane_awareness(): assert "MC_ENABLE_DEST_DEVICE_AFFINITY" not in env_vars +@pytest.mark.skipif( + os.environ.get("NEMO_RL_PY_EXECUTABLES_SYSTEM", "0") == "1", + reason="No venv is built when every PY_EXECUTABLES entry is sys.executable", +) def test_mcore_py_executable(): # The temporary directory is created within the project. # For some reason, creating a virtual environment outside of the project @@ -757,3 +763,54 @@ def test_default_port_ranges_ordered_and_below_ephemeral_floor(): # Avoid privileged ports (<1024). assert DEFAULT_GENERATION_ROUTER_PORT_RANGE_LOW > 1024 assert DEFAULT_MASTER_PORT_RANGE_LOW > 1024 + + +_REGISTRY_PROBE = """ +import json + +from nemo_rl.distributed.ray_actor_environment_registry import ( + ACTOR_ENVIRONMENT_REGISTRY, + get_actor_python_env, +) +from nemo_rl.distributed.virtual_cluster import PY_EXECUTABLES + +envs = {fqn: get_actor_python_env(fqn) for fqn in ACTOR_ENVIRONMENT_REGISTRY} +# Also assert on PY_EXECUTABLES directly: a constant with no registry entry +# is invisible to envs, so the class-level promise needs its own check. +constants = {n: getattr(PY_EXECUTABLES, n) for n in vars(PY_EXECUTABLES) if n.isupper()} +print( + json.dumps( + { + "all_system": set(envs.values()) | set(constants.values()) + == {PY_EXECUTABLES.SYSTEM}, + "envs": envs, + "constants": constants, + } + ) +) +""" + + +@pytest.mark.parametrize("use_system_executable", [False, True]) +def test_actor_registry_honors_system_flag(use_system_executable): + # The registry freezes its executable strings at import, so the flag can + # only be exercised in a fresh interpreter. + env = dict(os.environ) + env["NEMO_RL_PY_EXECUTABLES_SYSTEM"] = "1" if use_system_executable else "0" + result = subprocess.run( + [sys.executable, "-c", _REGISTRY_PROBE], + capture_output=True, + text=True, + env=env, + ) + assert result.returncode == 0, result.stderr + payload = json.loads(result.stdout.strip().splitlines()[-1]) + + if use_system_executable: + assert payload["all_system"], (payload["envs"], payload["constants"]) + else: + envs = payload["envs"] + assert envs[ + "nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker" + ].startswith("uv run") + assert envs["nemo_rl.environments.nemo_gym.NemoGym"].startswith("uv run")