`nemo_rl.environments.nemo_gym`#

Module Contents#

Classes#

`NemoGymCompatibleConfig`	Configuration fields required to select the NeMo Gym rollout path.
`ExternalServiceReadinessTargetConfig`	Health endpoint and required backend count for one external service.
`ExternalServiceReadinessConfig`	Startup gate for external services used by NeMo Gym rollouts.
`NemoGymConfig`
`NemoGym`	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.
`NemoGymShardSet`	The live actors behind one NeMo-Gym stack, sharded or not.

Functions#

`_count_structured_response_tokens`	Count tokens in Gym-extracted reasoning and answer text.
`_require_resolved_agent_refs`	Fail readably when Gym did not stamp an agent_ref onto every row.
`should_use_nemo_gym`	Determine whether NeMo Gym should handle rollouts and validation.
`_has_nan_generation_logprobs`	Return whether a postprocessed rollout contains NaN policy logprobs.
`_typed_gym_failure`	Map a NeMo-Gym HTTP failure onto a typed, PICKLABLE failure, or None if not one.
`get_nemo_gym_uv_cache_dir`	Return the uv cache directory inside a container, or None outside one.
`get_nemo_gym_venv_dir`	Return the NeMo Gym venv directory from NEMO_GYM_VENV_DIR, or None.
`_external_staging_backend`	Map the setup-derived generation backend to Gym’s capture backend.
`extract_external_service_readiness`	Remove and validate NeMo-RL’s external readiness block from Gym config.
`_probe_external_service`	Return a readiness problem for one service, or None when it is ready.
`_wait_for_external_services`	Wait for all external Gym dependencies under one shared deadline.
`_detect_invalid_tool_call_and_malformed_thinking`	Flag a NeMo-Gym output item as an invalid tool call / malformed thinking.
`get_pad_dynamic_image_shapes`	Return nemo_gym’s pad_dynamic_image_shapes from an env config, or False.
`extract_reward_components`	Return per-component rewards from a NeMo Gym verify result, or None.
`build_reward_component_columns`	Build `reward/<name>` batch columns from per-sample reward-component dicts.
`validate_reward_components_match_scalar`	Assert each multi-reward result sets `reward == sum(reward_components)`.
`setup_nemo_gym_config`
`build_nemo_gym_config`	Build the `NemoGymConfig` for a single, unsharded NeMo-Gym actor.
`_build_gym_actor_config`	Turn one already-resolved Gym config mapping into a `NemoGymConfig`.
`get_nemo_gym_route_name`	Return the entry name Gym uses to route a row.
`_shard_instances`	Expand shards into one (shard, replica_index) entry per actor.
`as_nemo_gym_shard_set`	Read the NeMo-Gym entry of `task_to_env` as a shard set either way.
`build_nemo_gym_actors`	Create and spin up every NeMo-Gym actor this job needs.
`_build_single_gym_actor`	The pre-sharding path: one actor, no placement group, no discovery.
`_build_sharded_gym_actors`
`_spinup_shards_concurrently`	Start every shard at once, naming the shard behind any failure.
`_discover_route_shard_map`	Ask one replica per shard what it spawned, then build routing metadata.
`spinup_nemo_gym_actor`	Spin up a single NeMo-Gym actor against the given generation server URLs.
`validate_dataset_agent_coverage`	Fail at setup if any row names a route no shard hosts.
`_iter_dataset_agent_names`	Collect the agent or task-source names a dataset’s rows reference.
`_load_agent_names_from_source`	Read a stable Gym source once per controller process.
`_get_agent_name`

Data#

`NEMO_GYM_ACTOR_FQN`
`NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S`
`GYM_SERVER_TYPE_KEYS`
`DEFAULT_SHARD_NAME`
`DEFAULT_SHARD_CPUS`
`DEFAULT_SHARD_PG_READY_TIMEOUT_SECONDS`
`DEFAULT_SHARD_SPINUP_TIMEOUT_SECONDS`
`DEFAULT_SHARD_DRAIN_TIMEOUT_SECONDS`
`_ROUTED_EXPERTS_DTYPES`
`DEFAULT_INVALID_TOOL_CALL_PATTERNS`
`DEFAULT_THINKING_TAGS`
`_POLICY_SERVER_NAME`
`_NG_ROLLOUT_ID_BODY_KEY`
`_TOKEN_CAPTURE_CONTROL_PREFIX`
`_TOKEN_CAPTURE_CONTROL_ENV`
`_POLICY_API_KEY_ENV`

API#

nemo_rl.environments.nemo_gym.NEMO_GYM_ACTOR_FQN#: ‘nemo_rl.environments.nemo_gym.NemoGym’

