`nemo_rl.environments.checkpointing.nemo_gym`#

NeMo Gym’s side of turn-level rollout snapshots.

Implements EnvironmentCheckpoint (nemo_rl.environments.checkpointing.interfaces) over a GymCheckpointCoordinator. Everything that knows Gym episodes, Gym participants, retire/forget, and gym_checkpoint.json lives here, so the single controller only pauses, snapshots, validates and restores through the interface.

Module Contents#

Classes#

`GymEnvironmentSnapshot`	Gym’s committed part of one rollout snapshot.
`GymEnvironmentCheckpoint`	Pause, commit, validate and restore Gym’s episodes for RL’s snapshots.

Functions#

`_summarize_gym_episodes`	Count episodes per instance and name a few, however many there are.
`build_gym_environment_checkpoint`	Build the Gym checkpoint and close Gym dispatch through its admission gate.

Data#

`_GYM_CHECKPOINT_DEADLINE_HOLDER`
`_GYM_EPISODE_SAMPLE_SIZE`

API#

nemo_rl.environments.checkpointing.nemo_gym._GYM_CHECKPOINT_DEADLINE_HOLDER#: ‘gym_checkpoint’

nemo_rl.environments.checkpointing.nemo_gym._GYM_EPISODE_SAMPLE_SIZE#: 5

nemo_rl.environments.checkpointing.nemo_gym._summarize_gym_episodes( episodes_by_instance: collections.abc.Mapping[str, collections.abc.Iterable[Any]], ) → str#: Count episodes per instance and name a few, however many there are.

nemo_rl.environments.checkpointing.nemo_gym.build_gym_environment_checkpoint( master_config: nemo_rl.algorithms.single_controller_utils.config.MasterConfig, *, env_handles: dict[str, Any], rollout_manager: nemo_rl.experience.rollout_manager.RolloutManager, dp_client: Any, ) → GymEnvironmentCheckpoint#: Build the Gym checkpoint and close Gym dispatch through its admission gate.

class nemo_rl.environments.checkpointing.nemo_gym.GymEnvironmentSnapshot#

Gym’s committed part of one rollout snapshot.

checkpoint: GymEnvironmentCheckpoint#: None

commit: nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointCommitResult#: None

checkpoint_path: pathlib.Path#: None

property restarted_rollouts: int#

async check_against_cut( cut: nemo_rl.algorithms.async_utils.replay_buffer.DataPlaneMutationCut, ) → None#

Require Gym ownership and referenced staging rows to match this cut.

Gym may have exported episodes RL dropped while it was prepared: it refuses their retire until resume, so they are still queued for one. They and their staged rows are left out of the saved manifest, since restore requires it to match RL’s inventory and RL no longer continues them; a dropped group may already have cleaned those rows up.

metrics(paused_seconds: float) → dict[str, float]#

class nemo_rl.environments.checkpointing.nemo_gym.GymEnvironmentCheckpoint( coordinator: nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointCoordinator, *, ledger: nemo_rl.experience.rollout_recovery.RolloutRecoveryLedger, rollout_manager: Any, list_staging_keys: collections.abc.Callable[[], collections.abc.Awaitable[collections.abc.Iterable[str]]], allow_restart_only_servers: bool = False, )#

Pause, commit, validate and restore Gym’s episodes for RL’s snapshots.

Initialization

Wrap coordinator.

Parameters:

coordinator – Talks to every Gym deployment’s checkpoint participants.
ledger – RL’s record of which Gym episodes it owns.
rollout_manager – Pauses rollout deadlines while Gym is parked.
list_staging_keys – Returns the token-capture staging rows TQ holds, to check the ones the committed episodes reference still exist.
allow_restart_only_servers – Accept Gym servers whose rollouts start over after a crash, with a warning, instead of refusing them.

property holds_rollout_state: bool#

is_unavailable(error: BaseException) → bool#

inventory() → dict[str, tuple[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode, ...]]#: Return the Gym-owned episode attempts of every live deployment.

_queued_retirements() → dict[str, frozenset[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode]]#: Return dropped episodes Gym may still hold, awaiting their retire.

async check_startup() → None#

Refuse, or warn about, Gym servers whose rollouts start over after a crash.

Such a server’s in-flight rollouts are reported as restarts at every checkpoint, so turn-level recovery would quietly become restart-from-task for them.

async _retire_dropped_episodes(checkpoint_id: str) → None#

Retire episodes RL dropped while Gym may still be running them.

Abandoning or discarding a group leaves an admitted episode running in Gym, which would export it at commit, outside the candidate set. Gym accepts a retire only while idle, so this runs before prepare. A drop recorded after it stays queued: this checkpoint commits it in scope but leaves it out of the saved manifest, and a later retire stops it.

async _forget_finished_rollouts(checkpoint_id: str) → None#

Let Gym drop its marks for rollouts RL will never send again.

Gym refuses a retired or restored rollout by ID until it is forgotten, so each finished one would otherwise stay on every participant for the rest of the run. Gym refuses a forget while a checkpoint is open, so this runs right after the retire. A deployment that refuses, for example because a retire has not finished freeing a rollout, gets the retire again and then the forget. Forget is best effort: a mark costs one small entry, so a deployment that still refuses keeps its rollouts queued for the next snapshot instead of failing this one.

async _drain_non_exported_candidates( candidates: dict[str, tuple[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode, ...]], commit: nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointCommitResult, ) → None#

Wait for on-wire replies that Gym did not retain in its checkpoint.

A restart is still running in Gym, parked until resume, so it is not waited for: the snapshot records it as starting over.

async snapshot( checkpoint_path: pathlib.Path, *, checkpoint_id: Optional[str] = None, ) → collections.abc.AsyncIterator[nemo_rl.environments.checkpointing.nemo_gym.GymEnvironmentSnapshot]#: Close Gym admission, commit owned episodes, and drain on-wire replies.

async restored( checkpoint_path: pathlib.Path, ) → collections.abc.AsyncIterator[nemo_rl.environments.checkpointing.interfaces.RestoredEnvironmentRollouts]#: Restore Gym’s saved episodes; on failure retire them before Gym reopens.

nemo_rl.environments.checkpointing.nemo_gym#