chore: unify setup nemogym

examples/nemo_gym/run_grpo_nemo_gym.py

@@ -41,9 +41,7 @@
 from nemo_rl.algorithms.utils import get_tokenizer
 from nemo_rl.data.utils import setup_response_data
 from nemo_rl.distributed.virtual_cluster import init_ray
-from nemo_rl.environments.nemo_gym import (
-    setup_nemo_gym_config,
-)
+from nemo_rl.environments.nemo_gym import setup_nemo_gym_config
 from nemo_rl.experience.rollouts import run_nemo_gym_rollout_sync
 from nemo_rl.models.generation import configure_generation_config
 from nemo_rl.utils.config import (

nemo_rl/algorithms/grpo.py

@@ -24,7 +24,6 @@
 import ray
 import torch
 from pydantic import BaseModel, Field, model_validator
-from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
 from torchdata.stateful_dataloader import StatefulDataLoader
 from transformers import AutoProcessor
 from transformers.tokenization_utils_base import PreTrainedTokenizerBase
@@ -80,12 +79,7 @@
     prepare_segment_topology,
 )
 from nemo_rl.environments.interfaces import EnvironmentInterface
-from nemo_rl.environments.nemo_gym import (
-    NemoGym,
-    NemoGymConfig,
-    get_nemo_gym_uv_cache_dir,
-    get_nemo_gym_venv_dir,
-)
+from nemo_rl.environments.nemo_gym import spinup_nemo_gym_actor
 from nemo_rl.experience.interfaces import (
     NEMO_GYM_TASK_INDEX_KEY,
     NEXT_NEMO_GYM_TASK_INDEX_KEY,
@@ -594,67 +588,24 @@ def init_train_dataloader(dataset, suffix: str = ""):
         master_config, enable_nemo_gym=enable_nemo_gym
     )
     nemo_gym_actor = None
-    if enable_nemo_gym:
-        nemo_gym_num_nodes = env_configs.get("nemo_gym", {}).get("num_gpu_nodes", 0)
-        ray_runtime_ctx = ray.get_runtime_context()
-        ray_cur_node_id = ray_runtime_ctx.get_node_id()
-    else:
-        nemo_gym_num_nodes = 0
-        ray_cur_node_id = None
 
     def _spinup_nemo_gym(base_urls, model_name):
         """Spin up the NeMo Gym actor against the given generation server URLs."""
         t0 = time.perf_counter()