nemo_rl.environments.nemo_gym.NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S#: 120

nemo_rl.environments.nemo_gym.GYM_SERVER_TYPE_KEYS#: (‘responses_api_agents’, ‘responses_api_models’, ‘resources_servers’)

nemo_rl.environments.nemo_gym.DEFAULT_SHARD_NAME#: ‘nemo_gym’

nemo_rl.environments.nemo_gym.DEFAULT_SHARD_CPUS#: 8

nemo_rl.environments.nemo_gym.DEFAULT_SHARD_PG_READY_TIMEOUT_SECONDS#: 180.0

nemo_rl.environments.nemo_gym.DEFAULT_SHARD_SPINUP_TIMEOUT_SECONDS#: 1800.0

nemo_rl.environments.nemo_gym.DEFAULT_SHARD_DRAIN_TIMEOUT_SECONDS#: None

nemo_rl.environments.nemo_gym._ROUTED_EXPERTS_DTYPES#: None

nemo_rl.environments.nemo_gym.DEFAULT_INVALID_TOOL_CALL_PATTERNS#: [‘<tool_call>’, ‘</tool_call>’, ‘<function_call>’, ‘</function_call>’]

nemo_rl.environments.nemo_gym.DEFAULT_THINKING_TAGS#: [’’, ‘’]

nemo_rl.environments.nemo_gym._count_structured_response_tokens( tokenizer: transformers.PreTrainedTokenizerBase, *, reasoning_text: str, answer_text: str, ) → tuple[int, int] | None#: Count tokens in Gym-extracted reasoning and answer text.

nemo_rl.environments.nemo_gym._require_resolved_agent_refs(nemo_gym_examples: list[dict]) → None#

Fail readably when Gym did not stamp an agent_ref onto every row.

run_examples resolves task_source to agent_ref in place before it returns, and every read after that point – this module’s counters, and Gym’s own dispatch, which posts to row["agent_ref"]["name"] – assumes it happened. Unguarded, a row that was not resolved surfaces as KeyError: 'agent_ref' inside a Ray TaskError inside an ExceptionGroup, forty lines from anything that names the cause.

The cause worth naming is a version skew rather than a bad row. task_source routing is new: an older Gym has no resolver, so a dataset prepared with a current Gym – which strips agent_ref and stamps task_source instead – arrives unroutable. That happens when the Gym actor’s venv is older than the checkout that prepared the data, which is what NRL_FORCE_REBUILD_VENVS=true exists to correct.

class nemo_rl.environments.nemo_gym.NemoGymCompatibleConfig#

Bases: typing.Protocol

Configuration fields required to select the NeMo Gym rollout path.

property env: dict[str, Any]#

property policy: nemo_rl.models.policy.PolicyConfig#

class nemo_rl.environments.nemo_gym.ExternalServiceReadinessTargetConfig#

Bases: pydantic.BaseModel

Health endpoint and required backend count for one external service.

url must be a vllm_pool_lb /health endpoint; see _probe_external_service for the reply it expects. A plain vLLM /health returns no backend counts and never reads as ready.

name: str#: ‘Field(…)’

url: str#: None

expected_backends: int#: ‘Field(…)’

classmethod validate_url(url: str) → str#: Require an absolute HTTP(S) health endpoint.

class nemo_rl.environments.nemo_gym.ExternalServiceReadinessConfig#

Bases: pydantic.BaseModel

Startup gate for external services used by NeMo Gym rollouts.

services: list[nemo_rl.environments.nemo_gym.ExternalServiceReadinessTargetConfig]#: ‘Field(…)’

timeout_seconds: float#: ‘Field(…)’

poll_interval_seconds: float#: ‘Field(…)’

request_timeout_seconds: float#: ‘Field(…)’

validate_unique_service_names() → Self#: Reject ambiguous duplicate service labels in diagnostics.

nemo_rl.environments.nemo_gym.should_use_nemo_gym( master_config: nemo_rl.environments.nemo_gym.NemoGymCompatibleConfig, ) → bool#: Determine whether NeMo Gym should handle rollouts and validation.

nemo_rl.environments.nemo_gym._has_nan_generation_logprobs(result: dict) → bool#: Return whether a postprocessed rollout contains NaN policy logprobs.

nemo_rl.environments.nemo_gym._typed_gym_failure(error: Exception) → Optional[Exception]#

Map a NeMo-Gym HTTP failure onto a typed, PICKLABLE failure, or None if not one.

Classification has to happen here, on the raising side, because run_rollouts runs inside the NemoGym Ray actor and the exception must survive the actor boundary to reach the retry policy on the driver.

