`nemo_rl.environments.gym_checkpoint_coordinator`#

NeMo RL coordination across independent Gym checkpoint deployments.

Module Contents#

Classes#

`GymCheckpointManifest`	Durable mapping from Gym actor instances to continued episodes.
`GymCheckpointCommitResult`	Outer manifest plus each Gym actor’s validated commit summary.
`GymCheckpointCoordinator`	Drive one Gym checkpoint operation across every NeMo-Gym actor.

Functions#

`is_transient_gym_checkpoint_failure`	Whether a failed Gym checkpoint may succeed if it is simply tried again.
`_require_schema_version`
`_episodes_by_instance`
`staging_keys_of`	Return the staged rows of `episodes`, the only rows a restore continues.
`load_gym_checkpoint_manifest`	Load and validate the Gym manifest from one committed RL snapshot.
`write_gym_checkpoint_manifest`	Atomically write a Gym manifest inside an unpublished RL snapshot.
`_cancel_remote`	Ask Ray to stop a control call RL stopped waiting for; best effort.

Data#

`GYM_CHECKPOINT_MANIFEST_FILENAME`
`GYM_CHECKPOINT_SCHEMA_VERSION`
`_MAX_REPLY_GRACE_S`
`_RESUME_ATTEMPTS`
`_RESUME_RETRY_PAUSE_S`

API#

nemo_rl.environments.gym_checkpoint_coordinator.GYM_CHECKPOINT_MANIFEST_FILENAME#: ‘gym_checkpoint.json’

nemo_rl.environments.gym_checkpoint_coordinator.GYM_CHECKPOINT_SCHEMA_VERSION#: 2

nemo_rl.environments.gym_checkpoint_coordinator._MAX_REPLY_GRACE_S#: 5.0

nemo_rl.environments.gym_checkpoint_coordinator._RESUME_ATTEMPTS#: 2

nemo_rl.environments.gym_checkpoint_coordinator._RESUME_RETRY_PAUSE_S#: 1.0

exception nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointOperationError( operation: str, failures: collections.abc.Mapping[str, BaseException], )#

Bases: RuntimeError

One or more independent Gym deployments failed one operation.

Initialization

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

__reduce__()#

exception nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointNotReady( blockers: collections.abc.Mapping[str, collections.abc.Mapping[str, tuple[str, ...]]], )#

Bases: RuntimeError

Gym participants could not all park before the checkpoint deadline.

Initialization

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

__reduce__()#

nemo_rl.environments.gym_checkpoint_coordinator.is_transient_gym_checkpoint_failure(error: BaseException) → bool#

Whether a failed Gym checkpoint may succeed if it is simply tried again.

Participants that could not park in time, calls that went unanswered, and actors or control planes that failed a call leave nothing behind: the coordinator resumes Gym, and Gym’s lease reopens it if that resume is lost too. Anything else, such as a malformed reply or an ownership or topology mismatch, is a broken invariant that a retry would only hide.

class nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointManifest#

Durable mapping from Gym actor instances to continued episodes.

restarts are the episodes RL still owned that Gym could not capture: a restore leaves them out and starts them over from their tasks.

schema_version: int#: None

checkpoint_id: str#: None

instances: dict[str, tuple[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode, ...]]#: None

staging_keys: dict[str, tuple[str, ...]]#: None

restarts: dict[str, tuple[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode, ...]]#: None

participants: dict[str, tuple[str, ...]]#: None

__post_init__() → None#

to_dict() → dict[str, Any]#

classmethod from_mapping( raw: collections.abc.Mapping[str, Any], ) → nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointManifest#: Parse a manifest RL wrote; post_init checks what it means.

nemo_rl.environments.gym_checkpoint_coordinator._require_schema_version(version: Any) → None#

nemo_rl.environments.gym_checkpoint_coordinator._episodes_by_instance( raw: collections.abc.Mapping[str, list[collections.abc.Mapping[str, Any]]], ) → dict[str, tuple[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode, ...]]#

class nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointCommitResult#

Outer manifest plus each Gym actor’s validated commit summary.

manifest: nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointManifest#: None

instances: dict[str, nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointCommitSummary]#: None

nemo_rl.environments.gym_checkpoint_coordinator.staging_keys_of( summary: nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointCommitSummary, episodes: collections.abc.Iterable[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode], ) → tuple[str, ...]#

Return the staged rows of episodes, the only rows a restore continues.

Gym attributes each row to the episodes whose records hold it, so a generation-cut row needs no parsing and a carried-over row stays while any attempt that uses it is kept.

