`nemo_rl.environments.checkpointing.interfaces`#

Save and restore rollout state that lives in the environment backend.

At turn-level recovery the environment backend holds part of every unfinished rollout: agent sessions, tool state, the turns already run. The single controller saves and restores that state only through EnvironmentCheckpoint and never names the backend. NeMo Gym implements it in nemo_rl.environments.checkpointing.nemo_gym; a backend that keeps no rollout state uses NoEnvironmentCheckpoint.

A rollout snapshot uses the interface in this order::

async with env_checkpoint.snapshot(path) as snapshot:  # backend paused
    async with data_plane_barrier.checkpoint() as cut:
        if snapshot is not None:
            await snapshot.check_against_cut(cut)
        ...  # RL saves TQ and its own state
# backend resumed

Module Contents#

Classes#

`EnvironmentEpisode`	One rollout attempt as the environment backend knows it.
`RestoredEnvironmentRollouts`	What an environment restore continued, so RL can match it.
`EnvironmentSnapshot`	The backend’s part of one rollout snapshot, while the backend is paused.
`EnvironmentCheckpoint`	Rollout state an environment backend keeps between turns.
`NoEnvironmentCheckpoint`	For a backend that keeps no rollout state between turns.

API#

class nemo_rl.environments.checkpointing.interfaces.EnvironmentEpisode[source]#

One rollout attempt as the environment backend knows it.

The recovery ledger names the environment’s rollouts this way when it asks the backend to retire, forget or continue them.

.. attribute:: rollout_id

The rollout’s base ID, the same for every attempt.

.. attribute:: attempt

How many times the backend restored this rollout. 0 is the first run; a restore continues it as the next attempt.

rollout_id: str#: None

attempt: int#: None

__post_init__() → None[source]#

property capture_key: str#

The attempt’s key on rollout rows: <rollout_id>-a<attempt>.

Attempt 0 uses the bare rollout ID. A backend must key its rows the same way.

next_attempt() → nemo_rl.environments.checkpointing.interfaces.EnvironmentEpisode[source]#: Return the attempt the backend creates when it restores this one.

class nemo_rl.environments.checkpointing.interfaces.RestoredEnvironmentRollouts[source]#

What an environment restore continued, so RL can match it.

.. attribute:: continued_attempts

(instance_id, rollout_id, attempt) of every rollout the backend restored. The recovery ledger continues exactly these and starts every other unfinished rollout over.

.. attribute:: staging_keys

Token-capture staging rows the restored rollouts still reference. Restore keeps them instead of clearing them as orphans.

continued_attempts: frozenset[tuple[str, str, int]]#: ‘frozenset(…)’

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

class nemo_rl.environments.checkpointing.interfaces.EnvironmentSnapshot[source]#

Bases: typing.Protocol

The backend’s part of one rollout snapshot, while the backend is paused.

property restarted_rollouts: int#: Rollouts the backend could not capture; they start over after a crash.

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

Under the data-plane cut, require the captured state to match RL’s.

Raises when the backend’s rollouts and RL’s ledger disagree, which would make the snapshot unrestorable.

metrics(paused_seconds: float) → dict[str, float][source]#: Backend metrics for this snapshot, logged with the snapshot timings.

class nemo_rl.environments.checkpointing.interfaces.EnvironmentCheckpoint[source]#

Bases: typing.Protocol

Rollout state an environment backend keeps between turns.

property holds_rollout_state: bool#

Whether rollout snapshots must capture this backend.

True means a trainer checkpoint, which never pauses the backend, does not stand in for a rollout snapshot.

async check_startup() → None[source]#: Refuse, or warn about, a deployment that cannot recover as configured.

snapshot( checkpoint_path: pathlib.Path, ) → contextlib.AbstractAsyncContextManager[Optional[nemo_rl.environments.checkpointing.interfaces.EnvironmentSnapshot]][source]#

Pause the backend and save its rollouts under checkpoint_path.

Yields None when there is nothing to save. The backend resumes, and accepts new rollouts again, when the context exits.

restored( checkpoint_path: pathlib.Path, ) → contextlib.AbstractAsyncContextManager[nemo_rl.environments.checkpointing.interfaces.RestoredEnvironmentRollouts][source]#

Restore the backend’s rollouts from checkpoint_path.

The restore is kept only if the context exits cleanly. On an exception the backend drops what it restored before it resumes.

is_unavailable(error: BaseException) → bool[source]#

Whether a failed snapshot only met a backend that was briefly unreachable.

Such a failure leaves nothing RL holds changed, so the next attempt can retry.

class nemo_rl.environments.checkpointing.interfaces.NoEnvironmentCheckpoint[source]#

For a backend that keeps no rollout state between turns.

property holds_rollout_state: bool#

async check_startup() → None[source]#

async snapshot( checkpoint_path: pathlib.Path, ) → collections.abc.AsyncIterator[Optional[nemo_rl.environments.checkpointing.interfaces.EnvironmentSnapshot]][source]#

async restored( checkpoint_path: pathlib.Path, ) → collections.abc.AsyncIterator[nemo_rl.environments.checkpointing.interfaces.RestoredEnvironmentRollouts][source]#

is_unavailable(error: BaseException) → bool[source]#

nemo_rl.environments.checkpointing.interfaces#

Module Contents#

Classes#

API#

`nemo_rl.environments.checkpointing.interfaces`#