It does not survive. aiohttp’s raise_for_status passes headers=self.headers, and those are a CIMultiDictProxy, which cloudpickle cannot serialize – so Ray drops the cause and the driver receives a bare RayTaskError with no type and no .status. Every gym HTTP failure then classified DATA, capping the gym path at max_data_attempts_per_prompt (2) and leaving max_attempts_per_prompt (5) unreachable on the very path whose dead-endpoint scenario motivates it. Two things made that the dominant case rather than a corner: Gym’s middleware turns inner-server failures into 500 – exactly the status the INFRA branch is for – and its transport layer retries disconnects in an uncapped loop, so those never arrive at all.

GymTransportError and RolloutDataFailure take a single str, so they pickle cleanly and classify_rollout_failure’s explicit-class fast path wins on the far side.

Returns None when the exception carries no HTTP status, leaving the caller to re-raise it untouched.

nemo_rl.environments.nemo_gym.get_nemo_gym_uv_cache_dir() → str | None#

Return the uv cache directory inside a container, or None outside one.

Inside a container (NRL_CONTAINER=1), returns the uv cache location so Gym stores its caches in the expected shared path. Returns None outside a container, meaning the caller should omit this arg and let Gym create the cache locally (the default when you may not be able to write to /opt).

nemo_rl.environments.nemo_gym.get_nemo_gym_venv_dir() → str | None#

Return the NeMo Gym venv directory from NEMO_GYM_VENV_DIR, or None.

Returns the value of NEMO_GYM_VENV_DIR if set, otherwise None. When None the caller should omit this arg and let Gym create venvs locally (the default when a container is not used since you may not be able to write to /opt).

class nemo_rl.environments.nemo_gym.NemoGymConfig#

Bases: typing.TypedDict

model_name: str#: None

base_urls: List[str]#: None

policy_api_key: NotRequired[Optional[str]]#: None

initial_global_config_dict: Dict[str, Any]#: None

port_range_low: NotRequired[int]#: None

port_range_high: NotRequired[int]#: None

invalid_tool_call_patterns: NotRequired[List[str] | None]#: None

thinking_tags: NotRequired[List[str] | None]#: None

require_routed_experts: NotRequired[bool]#: None

routed_experts_dtype: NotRequired[str]#: None

use_fastokens: NotRequired[bool]#: None

external_service_readiness: NotRequired[nemo_rl.environments.nemo_gym.ExternalServiceReadinessConfig | None]#: None

tokenizer_config: NotRequired[Optional[nemo_rl.models.policy.TokenizerConfig]]#: None

pad_dynamic_image_shapes: NotRequired[bool]#: None

token_capture: NotRequired[Dict[str, Any] | None]#: None

nemo_rl.environments.nemo_gym._POLICY_SERVER_NAME#: ‘policy_model’

nemo_rl.environments.nemo_gym._NG_ROLLOUT_ID_BODY_KEY#: ‘_ng_rollout_id’

nemo_rl.environments.nemo_gym._TOKEN_CAPTURE_CONTROL_PREFIX#: ‘/training-token-capture/control’

nemo_rl.environments.nemo_gym._TOKEN_CAPTURE_CONTROL_ENV#: ‘NEMO_GYM_TOKEN_CAPTURE_CONTROL_TOKEN’

nemo_rl.environments.nemo_gym._POLICY_API_KEY_ENV#: ‘NEMO_RL_POLICY_API_KEY’

nemo_rl.environments.nemo_gym._external_staging_backend( token_capture: Dict[str, Any], ) → str#: Map the setup-derived generation backend to Gym’s capture backend.

nemo_rl.environments.nemo_gym.extract_external_service_readiness( nemo_gym_dict: dict[str, Any], ) → nemo_rl.environments.nemo_gym.ExternalServiceReadinessConfig | None#: Remove and validate NeMo-RL’s external readiness block from Gym config.

nemo_rl.environments.nemo_gym._probe_external_service( service: nemo_rl.environments.nemo_gym.ExternalServiceReadinessTargetConfig, *, request_timeout_seconds: float, ) → str | None#

Return a readiness problem for one service, or None when it is ready.

The URL must be a vllm_pool_lb /health endpoint. This reads the status, healthy_backends and total_backends fields returned by LoadBalancer.handle_health in tools/external_gym_vllm/vllm_pool_lb.py.

nemo_rl.environments.nemo_gym._wait_for_external_services( config: nemo_rl.environments.nemo_gym.ExternalServiceReadinessConfig | None, ) → None#: Wait for all external Gym dependencies under one shared deadline.

