fix(runtime): honor system Python for all actor environments

nemo_rl/distributed/ray_actor_environment_registry.py

@@ -29,14 +29,25 @@
 TRTLLM_EXECUTABLE = (
     PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else PY_EXECUTABLES.TRTLLM
 )
+
+FSDP_EXECUTABLE = (
+    PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else PY_EXECUTABLES.FSDP
+)
+AUTOMODEL_EXECUTABLE = (
+    PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else PY_EXECUTABLES.AUTOMODEL
+)
+NEMO_GYM_EXECUTABLE = (
+    PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else PY_EXECUTABLES.NEMO_GYM
+)
+
 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.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.dtensor_policy_worker.DTensorPolicyWorker": FSDP_EXECUTABLE,
+    "nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2": AUTOMODEL_EXECUTABLE,
+    "nemo_rl.models.value.workers.dtensor_value_worker_v2.DTensorValueWorkerV2": AUTOMODEL_EXECUTABLE,
     "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker": MCORE_EXECUTABLE,
     "nemo_rl.models.value.workers.megatron_value_worker.MegatronValueWorker": MCORE_EXECUTABLE,
     "nemo_rl.models.generation.trtllm.trtllm_worker_async.TrtllmAsyncGenerationWorker": TRTLLM_EXECUTABLE,
@@ -48,18 +59,18 @@
     "nemo_rl.environments.code_jaccard_environment.CodeJaccardEnvironment": PY_EXECUTABLES.SYSTEM,
     "nemo_rl.environments.games.sliding_puzzle.SlidingPuzzleEnv": PY_EXECUTABLES.SYSTEM,
     # AsyncTrajectoryCollector needs vLLM environment to handle exceptions from VllmGenerationWorker
-    "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector": PY_EXECUTABLES.VLLM,
+    "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector": VLLM_EXECUTABLE,
     # ReplayBuffer needs vLLM environment to handle trajectory data from VllmGenerationWorker
-    "nemo_rl.algorithms.async_utils.ReplayBuffer": PY_EXECUTABLES.VLLM,
+    "nemo_rl.algorithms.async_utils.ReplayBuffer": VLLM_EXECUTABLE,
     # SyncRolloutActor doesn't import vllm directly — policy_generation is a
     # Ray actor handle. The VLLM env is needed because (1) transfer_queue is
     # bundled into the VLLM venv (and the policy training venvs), and the
     # actor writes flattened tensors to TQ via dp_client.put_samples;
     # (2) same-node colocation with VllmGenerationWorker avoids duplicate
     # venv caches.
-    "nemo_rl.experience.sync_rollout_actor.SyncRolloutActor": PY_EXECUTABLES.VLLM,
+    "nemo_rl.experience.sync_rollout_actor.SyncRolloutActor": VLLM_EXECUTABLE,
     "nemo_rl.environments.tools.retriever.RAGEnvironment": PY_EXECUTABLES.SYSTEM,
-    "nemo_rl.environments.nemo_gym.NemoGym": PY_EXECUTABLES.NEMO_GYM,
+    "nemo_rl.environments.nemo_gym.NemoGym": NEMO_GYM_EXECUTABLE,
 }
 
 from nemo_rl.modelopt.registry import MODELOPT_ACTOR_REGISTRY

tests/unit/distributed/test_ray_actor_environment_registry.py

@@ -0,0 +1,66 @@
+# Copyright (c) 2026, 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 importlib
+
+import nemo_rl.distributed.ray_actor_environment_registry as registry
+
+DEFAULT_EXECUTABLES = {
+    "nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker": registry.PY_EXECUTABLES.FSDP,
+    "nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2": registry.PY_EXECUTABLES.AUTOMODEL,
+    "nemo_rl.models.value.workers.dtensor_value_worker_v2.DTensorValueWorkerV2": registry.PY_EXECUTABLES.AUTOMODEL,
+    "nemo_rl.environments.nemo_gym.NemoGym": registry.PY_EXECUTABLES.NEMO_GYM,
+    "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector": registry.PY_EXECUTABLES.VLLM,
+    "nemo_rl.algorithms.async_utils.ReplayBuffer": registry.PY_EXECUTABLES.VLLM,
+    "nemo_rl.experience.sync_rollout_actor.SyncRolloutActor": registry.PY_EXECUTABLES.VLLM,
+}
+
+
+def _restore_actor_environment_registry(original_registry: dict[str, str]) -> None:
+    restored_registry = importlib.reload(registry).ACTOR_ENVIRONMENT_REGISTRY
+    original_registry.clear()
+    original_registry.update(restored_registry)
+    registry.ACTOR_ENVIRONMENT_REGISTRY = original_registry
+
+
+def test_actor_environments_use_default_executables(monkeypatch):
+    original_registry = registry.ACTOR_ENVIRONMENT_REGISTRY
+
+    try:
+        with monkeypatch.context() as context:
+            context.delenv("NEMO_RL_PY_EXECUTABLES_SYSTEM", raising=False)
+            importlib.reload(registry)
+
+            assert {
+                actor: registry.get_actor_python_env(actor)
+                for actor in DEFAULT_EXECUTABLES
+            } == DEFAULT_EXECUTABLES
+    finally:
+        _restore_actor_environment_registry(original_registry)
+
+
+def test_actor_environments_use_system_executable(monkeypatch):
+    original_registry = registry.ACTOR_ENVIRONMENT_REGISTRY
+
+    try:
+        with monkeypatch.context() as context:
+            context.setenv("NEMO_RL_PY_EXECUTABLES_SYSTEM", "1")
+            importlib.reload(registry)
+
+            assert {
+                actor: registry.get_actor_python_env(actor)
+                for actor in DEFAULT_EXECUTABLES
+            } == dict.fromkeys(DEFAULT_EXECUTABLES, registry.PY_EXECUTABLES.SYSTEM)
+    finally:
+        _restore_actor_environment_registry(original_registry)