From 9ff5657c3f3a975e3237527c01c2c41573be5567 Mon Sep 17 00:00:00 2001 From: Wenwen Gao Date: Wed, 17 Jun 2026 15:29:50 -0700 Subject: [PATCH 1/6] fix: prefer real NeMo-Gym package in actor Signed-off-by: Wenwen Gao --- examples/nemo_gym/prefetch_venvs.py | 5 +- nemo_rl/algorithms/grpo.py | 5 +- nemo_rl/environments/nemo_gym.py | 27 +++++++++++ .../unit/environments/test_nemo_gym_utils.py | 47 +++++++++++++++++++ 4 files changed, 80 insertions(+), 4 deletions(-) diff --git a/examples/nemo_gym/prefetch_venvs.py b/examples/nemo_gym/prefetch_venvs.py index 7128b7c92e..906c5e6c9e 100644 --- a/examples/nemo_gym/prefetch_venvs.py +++ b/examples/nemo_gym/prefetch_venvs.py @@ -55,6 +55,7 @@ def prefetch_nemo_gym_venvs(config_paths: list[str]) -> None: nemo_gym_py_exec = create_local_venv_on_each_node( nemo_gym_py_exec, "nemo_rl.environments.nemo_gym.NemoGym" ) + nemo_gym_venv = os.path.dirname(os.path.dirname(nemo_gym_py_exec)) succeeded = [] failed = [] @@ -96,8 +97,8 @@ def prefetch_nemo_gym_venvs(config_paths: list[str]) -> None: "py_executable": nemo_gym_py_exec, "env_vars": { **os.environ, - "VIRTUAL_ENV": nemo_gym_py_exec, - "UV_PROJECT_ENVIRONMENT": nemo_gym_py_exec, + "VIRTUAL_ENV": nemo_gym_venv, + "UV_PROJECT_ENVIRONMENT": nemo_gym_venv, }, }, } diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index ce90139e6c..9a0b5608e9 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -465,6 +465,7 @@ def _spinup_nemo_gym(base_urls, model_name): nemo_gym_py_exec = create_local_venv_on_each_node( nemo_gym_py_exec, "nemo_rl.environments.nemo_gym.NemoGym" ) + nemo_gym_venv = os.path.dirname(os.path.dirname(nemo_gym_py_exec)) nemo_gym_dict = env_configs["nemo_gym"] # NeMo-RL-side detection knobs are top-level NemoGymConfig fields # (where the detector reads them), not part of Gym's global config. @@ -497,8 +498,8 @@ def _spinup_nemo_gym(base_urls, model_name): "py_executable": nemo_gym_py_exec, "env_vars": { **os.environ, - "VIRTUAL_ENV": nemo_gym_py_exec, - "UV_PROJECT_ENVIRONMENT": nemo_gym_py_exec, + "VIRTUAL_ENV": nemo_gym_venv, + "UV_PROJECT_ENVIRONMENT": nemo_gym_venv, }, } actor = NemoGym.options(**nemo_gym_opts).remote(nemo_gym_cfg) diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index 611751af36..c85d755ab5 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -13,6 +13,7 @@ # limitations under the License. import os import subprocess +import sys from pathlib import Path from typing import Any, Dict, List, NotRequired, TypedDict @@ -38,6 +39,30 @@ DEFAULT_THINKING_TAGS = ["", ""] +def _ensure_nemo_gym_package_precedence() -> None: + """Prefer the third-party NeMo-Gym package over examples/nemo_gym.""" + repo_root = Path(__file__).resolve().parents[2] + gym_workspace = repo_root / "3rdparty" / "Gym-workspace" / "Gym" + gym_init = gym_workspace / "nemo_gym" / "__init__.py" + if not gym_init.exists(): + return + + gym_workspace_str = str(gym_workspace) + if sys.path[:1] != [gym_workspace_str]: + sys.path[:] = [p for p in sys.path if p != gym_workspace_str] + sys.path.insert(0, gym_workspace_str) + + shadowed_module = sys.modules.get("nemo_gym") + if ( + shadowed_module is not None + and getattr(shadowed_module, "__file__", None) is None + and not hasattr(shadowed_module, "PARENT_DIR") + ): + for module_name in list(sys.modules): + if module_name == "nemo_gym" or module_name.startswith("nemo_gym."): + del sys.modules[module_name] + + def get_nemo_gym_uv_cache_dir() -> str | None: """Return the uv cache directory inside a container, or None outside one. @@ -158,6 +183,8 @@ def _spinup(self) -> None: _gym_port_high = self.cfg.get("port_range_high", DEFAULT_GYM_PORT_RANGE_HIGH) self.head_server_port = _get_free_port_local(_gym_port_low, _gym_port_high) + _ensure_nemo_gym_package_precedence() + from nemo_gym.cli import GlobalConfigDictParserConfig, RunHelper from nemo_gym.rollout_collection import RolloutCollectionHelper from nemo_gym.server_utils import HEAD_SERVER_KEY_NAME, BaseServerConfig diff --git a/tests/unit/environments/test_nemo_gym_utils.py b/tests/unit/environments/test_nemo_gym_utils.py index fa60df05fb..3876d36d4c 100644 --- a/tests/unit/environments/test_nemo_gym_utils.py +++ b/tests/unit/environments/test_nemo_gym_utils.py @@ -17,10 +17,15 @@ (e.g. vllm) so the fast detector tests are not gated behind the nemo_gym extra. """ +import importlib +import sys +from pathlib import Path + import pytest from nemo_rl.environments import nemo_gym as nemo_gym_mod from nemo_rl.environments.nemo_gym import ( + _ensure_nemo_gym_package_precedence, _detect_invalid_tool_call_and_malformed_thinking, get_nemo_gym_uv_cache_dir, get_nemo_gym_venv_dir, @@ -97,3 +102,45 @@ def test_get_nemo_gym_uv_cache_dir_uses_uv_inside_container(monkeypatch): lambda *args, **kwargs: b" /root/.cache/uv\n", ) assert get_nemo_gym_uv_cache_dir() == "/root/.cache/uv" + + +def test_ensure_nemo_gym_package_precedence_recovers_from_examples_shadow(): + repo_root = Path(__file__).resolve().parents[3] + examples_dir = str(repo_root / "examples") + expected_gym_init = ( + repo_root + / "3rdparty" + / "Gym-workspace" + / "Gym" + / "nemo_gym" + / "__init__.py" + ).resolve() + + saved_path = list(sys.path) + saved_modules = { + name: module + for name, module in sys.modules.items() + if name == "nemo_gym" or name.startswith("nemo_gym.") + } + try: + for name in list(sys.modules): + if name == "nemo_gym" or name.startswith("nemo_gym."): + del sys.modules[name] + sys.path.insert(0, examples_dir) + + shadowed = importlib.import_module("nemo_gym") + assert getattr(shadowed, "__file__", None) is None + + _ensure_nemo_gym_package_precedence() + + actual = importlib.import_module("nemo_gym") + assert Path(actual.__file__).resolve() == expected_gym_init + from nemo_gym import PARENT_DIR # noqa: PLC0415 + + assert PARENT_DIR == expected_gym_init.parent.parent + finally: + sys.path[:] = saved_path + for name in list(sys.modules): + if name == "nemo_gym" or name.startswith("nemo_gym."): + del sys.modules[name] + sys.modules.update(saved_modules) From af874c4ea504aa6e8e8ed5ae4dd5baa975c318b2 Mon Sep 17 00:00:00 2001 From: Wenwen Gao Date: Wed, 24 Jun 2026 22:36:36 -0700 Subject: [PATCH 2/6] fix(slurm): stop Ray sidecars after driver exits Signed-off-by: Wenwen Gao --- ray.sub | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/ray.sub b/ray.sub index 62caadcbf7..89c3d2ac19 100644 --- a/ray.sub +++ b/ray.sub @@ -525,7 +525,12 @@ echo "All workers connected!" # This driver process is responsible for launching a job on the Ray cluster CONTAINER_CWD=$(scontrol show job $SLURM_JOB_ID | grep -oP 'WorkDir=\K[^ ]+' | head -1) if [[ -n "$COMMAND" ]]; then + set +e srun --no-container-mount-home --overlap --container-name=ray-head --container-workdir=$CONTAINER_CWD --nodes=1 --ntasks=1 -w "$head_node" -o $LOG_DIR/ray-driver.log bash -c "$COMMAND" + driver_exit_code=$? + set -e + touch "$LOG_DIR/ENDED" + exit "$driver_exit_code" else echo "[INFO]: Ray Cluster is idled, run this on the slurm head node to get a shell to the head node:" cat <$SLURM_SUBMIT_DIR/${SLURM_JOB_ID}-attach.sh From 826298012b610a659ff9cc60af5bec001ae40b4e Mon Sep 17 00:00:00 2001 From: Wenwen Gao Date: Fri, 19 Jun 2026 22:08:06 -0700 Subject: [PATCH 3/6] fix(nemo-gym): map overlong vllm prompts to 400 Signed-off-by: Wenwen Gao --- .../generation/vllm/vllm_worker_async.py | 15 +++-- .../models/generation/test_vllm_generation.py | 59 ++++++++++++++++++- 2 files changed, 69 insertions(+), 5 deletions(-) diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index c7e09455a4..1584879b77 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -722,15 +722,22 @@ async def create_chat_completion( generator = await openai_serving_chat.create_chat_completion( request, raw_request ) - except VLLMValidationError as e: + except (ValueError, VLLMValidationError) as e: # vLLM 0.20 raises VLLMValidationError for prompts exceeding # max_model_len during tokenization, instead of returning an - # ErrorResponse. Convert to HTTP 400 so the Gym proxy can - # detect context-length overflow and handle it gracefully. + # ErrorResponse. Our post-tokenization clamp can raise a local + # ValueError for the same condition after prefix replacement. + # Convert those cases to HTTP 400 so the Gym proxy can detect + # context-length overflow and handle it gracefully. + message = str(e) + if isinstance(e, ValueError) and not ( + "max_model_len" in message or "maximum context length" in message + ): + raise return JSONResponse( content={ "error": { - "message": str(e), + "message": message, "type": "invalid_request_error", "code": 400, } diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 7d4f789c28..bf19e0188d 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import importlib.util import json import os @@ -19,7 +20,7 @@ import types from copy import deepcopy from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest import ray @@ -344,6 +345,62 @@ def test_vllm_async_http_server_loads_reasoning_parser_plugin(monkeypatch): assert "reasoning_parser_plugin" not in openai_serving_chat.instances[0].kwargs +def _setup_fake_openai_chat_completion_route(monkeypatch): + ( + _tool_parser_manager, + _reasoning_parser_manager, + openai_serving_chat, + ) = _install_fake_vllm_openai_modules(monkeypatch) + + worker = VllmAsyncGenerationWorkerImpl.__new__(VllmAsyncGenerationWorkerImpl) + worker.cfg = { + "temperature": 1.0, + "top_p": 1.0, + "vllm_cfg": { + "http_server_serving_chat_kwargs": {}, + }, + } + worker.llm = MagicMock(model_config="model-config", renderer="renderer") + model_config = MagicMock(served_model_name="served-model", model="model-path") + worker.llm_async_engine_args = MagicMock() + worker.llm_async_engine_args.create_model_config.return_value = model_config + + app = _FakeFastAPIApp() + worker._setup_vllm_openai_api_server(app) + route = next(func for path, func in app.routes if path == "/v1/chat/completions") + return route, openai_serving_chat.instances[0] + + +def test_vllm_async_chat_completion_maps_context_value_error_to_400(monkeypatch): + route, openai_serving_chat = _setup_fake_openai_chat_completion_route(monkeypatch) + openai_serving_chat.create_chat_completion = AsyncMock( + side_effect=ValueError( + "Prompt length (8551) fills or exceeds max_model_len (8192). " + "No room for output tokens." + ) + ) + + request = types.SimpleNamespace(top_k=None, temperature=1.0, top_p=1.0) + response = asyncio.run(route(request, MagicMock())) + + assert response.status_code == 400 + body = json.loads(response.body.decode()) + assert body["error"]["code"] == 400 + assert body["error"]["type"] == "invalid_request_error" + assert "max_model_len" in body["error"]["message"] + + +def test_vllm_async_chat_completion_reraises_unrelated_value_error(monkeypatch): + route, openai_serving_chat = _setup_fake_openai_chat_completion_route(monkeypatch) + openai_serving_chat.create_chat_completion = AsyncMock( + side_effect=ValueError("unexpected internal validation failure") + ) + + request = types.SimpleNamespace(top_k=None, temperature=1.0, top_p=1.0) + with pytest.raises(ValueError, match="unexpected internal validation failure"): + asyncio.run(route(request, MagicMock())) + + def test_nano_v3_reasoning_parser_swaps_reasoning_when_thinking_disabled( monkeypatch, ): From cfc93f4cd0a7077f035ef78e0107930ad8abc4c7 Mon Sep 17 00:00:00 2001 From: Wenwen Gao Date: Fri, 26 Jun 2026 12:31:59 -0700 Subject: [PATCH 4/6] chore: remove test-only changes from nano-v3 PR --- .../unit/environments/test_nemo_gym_utils.py | 47 ---------- .../models/generation/test_vllm_generation.py | 88 +++++++------------ 2 files changed, 30 insertions(+), 105 deletions(-) diff --git a/tests/unit/environments/test_nemo_gym_utils.py b/tests/unit/environments/test_nemo_gym_utils.py index 3876d36d4c..fa60df05fb 100644 --- a/tests/unit/environments/test_nemo_gym_utils.py +++ b/tests/unit/environments/test_nemo_gym_utils.py @@ -17,15 +17,10 @@ (e.g. vllm) so the fast detector tests are not gated behind the nemo_gym extra. """ -import importlib -import sys -from pathlib import Path - import pytest from nemo_rl.environments import nemo_gym as nemo_gym_mod from nemo_rl.environments.nemo_gym import ( - _ensure_nemo_gym_package_precedence, _detect_invalid_tool_call_and_malformed_thinking, get_nemo_gym_uv_cache_dir, get_nemo_gym_venv_dir, @@ -102,45 +97,3 @@ def test_get_nemo_gym_uv_cache_dir_uses_uv_inside_container(monkeypatch): lambda *args, **kwargs: b" /root/.cache/uv\n", ) assert get_nemo_gym_uv_cache_dir() == "/root/.cache/uv" - - -def test_ensure_nemo_gym_package_precedence_recovers_from_examples_shadow(): - repo_root = Path(__file__).resolve().parents[3] - examples_dir = str(repo_root / "examples") - expected_gym_init = ( - repo_root - / "3rdparty" - / "Gym-workspace" - / "Gym" - / "nemo_gym" - / "__init__.py" - ).resolve() - - saved_path = list(sys.path) - saved_modules = { - name: module - for name, module in sys.modules.items() - if name == "nemo_gym" or name.startswith("nemo_gym.") - } - try: - for name in list(sys.modules): - if name == "nemo_gym" or name.startswith("nemo_gym."): - del sys.modules[name] - sys.path.insert(0, examples_dir) - - shadowed = importlib.import_module("nemo_gym") - assert getattr(shadowed, "__file__", None) is None - - _ensure_nemo_gym_package_precedence() - - actual = importlib.import_module("nemo_gym") - assert Path(actual.__file__).resolve() == expected_gym_init - from nemo_gym import PARENT_DIR # noqa: PLC0415 - - assert PARENT_DIR == expected_gym_init.parent.parent - finally: - sys.path[:] = saved_path - for name in list(sys.modules): - if name == "nemo_gym" or name.startswith("nemo_gym."): - del sys.modules[name] - sys.modules.update(saved_modules) diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index bf19e0188d..bb4f72feda 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import asyncio import importlib.util import json import os @@ -20,7 +19,7 @@ import types from copy import deepcopy from pathlib import Path -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import MagicMock import pytest import ray @@ -345,62 +344,6 @@ def test_vllm_async_http_server_loads_reasoning_parser_plugin(monkeypatch): assert "reasoning_parser_plugin" not in openai_serving_chat.instances[0].kwargs -def _setup_fake_openai_chat_completion_route(monkeypatch): - ( - _tool_parser_manager, - _reasoning_parser_manager, - openai_serving_chat, - ) = _install_fake_vllm_openai_modules(monkeypatch) - - worker = VllmAsyncGenerationWorkerImpl.__new__(VllmAsyncGenerationWorkerImpl) - worker.cfg = { - "temperature": 1.0, - "top_p": 1.0, - "vllm_cfg": { - "http_server_serving_chat_kwargs": {}, - }, - } - worker.llm = MagicMock(model_config="model-config", renderer="renderer") - model_config = MagicMock(served_model_name="served-model", model="model-path") - worker.llm_async_engine_args = MagicMock() - worker.llm_async_engine_args.create_model_config.return_value = model_config - - app = _FakeFastAPIApp() - worker._setup_vllm_openai_api_server(app) - route = next(func for path, func in app.routes if path == "/v1/chat/completions") - return route, openai_serving_chat.instances[0] - - -def test_vllm_async_chat_completion_maps_context_value_error_to_400(monkeypatch): - route, openai_serving_chat = _setup_fake_openai_chat_completion_route(monkeypatch) - openai_serving_chat.create_chat_completion = AsyncMock( - side_effect=ValueError( - "Prompt length (8551) fills or exceeds max_model_len (8192). " - "No room for output tokens." - ) - ) - - request = types.SimpleNamespace(top_k=None, temperature=1.0, top_p=1.0) - response = asyncio.run(route(request, MagicMock())) - - assert response.status_code == 400 - body = json.loads(response.body.decode()) - assert body["error"]["code"] == 400 - assert body["error"]["type"] == "invalid_request_error" - assert "max_model_len" in body["error"]["message"] - - -def test_vllm_async_chat_completion_reraises_unrelated_value_error(monkeypatch): - route, openai_serving_chat = _setup_fake_openai_chat_completion_route(monkeypatch) - openai_serving_chat.create_chat_completion = AsyncMock( - side_effect=ValueError("unexpected internal validation failure") - ) - - request = types.SimpleNamespace(top_k=None, temperature=1.0, top_p=1.0) - with pytest.raises(ValueError, match="unexpected internal validation failure"): - asyncio.run(route(request, MagicMock())) - - def test_nano_v3_reasoning_parser_swaps_reasoning_when_thinking_disabled( monkeypatch, ): @@ -526,6 +469,33 @@ def test_configure_generation_config_keeps_dummy_startup_weights_with_draft_refi assert configured["vllm_cfg"]["load_format"] == "dummy" +@pytest.mark.parametrize("method", ["deepseek_mtp", "mtp"]) +def test_configure_generation_config_keeps_dummy_startup_weights_for_mtp(method): + """MTP keeps dummy startup weights even without draft refit. + + The policy weights arrive via refit and only the MTP draft layer is loaded + from disk on the worker, so we must not force load_format="auto" (which would + read the full base-model checkpoint). + """ + vllm_config = deepcopy(basic_vllm_test_config) + vllm_config["vllm_kwargs"] = { + "speculative_config": { + "method": method, + "num_speculative_tokens": 1, + } + } + tokenizer = MagicMock(pad_token_id=0, eos_token_id=1) + + configured = configure_generation_config( + vllm_config, + tokenizer, + is_eval=False, + has_refit_draft_weights=False, + ) + + assert configured["vllm_cfg"]["load_format"] == "dummy" + + def get_basic_megatron_test_config( tp: int = 1, pp: int = 1, @@ -582,6 +552,7 @@ def get_basic_megatron_test_config( "bias_activation_fusion": True, "moe_per_layer_logging": False, "gradient_accumulation_fusion": False, + "use_fused_weighted_squared_relu": False, "train_iters": 100, # Required for Megatron training "optimizer": { "optimizer": "adam", @@ -2900,6 +2871,7 @@ def test_vllm_megatron_pipeline_parallel(cluster, tokenizer): vllm_config["model_name"] = model_name vllm_config["tokenizer"]["name"] = model_name vllm_config = configure_generation_config(vllm_config, test_tokenizer) + vllm_config["vllm_cfg"]["max_model_len"] = 128 megatron_config = get_basic_megatron_test_config( tp=1, From 875e702d8604c0f494c92fe73c965038d0b41a1c Mon Sep 17 00:00:00 2001 From: Wenwen Gao Date: Fri, 26 Jun 2026 12:32:40 -0700 Subject: [PATCH 5/6] chore: drop test changes from nano-v3 PR diff --- .../models/generation/test_vllm_generation.py | 29 ------------------- 1 file changed, 29 deletions(-) diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index bb4f72feda..7d4f789c28 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -469,33 +469,6 @@ def test_configure_generation_config_keeps_dummy_startup_weights_with_draft_refi assert configured["vllm_cfg"]["load_format"] == "dummy" -@pytest.mark.parametrize("method", ["deepseek_mtp", "mtp"]) -def test_configure_generation_config_keeps_dummy_startup_weights_for_mtp(method): - """MTP keeps dummy startup weights even without draft refit. - - The policy weights arrive via refit and only the MTP draft layer is loaded - from disk on the worker, so we must not force load_format="auto" (which would - read the full base-model checkpoint). - """ - vllm_config = deepcopy(basic_vllm_test_config) - vllm_config["vllm_kwargs"] = { - "speculative_config": { - "method": method, - "num_speculative_tokens": 1, - } - } - tokenizer = MagicMock(pad_token_id=0, eos_token_id=1) - - configured = configure_generation_config( - vllm_config, - tokenizer, - is_eval=False, - has_refit_draft_weights=False, - ) - - assert configured["vllm_cfg"]["load_format"] == "dummy" - - def get_basic_megatron_test_config( tp: int = 1, pp: int = 1, @@ -552,7 +525,6 @@ def get_basic_megatron_test_config( "bias_activation_fusion": True, "moe_per_layer_logging": False, "gradient_accumulation_fusion": False, - "use_fused_weighted_squared_relu": False, "train_iters": 100, # Required for Megatron training "optimizer": { "optimizer": "adam", @@ -2871,7 +2843,6 @@ def test_vllm_megatron_pipeline_parallel(cluster, tokenizer): vllm_config["model_name"] = model_name vllm_config["tokenizer"]["name"] = model_name vllm_config = configure_generation_config(vllm_config, test_tokenizer) - vllm_config["vllm_cfg"]["max_model_len"] = 128 megatron_config = get_basic_megatron_test_config( tp=1, From d35e2d5a28c74291f208f2567e6280d5dabcfffe Mon Sep 17 00:00:00 2001 From: Wenwen Gao Date: Fri, 26 Jun 2026 12:34:55 -0700 Subject: [PATCH 6/6] chore: drop prefetch venv change from nano-v3 PR --- examples/nemo_gym/prefetch_venvs.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/examples/nemo_gym/prefetch_venvs.py b/examples/nemo_gym/prefetch_venvs.py index 906c5e6c9e..7128b7c92e 100644 --- a/examples/nemo_gym/prefetch_venvs.py +++ b/examples/nemo_gym/prefetch_venvs.py @@ -55,7 +55,6 @@ def prefetch_nemo_gym_venvs(config_paths: list[str]) -> None: nemo_gym_py_exec = create_local_venv_on_each_node( nemo_gym_py_exec, "nemo_rl.environments.nemo_gym.NemoGym" ) - nemo_gym_venv = os.path.dirname(os.path.dirname(nemo_gym_py_exec)) succeeded = [] failed = [] @@ -97,8 +96,8 @@ def prefetch_nemo_gym_venvs(config_paths: list[str]) -> None: "py_executable": nemo_gym_py_exec, "env_vars": { **os.environ, - "VIRTUAL_ENV": nemo_gym_venv, - "UV_PROJECT_ENVIRONMENT": nemo_gym_venv, + "VIRTUAL_ENV": nemo_gym_py_exec, + "UV_PROJECT_ENVIRONMENT": nemo_gym_py_exec, }, }, }