nemo_rl.environments.nemo_gym._detect_invalid_tool_call_and_malformed_thinking( output_item_dict: dict[str, Any], invalid_tool_call_patterns: list[str] | None = None, thinking_tags: list[str] | None = None, ) → tuple[bool, bool]#

Flag a NeMo-Gym output item as an invalid tool call / malformed thinking.

Inspects the final output item of a model turn. For a final content message, any thinking tag is malformed (thinking should never leak into the answer); for a reasoning summary, only a repeated tag (count > 1) is malformed (a single pair is expected). A textual tool-call pattern in either indicates an invalid (unexecuted) tool call.

Returns:: (is_invalid_tool_call, has_malformed_thinking).

nemo_rl.environments.nemo_gym.get_pad_dynamic_image_shapes( env_config: collections.abc.Mapping[str, Any], ) → bool#

Return nemo_gym’s pad_dynamic_image_shapes from an env config, or False.

Takes master_config.env rather than the whole config: interpreting NeMo-Gym settings belongs with the environment, and callers outside it only need the resolved boolean.

The NemoGym actor reads the same key from its own config for the per-turn attach. The initial-payload attach runs in the driver instead, so it has to be read here and passed down, or multi-image prompts would be processed under different rules on the two paths.

class nemo_rl.environments.nemo_gym.NemoGym(cfg: nemo_rl.environments.nemo_gym.NemoGymConfig)#

Bases: nemo_rl.environments.interfaces.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.

Initialization

_require_spinup() → None#: Raise a diagnosable error if this instance never ran :meth:_spinup.

async health_check() → None#

Raise if the Gym head server or any subprocess server has died.

Thin wrapper over NeMo-Gym’s own RunHelper.poll, which is what gym env start calls every 60s from run_forever. NeMo-RL only calls rh.start, so without this the check Gym already implements never runs and a dead tool server surfaces as unexplained rollout timeouts instead of a named process.

Run the synchronous poll in a worker thread so this probe does not block concurrent rollouts on the actor’s event loop.

_spinup() → None#

Start the NeMo-Gym head server and rollout collection helper.

Deferred from init so the actor can be created cheaply (and scheduled onto reserved nodes) and spun up explicitly once the vLLM server URLs are available, overlapping with vLLM model loading.

set_tokenizer(tokenizer: transformers.PreTrainedTokenizerBase) → None#

Install the tokenizer run_rollouts postprocesses with.

Called once per actor at spinup. It used to be a run_rollouts argument, which meant Ray deserialized a tokenizer per prompt on this actor’s task-execution thread. That thread holds the GIL, so it blocked the actor’s event loop and no rollout could issue its first HTTP request until its own copy finished loading.

The cost is not marginal. A tokenizer of this shape measured 7.45 MB on the wire, 286 ms to serialize and 1052 ms to deserialize, so a SingleController recipe admitting ~1000 prompts per step spends roughly 18 minutes deserializing a single admission burst, against a step that should take minutes. Runs on that shape stalled without completing a rollout.

Measured on one CPU node with 1024 concurrent calls against one actor: 153 of 1024 prompts finished in 300 s passing the tokenizer per call, versus all 1024 in 16 s holding it here. Passing an ObjectRef instead does not help – Ray caches the object buffer, not the deserialized value, so it still pays per task.

_control_client()#: Gym ServerClient resolving servers by name from the head server.

async _control(method: str, path: str, **kwargs: Any) → dict#

list_entries() → Dict[str, List[str]]#

Report which config entries this actor actually spawned.

Returns {entry_name: [server_type_keys]} read from Gym’s resolved config, so entries that arrived via config_paths are included. The config NeMo RL passed in is not a substitute: it still holds config_paths as file paths and none of the entries they expand into, so reading it would miss every agent and judge loaded from a path.

Entries whose server config has no entrypoint are omitted because Gym does not start a process for them.

Callers compare these names across actors to build the agent->shard map and to catch an entry duplicated across shards; an unsharded job’s coverage check (validate_dataset_agent_coverage) reads them as the entries a dataset row may name. Names are all that is interpreted; what an entry means is Gym’s business.

async run_rollouts( nemo_gym_examples: list[dict], timer_prefix: str, deduplicate_multimodal_data: bool = False, per_prompt: bool = False, ) → collections.abc.AsyncGenerator[tuple[int, dict, dict, dict | None], None]#

Stream postprocessed rollouts as NeMo-Gym tasks complete.

A thin span-opening wrapper over :meth:_stream_rollouts, which holds the logic. Split so the span covers the whole stream without indenting the body under a with.

