fix: honor NEMO_RL_PY_EXECUTABLES_SYSTEM for all actor environments

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
 }

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__ <nemo_rl.distributed.worker_groups.RayWorkerBuilder.__call__>`. This provides flexibility for special use cases without modifying the core configurations.

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:

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

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.
     """

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

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",

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")