`nemo_rl.environments.gym_checkpoint_adapter`#

Thin NeMo RL adapter over one NeMo Gym v2 checkpoint coordinator.

Module Contents#

Classes#

`GymCheckpointInstance`	Stable identity and filesystem namespace for one logical Gym actor.
`GymCheckpointCommitSummary`	Actor-local Gym state selected by one checkpoint commit.
`GymCheckpointPrepareSummary`	Serializable result of preparing one actor’s Gym deployment.
`GymCheckpointAdapter`	Delegate checkpoint operations for exactly one Gym deployment to Gym.

Functions#

_reraise_as_gym_checkpoint_unavailable

Turn Gym’s unpicklable CoordinationError into GymCheckpointUnavailable.

Data#

GymCheckpointEpisode

API#

nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode#: None

class nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointInstance[source]#

Stable identity and filesystem namespace for one logical Gym actor.

shard_name: str#: None

replica_index: int#: None

__post_init__() → None[source]#

property instance_id: str#

live_capture_dir(capture_root: str | pathlib.Path) → pathlib.Path[source]#

checkpoint_dir(checkpoint_root: str | pathlib.Path) → pathlib.Path[source]#

class nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointCommitSummary[source]#

Actor-local Gym state selected by one checkpoint commit.

exported_episodes: tuple[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode, ...]#: None

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

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

staging_keys_by_episode: collections.abc.Mapping[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode, tuple[str, ...]]#: ‘field(…)’

__post_init__() → None[source]#

unowned_live_episodes: tuple[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode, ...]#: ()

class nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointPrepareSummary[source]#

Serializable result of preparing one actor’s Gym deployment.

restarts are in-flight episodes some participant cannot capture: Gym refuses a commit that names one, and they start over after a crash.

prepared: bool#: None

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

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

classmethod from_result( result: nemo_gym._checkpoint.coordination.PrepareResult, ) → nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointPrepareSummary[source]#

exception nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointUnavailable[source]#

Bases: RuntimeError

Gym’s control plane could not complete one checkpoint control call.

Raised in place of Gym’s CoordinationError, which cannot be pickled and so reaches the driver as an untyped UnserializableException. The participants’ failures arrive as text, so this says only that the call failed, not why; nothing RL holds changed, and a later checkpoint may succeed.

Initialization

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

nemo_rl.environments.gym_checkpoint_adapter._reraise_as_gym_checkpoint_unavailable() → collections.abc.Iterator[None][source]#

Turn Gym’s unpicklable CoordinationError into GymCheckpointUnavailable.

Every Gym coordination call goes through this, so its failure reaches the driver typed, with Gym’s message, instead of as UnserializableException.

class nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointAdapter( *, instance: nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointInstance, client: nemo_gym.server_utils.ServerClient, auth_token: str, )[source]#

Delegate checkpoint operations for exactly one Gym deployment to Gym.

Initialization

async discover() → None[source]#

async restart_only_servers() → list[str][source]#

Return the participants whose rollouts start over after a crash.

Gym reports their in-flight rollouts as restarts at every checkpoint, so turn-level recovery quietly becomes restart-from-task for them. Each participant declares this as its checkpoint mode, which discovery already returned.

_require_participants() → nemo_gym._checkpoint.coordination.Participants[source]#

async prepare( checkpoint_id: str, *, deadline_ts: float, ) → nemo_gym._checkpoint.coordination.PrepareResult[source]#

async renew(checkpoint_id: str, *, deadline_ts: float) → None[source]#

static _episode_ids( episodes: collections.abc.Iterable[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode], ) → list[Any][source]#: Convert the RL wire identity to Gym’s validated API model.

_commit_summary( checkpoint_id: str, episode_ids: list[Any], replies: dict[str, dict[str, Any]], ) → nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointCommitSummary[source]#: Convert participant replies into an actor-local ownership summary.

async retire( checkpoint_id: str, episodes: collections.abc.Iterable[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode], *, deadline_ts: float, ) → None[source]#

async forget( checkpoint_id: str, rollout_ids: collections.abc.Iterable[str], *, deadline_ts: float, ) → None[source]#

async commit( checkpoint_id: str, checkpoint_root: str | pathlib.Path, episodes: collections.abc.Iterable[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode], *, deadline_ts: float, ) → nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointCommitSummary[source]#

async restore( checkpoint_id: str, checkpoint_root: str | pathlib.Path, episodes: collections.abc.Iterable[nemo_rl.environments.gym_checkpoint_adapter.GymCheckpointEpisode], *, participants: collections.abc.Iterable[str], source_checkpoint_id: str, deadline_ts: float, ) → None[source]#

Restore episodes into this deployment.

participants are the server names that committed the checkpoint. Gym restores only live participants, so one missing here would have its state silently dropped; Gym itself fails a live participant that has no saved state.

async resume(checkpoint_id: str, *, deadline_ts: float) → None[source]#

nemo_rl.environments.gym_checkpoint_adapter#