The decorator parents this batch to the caller’s span, and everything below inherits it – including the HTTP requests Gym’s own client makes, which are instrumented at the library level (see

Func:: instrument_aiohttp_client) and read the ambient context rather than anything passed here.

An umbrella span either way, so it carries no rl.bucket: several of these are in flight at once, and their durations would sum past the wall clock they happened in.

Parameters:: per_prompt – Whether the caller dispatches this once per prompt. Decided by the caller because in_per_prompt_scope reads a ContextVar in the calling process, and this is a separate Ray actor. On the single-controller path a step issues one of these per prompt, so the span belongs in per_prompt rather than per_step, whose count is meant to scale with steps.
Yields:: One (rowidx, resolved_agent_ref, result, timing_metrics) tuple per completed task, in completion order rather than input order. rowidx echoes back the _rowidx the caller stamped on the example, which is how the caller maps a result to its slot. timing_metrics is None on every tuple but the last, which carries the batch totals.

async _stream_rollouts( nemo_gym_examples: list[dict], timer_prefix: str, deduplicate_multimodal_data: bool = False, ) → collections.abc.AsyncGenerator[tuple[int, dict, dict, dict | None], None]#: Body of :meth:run_rollouts; see there for the tracing wrapper.

async _postprocess_receipt_mode( nemo_gym_row: dict, nemo_gym_result: dict, ) → dict#

Fetch the ledger manifest and assemble the receipt locally.

The legacy token walk (and its contiguity assert) does not run: the capture ledger owns lineage, output items carry no token arrays, and the canonical row is rebuilt by the finalizer from staged deltas. The Ray return carries only the receipt (~100 B/call) beside the agent-level result.

static _assemble_receipt( rollout_id: str, manifest: dict, *, terminal_response_id: Optional[str], scored_response: Optional[dict] = None, reward: float, ) → dict#

Build the token-free RolloutReceipt payload from a ledger manifest.

terminal_response_id is the served response envelope id the harness reports for the completion it kept (Gym’s terminal_response_id result key), forwarded to resolve_terminal as declared_response_id. It is not the logical-request header.

Terminal selection is staged, fail-closed at every stage:

Witness attribution (resolve_terminal): the declared response id (terminal_response_id), the scored response’s envelope id, and its content fingerprints each independently name a manifest row through CallRecord.response_id and the recorded fingerprints. Agreeing witnesses attribute; a declared id that matches no row masks and never falls back; disagreeing witnesses attribute nothing.
Heuristic fallback (select_terminal_call): with no witness, the manifest’s explicit parent links infer the terminal (fail-closed: ambiguity masks).

The receipt records the resolving stage in terminal_selection (declared/response_id/content/heuristic — failed selections stamp the last stage attempted; None when no stage ran because a manifest row failed to parse) and the witness trail in terminal_attribution_reason. A row that fails CallRecord validation masks the rollout (invalid_manifest_row) and is dropped from the shipped manifest, so the finalizer can still enumerate and clean the rows that did parse. Retry duplicates are dead-branch rows: they stay in the manifest (their staged rows are fetched, verified, and cleaned) but never join the terminal chain — verify_and_linearize tolerates rows unreferenced by the terminal chain.

Poisoning is fail-closed with one carve-out. A failure row whose reason is request_finished_without_staged_coordinates is a call that never returned a completion (the ledger commit precedes the response leaving the server) and can never be a lineage parent (an uncommitted call has no row to resolve against) — e.g. the doomed final call of a rollout that exhausted the model’s context window. Such rows are structurally off-chain and do not poison; if the terminal request itself died this way, the missing-terminal-row check below still masks the rollout. Every other failure reason (for example worker_capture_failed, invalid_worker_commit_coordinates, or unresolved_parent; a reason-less failure row poisons as capture_failed) marks a call whose completion WAS served — a hole in the chain — and poisons.

_postprocess_nemo_gym_to_nemo_rl_result( nemo_gym_row: dict, nemo_gym_result: dict, tokenizer: transformers.PreTrainedTokenizerBase, *, include_initial_multimodal_data: bool = True, ) → dict#

shutdown() → None#

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.

abstractmethod step(message_log_batch, metadata)#

abstractmethod global_post_process_and_metrics(batch)#

nemo_rl.environments.nemo_gym.extract_reward_components( nemo_gym_result: dict, ) → Dict[str, float] | None#

Return per-component rewards from a NeMo Gym verify result, or None.

