diff --git a/examples/nemo_gym/run_grpo_nemo_gym.py b/examples/nemo_gym/run_grpo_nemo_gym.py index 7210879a725..4d1aa3eefcb 100644 --- a/examples/nemo_gym/run_grpo_nemo_gym.py +++ b/examples/nemo_gym/run_grpo_nemo_gym.py @@ -36,7 +36,6 @@ grpo_train, refit_policy_generation, setup, - shutdown_environments, ) from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.data.utils import setup_response_data @@ -45,6 +44,7 @@ setup_nemo_gym_config, should_use_nemo_gym, ) +from nemo_rl.environments.utils import shutdown_environments from nemo_rl.experience.rollouts import run_nemo_gym_rollout_sync from nemo_rl.models.generation import configure_generation_config from nemo_rl.models.generation.vllm.config import materialize_vllm_video_config @@ -271,6 +271,7 @@ def main() -> None: task_to_env = {"nemo_gym": nemo_gym} val_task_to_env = task_to_env + trainer_owns_environment_shutdown = False try: if is_trajectory_collection: collect_trajectories( @@ -331,6 +332,7 @@ def main() -> None: print("🚀 Running synchronous GRPO training") # Run standard GRPO training + trainer_owns_environment_shutdown = True grpo_train( policy, policy_generation, @@ -347,7 +349,8 @@ def main() -> None: processor=processor, ) finally: - shutdown_environments(task_to_env, val_task_to_env) + if not trainer_owns_environment_shutdown: + shutdown_environments(task_to_env, val_task_to_env) try: policy_generation.shutdown() except Exception as error: diff --git a/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh b/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh index b12905694dc..6e600ec023f 100755 --- a/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh +++ b/examples/nemo_gym/run_nemo_gym_single_node_sanity_tests.sh @@ -41,7 +41,16 @@ uv run python -c "import ray; ray.shutdown()" ./tests/run_unit.sh unit/environments/test_math_environment.py::test_math_env_step_basic # NeMo Gym integrates directly into NeMo RL as an Environment since that is the cleanest way. This tests the NeMo Gym integration logic and correctness. -./tests/run_unit.sh unit/environments/test_nemo_gym.py::test_nemo_gym_sanity +./tests/run_unit.sh \ + unit/environments/test_nemo_gym.py::test_nemo_gym_sanity \ + --nemo-gym-only # NeMo Gym uses a separate rollout loop inside grpo_train in NeMo RL. This tests the e2e rollout functionality and correctness. -./tests/run_unit.sh unit/experience/test_rollouts.py::test_run_async_nemo_gym_rollout +./tests/run_unit.sh \ + unit/experience/test_rollouts.py::test_run_async_nemo_gym_rollout \ + --nemo-gym-only + +# Exercise one pinned example from each L0 Gym environment through the real NeMo RL rollout boundary. +./tests/run_unit.sh \ + unit/environments/test_nemo_gym_rollout.py::test_l0_gym_environments_roll_out_through_nemo_rl \ + --nemo-gym-only diff --git a/examples/run_grpo.py b/examples/run_grpo.py index 51b1f08698e..3670bc82019 100644 --- a/examples/run_grpo.py +++ b/examples/run_grpo.py @@ -23,11 +23,11 @@ MasterConfig, grpo_train, setup, - shutdown_environments, ) 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.utils import shutdown_environments from nemo_rl.models.generation import configure_generation_config from nemo_rl.utils.config import ( load_config, @@ -180,6 +180,7 @@ def _make_policy(**kwargs): print("=" * 60 + "\n", flush=True) try: + trainer_owns_environment_shutdown = False # Check if async mode is enabled if config.grpo.async_grpo.enabled: # Async GRPO does not support dynamic sampling, reward scaling, or reward shaping (DAPO features) @@ -228,6 +229,7 @@ def _make_policy(**kwargs): # Two parallel synchronous trainers (verl-style — main_ppo.py vs # main_ppo_sync.py). data_plane.enabled selects which one runs. trainer = _select_trainer(master_config) + trainer_owns_environment_shutdown = trainer is grpo_train # grpo_train_sync defers checkpoint finalization to the checkpointer's # background threads; the context manager guarantees they are flushed on # exit. (grpo_train also flushes internally; shutdown() is idempotent.) @@ -247,7 +249,8 @@ def _make_policy(**kwargs): master_config, ) finally: - shutdown_environments(task_to_env, val_task_to_env) + if not trainer_owns_environment_shutdown: + shutdown_environments(task_to_env, val_task_to_env) try: policy_generation.shutdown() except Exception as error: diff --git a/nemo_rl/algorithms/distillation.py b/nemo_rl/algorithms/distillation.py index 91e00346576..01c49d3e30b 100644 --- a/nemo_rl/algorithms/distillation.py +++ b/nemo_rl/algorithms/distillation.py @@ -56,6 +56,7 @@ get_nemo_gym_venv_dir, should_use_nemo_gym, ) +from nemo_rl.environments.utils import shutdown_environments from nemo_rl.experience.rollouts import ( run_async_multi_turn_rollout, run_multi_turn_rollout, @@ -674,7 +675,7 @@ def init_nemo_gym(): # =============================================================================== -def distillation_train( +def _distillation_train_impl( student_policy: ColocatablePolicyInterface, teacher_policy: ColocatablePolicyInterface, student_generation: Optional[GenerationInterface], @@ -1172,6 +1173,43 @@ def distillation_train( checkpointer.shutdown() +@trace_fn(RLSpanGroup.JOB, "rl.distillation.job") +def distillation_train( + student_policy: ColocatablePolicyInterface, + teacher_policy: ColocatablePolicyInterface, + student_generation: Optional[GenerationInterface], + dataloader: StatefulDataLoader, + val_dataloader: Optional[StatefulDataLoader], + tokenizer: TokenizerType, + loss_fn: DistillationLossFn, + task_to_env: dict[str, EnvironmentInterface], + val_task_to_env: Optional[dict[str, EnvironmentInterface]], + logger: Logger, + checkpointer: CheckpointManager, + distillation_save_state: DistillationSaveState, + master_config: MasterConfig, +) -> None: + """Run distillation training and always tear down its environments.""" + try: + _distillation_train_impl( + student_policy=student_policy, + teacher_policy=teacher_policy, + student_generation=student_generation, + dataloader=dataloader, + val_dataloader=val_dataloader, + tokenizer=tokenizer, + loss_fn=loss_fn, + task_to_env=task_to_env, + val_task_to_env=val_task_to_env, + logger=logger, + checkpointer=checkpointer, + distillation_save_state=distillation_save_state, + master_config=master_config, + ) + finally: + shutdown_environments(task_to_env, val_task_to_env) + + def validate( policy_generation: GenerationInterface, val_dataloader: Optional[StatefulDataLoader], diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 2370f2459d8..c1790dd78f2 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -86,6 +86,7 @@ ) from nemo_rl.environments.interfaces import EnvironmentInterface from nemo_rl.environments.nemo_gym import should_use_nemo_gym, spinup_nemo_gym_actor +from nemo_rl.environments.utils import shutdown_environments from nemo_rl.experience.interfaces import ( NEXT_NEMO_GYM_TASK_INDEX_KEY, ) @@ -402,35 +403,6 @@ def _needs_hf_refit_handshake( return not (nccl_reshard_refit_enabled and not colocated_inference) -def shutdown_environments( - task_to_env: dict[str, EnvironmentInterface] | None, - val_task_to_env: dict[str, EnvironmentInterface] | None, -) -> None: - """Shut down each unique environment actor before generation stops.""" - seen_environment_handles: set[int] = set() - for environment_map in (task_to_env, val_task_to_env): - if environment_map is None: - continue - for task_name, environment in environment_map.items(): - handle_id = id(environment) - if handle_id in seen_environment_handles: - continue - seen_environment_handles.add(handle_id) - - print(f"🛑 Shutting down environment {task_name}...") - try: - ray.get(environment.shutdown.remote(), timeout=10) - except Exception as shutdown_error: - print( - f"Environment {task_name} graceful shutdown failed: " - f"{shutdown_error}" - ) - try: - ray.kill(environment) - except Exception as kill_error: - print(f"Error stopping environment {task_name}: {kill_error}") - - def setup( master_config: MasterConfig, tokenizer: TokenizerType, @@ -2698,7 +2670,7 @@ def _validation_early_stop_message( ) -def grpo_train( +def _grpo_train_impl( policy: ColocatablePolicyInterface, policy_generation: Optional[GenerationInterface], wrapped_dataloader: StatefulDataLoader | MultipleDataloaderWrapper, @@ -3838,6 +3810,43 @@ def grpo_train( checkpointer.shutdown() +@trace_fn(RLSpanGroup.JOB, "rl.grpo.job") +def grpo_train( + policy: ColocatablePolicyInterface, + policy_generation: Optional[GenerationInterface], + wrapped_dataloader: StatefulDataLoader | MultipleDataloaderWrapper, + val_dataloader: Optional[StatefulDataLoader], + tokenizer: TokenizerType, + loss_fn: LossFunction, + task_to_env: dict[str, EnvironmentInterface], + val_task_to_env: Optional[dict[str, EnvironmentInterface]], + logger: Logger, + checkpointer: CheckpointManager, + grpo_save_state: GRPOSaveState, + master_config: MasterConfig, + processor: Optional[AutoProcessor] = None, +) -> None: + """Run GRPO training and always tear down its environments.""" + try: + _grpo_train_impl( + policy=policy, + policy_generation=policy_generation, + wrapped_dataloader=wrapped_dataloader, + val_dataloader=val_dataloader, + tokenizer=tokenizer, + loss_fn=loss_fn, + task_to_env=task_to_env, + val_task_to_env=val_task_to_env, + logger=logger, + checkpointer=checkpointer, + grpo_save_state=grpo_save_state, + master_config=master_config, + processor=processor, + ) + finally: + shutdown_environments(task_to_env, val_task_to_env) + + def validate( policy_generation: GenerationInterface, val_dataloader: Optional[StatefulDataLoader], diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index c4ee8c07c19..f6957ef053f 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -42,6 +42,11 @@ _get_node_ip_local, ) from nemo_rl.environments.interfaces import EnvironmentInterface +from nemo_rl.environments.nemo_gym_shards import ( + SHARDING_CONFIG_KEYS, + ShardConfigError, + parse_shard_plan, +) from nemo_rl.experience.failures import ( GymTransportError, RolloutDataFailure, @@ -53,6 +58,8 @@ from nemo_rl.utils.timer import Timer from nemo_rl.utils.venvs import create_local_venv_on_each_node +NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S = 120 + # Kept local so the Gym actor does not depend on model-config dtype resolution. # Must cover every name resolve_routed_experts_dtype can produce. _ROUTED_EXPERTS_DTYPES = { @@ -485,7 +492,11 @@ def _attach_multimodal_data_to_user_message( ) -@ray.remote(max_restarts=-1, max_task_retries=-1) # pragma: no cover +# Fail fast rather than restart. The servers this actor owns are started in +# _spinup, which Ray does not re-run after a restart, so a restarted actor is +# permanently broken: _require_spinup() rejects every later rollout call, and +# the caller never sees the RayActorError it is waiting for. +@ray.remote(max_restarts=0, max_task_retries=0) # pragma: no cover class NemoGym(EnvironmentInterface): """This environment class isn't really used for training. It's really meant as an integration wrapper around NeMo-Gym that hooks into the existing NeMo RL resource management via ray. So there is still one source of truth for resource management in NeMo RL.""" @@ -1029,13 +1040,16 @@ def _postprocess_nemo_gym_to_nemo_rl_result( return result def shutdown(self) -> None: - # Teardown runs in a finally block, so it must not turn a real training error - # into a confusing AttributeError from a never-spun-up (e.g. restarted) actor. - if self.rh is None: - return - run_helper = self.rh - self.rh = None - run_helper.shutdown() + """Stop the Gym servers. Safe to call more than once, and before spinup. + + Teardown runs in a finally block and may be requested more than once. + RunHelper.shutdown() is not idempotent, so the handle is cleared before + it is used. A failure therefore cannot leave a live handle that a later + cleanup attempt invokes again. + """ + rh, self.rh = self.rh, None + if rh is not None: + rh.shutdown() def step(self, message_log_batch, metadata): # This is not used since NeMo-Gym will handle the rollouts entirely. @@ -1301,11 +1315,34 @@ def spinup_nemo_gym_actor( """ nemo_gym_dict = dict(env_configs["nemo_gym"]) + # Validate the shards block even though only single-actor creation is wired + # up so far, so a malformed or premature sharded config fails at setup with + # a precise message instead of silently running unsharded. + shard_plan = parse_shard_plan(nemo_gym_dict) + if shard_plan is not None: + raise ShardConfigError( + f"env.nemo_gym.shards defines {len(shard_plan.shards)} shards " + f"({', '.join(s.name for s in shard_plan.shards)}), but multi-actor " + f"creation is not wired up yet. Remove 'shards' and use " + f"'config_paths' to run this job on a single actor." + ) + + # NeMo-RL-only keys are consumed here and must never reach Gym: the merged + # config is serialized into every Gym child process, and unrecognized + # dict-shaped top-level keys are parsed as server instance configs. + for key in SHARDING_CONFIG_KEYS: + nemo_gym_dict.pop(key, None) + # 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) thinking_tags = nemo_gym_dict.pop("thinking_tags", None) tokenizer_config = nemo_gym_dict.pop("tokenizer_config", None) + port_range = { + key: value + for key in ("port_range_low", "port_range_high") + if (value := nemo_gym_dict.pop(key, None)) is not None + } # Same treatment for the multimodal knobs: NemoGymConfig declares them as # top-level fields, so populate them here instead of leaving the actor to # read them back out of Gym's global config dict. @@ -1340,6 +1377,7 @@ def spinup_nemo_gym_actor( routed_experts_dtype=routed_experts_dtype, use_fastokens=use_fastokens, initial_global_config_dict=nemo_gym_dict, + **port_range, **multimodal_flags, ) @@ -1386,6 +1424,27 @@ def spinup_nemo_gym_actor( flush=True, ) actor = NemoGym.options(**nemo_gym_opts).remote(nemo_gym_cfg) - ray.get(actor._spinup.remote()) - ray.get(actor.set_tokenizer.remote(tokenizer)) + try: + ray.get(actor._spinup.remote()) + ray.get(actor.set_tokenizer.remote(tokenizer)) + except Exception: + # _spinup can fail after RunHelper has started some Gym subprocesses. + # Ask the actor to reap anything it owns, then force-stop the actor as a + # final safety net. Cleanup errors must not hide the startup failure. + try: + ray.get( + actor.shutdown.remote(), + timeout=NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S, + ) + except Exception as cleanup_error: + print( + f"Warning: NeMo-Gym actor cleanup after startup failure failed: {cleanup_error}" + ) + try: + ray.kill(actor) + except Exception as kill_error: + print( + f"Warning: NeMo-Gym actor kill after startup failure failed: {kill_error}" + ) + raise return actor diff --git a/nemo_rl/environments/nemo_gym_shards.py b/nemo_rl/environments/nemo_gym_shards.py new file mode 100644 index 00000000000..454c099c0ad --- /dev/null +++ b/nemo_rl/environments/nemo_gym_shards.py @@ -0,0 +1,570 @@ +# 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. +"""Schema for splitting one NeMo-Gym stack across several actors. + +A single ``NemoGym`` actor runs every environment in the job on one node, which +caps a run at what that node can host. Sharding partitions the Gym config into +groups that each get their own actor on their own node. + +NeMo RL composes shard configs; it never interprets them. Everything forwarded +into a shard's merge is opaque Gym config, and validation compares entry +*names*, never meanings. +""" + +import math +import re +from dataclasses import dataclass, field, fields, replace +from pathlib import PurePosixPath +from typing import Any, Mapping + +from omegaconf import OmegaConf + +# Keys under ``env.nemo_gym`` that NeMo RL consumes itself. They must never +# reach NeMo Gym: the merged config is serialized into NEMO_GYM_CONFIG_DICT for +# every one of the ~60 child processes, so anything left here is shipped to all +# of them. + +# Dict-shaped NeMo RL integration settings that are not Gym server entries. +# The actor factory or rollout code consumes these separately. +NEMO_RL_DICT_CONFIG_KEYS = frozenset({"effort_levels", "tokenizer_config"}) + +# Ray placement-group strategies a shard plan may ask for. STRICT_SPREAD is the +# point of sharding -- one actor per node -- and anything else colocates shards +# and gives up the capacity isolation. PACK exists so the mechanism can be +# exercised on a single machine. +PLACEMENT_STRATEGIES = frozenset({"STRICT_SPREAD", "SPREAD", "STRICT_PACK", "PACK"}) +DEFAULT_PLACEMENT_STRATEGY = "STRICT_SPREAD" +DEFAULT_REPLICAS = 1 +SHARD_NAME_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*$") + +# Gym's log-directory key (``nemo_gym.global_config.NEMO_GYM_LOG_DIR_KEY_NAME``). +# Restated rather than imported: this module is imported on the driver, and +# nemo_gym is installed only in the actor's venv. +GYM_LOG_DIR_KEY = "nemo_gym_log_dir" + +# Gym's server-type keys that make an entry a rollout destination. +GYM_ROUTABLE_KEYS = frozenset({"responses_api_agents", "resources_servers"}) + + +class ShardConfigError(ValueError): + """Raised for a malformed ``env.nemo_gym.shards`` block.""" + + +class ShardSetupError(RuntimeError): + """Raised when a sharded stack fails to start or fails its startup checks.""" + + +@dataclass(frozen=True) +class ShardSpec: + """One shard: a slice of the Gym config that gets its own actor. + + ``replicas`` instances are created from a single merge, so replicas are + identical apart from the node they land on. + """ + + name: str + config_paths: list[str] + inherited_overlays: frozenset[str] = frozenset() + overrides: dict[str, Any] = field(default_factory=dict) + replicas: int = DEFAULT_REPLICAS + actor_cpus: float | None = None + port_range_low: int | None = None + port_range_high: int | None = None + + +# Per-shard keys consumed by NeMo RL rather than forwarded to Gym. +SHARD_SPEC_KEYS = frozenset(field.name for field in fields(ShardSpec)) + + +@dataclass(frozen=True) +class ShardPlan: + """The parsed ``shards`` block plus the settings that apply to all shards.""" + + shards: list[ShardSpec] + common_inherited_overlays: frozenset[str] = frozenset() + common_overrides: dict[str, Any] = field(default_factory=dict) + allowed_duplicate_entries: frozenset[str] = frozenset() + placement_strategy: str = DEFAULT_PLACEMENT_STRATEGY + + +SHARDING_CONFIG_KEYS = frozenset(field.name for field in fields(ShardPlan)) + + +def _as_plain_dict(value: Any, *, context: str) -> dict[str, Any]: + """Coerce a mapping (possibly OmegaConf-backed) to a plain dict.""" + if value is None: + return {} + if OmegaConf.is_config(value): + value = OmegaConf.to_container(value, resolve=True) + if not isinstance(value, Mapping): + raise ShardConfigError( + f"{context} must be a mapping, got {type(value).__name__}" + ) + plain: dict[str, Any] = {} + for key, item in value.items(): + if not isinstance(key, str): + raise ShardConfigError(f"{context} keys must be strings") + plain[key] = item + return plain + + +def _as_string_set(value: Any, *, context: str) -> frozenset[str]: + """Coerce a list of unique strings to a set.""" + if value is None: + return frozenset() + if OmegaConf.is_config(value): + value = OmegaConf.to_container(value, resolve=True) + if not isinstance(value, list) or not all( + isinstance(entry, str) and entry for entry in value + ): + raise ShardConfigError(f"{context} must be a list of non-empty strings") + if len(value) != len(set(value)): + raise ShardConfigError(f"{context} contains duplicate keys") + return frozenset(value) + + +def _parse_shard(raw: Any, index: int) -> ShardSpec: + entry = _as_plain_dict(raw, context=f"shards[{index}]") + + unknown = set(entry) - SHARD_SPEC_KEYS + if unknown: + raise ShardConfigError( + f"shards[{index}] has unrecognized keys {sorted(unknown)}. " + f"Gym config overlays belong under this shard's 'overrides'." + ) + + name = entry.get("name") + if not isinstance(name, str) or not name.strip(): + raise ShardConfigError(f"shards[{index}] needs a non-empty string 'name'") + if not SHARD_NAME_PATTERN.fullmatch(name) or name in {".", ".."}: + raise ShardConfigError( + f"shard name {name!r} must be one safe path component containing " + "only letters, numbers, dots, underscores, and hyphens" + ) + + config_paths = entry.get("config_paths") + if OmegaConf.is_config(config_paths): + config_paths = OmegaConf.to_container(config_paths, resolve=True) + if not isinstance(config_paths, list) or not config_paths: + raise ShardConfigError( + f"shard '{name}' needs at least one entry in 'config_paths'" + ) + if not all(isinstance(path, str) and path for path in config_paths): + raise ShardConfigError(f"shard '{name}' config_paths must be non-empty strings") + + replicas = entry.get("replicas", DEFAULT_REPLICAS) + if not isinstance(replicas, int) or isinstance(replicas, bool) or replicas < 1: + raise ShardConfigError( + f"shard '{name}' has replicas={replicas!r}; must be an integer >= 1" + ) + + port_low = entry.get("port_range_low") + port_high = entry.get("port_range_high") + if (port_low is None) != (port_high is None): + raise ShardConfigError( + f"shard '{name}' must set both port_range_low and port_range_high, or neither" + ) + if port_low is not None: + if ( + not isinstance(port_low, int) + or isinstance(port_low, bool) + or not isinstance(port_high, int) + or isinstance(port_high, bool) + or port_low < 1 + or port_high > 65_535 + ): + raise ShardConfigError( + f"shard '{name}' port range must use integers within [1, 65535]" + ) + if port_low >= port_high: + raise ShardConfigError( + f"shard '{name}' must reserve at least two ports in the inclusive " + f"range [{port_low}, {port_high}]" + ) + + actor_cpus = entry.get("actor_cpus") + if actor_cpus is not None and ( + not isinstance(actor_cpus, (int, float)) + or isinstance(actor_cpus, bool) + or not math.isfinite(actor_cpus) + or actor_cpus <= 0 + ): + raise ShardConfigError( + f"shard '{name}' has actor_cpus={actor_cpus!r}; " + "must be a positive finite number" + ) + + return ShardSpec( + name=name, + config_paths=list(config_paths), + inherited_overlays=_as_string_set( + entry.get("inherited_overlays"), + context=f"shard '{name}' inherited_overlays", + ), + overrides=_as_plain_dict( + entry.get("overrides"), context=f"shard '{name}' overrides" + ), + replicas=replicas, + actor_cpus=float(actor_cpus) if actor_cpus is not None else None, + port_range_low=port_low, + port_range_high=port_high, + ) + + +def find_gym_config_entries(nemo_gym_config: Mapping[str, Any]) -> list[str]: + """Return top-level keys that Gym would read as server instance configs. + + Deliberately broader than Gym's own rule + (``filter_for_server_instance_configs``), which also requires the value to + validate as a ``ServerInstanceConfig``: this module cannot import Gym, so + any dict-shaped top-level key that is not a known setting is treated as an + entry overlay. Under sharding these values are inherited from the resolved + parent config. A shard claims them by key so their YAML bodies do not need + to be copied. + """ + return sorted( + key + for key, value in nemo_gym_config.items() + if key not in SHARDING_CONFIG_KEYS + and key not in NEMO_RL_DICT_CONFIG_KEYS + and (isinstance(value, Mapping) or OmegaConf.is_dict(value)) + ) + + +def parse_shard_plan(nemo_gym_config: Mapping[str, Any]) -> ShardPlan | None: + """Parse ``env.nemo_gym.shards``, or return None when the job is unsharded. + + Returning None is the backward-compatible path: a config with + ``config_paths`` and no ``shards`` behaves exactly as it did before + sharding existed. + + Raises: + ShardConfigError: the shards block is malformed, or it is combined with + config that has no unambiguous shard to belong to. + """ + if "shards" not in nemo_gym_config: + stray = sorted( + key + for key in SHARDING_CONFIG_KEYS + if key != "shards" and key in nemo_gym_config + ) + if stray: + raise ShardConfigError( + f"env.nemo_gym has sharding keys {stray} but no 'shards' block. " + "Add a 'shards' block, or remove these keys." + ) + return None + + raw_shards = nemo_gym_config["shards"] + if OmegaConf.is_config(raw_shards): + raw_shards = OmegaConf.to_container(raw_shards, resolve=True) + if not isinstance(raw_shards, list) or not raw_shards: + raise ShardConfigError("env.nemo_gym.shards must be a non-empty list") + + # A top-level config_paths could plausibly belong to one shard or all of + # them, so its scope must remain explicit. A null reads as absent because + # Hydra can blank an inherited key but cannot delete it. + if nemo_gym_config.get("config_paths") is not None: + raise ShardConfigError( + "env.nemo_gym.config_paths cannot be combined with 'shards'. " + "Move each path into the shard that should host it, and set the " + "inherited config_paths to null." + ) + shards = [_parse_shard(raw, index) for index, raw in enumerate(raw_shards)] + + duplicate_names = sorted( + { + shard.name + for shard in shards + if [s.name for s in shards].count(shard.name) > 1 + } + ) + if duplicate_names: + raise ShardConfigError(f"Duplicate shard names: {duplicate_names}") + + inherited_inventory = frozenset(find_gym_config_entries(nemo_gym_config)) + common_inherited = _as_string_set( + nemo_gym_config.get("common_inherited_overlays"), + context="env.nemo_gym.common_inherited_overlays", + ) + claims: dict[str, list[str]] = {} + for key in common_inherited: + claims.setdefault(key, []).append("common_inherited_overlays") + for shard in shards: + for key in shard.inherited_overlays: + claims.setdefault(key, []).append(f"shard '{shard.name}'") + + unknown = sorted(set(claims) - inherited_inventory) + if unknown: + raise ShardConfigError( + f"Inherited overlay keys {unknown} do not exist in the resolved " + f"env.nemo_gym config. Available overlays: {sorted(inherited_inventory)}." + ) + + duplicated = { + key: owners for key, owners in sorted(claims.items()) if len(owners) > 1 + } + if duplicated: + raise ShardConfigError( + "Each inherited Gym overlay must be claimed once; duplicate claims: " + f"{duplicated}." + ) + + unclaimed = inherited_inventory - set(claims) + if len(shards) == 1: + if unclaimed: + shards[0] = replace( + shards[0], + inherited_overlays=shards[0].inherited_overlays | unclaimed, + ) + elif unclaimed: + raise ShardConfigError( + f"Inherited Gym overlays {sorted(unclaimed)} have no owner. List each " + f"key under one shard's 'inherited_overlays' or under " + f"'common_inherited_overlays'. Complete overlay inventory: " + f"{sorted(inherited_inventory)}." + ) + + configured_strategy = nemo_gym_config.get("placement_strategy") + placement_strategy = ( + DEFAULT_PLACEMENT_STRATEGY + if configured_strategy is None + else configured_strategy + ) + if placement_strategy not in PLACEMENT_STRATEGIES: + raise ShardConfigError( + f"env.nemo_gym.placement_strategy must be one of " + f"{sorted(PLACEMENT_STRATEGIES)}, got {placement_strategy!r}" + ) + + # Replicas are stamped from one merge, so they share that shard's port + # range and no per-replica override exists to give them separate ones. + # Only STRICT_SPREAD guarantees each one its own node, and therefore its + # own port space; under any other strategy two replicas can land together + # and race for the same ports at spinup. + replicated = [shard.name for shard in shards if shard.replicas > 1] + if replicated and placement_strategy != DEFAULT_PLACEMENT_STRATEGY: + raise ShardConfigError( + f"Shards {replicated} declare replicas, which requires " + f"placement_strategy {DEFAULT_PLACEMENT_STRATEGY} so that each " + f"replica gets its own node and port space; got " + f"{placement_strategy!r}. Replicas share one merged config, so " + f"they cannot be given separate port ranges. Either drop the " + f"replicas or split them into separate shards with their own " + f"port_range_low/high." + ) + + if placement_strategy != DEFAULT_PLACEMENT_STRATEGY: + missing_ranges = [ + shard.name for shard in shards if shard.port_range_low is None + ] + if missing_ranges: + raise ShardConfigError( + f"Shards {missing_ranges} require explicit port_range_low/high " + f"with placement_strategy {placement_strategy!r}, because relaxed " + "placement can colocate shard stacks on one node." + ) + ranges: list[tuple[int, int, str]] = [] + for shard in shards: + assert shard.port_range_low is not None + assert shard.port_range_high is not None + ranges.append((shard.port_range_low, shard.port_range_high, shard.name)) + ranges.sort() + for (_, previous_high, previous_name), ( + current_low, + _, + current_name, + ) in zip(ranges, ranges[1:]): + if current_low <= previous_high: + raise ShardConfigError( + f"Shards '{previous_name}' and '{current_name}' have " + "overlapping port ranges under relaxed placement" + ) + + common_overrides = _as_plain_dict( + nemo_gym_config.get("common_overrides"), + context="env.nemo_gym.common_overrides", + ) + if "config_paths" in common_overrides: + raise ShardConfigError( + "env.nemo_gym.common_overrides cannot set config_paths; " + "declare paths on each shard" + ) + wrongly_common = sorted( + key + for key in common_overrides + if key in inherited_inventory and key not in common_inherited + ) + if wrongly_common: + raise ShardConfigError( + f"common_overrides cannot override inherited entries {wrongly_common} " + "because they are not claimed by common_inherited_overlays" + ) + for shard in shards: + if "config_paths" in shard.overrides: + raise ShardConfigError( + f"shard '{shard.name}' overrides cannot set config_paths; " + "use the shard's config_paths field" + ) + wrongly_owned = sorted( + key + for key in shard.overrides + if key in inherited_inventory + and key not in common_inherited + and key not in shard.inherited_overlays + ) + if wrongly_owned: + raise ShardConfigError( + f"shard '{shard.name}' overrides inherited entries {wrongly_owned} " + "that are claimed by another shard" + ) + + raw_allowed = _as_string_set( + nemo_gym_config.get("allowed_duplicate_entries"), + context="env.nemo_gym.allowed_duplicate_entries", + ) + + return ShardPlan( + shards=shards, + common_inherited_overlays=common_inherited, + common_overrides=common_overrides, + allowed_duplicate_entries=raw_allowed, + placement_strategy=placement_strategy, + ) + + +def apply_shard_overlay( + base_config: dict[str, Any], plan: ShardPlan, shard: ShardSpec +) -> dict[str, Any]: + """Build one shard's Gym config from shared settings and its claimed overlays. + + The ``config_paths`` merge itself still happens inside Gym, exactly as for + an unsharded job. Dict-shaped entries inherited from a parent recipe are + removed from the shared base and restored only for the shards that claim + them. Explicit overrides win over inherited values. + + Top-level keys left null are dropped rather than forwarded. Blanking a key + is the only way a Hydra override can retract one it inherited (see + ``parse_shard_plan``), and Gym should not receive a null where it expects + an entry. + """ + inherited_entries = set(find_gym_config_entries(base_config)) + shared_base = { + key: value + for key, value in base_config.items() + if key not in inherited_entries and key not in SHARDING_CONFIG_KEYS + } + selected_inherited = { + key: base_config[key] + for key in plan.common_inherited_overlays | shard.inherited_overlays + } + merged = _as_plain_dict( + OmegaConf.to_container( + OmegaConf.merge( + OmegaConf.create(shared_base), + OmegaConf.create(selected_inherited), + OmegaConf.create(plan.common_overrides), + OmegaConf.create(shard.overrides), + ), + resolve=True, + ), + context=f"merged config for shard '{shard.name}'", + ) + merged = {key: value for key, value in merged.items() if value is not None} + merged["config_paths"] = list(shard.config_paths) + if shard.port_range_low is not None: + merged["port_range_low"] = shard.port_range_low + merged["port_range_high"] = shard.port_range_high + return merged + + +def build_route_shard_map( + entries_by_shard: Mapping[str, Mapping[str, list[str]]], + allowed_duplicate_entries: frozenset[str] | set[str] = frozenset(), +) -> dict[str, str]: + """Map each routable entry to its shard, rejecting ambiguous routes. + + Takes what each shard reported from ``NemoGym.list_entries()`` and returns + ``{route_name: shard_name}``, the lookup the router dispatches on. Current + Gym rows route either by ``agent_ref.name`` or by ``task_source``. The + latter names the agent or resources-server entry that declared the dataset, + so both entry types must be included. + + Two failures are caught here rather than at first dispatch. A routable + entry hosted by two shards is always an error: rows naming it could go to + either, so routing would be silently nondeterministic. Any other entry in + two shards has to be allowlisted, because duplication is usually accidental + — a shared YAML dropped into two shards' path lists quietly brings its judge + along and doubles that judge's GPU claim. + + Only names are compared. What an entry means is Gym's business. + """ + route_to_shard: dict[str, str] = {} + for shard_name, entries in entries_by_shard.items(): + for entry, types in entries.items(): + if not GYM_ROUTABLE_KEYS.intersection(types): + continue + if entry in route_to_shard: + raise ShardSetupError( + f"Routable Gym entry '{entry}' is hosted by both shard " + f"'{route_to_shard[entry]}' and shard '{shard_name}'. An " + f"agent or resources server must live in exactly one shard " + f"so rows naming it have one destination." + ) + route_to_shard[entry] = shard_name + + hosting_shard: dict[str, str] = {} + for shard_name, entries in entries_by_shard.items(): + for entry in entries: + if entry in hosting_shard and entry not in allowed_duplicate_entries: + raise ShardSetupError( + f"Config entry '{entry}' appears in shard " + f"'{hosting_shard[entry]}' and shard '{shard_name}' but is " + f"not listed in allowed_duplicate_entries. If this entry " + f"starts a model engine, duplicating it doubles its GPU " + f"claim; if the duplication is intended, allowlist it." + ) + hosting_shard.setdefault(entry, shard_name) + + return route_to_shard + + +def apply_shard_log_dir( + gym_config: Mapping[str, Any], + shard_name: str, + *, + replica_index: int | None = None, +) -> dict[str, Any]: + """Give one shard instance its own subdirectory of the configured log dir. + + Gym names each server's log file after the server and appends to it with + ``tee -a`` (``cli/setup_command.py``). One log dir shared by every instance + therefore interleaves output whenever two instances host the same server + name — which replicas always do, since they are stamped from one merge — and + on shared storage that happens across nodes without any error. + + Returns the config unchanged when no log dir is set, which is the default. + """ + config = dict(gym_config) + log_dir = config.get(GYM_LOG_DIR_KEY) + if not log_dir: + return config + + # Gym applies the same sanitization to server names before using them as + # path components. + suffix = shard_name.replace("/", "_") + if replica_index is not None: + suffix = f"{suffix}/{replica_index}" + config[GYM_LOG_DIR_KEY] = str(PurePosixPath(str(log_dir)) / suffix) + return config diff --git a/nemo_rl/environments/utils.py b/nemo_rl/environments/utils.py index 1dbd691a583..20702259a0d 100644 --- a/nemo_rl/environments/utils.py +++ b/nemo_rl/environments/utils.py @@ -12,14 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. import os -from typing import Any, Dict, NotRequired, TypedDict +from typing import Any, Dict, NotRequired, Optional, TypedDict +import ray from hydra.utils import get_object from nemo_rl.distributed.ray_actor_environment_registry import get_actor_python_env from nemo_rl.environments.interfaces import EnvironmentInterface from nemo_rl.utils.venvs import create_local_venv_on_each_node +DEFAULT_ENV_SHUTDOWN_TIMEOUT_SECONDS = 10.0 + # Environment registry entry schema. class EnvRegistryEntry(TypedDict, total=False): @@ -137,3 +140,43 @@ def register_env(env_name: str, actor_class_fqn: str) -> None: raise ValueError(f"Env name {env_name} already registered") ENV_REGISTRY[env_name] = {"actor_class_fqn": actor_class_fqn} + + +def shutdown_environments( + *env_maps: Optional[Dict[str, EnvironmentInterface]], + timeout: float = DEFAULT_ENV_SHUTDOWN_TIMEOUT_SECONDS, +) -> None: + """Gracefully shut down every distinct environment actor in the given maps. + + Runners commonly bind the same actor — often the very same mapping — to + both training and validation, so each handle is asked to stop once no + matter how many maps contain it. + + Environments must be torn down before generation workers: they may have + in-flight HTTP requests to the vLLM endpoints, and killing generation first + leaves them retrying dead connections. + + Args: + env_maps: Task-name to environment mappings. ``None`` and empty + mappings are skipped. + timeout: Seconds to wait for each actor's ``shutdown()`` before killing it. + """ + seen: set[int] = set() + for env_map in env_maps: + if not env_map: + continue + for task_name, env in env_map.items(): + if id(env) in seen: + continue + seen.add(id(env)) + print(f"🛑 Shutting down environment {task_name}...") + try: + ray.get(env.shutdown.remote(), timeout=timeout) + except Exception as e: + print(f"Graceful shutdown of environment {task_name} failed: {e}") + try: + ray.kill(env) + except Exception as kill_error: + print( + f"Error killing environment {task_name}: {kill_error}", + ) diff --git a/pyrefly.toml b/pyrefly.toml index aee1da0635f..de906d6c955 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -140,6 +140,7 @@ project-includes = [ "nemo_rl/environments/math_environment.py", "nemo_rl/environments/metrics.py", "nemo_rl/environments/nemo_gym_video.py", + "nemo_rl/environments/nemo_gym_shards.py", "nemo_rl/environments/nemotron_utils.py", "nemo_rl/environments/rewards.py", "nemo_rl/environments/utils.py", diff --git a/tests/unit/algorithms/test_distillation.py b/tests/unit/algorithms/test_distillation.py index 8fb307fd0fc..97fb13c72ce 100644 --- a/tests/unit/algorithms/test_distillation.py +++ b/tests/unit/algorithms/test_distillation.py @@ -917,6 +917,64 @@ def test_noncolocated_inference_requires_explicit_gpus_per_node_single_node(): setup(master_config, tokenizer, dataset, None) +def test_distillation_train_shuts_down_environments_after_failure(): + task_to_env = {"nemo_gym": MagicMock()} + val_task_to_env = task_to_env + + with ( + patch.object( + distil_mod, + "_distillation_train_impl", + side_effect=RuntimeError("rollout failed"), + ), + patch.object(distil_mod, "shutdown_environments") as shutdown, + pytest.raises(RuntimeError, match="rollout failed"), + ): + distillation_train( + student_policy=MagicMock(), + teacher_policy=MagicMock(), + student_generation=MagicMock(), + dataloader=MagicMock(), + val_dataloader=None, + tokenizer=MagicMock(), + loss_fn=MagicMock(), + task_to_env=task_to_env, + val_task_to_env=val_task_to_env, + logger=MagicMock(), + checkpointer=MagicMock(), + distillation_save_state=MagicMock(), + master_config=MagicMock(), + ) + + shutdown.assert_called_once_with(task_to_env, val_task_to_env) + + +def test_distillation_train_shuts_down_environments_after_success(): + task_to_env = {"nemo_gym": MagicMock()} + + with ( + patch.object(distil_mod, "_distillation_train_impl"), + patch.object(distil_mod, "shutdown_environments") as shutdown, + ): + distillation_train( + student_policy=MagicMock(), + teacher_policy=MagicMock(), + student_generation=MagicMock(), + dataloader=MagicMock(), + val_dataloader=None, + tokenizer=MagicMock(), + loss_fn=MagicMock(), + task_to_env=task_to_env, + val_task_to_env=task_to_env, + logger=MagicMock(), + checkpointer=MagicMock(), + distillation_save_state=MagicMock(), + master_config=MagicMock(), + ) + + shutdown.assert_called_once_with(task_to_env, task_to_env) + + @pytest.mark.parametrize("refit_transport", [None, "nixl"]) def test_distillation_setup_non_colocated_smoke(monkeypatch, refit_transport): """Smoke test: calling setup with a non-colocated config should succeed.""" diff --git a/tests/unit/algorithms/test_grpo.py b/tests/unit/algorithms/test_grpo.py index 403bf2c243d..5e2a4f6fe9b 100644 --- a/tests/unit/algorithms/test_grpo.py +++ b/tests/unit/algorithms/test_grpo.py @@ -55,7 +55,6 @@ grpo_train, refit_policy_generation, setup, - shutdown_environments, validate, ) from nemo_rl.algorithms.grpo_sync import _train_fields_for_step, grpo_train_sync @@ -1446,35 +1445,6 @@ def test_async_grpo_awaits_resume_after_refit_failure(mock_grpo_components) -> N ) -def test_shutdown_environments_drains_unique_actors_before_kill() -> None: - shared_environment = MagicMock() - failing_environment = MagicMock() - shared_shutdown_ref = object() - failing_shutdown_ref = object() - shared_environment.shutdown.remote.return_value = shared_shutdown_ref - failing_environment.shutdown.remote.return_value = failing_shutdown_ref - - def get_or_fail(ref, timeout=None): - assert timeout == 10 - if ref is failing_shutdown_ref: - raise RuntimeError("environment shutdown failed") - assert ref is shared_shutdown_ref - return True - - with ( - patch("nemo_rl.algorithms.grpo.ray.get", side_effect=get_or_fail), - patch("nemo_rl.algorithms.grpo.ray.kill") as ray_kill, - ): - shutdown_environments( - {"train": shared_environment, "failing": failing_environment}, - {"validation": shared_environment}, - ) - - shared_environment.shutdown.remote.assert_called_once_with() - failing_environment.shutdown.remote.assert_called_once_with() - ray_kill.assert_called_once_with(failing_environment) - - def test_should_use_nemo_gym_requires_dynamo_token_wrapper() -> None: master_config = MagicMock() master_config.env = {"should_use_nemo_gym": True} @@ -5297,3 +5267,58 @@ def test_train_fields_for_step(skip_prev_logprobs, expect_prev): ) def test_needs_hf_refit_handshake(backend, nccl_reshard, colocated, expected): assert _needs_hf_refit_handshake(backend, nccl_reshard, colocated) is expected + + +def test_grpo_train_shuts_down_environments_after_failure(): + task_to_env = {"nemo_gym": MagicMock()} + val_task_to_env = task_to_env + + with ( + patch( + "nemo_rl.algorithms.grpo._grpo_train_impl", + side_effect=RuntimeError("rollout failed"), + ), + patch("nemo_rl.algorithms.grpo.shutdown_environments") as shutdown, + pytest.raises(RuntimeError, match="rollout failed"), + ): + grpo_train( + policy=MagicMock(), + policy_generation=MagicMock(), + wrapped_dataloader=MagicMock(), + val_dataloader=None, + tokenizer=MagicMock(), + loss_fn=MagicMock(), + task_to_env=task_to_env, + val_task_to_env=val_task_to_env, + logger=MagicMock(), + checkpointer=MagicMock(), + grpo_save_state=MagicMock(), + master_config=MagicMock(), + ) + + shutdown.assert_called_once_with(task_to_env, val_task_to_env) + + +def test_grpo_train_shuts_down_environments_after_success(): + task_to_env = {"nemo_gym": MagicMock()} + + with ( + patch("nemo_rl.algorithms.grpo._grpo_train_impl"), + patch("nemo_rl.algorithms.grpo.shutdown_environments") as shutdown, + ): + grpo_train( + policy=MagicMock(), + policy_generation=MagicMock(), + wrapped_dataloader=MagicMock(), + val_dataloader=None, + tokenizer=MagicMock(), + loss_fn=MagicMock(), + task_to_env=task_to_env, + val_task_to_env=task_to_env, + logger=MagicMock(), + checkpointer=MagicMock(), + grpo_save_state=MagicMock(), + master_config=MagicMock(), + ) + + shutdown.assert_called_once_with(task_to_env, task_to_env) diff --git a/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml b/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml new file mode 100644 index 00000000000..a852a6f3ff8 --- /dev/null +++ b/tests/unit/environments/nemo_gym_test_data/l0_rollout_acceptance.yaml @@ -0,0 +1,158 @@ +cases: + - name: math_with_judge + config_path: resources_servers/math_with_judge/configs/math_with_judge.yaml + data_path: resources_servers/math_with_judge/data/example.jsonl + example_index: 0 + example_sha256: 5e8d77674814a6b1fdcd0f2dbef0c8597cc109814763056070ff2eeb0c9d70ef # pragma: allowlist secret + agent_ref: + type: responses_api_agents + name: math_with_judge_simple_agent + expected_generations: + - '\boxed{32}' + expected_prompt_fragment: 1000 digit numbers + rejected_generations: + - '\boxed{999999}' + expected_reward: 1.0 + expected_result: + extracted_answer: "32" + library_reward: 1.0 + + - name: code_gen + config_path: resources_servers/code_gen/configs/code_gen.yaml + data_path: resources_servers/code_gen/data/example.jsonl + example_index: 4 + example_sha256: ceea8ed47c6615783bc9df8140110b2aa86729438448ee04e74c5d4850ebff24 # pragma: allowlist secret + agent_ref: + type: responses_api_agents + name: code_gen_simple_agent + expected_generations: + - |- + ```python + import sys + + values = iter(map(int, sys.stdin.buffer.read().split())) + test_count = next(values) + answers = [] + for _ in range(test_count): + n = next(values) + positions = [0] * (n + 1) + for index in range(n): + positions[next(values)] = index + left = right = positions[1] + bits = [] + for value in range(1, n + 1): + left = min(left, positions[value]) + right = max(right, positions[value]) + bits.append("1" if right - left + 1 == value else "0") + answers.append("".join(bits)) + sys.stdout.write("\n".join(answers)) + ``` + expected_prompt_fragment: beautiful number + rejected_generations: + - |- + ```python + print(0) + ``` + expected_reward: 1.0 + expected_result: + result: + - true + - true + - true + + - name: single_step_tool_use_with_argument_comparison + config_path: resources_servers/single_step_tool_use_with_argument_comparison/configs/single_step_tool_use_with_argument_comparison.yaml + data_path: resources_servers/single_step_tool_use_with_argument_comparison/data/example.jsonl + example_index: 4 + example_sha256: 2e24731dfeda02e16351dd5abd5c1e3fd8307554f4824cf4fb3630672c63f5ca # pragma: allowlist secret + agent_ref: + type: responses_api_agents + name: single_step_tool_use_with_argument_comparison_agent + expected_generations: + - '{"name":"check_seat_availability","arguments":{"event_id":"SHOW24","section":"Medical Zone"}}' + expected_prompt_fragment: SHOW24 + rejected_generations: + - I cannot call the requested tool. + expected_reward: 1.0 + expected_result: + category: A tool call that matches the expected tool call was found + expected_action: + type: function_call + name: check_seat_availability + arguments: '{"event_id": "SHOW24", "section": "Medical Zone"}' + + - name: mcqa + config_path: resources_servers/mcqa/configs/mcqa.yaml + data_path: resources_servers/mcqa/data/example.jsonl + example_index: 0 + example_sha256: 3e3bd43ed3707cc35f84edf0a57357d0cc0f8dfd7692991d0daf05ad421d0846 # pragma: allowlist secret + agent_ref: + type: responses_api_agents + name: mcqa_simple_agent + expected_generations: + - '\boxed{B}' + expected_prompt_fragment: cystic fibrosis + rejected_generations: + - '\boxed{A}' + expected_reward: 1.0 + expected_result: + expected_answer: B + extracted_answer: B + + - name: equivalence_llm_judge + config_path: resources_servers/equivalence_llm_judge/configs/equivalence_llm_judge.yaml + data_path: resources_servers/equivalence_llm_judge/data/example.jsonl + example_index: 0 + example_sha256: efabdc976b384f6f48ee5e42cf855a66e6fba84feedd1b15fd30f256f56ecd9b # pragma: allowlist secret + agent_ref: + type: responses_api_agents + name: equivalence_llm_judge_simple_agent + expected_generations: + - '\boxed{Charles Darwin}' + expected_prompt_fragment: theory of evolution by natural selection + rejected_generations: + - '\boxed{Ada Lovelace}' + rejected_judge_verdicts: + - "[[A!=B]]" + expected_reward: 1.0 + expected_result: + expected_answer: darwin + expected_judge_verdicts: + - '[[A=B]]' + + - name: structured_outputs_v4 + config_path: resources_servers/structured_outputs/configs/structured_outputs_v4.yaml + data_path: resources_servers/structured_outputs/data/structured_outputs_v4_example.jsonl + example_index: 0 + example_sha256: 138b2f4383ec722989fdefd60ea0f1724712037e46e9d62ca2df0cc14308e1a4 # pragma: allowlist secret + agent_ref: + type: responses_api_agents + name: structured_outputs_v4_simple_agent + expected_generations: + - '{"name":"response_tool_8","arguments":{"extraction":{"name":"Dizer Kola","native_name":"ديزركلا","romanized_name":"Dizer Kola","settlement_type":"village","country":"Iran","province":"Mazandaran","county":"Nowshahr","bakhsh":"Central","rural_district":"Baladeh Kojur","coordinates":{"latitude":36.55694,"longitude":51.79389},"population_total":250,"population_year":2006,"number_of_families":64,"timezone_standard":"UTC+3:30 (IRST)","timezone_dst":"UTC+4:30 (IRDT)"}}}' + expected_prompt_fragment: Dizer Kola + rejected_generations: + - I cannot extract that. + expected_reward: 1.0 + expected_result: + schema_type: json + error_type: null + + - name: workplace_assistant + config_path: resources_servers/workplace_assistant/configs/workplace_assistant.yaml + data_path: resources_servers/workplace_assistant/data/example.jsonl + example_index: 0 + example_sha256: c7a80ae027d7bf89dea4700acc00a20d37400a91a19dbe9908d4b945367b657a # pragma: allowlist secret + agent_ref: + type: responses_api_agents + name: workplace_assistant_simple_agent + expected_generations: + - '{"name":"email_reply_email","arguments":{"email_id":"00000057","body":"Thanks for the update - I will get back to you tomorrow."}}' + - Done + expected_prompt_fragment: Task Update on Develop prototype + rejected_generations: + - Done + expected_reward: 1.0 + expected_result: {} + expected_tool_outputs: + - '{"output":"Email replied successfully."}' diff --git a/tests/unit/environments/test_environment_utils.py b/tests/unit/environments/test_environment_utils.py index 37f36589502..f5c7195a90b 100644 --- a/tests/unit/environments/test_environment_utils.py +++ b/tests/unit/environments/test_environment_utils.py @@ -11,9 +11,15 @@ # 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. +from unittest.mock import MagicMock, patch + import pytest -from nemo_rl.environments.utils import ENV_REGISTRY, register_env +from nemo_rl.environments.utils import ( + ENV_REGISTRY, + register_env, + shutdown_environments, +) def test_register_new_env_success(): @@ -50,3 +56,74 @@ def test_register_env_duplicate_raises_error(): # Restore original registry state ENV_REGISTRY.clear() ENV_REGISTRY.update(original_registry) + + +def _fake_env(name): + env = MagicMock(name=name) + env.shutdown.remote.return_value = f"{name}-ref" + return env + + +def test_shutdown_environments_dedupes_shared_handles(): + """Runners can bind one actor to train and val; it must be stopped once.""" + env = _fake_env("gym") + task_to_env = {"nemo_gym": env} + val_task_to_env = task_to_env # the aliasing real runners use + + with patch("nemo_rl.environments.utils.ray") as mock_ray: + shutdown_environments(task_to_env, val_task_to_env) + + env.shutdown.remote.assert_called_once_with() + mock_ray.get.assert_called_once_with("gym-ref", timeout=10.0) + mock_ray.kill.assert_not_called() + + +def test_shutdown_environments_stops_each_distinct_actor(): + train_env, val_env = _fake_env("train"), _fake_env("val") + + with patch("nemo_rl.environments.utils.ray") as mock_ray: + shutdown_environments({"a": train_env}, {"a": val_env}, None) + + train_env.shutdown.remote.assert_called_once_with() + val_env.shutdown.remote.assert_called_once_with() + assert mock_ray.get.call_count == 2 + mock_ray.kill.assert_not_called() + + +def test_shutdown_environments_continues_after_a_failure(): + wedged, healthy = _fake_env("wedged"), _fake_env("healthy") + + with patch("nemo_rl.environments.utils.ray") as mock_ray: + mock_ray.get.side_effect = ( + lambda ref, timeout: None + if ref == "healthy-ref" + else (_ for _ in ()).throw(TimeoutError("no response")) + ) + shutdown_environments({"wedged": wedged, "healthy": healthy}) + + mock_ray.kill.assert_called_once_with(wedged) + healthy.shutdown.remote.assert_called_once_with() + + +def test_shutdown_environments_continues_after_a_failed_kill(): + wedged, healthy = _fake_env("wedged"), _fake_env("healthy") + + with patch("nemo_rl.environments.utils.ray") as mock_ray: + mock_ray.get.side_effect = ( + lambda ref, timeout: None + if ref == "healthy-ref" + else (_ for _ in ()).throw(TimeoutError("no response")) + ) + mock_ray.kill.side_effect = RuntimeError("already gone") + shutdown_environments({"wedged": wedged, "healthy": healthy}) + + mock_ray.kill.assert_called_once_with(wedged) + healthy.shutdown.remote.assert_called_once_with() + + +def test_shutdown_environments_tolerates_empty_input(): + with patch("nemo_rl.environments.utils.ray") as mock_ray: + shutdown_environments(None, {}) + + mock_ray.get.assert_not_called() + mock_ray.kill.assert_not_called() diff --git a/tests/unit/environments/test_nemo_gym_rollout.py b/tests/unit/environments/test_nemo_gym_rollout.py new file mode 100644 index 00000000000..6ed806f5dd3 --- /dev/null +++ b/tests/unit/environments/test_nemo_gym_rollout.py @@ -0,0 +1,678 @@ +# 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 json +import math +from copy import deepcopy +from hashlib import sha256 +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from threading import Thread +from typing import Any + +import pytest +import ray +import torch +from transformers import ByT5Tokenizer +from yaml import safe_load + +from nemo_rl.data.collate_fn import rl_collate_fn +from nemo_rl.data.datasets.response_datasets import NemoGymDataset +from nemo_rl.data.interfaces import DatumSpec +from nemo_rl.data.processors import nemo_gym_data_processor +from nemo_rl.distributed.virtual_cluster import _get_node_ip_local +from nemo_rl.environments.nemo_gym import ( + NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S, + spinup_nemo_gym_actor, +) +from nemo_rl.experience.rollouts import run_nemo_gym_rollout_sync + +_REPO_ROOT = Path(__file__).parents[3] +_GYM_ROOT = _REPO_ROOT / "3rdparty/Gym-workspace/Gym" +_CASES_PATH = Path(__file__).parent / "nemo_gym_test_data/l0_rollout_acceptance.yaml" +# "L0" names the RL CI tier chosen for this test. The exercised boundary is +# the conceptual L2 contract: a real Gym environment rolls out through NeMo RL. +_POLICY_MODEL_CONFIG = ( + "responses_api_models/vllm_model/configs/vllm_model_for_training.yaml" +) +_REQUIRED_L0_CASES = { + "code_gen", + "equivalence_llm_judge", + "math_with_judge", + "mcqa", + "single_step_tool_use_with_argument_comparison", + "structured_outputs_v4", + "workplace_assistant", +} +_GENERATION_CONFIG = { + "backend": "test", + "max_new_tokens": 1024, + # Tool-heavy environments such as workplace_assistant serialize more than + # 64K byte-level tokens before a continuation generation. + "max_total_sequence_length": 131072, + "temperature": 0.0, + "top_p": 1.0, + "top_k": None, + "stop_token_ids": None, + "stop_strings": None, +} +_CODE_GEN_COMPLETION = """```python +import sys + +values = iter(map(int, sys.stdin.buffer.read().split())) +test_count = next(values) +answers = [] +for _ in range(test_count): + n = next(values) + positions = [0] * (n + 1) + for index in range(n): + positions[next(values)] = index + left = right = positions[1] + bits = [] + for value in range(1, n + 1): + left = min(left, positions[value]) + right = max(right, positions[value]) + bits.append("1" if right - left + 1 == value else "0") + answers.append("".join(bits)) +sys.stdout.write("\\n".join(answers)) +```""" +_TOOL_CALL_ARGUMENTS = '{"event_id": "SHOW24", "section": "Medical Zone"}' +_TOOL_CALL_GENERATION = ( + '{"name":"check_seat_availability","arguments":' + '{"event_id":"SHOW24","section":"Medical Zone"}}' +) +_MCQA_COMPLETION = r"\boxed{B}" +_EQUIVALENCE_COMPLETION = r"\boxed{Charles Darwin}" +_EQUIVALENCE_JUDGE_COMPLETION = ( + "The candidate identifies Darwin, so it matches the reference.\n\n" + "[[A=B]] they are equivalent" +) +_STRUCTURED_OUTPUTS_ARGUMENTS = { + "name": "Dizer Kola", + "native_name": "ديزركلا", + "romanized_name": "Dizer Kola", + "settlement_type": "village", + "country": "Iran", + "province": "Mazandaran", + "county": "Nowshahr", + "bakhsh": "Central", + "rural_district": "Baladeh Kojur", + "coordinates": {"latitude": 36.55694, "longitude": 51.79389}, + "population_total": 250, + "population_year": 2006, + "number_of_families": 64, + "timezone_standard": "UTC+3:30 (IRST)", + "timezone_dst": "UTC+4:30 (IRDT)", +} +_WORKPLACE_ARGUMENTS = { + "email_id": "00000057", + "body": "Thanks for the update - I will get back to you tomorrow.", +} +_BYTE_TOKEN_OFFSET = 3 +_CONTINUATION_MARKER = "\x02nemo-rl-continuation\x00" +_REJECTED_ROLLOUT_MARKER = "[NEMO_RL_L0_REJECTED_ROLLOUT]" + + +def _tool_call_message( + name: str, arguments: dict[str, Any], call_id: str +) -> tuple[dict[str, Any], str]: + message = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": call_id, + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(arguments, ensure_ascii=False), + }, + } + ], + } + generation = json.dumps( + {"name": name, "arguments": arguments}, + ensure_ascii=False, + separators=(",", ":"), + ) + return message, generation + + +def _text_token_ids(text: str) -> list[int]: + """Match ByT5's byte vocabulary without deriving goldens from the tokenizer.""" + return [byte + _BYTE_TOKEN_OFFSET for byte in text.encode()] + + +def _prompt_token_ids(body: dict[str, Any]) -> list[int]: + serialized_request = json.dumps(body, sort_keys=True, separators=(",", ":")) + prompt_token_ids = _text_token_ids(serialized_request) + + messages = body.get("messages", []) + if not any(message.get("role") == "tool" for message in messages): + return prompt_token_ids + + # The RL postprocessor requires each later prompt to start with every token + # already observed. Reconstruct the first workplace request, then append the + # first generation and a marked encoding of the continuation request. + first_assistant_index = next( + index + for index, message in enumerate(messages) + if message.get("role") == "assistant" + ) + initial_body = deepcopy(body) + initial_body["messages"] = messages[:first_assistant_index] + initial_request = json.dumps(initial_body, sort_keys=True, separators=(",", ":")) + _, first_generation = _tool_call_message( + "email_reply_email", + _WORKPLACE_ARGUMENTS, + "call-l0-workplace", + ) + return ( + _text_token_ids(initial_request) + + _text_token_ids(first_generation) + + _text_token_ids(_CONTINUATION_MARKER + serialized_request) + ) + + +class _ScriptedPolicyGeneration: + """Minimum policy-generation surface consumed by the Gym rollout path.""" + + cfg = _GENERATION_CONFIG + + +class _ScriptedOpenAIHandler(BaseHTTPRequestHandler): + """Return deterministic policy outputs to real Gym model and agent services.""" + + protocol_version = "HTTP/1.1" + + @staticmethod + def _assert_request_contract(body: dict[str, Any]) -> None: + serialized_body = json.dumps(body) + is_judge_request = ( + "GOLD:" in serialized_body and "CANDIDATE:" in serialized_body + ) + + assert body["model"] == "scripted-model" + assert body["logprobs"] is True + assert body["top_logprobs"] == 0 + assert body["return_tokens_as_token_ids"] is True + + if is_judge_request: + assert "temperature" not in body + assert "top_p" not in body + assert "max_tokens" not in body + else: + assert body["temperature"] == _GENERATION_CONFIG["temperature"] + assert body["top_p"] == _GENERATION_CONFIG["top_p"] + assert body["max_tokens"] == _GENERATION_CONFIG["max_new_tokens"] + + expected_tool_name = None + if "SHOW24" in serialized_body and "Medical Zone" in serialized_body: + expected_tool_name = "check_seat_availability" + elif "Dizer Kola" in serialized_body: + expected_tool_name = "response_tool_8" + elif "Task Update on Develop prototype" in serialized_body: + expected_tool_name = "email_reply_email" + + if expected_tool_name is not None: + tool_names = { + tool["function"]["name"] + for tool in body.get("tools", []) + if tool.get("type") == "function" + } + assert expected_tool_name in tool_names + + @staticmethod + def _message_for(body: dict[str, Any]) -> tuple[dict[str, Any], str]: + serialized_body = json.dumps(body) + if "GOLD:" in serialized_body and "CANDIDATE:" in serialized_body: + if r"\\boxed{Ada Lovelace}" in serialized_body: + content = "The candidate does not match the reference.\n\n[[A!=B]]" + return {"role": "assistant", "content": content}, content + return { + "role": "assistant", + "content": _EQUIVALENCE_JUDGE_COMPLETION, + }, _EQUIVALENCE_JUDGE_COMPLETION + if _REJECTED_ROLLOUT_MARKER in serialized_body: + if "1000 digit numbers" in serialized_body: + content = r"\boxed{999999}" + elif "python programming language only" in serialized_body: + content = "```python\nprint(0)\n```" + elif "SHOW24" in serialized_body and "Medical Zone" in serialized_body: + content = "I cannot call the requested tool." + elif "cystic fibrosis" in serialized_body: + content = r"\boxed{A}" + elif "theory of evolution by natural selection" in serialized_body: + content = r"\boxed{Ada Lovelace}" + elif "Dizer Kola" in serialized_body: + content = "I cannot extract that." + elif "Task Update on Develop prototype" in serialized_body: + content = "Done" + else: + raise AssertionError( + "L0 rejected policy received an unrecognized prompt" + ) + return {"role": "assistant", "content": content}, content + if "1000 digit numbers" in serialized_body: + content = r"\boxed{32}" + return {"role": "assistant", "content": content}, content + if "python programming language only" in serialized_body: + return { + "role": "assistant", + "content": _CODE_GEN_COMPLETION, + }, _CODE_GEN_COMPLETION + if "SHOW24" in serialized_body and "Medical Zone" in serialized_body: + message, _ = _tool_call_message( + "check_seat_availability", + json.loads(_TOOL_CALL_ARGUMENTS), + "call-l0-tool-use", + ) + return message, _TOOL_CALL_GENERATION + if "cystic fibrosis" in serialized_body: + return { + "role": "assistant", + "content": _MCQA_COMPLETION, + }, _MCQA_COMPLETION + if "theory of evolution by natural selection" in serialized_body: + return { + "role": "assistant", + "content": _EQUIVALENCE_COMPLETION, + }, _EQUIVALENCE_COMPLETION + if "Dizer Kola" in serialized_body: + return _tool_call_message( + "response_tool_8", + {"extraction": _STRUCTURED_OUTPUTS_ARGUMENTS}, + "call-l0-structured-output", + ) + if "Task Update on Develop prototype" in serialized_body: + if any( + message.get("role") == "tool" for message in body.get("messages", []) + ): + content = "Done" + return {"role": "assistant", "content": content}, content + return _tool_call_message( + "email_reply_email", + _WORKPLACE_ARGUMENTS, + "call-l0-workplace", + ) + raise AssertionError("L0 scripted policy received an unrecognized prompt") + + def do_POST(self) -> None: # noqa: N802 + if self.path != "/v1/chat/completions": + self.send_error(404) + return + + content_length = int(self.headers.get("Content-Length", "0")) + body = json.loads(self.rfile.read(content_length)) + self._assert_request_contract(body) + message, generation_text = self._message_for(body) + prompt_token_ids = _prompt_token_ids(body) + generation_token_ids = _text_token_ids(generation_text) + assert ( + len(prompt_token_ids) + len(generation_token_ids) + <= _GENERATION_CONFIG["max_total_sequence_length"] + ) + message.update( + { + "prompt_token_ids": prompt_token_ids, + "generation_token_ids": generation_token_ids, + "generation_log_probs": [-0.1] * len(generation_token_ids), + } + ) + response = { + "id": "chatcmpl-l0-acceptance", + "object": "chat.completion", + "created": 0, + "model": body.get("model", "scripted-model"), + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls" + if message.get("tool_calls") + else "stop", + "message": message, + } + ], + "usage": { + "prompt_tokens": len(prompt_token_ids), + "completion_tokens": len(generation_token_ids), + "total_tokens": len(prompt_token_ids) + len(generation_token_ids), + }, + } + payload = json.dumps(response).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, format: str, *args: Any) -> None: + pass + + +def _load_acceptance_cases() -> list[dict[str, Any]]: + with _CASES_PATH.open() as case_file: + cases = safe_load(case_file)["cases"] + + assert cases, "L0 Gym rollout acceptance matrix must not be empty" + required_fields = { + "name", + "config_path", + "data_path", + "example_index", + "example_sha256", + "agent_ref", + "expected_generations", + "rejected_generations", + "expected_prompt_fragment", + "expected_result", + "expected_reward", + } + for case in cases: + assert required_fields <= case.keys(), ( + f"acceptance case is missing fields: {required_fields - case.keys()}" + ) + assert (_GYM_ROOT / case["config_path"]).is_file(), ( + f"{case['name']}: missing Gym config {case['config_path']}" + ) + data_path = _GYM_ROOT / case["data_path"] + assert data_path.is_file(), ( + f"{case['name']}: missing Gym data {case['data_path']}" + ) + with data_path.open("rb") as data_file: + examples = [line.rstrip(b"\r\n") for line in data_file if line.strip()] + assert 0 <= case["example_index"] < len(examples), ( + f"{case['name']}: example_index {case['example_index']} is outside a {len(examples)}-row dataset" + ) + actual_sha256 = sha256(examples[case["example_index"]]).hexdigest() + assert actual_sha256 == case["example_sha256"], ( + f"{case['name']}: pinned example changed; review the row and update its golden values" + ) + assert case["agent_ref"].keys() >= {"type", "name"} + assert case["expected_generations"] + assert case["rejected_generations"] + assert all( + isinstance(generation, str) + for key in ("expected_generations", "rejected_generations") + for generation in case[key] + ) + assert math.isfinite(case["expected_reward"]) + + names = [case["name"] for case in cases] + assert len(names) == len(set(names)), "acceptance case names must be unique" + assert set(names) == _REQUIRED_L0_CASES, ( + f"L0 matrix must contain exactly {_REQUIRED_L0_CASES}, got {set(names)}" + ) + return cases + + +_L0_CASES = _load_acceptance_cases() + + +@pytest.mark.nemo_gym +def test_l0_scripted_multiturn_tokens_are_contiguous(): + initial_body = { + "model": "scripted-model", + "messages": [ + {"role": "system", "content": "system"}, + { + "role": "user", + "content": "Task Update on Develop prototype", + }, + ], + } + continuation_body = deepcopy(initial_body) + continuation_body["messages"].extend( + [ + {"role": "assistant", "content": None, "tool_calls": []}, + { + "role": "tool", + "tool_call_id": "call-l0-workplace", + "content": '{"output":"Email replied successfully."}', + }, + ] + ) + _, first_generation = _tool_call_message( + "email_reply_email", + _WORKPLACE_ARGUMENTS, + "call-l0-workplace", + ) + seen_token_ids = _prompt_token_ids(initial_body) + _text_token_ids(first_generation) + + continuation_prompt = _prompt_token_ids(continuation_body) + + assert continuation_prompt[: len(seen_token_ids)] == seen_token_ids + decoded_continuation = ByT5Tokenizer().batch_decode([continuation_prompt])[0] + assert _CONTINUATION_MARKER in decoded_continuation + assert json.loads(decoded_continuation.rpartition(_CONTINUATION_MARKER)[2]) == ( + continuation_body + ) + + +def _load_case_datum(case: dict[str, Any], *, accepted: bool) -> DatumSpec: + dataset = NemoGymDataset(str(_GYM_ROOT / case["data_path"])) + datum = nemo_gym_data_processor( + dataset.dataset[case["example_index"]], None, None, None, 0 + ) + datum["extra_env_info"]["agent_ref"] = deepcopy(case["agent_ref"]) + if not accepted: + input_items = datum["extra_env_info"]["responses_create_params"]["input"] + final_user_item = next( + item for item in reversed(input_items) if item.get("role") == "user" + ) + final_user_item["content"] += f"\n\n{_REJECTED_ROLLOUT_MARKER}" + return datum + + +def _assert_contract_preserved( + actual: Any, expected: Any, *, path: str = "responses_create_params" +) -> None: + """Require every source value while allowing Gym to add schema defaults.""" + if isinstance(expected, dict): + assert isinstance(actual, dict), ( + f"{path}: expected a mapping, got {type(actual).__name__}" + ) + missing_keys = expected.keys() - actual.keys() + assert not missing_keys, f"{path}: missing keys {sorted(missing_keys)}" + for key, expected_value in expected.items(): + _assert_contract_preserved( + actual[key], expected_value, path=f"{path}.{key}" + ) + return + + if isinstance(expected, list): + assert isinstance(actual, list), ( + f"{path}: expected a list, got {type(actual).__name__}" + ) + assert len(actual) == len(expected), ( + f"{path}: expected {len(expected)} items, got {len(actual)}" + ) + for index, (actual_value, expected_value) in enumerate( + zip(actual, expected, strict=True) + ): + _assert_contract_preserved( + actual_value, expected_value, path=f"{path}[{index}]" + ) + return + + assert actual == expected, f"{path}: expected {expected!r}, got {actual!r}" + + +@pytest.fixture(scope="module") +def scripted_openai_base_url(): + server = ThreadingHTTPServer(("0.0.0.0", 0), _ScriptedOpenAIHandler) + server_thread = Thread(target=server.serve_forever, daemon=True) + server_thread.start() + try: + yield f"http://{_get_node_ip_local()}:{server.server_port}/v1" + finally: + server.shutdown() + server.server_close() + server_thread.join(timeout=5) + + +@pytest.fixture(scope="module") +def l0_nemo_gym(scripted_openai_base_url): + """Start all acceptance environments in one real Gym actor.""" + config_paths = [_POLICY_MODEL_CONFIG] + [case["config_path"] for case in _L0_CASES] + tokenizer = ByT5Tokenizer() + env = spinup_nemo_gym_actor( + { + "nemo_gym": { + "config_paths": config_paths, + "skip_venv_if_present": True, + } + }, + base_urls=[scripted_openai_base_url], + model_name="scripted-model", + tokenizer=tokenizer, + enable_router_replay=False, + use_fastokens=False, + ) + try: + yield env + finally: + try: + # The shared actor owns many Gym subprocesses, which are reaped + # sequentially during graceful shutdown. + ray.get( + env.shutdown.remote(), + timeout=NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S, + ) + finally: + ray.kill(env) + + +@pytest.mark.nemo_gym +@pytest.mark.timeout(900) +@pytest.mark.parametrize("case", _L0_CASES, ids=[case["name"] for case in _L0_CASES]) +@pytest.mark.parametrize("accepted", [True, False], ids=["accepted", "rejected"]) +def test_l0_gym_environments_roll_out_through_nemo_rl(l0_nemo_gym, case, accepted): + """A pinned Gym example must preserve its contract across the NeMo RL boundary.""" + tokenizer = ByT5Tokenizer() + expected_generations = case[ + "expected_generations" if accepted else "rejected_generations" + ] + expected_reward = case["expected_reward"] if accepted else 0.0 + datum = _load_case_datum(case, accepted=accepted) + extra_env_info = datum["extra_env_info"] + expected_responses_create_params = deepcopy( + extra_env_info["responses_create_params"] + ) + expected_responses_create_params["temperature"] = _GENERATION_CONFIG["temperature"] + expected_responses_create_params["top_p"] = _GENERATION_CONFIG["top_p"] + row_max_output_tokens = expected_responses_create_params.get("max_output_tokens") + expected_responses_create_params["max_output_tokens"] = ( + min(row_max_output_tokens, _GENERATION_CONFIG["max_new_tokens"]) + if row_max_output_tokens is not None + else _GENERATION_CONFIG["max_new_tokens"] + ) + assert case["expected_prompt_fragment"] in json.dumps( + extra_env_info["responses_create_params"] + ) + result = run_nemo_gym_rollout_sync( + policy_generation=_ScriptedPolicyGeneration(), + input_batch=rl_collate_fn([datum]), + tokenizer=tokenizer, + task_to_env={"nemo_gym": l0_nemo_gym}, + max_seq_len=_GENERATION_CONFIG["max_total_sequence_length"], + generation_config=deepcopy(_GENERATION_CONFIG), + log_full_result_tables=True, + ) + + final_batch = result.final_batch + assert final_batch.size == 1 + assert final_batch["agent_ref"] == [case["agent_ref"]] + + reward = final_batch["total_reward"].item() + assert math.isfinite(reward) + assert reward == pytest.approx(expected_reward) + assert final_batch["length"].item() > 0 + + assistant_messages = [ + message + for message in final_batch["message_log"][0] + if message["role"] == "assistant" + ] + assert assistant_messages + assert len(assistant_messages) == len(expected_generations) + for message, generation in zip( + assistant_messages, expected_generations, strict=True + ): + assert isinstance(message["token_ids"], torch.Tensor) + assert isinstance(message["generation_logprobs"], torch.Tensor) + expected_token_ids = torch.tensor( + [byte + _BYTE_TOKEN_OFFSET for byte in generation.encode()] + ) + torch.testing.assert_close(message["token_ids"], expected_token_ids) + torch.testing.assert_close( + message["generation_logprobs"], + torch.full((len(expected_token_ids),), -0.1), + ) + + metric_prefix = f"{case['agent_ref']['name']}/reward/" + assert any(key.startswith(metric_prefix) for key in result.rollout_metrics) + + full_result_key = f"{case['agent_ref']['name']}/full_result" + assert full_result_key in result.rollout_metrics + full_result_table = result.rollout_metrics[full_result_key] + assert len(full_result_table.data) == 1 + full_result = json.loads(full_result_table.data[0][0]) + _assert_contract_preserved( + full_result["responses_create_params"], + expected_responses_create_params, + ) + if accepted: + for field, expected_value in case["expected_result"].items(): + assert full_result[field] == expected_value + assert full_result["response"]["output"] + if accepted and "expected_tool_outputs" in case: + tool_outputs = [ + item["output"] + for item in full_result["response"]["output"] + if item["type"] == "function_call_output" + ] + assert tool_outputs == case["expected_tool_outputs"] + expected_judge_verdicts = case.get( + "expected_judge_verdicts" if accepted else "rejected_judge_verdicts" + ) + if expected_judge_verdicts is not None: + verdicts = [ + evaluation["verdict_label"] + for evaluation in full_result["judge_evaluations"] + ] + assert verdicts == expected_judge_verdicts + generation_strings = [ + item["generation_str"] + for item in full_result["response"]["output"] + if "generation_str" in item + ] + assert generation_strings == expected_generations + prompt_strings = [ + item["prompt_str"] + for item in full_result["response"]["output"] + if "prompt_str" in item + ] + assert len(prompt_strings) == len(expected_generations) + assert all(case["expected_prompt_fragment"] in prompt for prompt in prompt_strings) + if accepted and case["name"] == "workplace_assistant": + assert _CONTINUATION_MARKER in prompt_strings[1] + continuation_json = prompt_strings[1].rpartition(_CONTINUATION_MARKER)[2] + continuation_messages = json.loads(continuation_json)["messages"] + assert any(message.get("tool_calls") for message in continuation_messages) + assert any( + message.get("role") == "tool" + and "Email replied successfully" in message.get("content", "") + for message in continuation_messages + ) diff --git a/tests/unit/environments/test_nemo_gym_shards.py b/tests/unit/environments/test_nemo_gym_shards.py new file mode 100644 index 00000000000..102187e51ed --- /dev/null +++ b/tests/unit/environments/test_nemo_gym_shards.py @@ -0,0 +1,698 @@ +# 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. +"""Unit tests for the NeMo-Gym shard schema (pure Python, no Gym import).""" + +import pytest +from omegaconf import OmegaConf + +from nemo_rl.environments.nemo_gym_shards import ( + GYM_LOG_DIR_KEY, + ShardConfigError, + ShardPlan, + ShardSetupError, + ShardSpec, + apply_shard_log_dir, + apply_shard_overlay, + build_route_shard_map, + find_gym_config_entries, + parse_shard_plan, +) + + +def _sharded_config(**overrides): + config = { + "uv_venv_dir": "/opt/gym_venvs", + "num_gpu_nodes": 2, + "shards": [ + {"name": "judged", "config_paths": ["judge.yaml"]}, + {"name": "tools", "config_paths": ["proxy.yaml", "tools.yaml"]}, + ], + } + config.update(overrides) + return config + + +def test_unsharded_config_returns_no_plan(): + """config_paths without shards is the pre-sharding path and must stay working.""" + assert parse_shard_plan({"config_paths": ["gym.yaml"], "num_gpu_nodes": 1}) is None + + +@pytest.mark.parametrize( + "key", + [ + "common_inherited_overlays", + "common_overrides", + "allowed_duplicate_entries", + ], +) +def test_unsharded_config_rejects_stray_sharding_keys(key): + with pytest.raises(ShardConfigError, match=f"sharding keys.*{key}.*no 'shards'"): + parse_shard_plan({"config_paths": ["gym.yaml"], key: []}) + + +def test_parse_shard_plan_reads_shards_and_common_settings(): + plan = parse_shard_plan( + _sharded_config( + common_overrides={"policy_model": {"num_workers": 4}}, + allowed_duplicate_entries=["policy_model"], + ) + ) + + assert [shard.name for shard in plan.shards] == ["judged", "tools"] + assert plan.shards[1].config_paths == ["proxy.yaml", "tools.yaml"] + assert plan.common_overrides == {"policy_model": {"num_workers": 4}} + assert plan.allowed_duplicate_entries == frozenset({"policy_model"}) + + +def test_parse_shard_plan_accepts_omegaconf_input(): + """Job configs arrive as OmegaConf nodes, not plain dicts.""" + plan = parse_shard_plan(OmegaConf.create(_sharded_config())) + + assert [shard.name for shard in plan.shards] == ["judged", "tools"] + assert isinstance(plan.shards[0].config_paths, list) + + +def test_top_level_config_paths_conflicts_with_shards(): + with pytest.raises(ShardConfigError, match="cannot be combined with 'shards'"): + parse_shard_plan(_sharded_config(config_paths=["gym.yaml"])) + + +def test_an_inherited_config_paths_can_be_retracted_with_null(): + """A Hydra override can blank a key it inherited but cannot delete it, so + null has to read as absent or no recipe could ever be sharded.""" + plan = parse_shard_plan(_sharded_config(config_paths=None)) + + assert plan is not None + assert [shard.name for shard in plan.shards] == ["judged", "tools"] + + +def test_a_nulled_entry_overlay_is_not_forwarded_to_gym(): + """Gym should get the key gone, not the key set to null.""" + plan = parse_shard_plan(_sharded_config(policy_model=None)) + assert plan is not None + + merged = apply_shard_overlay({"policy_model": None}, plan, plan.shards[0]) + + assert "policy_model" not in merged + + +def test_multi_shard_plan_rejects_an_unclaimed_inherited_overlay(): + with pytest.raises(ShardConfigError, match="have no owner"): + parse_shard_plan( + _sharded_config( + nl2bash_judge_model={"responses_api_models": {"local_vllm_model": {}}} + ) + ) + + +def test_single_shard_plan_automatically_claims_every_inherited_overlay(): + plan = parse_shard_plan( + _sharded_config( + shards=[{"name": "only", "config_paths": ["gym.yaml"]}], + policy_model={"responses_api_models": {"vllm_model": {"num_workers": 4}}}, + math={"resources_servers": {"math": {"workers": 8}}}, + ) + ) + + assert plan.shards[0].inherited_overlays == frozenset({"math", "policy_model"}) + merged = apply_shard_overlay( + _sharded_config( + shards=[], + policy_model={"responses_api_models": {"vllm_model": {"num_workers": 4}}}, + math={"resources_servers": {"math": {"workers": 8}}}, + ), + plan, + plan.shards[0], + ) + assert ( + merged["policy_model"]["responses_api_models"]["vllm_model"]["num_workers"] == 4 + ) + assert merged["math"]["resources_servers"]["math"]["workers"] == 8 + + +def test_multi_shard_plan_routes_inherited_overlays_by_key(): + policy_model = {"responses_api_models": {"vllm_model": {"num_workers": 4}}} + judged = {"resources_servers": {"judge": {"workers": 8}}} + tools = {"resources_servers": {"tools": {"workers": 16}}} + config = _sharded_config( + shards=[ + { + "name": "judged", + "config_paths": ["judge.yaml"], + "inherited_overlays": ["judged"], + }, + { + "name": "tools", + "config_paths": ["tools.yaml"], + "inherited_overlays": ["tools"], + }, + ], + common_inherited_overlays=["policy_model"], + policy_model=policy_model, + judged=judged, + tools=tools, + ) + + plan = parse_shard_plan(config) + judged_config = apply_shard_overlay(config, plan, plan.shards[0]) + tools_config = apply_shard_overlay(config, plan, plan.shards[1]) + + assert judged_config["policy_model"] == policy_model + assert judged_config["judged"] == judged + assert "tools" not in judged_config + assert tools_config["policy_model"] == policy_model + assert tools_config["tools"] == tools + assert "judged" not in tools_config + for merged in (judged_config, tools_config): + assert "shards" not in merged + assert "common_inherited_overlays" not in merged + + +def test_inherited_overlay_claims_must_be_known_and_unique(): + config = _sharded_config( + shards=[ + { + "name": "judged", + "config_paths": ["judge.yaml"], + "inherited_overlays": ["missing"], + }, + {"name": "tools", "config_paths": ["tools.yaml"]}, + ], + policy_model={"responses_api_models": {}}, + ) + with pytest.raises(ShardConfigError, match="Available overlays.*policy_model"): + parse_shard_plan(config) + + config["shards"][0]["inherited_overlays"] = ["policy_model"] + config["shards"][1]["inherited_overlays"] = ["policy_model"] + with pytest.raises(ShardConfigError, match="duplicate claims.*policy_model"): + parse_shard_plan(config) + + +def test_inherited_overlay_lists_require_unique_non_empty_strings(): + with pytest.raises(ShardConfigError, match="list of non-empty strings"): + parse_shard_plan( + _sharded_config( + shards=[ + { + "name": "only", + "config_paths": ["gym.yaml"], + "inherited_overlays": [""], + } + ] + ) + ) + + with pytest.raises(ShardConfigError, match="contains duplicate keys"): + parse_shard_plan( + _sharded_config( + shards=[{"name": "only", "config_paths": ["gym.yaml"]}], + common_inherited_overlays=["policy_model", "policy_model"], + ) + ) + + +@pytest.mark.parametrize( + ("shards", "expected"), + [ + ([], "non-empty list"), + ([{"config_paths": ["a.yaml"]}], "non-empty string 'name'"), + ([{"name": " ", "config_paths": ["a.yaml"]}], "non-empty string 'name'"), + ([{"name": "a", "config_paths": []}], "at least one entry in 'config_paths'"), + ([{"name": "a"}], "at least one entry in 'config_paths'"), + ( + [{"name": "a", "config_paths": ["a.yaml"], "replicas": 0}], + "must be an integer >= 1", + ), + ( + [{"name": "a", "config_paths": ["a.yaml"], "port_range_low": 5000}], + "both port_range_low and port_range_high", + ), + ( + [ + { + "name": "a", + "config_paths": ["a.yaml"], + "port_range_low": 5200, + "port_range_high": 5200, + } + ], + "must reserve at least two ports", + ), + ( + [ + { + "name": "a", + "config_paths": ["a.yaml"], + "port_range_low": 65000, + "port_range_high": 65536, + } + ], + r"within \[1, 65535\]", + ), + ( + [{"name": "a", "config_paths": ["a.yaml"], "judge_model": {}}], + "unrecognized keys", + ), + ( + [ + {"name": "dup", "config_paths": ["a.yaml"]}, + {"name": "dup", "config_paths": ["b.yaml"]}, + ], + "Duplicate shard names", + ), + ], +) +def test_malformed_shards_are_rejected(shards, expected): + with pytest.raises(ShardConfigError, match=expected): + parse_shard_plan({"shards": shards}) + + +@pytest.mark.parametrize("name", ["../judged", "team/judged", ".", ".."]) +def test_shard_names_must_be_safe_log_path_components(name): + with pytest.raises(ShardConfigError, match="one safe path component"): + parse_shard_plan( + _sharded_config(shards=[{"name": name, "config_paths": ["judge.yaml"]}]) + ) + + +@pytest.mark.parametrize("actor_cpus", [0, -1, float("nan"), float("inf"), "8"]) +def test_actor_cpus_must_be_positive_and_finite(actor_cpus): + with pytest.raises(ShardConfigError, match="positive finite number"): + parse_shard_plan( + _sharded_config( + shards=[ + { + "name": "judged", + "config_paths": ["judge.yaml"], + "actor_cpus": actor_cpus, + } + ] + ) + ) + + +def test_config_paths_and_port_types_are_validated_at_parse_time(): + with pytest.raises( + ShardConfigError, match="config_paths must be non-empty strings" + ): + parse_shard_plan( + _sharded_config(shards=[{"name": "judged", "config_paths": [123]}]) + ) + + with pytest.raises(ShardConfigError, match="port range must use integers"): + parse_shard_plan( + _sharded_config( + shards=[ + { + "name": "judged", + "config_paths": ["judge.yaml"], + "port_range_low": "5000", + "port_range_high": 5500, + } + ] + ) + ) + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ([{"nope": 1}], "list of non-empty strings"), + ([""], "list of non-empty strings"), + (["policy_model", "policy_model"], "contains duplicate keys"), + ({}, "list of non-empty strings"), + ], +) +def test_allowed_duplicate_entries_requires_unique_non_empty_strings(value, expected): + with pytest.raises(ShardConfigError, match=expected): + parse_shard_plan(_sharded_config(allowed_duplicate_entries=value)) + + +@pytest.mark.parametrize("layer", ["common", "shard"]) +def test_override_layers_cannot_set_config_paths(layer): + config = _sharded_config() + if layer == "common": + config["common_overrides"] = {"config_paths": ["wrong.yaml"]} + else: + config["shards"][0]["overrides"] = {"config_paths": ["wrong.yaml"]} + + with pytest.raises(ShardConfigError, match="cannot set config_paths"): + parse_shard_plan(config) + + +def test_override_layers_cannot_override_an_entry_owned_by_another_shard(): + config = _sharded_config( + shards=[ + { + "name": "judged", + "config_paths": ["judge.yaml"], + "inherited_overlays": ["judge_model"], + }, + { + "name": "tools", + "config_paths": ["tools.yaml"], + "overrides": {"judge_model": {"responses_api_models": {}}}, + }, + ], + judge_model={"responses_api_models": {"vllm_model": {}}}, + ) + + with pytest.raises(ShardConfigError, match="claimed by another shard"): + parse_shard_plan(config) + + +def test_common_overrides_only_override_common_inherited_entries(): + config = _sharded_config( + shards=[ + { + "name": "judged", + "config_paths": ["judge.yaml"], + "inherited_overlays": ["judge_model"], + }, + {"name": "tools", "config_paths": ["tools.yaml"]}, + ], + common_overrides={"judge_model": {"responses_api_models": {}}}, + judge_model={"responses_api_models": {"vllm_model": {}}}, + ) + + with pytest.raises(ShardConfigError, match="not claimed by common"): + parse_shard_plan(config) + + +def test_shards_land_on_distinct_nodes_unless_told_otherwise(): + assert parse_shard_plan(_sharded_config()).placement_strategy == "STRICT_SPREAD" + + +def test_placement_strategy_can_be_relaxed_to_run_shards_on_one_machine(): + plan = parse_shard_plan( + _sharded_config( + placement_strategy="PACK", + shards=[ + { + "name": "judged", + "config_paths": ["judge.yaml"], + "port_range_low": 5000, + "port_range_high": 5500, + }, + { + "name": "tools", + "config_paths": ["tools.yaml"], + "port_range_low": 5501, + "port_range_high": 6000, + }, + ], + ) + ) + + assert plan.placement_strategy == "PACK" + + +def test_relaxed_placement_requires_disjoint_explicit_port_ranges(): + with pytest.raises(ShardConfigError, match="require explicit port_range"): + parse_shard_plan(_sharded_config(placement_strategy="PACK")) + + with pytest.raises(ShardConfigError, match="overlapping port ranges"): + parse_shard_plan( + _sharded_config( + placement_strategy="PACK", + shards=[ + { + "name": "judged", + "config_paths": ["judge.yaml"], + "port_range_low": 5000, + "port_range_high": 5600, + }, + { + "name": "tools", + "config_paths": ["tools.yaml"], + "port_range_low": 5600, + "port_range_high": 6000, + }, + ], + ) + ) + + +def test_an_unknown_placement_strategy_is_rejected(): + with pytest.raises(ShardConfigError, match="placement_strategy must be one of"): + parse_shard_plan(_sharded_config(placement_strategy="SPRED")) + + +def test_replicas_that_could_share_a_node_would_share_a_port_range(): + """Replicas come from one merge, so they cannot be given separate ranges.""" + config = _sharded_config(placement_strategy="PACK") + config["shards"][1]["replicas"] = 2 + + with pytest.raises(ShardConfigError, match=r"\['tools'\] declare replicas"): + parse_shard_plan(config) + + +def test_replicas_are_fine_on_the_default_strategy_that_spreads_them(): + config = _sharded_config() + config["shards"][1]["replicas"] = 2 + + assert [shard.replicas for shard in parse_shard_plan(config).shards] == [1, 2] + + +def test_apply_shard_overlay_layers_shard_over_common(): + plan = ShardPlan( + shards=[], + common_overrides={ + "policy_model": {"num_workers": 4, "timeout": 30}, + "shared_judge": {"replicas": 1}, + }, + ) + shard = ShardSpec( + name="judged", + config_paths=["judge.yaml"], + overrides={"policy_model": {"num_workers": 16}}, + ) + + merged = apply_shard_overlay({"default_host": "10.0.0.1"}, plan, shard) + + assert merged["config_paths"] == ["judge.yaml"] + assert merged["default_host"] == "10.0.0.1" + # Shard wins on conflict, and sibling keys from common survive the merge. + assert merged["policy_model"] == {"num_workers": 16, "timeout": 30} + assert merged["shared_judge"] == {"replicas": 1} + + +def test_explicit_overrides_layer_over_inherited_values(): + base = { + "policy_model": { + "responses_api_models": {"vllm_model": {"num_workers": 2, "timeout": 30}} + } + } + plan = ShardPlan( + shards=[], + common_inherited_overlays=frozenset({"policy_model"}), + common_overrides={ + "policy_model": {"responses_api_models": {"vllm_model": {"num_workers": 4}}} + }, + ) + shard = ShardSpec( + name="judged", + config_paths=["judge.yaml"], + overrides={ + "policy_model": {"responses_api_models": {"vllm_model": {"num_workers": 8}}} + }, + ) + + merged = apply_shard_overlay(base, plan, shard) + + assert merged["policy_model"]["responses_api_models"]["vllm_model"] == { + "num_workers": 8, + "timeout": 30, + } + + +def test_apply_shard_overlay_applies_per_shard_port_range(): + shard = ShardSpec( + name="ci", + config_paths=["a.yaml"], + port_range_low=6000, + port_range_high=6099, + ) + + merged = apply_shard_overlay( + {"port_range_low": 5000, "port_range_high": 5999}, ShardPlan(shards=[]), shard + ) + + assert (merged["port_range_low"], merged["port_range_high"]) == (6000, 6099) + + +def test_apply_shard_overlay_keeps_global_ports_by_default(): + merged = apply_shard_overlay( + {"port_range_low": 5000, "port_range_high": 5999}, + ShardPlan(shards=[]), + ShardSpec(name="prod", config_paths=["a.yaml"]), + ) + + assert (merged["port_range_low"], merged["port_range_high"]) == (5000, 5999) + + +def test_build_route_shard_map_routes_agents_and_task_sources(): + route_to_shard = build_route_shard_map( + { + "judged": { + "math_agent": ["responses_api_agents"], + "math_env": ["resources_servers"], + }, + "tools": { + "bash_agent": ["responses_api_agents"], + "bash_tools": ["resources_servers"], + }, + } + ) + + assert route_to_shard == { + "math_agent": "judged", + "math_env": "judged", + "bash_agent": "tools", + "bash_tools": "tools", + } + + +def test_build_route_shard_map_rejects_an_agent_in_two_shards(): + """Rows naming the agent could go to either shard, so routing is undefined.""" + with pytest.raises(ShardSetupError, match="hosted by both shard"): + build_route_shard_map( + { + "judged": {"math_agent": ["responses_api_agents"]}, + "tools": {"math_agent": ["responses_api_agents"]}, + } + ) + + +def test_agents_are_never_allowlisted_for_duplication(): + """The allowlist covers shared support entries, not routing ambiguity.""" + with pytest.raises(ShardSetupError, match="hosted by both shard"): + build_route_shard_map( + { + "a": {"math_agent": ["responses_api_agents"]}, + "b": {"math_agent": ["responses_api_agents"]}, + }, + allowed_duplicate_entries={"math_agent"}, + ) + + +def test_task_sources_are_never_allowlisted_for_duplication(): + with pytest.raises(ShardSetupError, match="hosted by both shard"): + build_route_shard_map( + { + "a": {"math_env": ["resources_servers"]}, + "b": {"math_env": ["resources_servers"]}, + }, + allowed_duplicate_entries={"math_env"}, + ) + + +def test_build_route_shard_map_rejects_an_unlisted_duplicate_entry(): + with pytest.raises(ShardSetupError, match="allowed_duplicate_entries"): + build_route_shard_map( + { + "judged": {"shared_judge": ["responses_api_models"]}, + "tools": {"shared_judge": ["responses_api_models"]}, + } + ) + + +def test_build_route_shard_map_allows_a_listed_duplicate_entry(): + """Policy proxies are copied into every shard on purpose.""" + route_to_shard = build_route_shard_map( + { + "judged": { + "math_agent": ["responses_api_agents"], + "policy_model": ["responses_api_models"], + }, + "tools": { + "bash_agent": ["responses_api_agents"], + "policy_model": ["responses_api_models"], + }, + }, + allowed_duplicate_entries={"policy_model"}, + ) + + assert route_to_shard == {"math_agent": "judged", "bash_agent": "tools"} + + +def test_a_routable_name_is_still_a_duplicate_when_it_is_a_model_elsewhere(): + """Reusing a routable name for a model engine doubles that engine's GPU claim.""" + with pytest.raises(ShardSetupError, match="allowed_duplicate_entries"): + build_route_shard_map( + { + "judged": {"math": ["responses_api_agents"]}, + "tools": { + "math": ["responses_api_models"], + "bash": ["responses_api_agents"], + }, + } + ) + + +def test_apply_shard_log_dir_gives_each_shard_its_own_directory(): + config = {GYM_LOG_DIR_KEY: "/logs/gym", "default_host": "10.0.0.1"} + + judged = apply_shard_log_dir(config, "judged") + tools = apply_shard_log_dir(config, "tools") + + assert judged[GYM_LOG_DIR_KEY] == "/logs/gym/judged" + assert tools[GYM_LOG_DIR_KEY] == "/logs/gym/tools" + # Unrelated settings survive, and the input is not mutated. + assert judged["default_host"] == "10.0.0.1" + assert config[GYM_LOG_DIR_KEY] == "/logs/gym" + + +def test_apply_shard_log_dir_separates_replicas(): + """Replicas are stamped from one merge, so they host identical server names.""" + config = {GYM_LOG_DIR_KEY: "/logs/gym"} + + first = apply_shard_log_dir(config, "tools", replica_index=0) + second = apply_shard_log_dir(config, "tools", replica_index=1) + + assert first[GYM_LOG_DIR_KEY] == "/logs/gym/tools/0" + assert second[GYM_LOG_DIR_KEY] == "/logs/gym/tools/1" + + +def test_apply_shard_log_dir_is_a_noop_without_a_configured_log_dir(): + """Gym only writes log files when the key is set, which is not the default.""" + assert apply_shard_log_dir({"default_host": "10.0.0.1"}, "judged") == { + "default_host": "10.0.0.1" + } + + +def test_apply_shard_log_dir_sanitizes_the_shard_name(): + config = apply_shard_log_dir({GYM_LOG_DIR_KEY: "/logs"}, "team/judged") + + assert config[GYM_LOG_DIR_KEY] == "/logs/team_judged" + + +def test_find_gym_config_entries_ignores_scalars_and_known_keys(): + entries = find_gym_config_entries( + { + "uv_venv_dir": "/opt/venvs", + "num_gpu_nodes": 2, + "shards": [{"name": "a"}], + "common_inherited_overlays": ["genrm_model"], + "common_overrides": {"x": 1}, + "allowed_duplicate_entries": ["policy_model"], + "effort_levels": {"low_weight": 0.0}, + "tokenizer_config": {"name": "policy"}, + "genrm_model": {"responses_api_models": {}}, + "safety_judge_model": {"responses_api_models": {}}, + } + ) + + assert entries == ["genrm_model", "safety_judge_model"]