`nemo_rl.environments.nemo_gym_shards`#

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.

Module Contents#

Classes#

`ShardSpec`	One shard: a slice of the Gym config that gets its own actor.
`ShardPlan`	The parsed `shards` block plus the settings that apply to all shards.

Functions#

`_as_plain_dict`	Coerce a mapping (possibly OmegaConf-backed) to a plain dict.
`_as_string_set`	Coerce a list of unique strings to a set.
`_parse_shard`
`find_gym_config_entries`	Return top-level keys that Gym would read as server instance configs.
`parse_shard_plan`	Parse `env.nemo_gym.shards`, or return None when the job is unsharded.
`apply_shard_overlay`	Build one shard’s Gym config from shared settings and its claimed overlays.
`build_route_shard_map`	Map each routable entry to its shard, rejecting ambiguous routes.
`apply_shard_log_dir`	Give one shard instance its own subdirectory of the configured log dir.

Data#

`NEMO_RL_DICT_CONFIG_KEYS`
`DEFAULT_REPLICAS`
`SHARD_NAME_PATTERN`
`GYM_LOG_DIR_KEY`
`GYM_ROUTABLE_KEYS`
`SHARD_SPEC_KEYS`
`SHARDING_CONFIG_KEYS`

API#

nemo_rl.environments.nemo_gym_shards.NEMO_RL_DICT_CONFIG_KEYS#: ‘frozenset(…)’

nemo_rl.environments.nemo_gym_shards.DEFAULT_REPLICAS#: 1

nemo_rl.environments.nemo_gym_shards.SHARD_NAME_PATTERN#: ‘compile(…)’

nemo_rl.environments.nemo_gym_shards.GYM_LOG_DIR_KEY#: ‘nemo_gym_log_dir’

nemo_rl.environments.nemo_gym_shards.GYM_ROUTABLE_KEYS#: ‘frozenset(…)’

exception nemo_rl.environments.nemo_gym_shards.ShardConfigError#

Bases: ValueError

Raised for a malformed env.nemo_gym.shards block.

Initialization

Initialize self. See help(type(self)) for accurate signature.

exception nemo_rl.environments.nemo_gym_shards.ShardSetupError#

Bases: RuntimeError

Raised when a sharded stack fails to start or fails its startup checks.

Initialization

Initialize self. See help(type(self)) for accurate signature.

class nemo_rl.environments.nemo_gym_shards.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#: None

config_paths: list[str]#: None

inherited_overlays: frozenset[str]#: ‘frozenset(…)’

overrides: dict[str, Any]#: ‘field(…)’

replicas: int#: None

actor_cpus: float | None#: None

port_range_low: int | None#: None

port_range_high: int | None#: None

nemo_rl.environments.nemo_gym_shards.SHARD_SPEC_KEYS#: ‘frozenset(…)’

class nemo_rl.environments.nemo_gym_shards.ShardPlan#

The parsed shards block plus the settings that apply to all shards.

shards: list[nemo_rl.environments.nemo_gym_shards.ShardSpec]#: None

common_inherited_overlays: frozenset[str]#: ‘frozenset(…)’

common_overrides: dict[str, Any]#: ‘field(…)’

allowed_duplicate_entries: frozenset[str]#: ‘frozenset(…)’

nemo_rl.environments.nemo_gym_shards.SHARDING_CONFIG_KEYS#: ‘frozenset(…)’

nemo_rl.environments.nemo_gym_shards._as_plain_dict( value: Any, *, context: str, ) → dict[str, Any]#: Coerce a mapping (possibly OmegaConf-backed) to a plain dict.

nemo_rl.environments.nemo_gym_shards._as_string_set(value: Any, *, context: str) → frozenset[str]#: Coerce a list of unique strings to a set.

nemo_rl.environments.nemo_gym_shards._parse_shard( raw: Any, index: int, ) → nemo_rl.environments.nemo_gym_shards.ShardSpec#

nemo_rl.environments.nemo_gym_shards.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.

nemo_rl.environments.nemo_gym_shards.parse_shard_plan( nemo_gym_config: Mapping[str, Any], ) → nemo_rl.environments.nemo_gym_shards.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.

nemo_rl.environments.nemo_gym_shards.apply_shard_overlay( base_config: dict[str, Any], plan: nemo_rl.environments.nemo_gym_shards.ShardPlan, shard: nemo_rl.environments.nemo_gym_shards.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.

nemo_rl.environments.nemo_gym_shards.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.

nemo_rl.environments.nemo_gym_shards.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.

nemo_rl.environments.nemo_gym_shards#