Single-reward NeMo Gym environments return only a scalar reward. Multi-reward environments additionally return reward_components: a mapping of component-name -> score. These are surfaced as reward/<name> batch keys and consumed by GDPO (see nemo_rl.algorithms.advantage_estimator.GDPOAdvantageEstimator).

Returns None when the environment is single-reward (no reward_components), so callers fall back to the scalar reward path unchanged.

nemo_rl.environments.nemo_gym.build_reward_component_columns( component_dicts: List[Dict[str, float] | None], ) → Dict[str, torch.Tensor]#

Build reward/<name> batch columns from per-sample reward-component dicts.

Takes the union of component names across the batch in sorted (deterministic) order and, for each, emits a reward/<name> tensor with one entry per sample. A component absent on a given sample is filled with 0.0 so every column covers all samples (the per-prompt baseline requires each component present for all responses).

Keys are prefixed reward/ so they are exactly what nemo_rl.algorithms.utils.get_gdpo_reward_component_keys selects (it matches startswith("reward/") and sorts by name); the name carries the component identity, so no positional index is needed. Returns an empty dict when no sample has components.

nemo_rl.environments.nemo_gym.validate_reward_components_match_scalar( nemo_gym_results: List[dict], ) → None#

Assert each multi-reward result sets reward == sum(reward_components).

A multi-reward verifier must set the scalar reward to the sum of its reward_components so single-reward (GRPO) consumers and GDPO read the same aggregate. We keep the verifier’s scalar reward as total_reward rather than silently overwriting it with the component sum, so a verifier that violates this contract must be surfaced here instead of masked.

Raises ValueError on the first violating result. A no-op for single-reward results (those without reward_components).

nemo_rl.environments.nemo_gym.setup_nemo_gym_config(config, tokenizer) → None#

nemo_rl.environments.nemo_gym.build_nemo_gym_config( env_configs: dict[str, Any], *, base_urls: list[str], model_name: str, enable_router_replay: bool, use_fastokens: bool, token_capture: Optional[dict[str, Any]] = None, policy_api_key: Optional[str] = None, ) → nemo_rl.environments.nemo_gym.NemoGymConfig#

Build the NemoGymConfig for a single, unsharded NeMo-Gym actor.

Splits env_configs["nemo_gym"] into the NeMo-RL-side fields the actor reads directly and the remainder, which is forwarded verbatim as NeMo-Gym’s initial global config.

Parameters:

env_configs – The master_config.env mapping; env_configs[“nemo_gym”] supplies the Gym global config plus NeMo-RL detection knobs (invalid_tool_call_patterns, thinking_tags, num_gpu_nodes).
base_urls – Per-DP-rank OpenAI-compatible server base URLs from the generation backend.
model_name – Served model name the Gym rollouts should target.
enable_router_replay – Sets require_routed_experts and selects the routed-experts carry dtype (“int8”/”int16”/”int32”) for the model.
use_fastokens – Forwarded from policy.tokenizer.use_fastokens so the actor patches its tokenizer the same way the driver does.
policy_api_key – The bearer token the generation server requires on its OpenAI routes (VllmGeneration.http_server_api_key), or None when it checks none.

Returns:

A NemoGymConfig with NeMo-RL fields at the top level and the remaining env_configs["nemo_gym"] keys under initial_global_config_dict. The caller’s env_configs is not mutated.

Raises:

ShardConfigError – The config is sharded. One config cannot describe several shards; use :func:build_nemo_gym_actors.

nemo_rl.environments.nemo_gym._build_gym_actor_config( nemo_gym_dict: dict[str, Any], *, base_urls: list[str], model_name: str, enable_router_replay: bool, use_fastokens: bool, token_capture: Optional[dict[str, Any]] = None, policy_api_key: Optional[str] = None, ) → nemo_rl.environments.nemo_gym.NemoGymConfig#

Turn one already-resolved Gym config mapping into a NemoGymConfig.

Shared by the unsharded path and by each shard, so every actor gets the same treatment of NeMo-RL-side keys regardless of how it was composed.

nemo_rl.environments.nemo_gym.get_nemo_gym_route_name( row: collections.abc.Mapping[str, Any], ) → str#: Return the entry name Gym uses to route a row.

class nemo_rl.environments.nemo_gym.NemoGymShardSet#

The live actors behind one NeMo-Gym stack, sharded or not.

An unsharded job is the one-shard, one-replica case, so callers do not need a separate code path for it.

.. attribute:: handles

Shard name to its replica handles, in replica order.

.. attribute:: route_to_shard

Agent or task-source entry name to the shard hosting it. Empty when unsharded, where every row goes to the only actor.

.. attribute:: placement_group

The STRICT_SPREAD group pinning shards to distinct nodes, or None when unsharded.