-        nemo_gym_py_exec = get_actor_python_env("nemo_rl.environments.nemo_gym.NemoGym")
-        if nemo_gym_py_exec.startswith("uv"):
-            nemo_gym_py_exec = create_local_venv_on_each_node(
-                nemo_gym_py_exec, "nemo_rl.environments.nemo_gym.NemoGym"
-            )
-        nemo_gym_dict = 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.
-        invalid_tool_call_patterns = nemo_gym_dict.pop(
-            "invalid_tool_call_patterns", None
+        enable_router_replay = router_replay_enabled(policy_config)
+        routed_experts_dtype = (
+            resolve_routed_experts_dtype_name_for_model(model_name)
+            if enable_router_replay
+            else "int16"
         )
-        thinking_tags = nemo_gym_dict.pop("thinking_tags", None)
-        # Pass prebuilt cache + venv dirs through the global config so the gym reuses
-        # image-baked venvs instead of rebuilding them.
-        uv_cache_dir = get_nemo_gym_uv_cache_dir()
-        if uv_cache_dir is not None:
-            nemo_gym_dict.setdefault("uv_cache_dir", uv_cache_dir)
-        uv_venv_dir = get_nemo_gym_venv_dir()
-        if uv_venv_dir is not None:
-            nemo_gym_dict.setdefault("uv_venv_dir", uv_venv_dir)
-        nemo_gym_cfg = NemoGymConfig(
-            model_name=model_name,
+        actor = spinup_nemo_gym_actor(
+            env_configs=env_configs,
             base_urls=base_urls,
-            invalid_tool_call_patterns=invalid_tool_call_patterns,
-            thinking_tags=thinking_tags,
-            require_routed_experts=router_replay_enabled(policy_config),
-            routed_experts_dtype=(
-                resolve_routed_experts_dtype_name_for_model(model_name)
-                if router_replay_enabled(policy_config)
-                else "int16"
-            ),
+            model_name=model_name,
+            enable_router_replay=enable_router_replay,
+            routed_experts_dtype=routed_experts_dtype,
             use_fastokens=bool(policy_config["tokenizer"].get("use_fastokens")),
-            initial_global_config_dict=nemo_gym_dict,
         )
-        nemo_gym_opts = {}
-        if nemo_gym_num_nodes:
-            nemo_gym_opts["scheduling_strategy"] = NodeAffinitySchedulingStrategy(
-                node_id=ray_cur_node_id,
-                soft=True,
-            )
-        nemo_gym_opts["runtime_env"] = {
-            "py_executable": nemo_gym_py_exec,
-            "env_vars": {
-                **os.environ,
-                "VIRTUAL_ENV": nemo_gym_py_exec,
-                "UV_PROJECT_ENVIRONMENT": nemo_gym_py_exec,
-            },
-        }
-        actor = NemoGym.options(**nemo_gym_opts).remote(nemo_gym_cfg)
-        ray.get(actor._spinup.remote())
         return actor, time.perf_counter() - t0
 
     total_nodes = cluster_config["num_nodes"]

nemo_rl/algorithms/single_controller_utils/setup.py

@@ -30,9 +30,7 @@
 from transformers.tokenization_utils_base import PreTrainedTokenizerBase
 
 from nemo_rl.algorithms.async_utils.replay_buffer import TQReplayBuffer
-from nemo_rl.algorithms.grpo import (
-    MasterConfig as GrpoMasterConfig,
-)
+from nemo_rl.algorithms.grpo import MasterConfig as GrpoMasterConfig
 from nemo_rl.algorithms.grpo import (
     _create_advantage_estimator,
     _should_use_nemo_gym,
@@ -51,6 +49,9 @@
 from nemo_rl.environments.interfaces import EnvironmentInterface
 from nemo_rl.environments.nemo_gym import spinup_nemo_gym_actor
 from nemo_rl.experience.rollout_manager import RolloutManager
+from nemo_rl.models.generation.interfaces import (
+    resolve_routed_experts_dtype_name_for_model,
+)
 from nemo_rl.models.generation.sglang.config import SGLangConfig
 from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration
 from nemo_rl.models.generation.vllm import VllmGeneration
@@ -384,12 +385,19 @@ def setup_single_controller(
     if use_nemo_gym:
         # TODO(#2625): Mirror GRPO's deferred vLLM load so NeMo-Gym spinup
         # overlaps model loading instead of running serially afterward.
-        enable_router_replay = router_replay_enabled(master_config.policy)
+        enable_router_replay = router_replay_enabled(policy_config)
+        routed_experts_dtype = (
+            resolve_routed_experts_dtype_name_for_model(generation_config["model_name"])
+            if enable_router_replay
+            else "int16"
+        )
         env_handles["nemo_gym"] = spinup_nemo_gym_actor(
             env_configs=master_config.env,
             base_urls=generation.dp_openai_server_base_urls,
             model_name=generation_config["model_name"],
             enable_router_replay=enable_router_replay,
+            routed_experts_dtype=routed_experts_dtype,
+            use_fastokens=bool(policy_config["tokenizer"].get("use_fastokens")),
         )
 
     # ==========================

nemo_rl/environments/nemo_gym.py

@@ -613,26 +613,31 @@ def spinup_nemo_gym_actor(
     env_configs: dict[str, Any],
     base_urls: list[Optional[str]],
     model_name: str,
+    *,
     enable_router_replay: bool,
+    routed_experts_dtype: str,
+    use_fastokens: bool,
 ) -> Any:
     """Spin up the NeMo-Gym actor against the given generation server URLs.
 
-    When ``env_configs["nemo_gym"]["num_gpu_nodes"] > 0``, the actor is
-    scheduled with soft NodeAffinity to the current Ray node so its colocated
-    GPU resources land where the caller expects.
+    When env_configs["nemo_gym"]["num_gpu_nodes"] > 0, the actor is scheduled
+    with soft NodeAffinity to the current Ray node so its colocated GPU
+    resources land where the caller expects.
 
     Args:
-        env_configs: The ``master_config.env`` mapping; ``env_configs["nemo_gym"]``
-            supplies the Gym global config plus NeMo-RL detection knobs
-            (``invalid_tool_call_patterns``, ``thinking_tags``, ``num_gpu_nodes``).
-        base_urls: Per-DP-rank OpenAI-compatible server base URLs from the
-            generation backend.
+        env_configs: The master_config.env mapping; env_configs["nemo_gym"] supplies
+            the Gym global config plus NeMo-RL detection knobs (invalid_tool_call_patterns,
+            thinking_tags, num_gpu_nodes).
+        base_urls: Per-DP-rank OpenAI-compatible server base URLs from the generation backend.
         model_name: Served model name the Gym rollouts should target.
-        enable_router_replay: Sets ``require_routed_experts`` on the
-            ``NemoGymConfig``.
+        enable_router_replay: Sets require_routed_experts on the NemoGymConfig.
+        routed_experts_dtype: Dtype name for R3 routed_experts tensors ("int8"/"int16"/"int32"),
+            resolved by the caller from the model's expert count.
+        use_fastokens: Forwarded from policy.tokenizer.use_fastokens so the rollout actor
+            patches its tokenizer consistently with the driver.
 
     Returns:
-        The spun-up ``NemoGym`` Ray actor handle (``_spinup`` already awaited).
+        The spun-up NemoGym Ray actor handle (_spinup already awaited).
     """
     nemo_gym_dict = dict(env_configs["nemo_gym"])
 
@@ -656,6 +661,8 @@ def spinup_nemo_gym_actor(
         invalid_tool_call_patterns=invalid_tool_call_patterns,
         thinking_tags=thinking_tags,
         require_routed_experts=enable_router_replay,
+        routed_experts_dtype=routed_experts_dtype,
+        use_fastokens=use_fastokens,
         initial_global_config_dict=nemo_gym_dict,
     )
 

tests/unit/single_controller/test_single_controller_setup.py

@@ -70,6 +70,7 @@ def _make_master_config(
         policy={
             "train_global_batch_size": num_prompts_per_step * 2,
             "max_total_sequence_length": 32,
+            "tokenizer": {"use_fastokens": False},
             "megatron_cfg": {"enabled": megatron_enabled},
             "generation": {
                 "backend": backend,
@@ -421,6 +422,8 @@ def test_nemo_gym_wires_env_handle(self, patched_factories):
             ].return_value.dp_openai_server_base_urls,
             model_name="test-model",
             enable_router_replay=False,
+            routed_experts_dtype="int16",
+            use_fastokens=False,
         )
         assert actor_args.env_handles["nemo_gym"] is fake_gym_actor