nemo_rl.environments.gym_checkpoint_coordinator.load_gym_checkpoint_manifest( checkpoint_root: str | pathlib.Path, ) → nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointManifest#: Load and validate the Gym manifest from one committed RL snapshot.

nemo_rl.environments.gym_checkpoint_coordinator.write_gym_checkpoint_manifest( checkpoint_root: str | pathlib.Path, manifest: nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointManifest, ) → None#: Atomically write a Gym manifest inside an unpublished RL snapshot.

class nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointCoordinator( shard_set: nemo_rl.environments.nemo_gym.NemoGymShardSet, *, control_timeout_s: float, prepare_timeout_s: float, )#

Drive one Gym checkpoint operation across every NeMo-Gym actor.

Initialization

property instance_ids: frozenset[str]#

property control_timeout_s: float#: Maximum wall-clock budget for one coordinated control call but prepare.

_deadline() → float#

async _collect( operation: str, references: collections.abc.Mapping[str, Any], *, deadline_ts: float, ) → dict[str, Any]#

Wait for every actor’s reply, but never past the call’s deadline.

Gym enforces deadline_ts only once it runs the call. An actor whose event loop is stuck, or a call Ray never schedules, would otherwise be waited on forever. Giving up is safe: a Gym deployment left parked reopens by itself when its checkpoint lease expires.

async prepare( checkpoint_id: str, ) → dict[str, nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointPrepareSummary]#

async renew(checkpoint_id: str) → None#

async restart_only_servers() → dict[str, list[str]]#: Return each instance’s Gym participants whose rollouts start over after a crash.

_validate_inventory( episodes_by_instance: collections.abc.Mapping[str, tuple[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode, ...]], ) → dict[str, tuple[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode, ...]]#

async commit( checkpoint_id: str, checkpoint_root: str | pathlib.Path, episodes_by_instance: collections.abc.Mapping[str, tuple[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode, ...]], *, restarts: collections.abc.Mapping[str, tuple[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode, ...]], ) → nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointCommitResult#

Save episodes_by_instance and record restarts as starting over.

Gym refuses a commit whose scope names an episode it reported as a restart, so the two must not overlap.

async retire( checkpoint_id: str, episodes_by_instance: collections.abc.Mapping[str, tuple[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode, ...]], ) → None#: Retire selected physical attempts from every Gym deployment.

async forget( checkpoint_id: str, episodes_by_instance: collections.abc.Mapping[str, tuple[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode, ...]], ) → None#

Let Gym stop refusing these rollouts: RL will never send them again.

One call per deployment with anything to forget, naming each rollout once.

async retire_restored( restore_id: str, manifest: nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointManifest, ) → None#: Retire the replacement attempts installed by a successful restore.

async restore( restore_id: str, checkpoint_root: str | pathlib.Path, manifest: nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointManifest, ) → None#

async _discard_restore( restore_id: str, manifest: nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointManifest, error: BaseException, ) → None#: Retire the restored attempts, then reopen Gym, noting failures on error.

async restored( restore_id: str, checkpoint_root: str | pathlib.Path, manifest: nemo_rl.environments.gym_checkpoint_coordinator.GymCheckpointManifest, ) → collections.abc.AsyncIterator[None]#

Restore every deployment for the body, then always reopen it once.

A failed restore cleans up and reopens Gym itself. A body that fails could not take over the restored attempts, so they are retired before Gym reopens and never continue.

async resume(checkpoint_id: str) → None#: Reopen every deployment, retrying the ones a transient failure left closed.

async _renew_until_stopped(checkpoint_id: str, stop: asyncio.Event) → None#: Keep a prepared Gym deployment parked during a slow outer save.

async prepared( checkpoint_id: str, ) → collections.abc.AsyncIterator[dict[str, frozenset[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode]]]#

Prepare every deployment, renew its lease, and always release it.

Yields each instance’s restarts: the episodes a commit must leave out. Gym resumes when the body exits, before the caller publishes, so new rollouts wait only for Gym’s own commit. Once the body succeeded every commit is written and resume only reopens admission: a resume that still fails is reported, not raised, since Gym’s lease reopens that deployment on its own and the checkpoint stays valid.

nemo_rl.environments.gym_checkpoint_coordinator._cancel_remote(reference: Any) → None#: Ask Ray to stop a control call RL stopped waiting for; best effort.

nemo_rl.environments.gym_checkpoint_coordinator#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.environments.gym_checkpoint_coordinator`#