handles: Dict[str, List[ray.actor.ActorHandle]]#: None

route_to_shard: Dict[str, str]#: ‘field(…)’

placement_group: Optional[ray.util.placement_group.PlacementGroup]#: None

_next_replica: Dict[str, int]#: ‘field(…)’

_replica_lock: threading.Lock#: ‘field(…)’

__getstate__() → dict[str, Any]#

__setstate__(state: dict[str, Any]) → None#

property is_sharded: bool#

property all_handles: List[Any]#

property hosted_routes: frozenset[str]#: Agent and task-source entry names this set can route to.

shard_for_route(route_name: str) → str#

Name the shard hosting an agent or task source.

Unsharded jobs have one actor and no map, so every route resolves to it; nothing was discovered because nothing could have conflicted.

Raises:: ShardSetupError – No shard hosts the route, so its rows have nowhere to go.

pick_handle(route_name: str) → Any#

Choose the actor instance to serve a route’s next prompt group.

The shard is fixed by the data; the replica rotates round-robin. Round robin is deterministic and easy to reason about, which matters more than adaptivity here: within a synchronous step there is no completion feedback to adapt on, so an even split is the best available policy. A least-in-flight policy would require callers to release a lease when each dispatch completes. That lifecycle is intentionally outside this round-robin implementation.

instance_label(handle: Any) → str#

Name one actor, for error messages and metric keys.

A shard with one replica is named by the shard alone, so the common case reads as it did before replicas existed; a replicated shard adds the replica index. This is the same rule the per-shard log directories follow, so a metric and its logs carry the same name.

sole_handle() → Any#

The only actor, for callers that predate routing.

Raises:: ShardSetupError – There is more than one actor, so picking one would silently drop the rest.

shutdown( *, timeout: float | None = NEMO_GYM_GRACEFUL_SHUTDOWN_TIMEOUT_S, force_kill: bool = True, ) → None#: Stop every actor, then release the bundles they were pinned to.

nemo_rl.environments.nemo_gym._shard_instances( plan: nemo_rl.environments.nemo_gym_shards.ShardPlan, ) → List[tuple[nemo_rl.environments.nemo_gym_shards.ShardSpec, int]]#: Expand shards into one (shard, replica_index) entry per actor.

nemo_rl.environments.nemo_gym.as_nemo_gym_shard_set( environment: Any, ) → nemo_rl.environments.nemo_gym.NemoGymShardSet#

Read the NeMo-Gym entry of task_to_env as a shard set either way.

Call sites that predate sharding put a bare actor handle there. Rather than make every one of them build a set first, treat a lone handle as the one-shard, one-replica case it already is, so the routing path is identical whether or not the job is sharded.

nemo_rl.environments.nemo_gym.build_nemo_gym_actors( env_configs: dict[str, Any], *, base_urls: list[str], model_name: str, tokenizer: transformers.PreTrainedTokenizerBase, enable_router_replay: bool, use_fastokens: bool, token_capture: Optional[dict[str, Any]] = None, policy_api_key: Optional[str] = None, pg_ready_timeout: float = DEFAULT_SHARD_PG_READY_TIMEOUT_SECONDS, spinup_timeout: float = DEFAULT_SHARD_SPINUP_TIMEOUT_SECONDS, ) → nemo_rl.environments.nemo_gym.NemoGymShardSet#

Create and spin up every NeMo-Gym actor this job needs.

Without shards this makes exactly one actor, scheduled as before. With shards it makes one actor per replica, each pinned by a STRICT_SPREAD placement group to a distinct node, each holding its own complete Gym stack. Actors are spun up concurrently, so the wall-clock cost is roughly the slowest shard rather than the sum.

Parameters:: tokenizer – Installed on every actor once it is up, rather than passed per rollout call. See NemoGym.set_tokenizer for why.
Returns:: A :class:NemoGymShardSet whose actors are all running and validated.
Raises:: ShardSetupError – The bundles could not be placed, a shard failed to start, or the shards’ entries did not pass the startup checks. Any actors already created are torn down first.

nemo_rl.environments.nemo_gym._build_single_gym_actor( nemo_gym_dict: dict[str, Any], *, base_urls: list[str], model_name: str, tokenizer: transformers.PreTrainedTokenizerBase, enable_router_replay: bool, use_fastokens: bool, token_capture: Optional[dict[str, Any]], policy_api_key: Optional[str] = None, ) → nemo_rl.environments.nemo_gym.NemoGymShardSet#

The pre-sharding path: one actor, no placement group, no discovery.

Discovery is skipped rather than merely unused. Its checks compare entry names between shards, so with one shard there is nothing they could find.

nemo_rl.environments.nemo_gym._build_sharded_gym_actors( nemo_gym_dict: dict[str, Any], plan: nemo_rl.environments.nemo_gym_shards.ShardPlan, *, base_urls: list[str], model_name: str, tokenizer: transformers.PreTrainedTokenizerBase, enable_router_replay: bool, use_fastokens: bool, token_capture: Optional[dict[str, Any]], policy_api_key: Optional[str] = None, pg_ready_timeout: float, spinup_timeout: float, ) → nemo_rl.environments.nemo_gym.NemoGymShardSet#

nemo_rl.environments.nemo_gym._spinup_shards_concurrently( shard_set: nemo_rl.environments.nemo_gym.NemoGymShardSet, spinup_timeout: float, *, tokenizer: transformers.PreTrainedTokenizerBase, ) → None#

Start every shard at once, naming the shard behind any failure.

Config faults surface inside _spinup – a server ref pointing at an entry that is missing from this shard’s slice raises Gym’s ServerRefNotFoundError there. Gym names the entry and field but has no concept of a shard, so the shard name is added here.

The tokenizer install is a second pass rather than a call queued behind each _spinup. Queuing both up front would mean waiting on the install to learn the spinup failed, which loses the message above.

nemo_rl.environments.nemo_gym._discover_route_shard_map( shard_set: nemo_rl.environments.nemo_gym.NemoGymShardSet, plan: nemo_rl.environments.nemo_gym_shards.ShardPlan, ) → Dict[str, str]#

Ask one replica per shard what it spawned, then build routing metadata.

Replicas of a shard are stamped from one merge, so they host identical entries and only the first needs to be asked.

nemo_rl.environments.nemo_gym.spinup_nemo_gym_actor( env_configs: dict[str, Any], *, base_urls: list[str], model_name: str, tokenizer: transformers.PreTrainedTokenizerBase, enable_router_replay: bool, use_fastokens: bool, token_capture: Optional[dict[str, Any]] = None, policy_api_key: Optional[str] = None, ) → Any#

Spin up a single NeMo-Gym actor against the given generation server URLs.

When env_configs["nemo_gym"]["num_gpu_nodes"] > 0, the actor is scheduled with soft NodeAffinity to the caller’s Ray node so its colocated GPU resources land where the caller expects.

Parameters:

tokenizer – Installed on the actor once, here, rather than passed per rollout call. See NemoGym.set_tokenizer for why that distinction is the difference between a working run and a stalled one.
token_capture – Dumped TokenCaptureConfig when ledger-authoritative token capture is enabled, else None. Forwarded to build_nemo_gym_config.
policy_api_key – The bearer token the generation server requires on its OpenAI routes (VllmGeneration.http_server_api_key), or None when it checks none.

Returns:

The spun-up NemoGym Ray actor handle (_spinup already awaited).

Raises:

ShardConfigError – The config is sharded. Callers that dispatch to a single handle cannot serve several shards; they need the router.

nemo_rl.environments.nemo_gym.validate_dataset_agent_coverage( shard_set: nemo_rl.environments.nemo_gym.NemoGymShardSet, datasets: collections.abc.Mapping[str, Any], ) → None#

Fail at setup if any row names a route no shard hosts.

Rows can name a legacy agent_ref or a current Gym task_source. Without this scan, a rare route can sit unseen for hours of training before its first dispatch fails.

A sharded job checks against the route map built at setup. An unsharded job has one actor and no route map, so it asks that actor which entries its resolved Gym config spawned and keeps the routable ones, agents and resources servers, exactly as the sharded route map does. A model or environment server is never a route, on one actor or on many.

Parameters:

shard_set – The running actors. A sharded set carries the route map built at setup; an unsharded set is asked for its entries here.
datasets – Split name to dataset, for the error message. None values and datasets without gym rows are skipped.

Raises:

ShardSetupError – A split references routes no shard hosts.

nemo_rl.environments.nemo_gym._iter_dataset_agent_names(dataset: Any) → set[str]#

Collect the agent or task-source names a dataset’s rows reference.

A stable source file is scanned once and its names cached; a custom or changed source falls back to scanning the loaded rows.

nemo_rl.environments.nemo_gym._load_agent_names_from_source( source: nemo_rl.data.interfaces.NemoGymSourceIdentity, ) → frozenset[str] | None#: Read a stable Gym source once per controller process.

nemo_rl.environments.nemo_gym._get_agent_name(row: object) → str | None#

nemo_rl.environments.nemo_gym#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.environments.nemo_gym`#