Partial rollout checkpointing for long-running Gym executions #3024

Current status

The design changed after this issue was filed: checkpointing is now built around environment servers, with each server kind (policy model, environment, resources, agent) taking part in a checkpoint through a shared control plane.
It lands as a stack of pull requests, one commit each.

Merged:

• feat(checkpoint): add the participant control plane, episode steps, and coordination #3882: participant control plane, episode steps, and coordination
• feat(checkpoint): make policy model servers checkpoint participants with generation cuts #3883: policy model servers as checkpoint participants, with generation cuts
• feat(token-capture): add worker staging helpers for generation cuts #3884: worker staging helpers for generation cuts
• feat(checkpoint): continue environment server episodes from their boundaries #3885: environment server episodes continue from their boundaries
• feat(checkpoint): add the resources server participant with asynchronous session hooks #3886: resources server participant
• feat(checkpoint): declare twenty training verifiers stateless with replayable verification #3887: training verifiers declared stateless with replayable verification
• feat(checkpoint): add the agent session participant and Simple Agent continuation #3888: agent session participant and Simple Agent continuation
• test(checkpoint): add a process-level end-to-end suite driven by coordination #3889: process-level end-to-end suite
• fix(checkpoint): reply idle to a prepare whose checkpoint ended while it waited #4660: a prepare whose checkpoint ended while it waited replies idle
• test(checkpoint): expect agent sessions through close_agent in the real-model end-to-end test #4694: the real-model end-to-end test expects agent sessions through close_agent

In review:

• feat(checkpoint): checkpoint evaluation runs from rollout collection #3893: checkpointing evaluation runs from rollout collection; targets main on its own.
• feat(checkpoint): spans and metrics for partial-rollout checkpoints #3903, then feat(checkpoint): resources, agent, and environment servers with several workers #3909, then test(checkpoint): check checkpoint invariants over random sequences of operations #4197, as one stack:
  • feat(checkpoint): spans and metrics for partial-rollout checkpoints #3903: spans and metrics;
  • feat(checkpoint): resources, agent, and environment servers with several workers #3909: resources, agent, and environment servers with several workers;
  • test(checkpoint): check checkpoint invariants over random sequences of operations #4197: checkpoint invariants over random sequences of operations.
• fix(server): route a session's requests to the worker that holds it #3905: session routing for servers with several workers; targets main on its own and lands before feat(checkpoint): resources, agent, and environment servers with several workers #3909.
• feat(checkpoint): name the episodes an agent's own /run drives in the agent commit reply #4690: the agent commit reply names the episodes an agent's own /run drives.

Resources server ports, each on main:

• feat(checkpoint): checkpoint Workplace Assistant sessions #3894: Workplace Assistant;
• fix(gymnasium): build Gymnasium servers on the shared resources server setup #3895: Gymnasium, then feat(checkpoint): export Blackjack game state for partial-rollout checkpoints #3896: Blackjack on top of it;
• feat(checkpoint): export indirect prompt injection session state #3897: indirect prompt injection (draft);
• feat(checkpoint): continue proof refinement agent turns from their boundaries #3898: proof refinement (draft).

The first prototype stack, #2939 to #2946, is closed in favor of this one.
The sub-issues filed for the original design are closed as superseded.
The original proposal follows for reference.

Use cases, pain points, and background

Long-running Gym rollouts can lose all progress when a worker, server process, node, or training job fails. Restarting every unfinished rollout is safe for short stateless tasks, but it wastes completed turns in multi-turn agents, stateful environments, and long tool-using tasks.

The recoverable state is distributed:

• The calling framework knows which rollout inputs are queued, running, completed, or already consumed.
• A Gym agent server owns the active /run loop, conversation history, usage, cookies, and control-flow position.
• A Gym resources server may own mutable per-session environment state.
• A Gym policy model server owns model-call identity and token-custody lineage.
• The configured generation and token-storage implementation owns token IDs, masks, log probabilities, and policy-version metadata.
• An opaque agent may additionally own a process tree, filesystem, sandbox, and requests that are still in flight.

Copying these states independently does not produce a valid checkpoint. An agent boundary can refer to a model call or resources mutation that is absent from another component's snapshot. File digests detect corruption but cannot prove that independently copied files describe the same execution boundary.

Gym therefore needs stable execution identity, explicit component checkpoint APIs, and a protocol that lets an external coordinator publish a checkpoint only after every required component has committed compatible state.

Description

Add partial rollout checkpointing as an additive Gym capability that any evaluation or training framework can adopt. Gym defines execution identity, component capabilities, admission control, serializable state, manifest publication, restore, and stale-attempt fencing. The integrating framework owns the global checkpoint transaction, rollout scheduling, result consumption, and final recovery decision.

The first milestone preserves completed training-ready data and safely restarts unfinished stateless rollouts. Later milestones resume supported whitebox agents at completed operation boundaries, restore stateful resources sessions, and eventually support opaque blackbox runtimes through provider snapshots.

The design must fail closed. If the integrating framework cannot prove that an agent boundary, resources revision, model-call lineage, and external token-store snapshot agree, the rollout must start a replacement attempt from its original input or the global checkpoint must be rejected. It must never resume from an incomplete combination of component states.

Proposed design

Framework-neutral ownership

• The external checkpoint coordinator owns the global transaction and publishes the final checkpoint.
• A framework-specific Gym integration owns rollout dispatch, a live-execution registry, result-handoff acknowledgement, and orchestration across Gym servers.
• Each Gym server owns admission and serialization for the state held in that service.
• A framework-provided token sink, source, and lineage store own durable training-token data. Gym checkpoint artifacts contain references and integrity metadata rather than token tensors.

flowchart LR
    C[Framework checkpoint coordinator] --> I[Framework Gym integration]
    I --> A[Gym agent server]
    I --> M[Gym policy model server]
    I --> R[Gym resources server]
    M --> G[Generation backend]
    G --> TS[Framework token store]

    A --> AD[Agent boundary manifest]
    M --> MD[Model lineage manifest]
    R --> RD[Resources state manifest]
    TS --> TD[Token-store snapshot]

    AD --> P[Global checkpoint publication]
    MD --> P
    RD --> P
    TD --> P
    I --> P

The Gym APIs do not depend on a particular controller, actor system, distributed runtime, or token store. For example, NeMo-RL can implement the external coordinator with SingleController, the Gym integration with its NemoGym actor, and token storage with TransferQueue. Another framework can provide different implementations while preserving the same Gym contracts.

Stable execution identity

The integrating framework assigns rollout_id before dispatch. The ID names one logical rollout across retries and restore. _ng_attempt_index names a physical execution of that rollout. Every agent, model, resources, result, checkpoint, and retry artifact carries both values independently of observability or token-capture settings.

Restore creates attempt N+1 and tombstones attempt N. Delayed requests from N receive a stale-attempt error and cannot modify restored state or publish a winning model-call root.

Framework-facing checkpoint lifecycle

sequenceDiagram
    participant C as Framework coordinator
    participant I as Gym integration
    participant A as Agent server
    participant M as Policy model server
    participant R as Resources server
    participant TS as Token store
    participant D as Coordinator-provided directory

    C->>I: prepare(checkpoint_id, deadline)
    I->>I: Freeze dispatch membership
    I->>A: Request parking at safe agent boundaries
    A-->>I: Parked and terminal execution report
    I->>M: Close policy-model admission and drain accepted calls
    I->>R: Close mutation admission and export locked sessions
    M-->>I: Paused with durable model custody
    R-->>I: Prepared resources snapshots

    C->>I: commit(checkpoint_id, D)
    I->>A: Commit agent boundaries under D
    I->>M: Commit model lineage under D
    I->>R: Commit resources state under D
    A-->>I: Manifest path and digest
    M-->>I: Manifest path and digest
    R-->>I: Manifest path and digest
    I->>D: Publish aggregate Gym manifest last

    C->>TS: Save a compatible token-store snapshot
    C->>D: Publish global checkpoint last
    C->>I: resume(checkpoint_id)

Prepare, commit, and restore leave Gym paused. Resume is explicit and occurs only after generation, token storage, Gym participants, and framework state are ready. The integration reopens dependencies before releasing parked agent loops and new dispatches.

Shared server contracts

Every checkpoint-capable Gym server exposes an authenticated capabilities route and fenced control operations carrying checkpoint_id plus an absolute deadline. Repeated operations for the same checkpoint are idempotent. Stale or concurrent checkpoint IDs cannot commit state or reopen traffic.

The capabilities response declares:

• component type and instance role;
• checkpoint mode and schema version;
• supported admission states;
• single-process, coordinated multi-process, or unmanaged operation;
• environment concurrency contract;
• active checkpoint ID and current phase.

Policy model server

The policy model participant must:

• close new policy-generation admission;
• aggregate admission and in-flight state across every FastAPI worker process;
• let already accepted calls finish durable token staging and Gym custody recording;
• commit and restore model lineage without embedding token tensors in public manifests;
• retain tombstones for abandoned attempts;
• remain paused after restore.

The external token implementation satisfies Gym's TokenSink, TokenSource, and LineageStore contracts. It may use a queue, object store, database, shared filesystem, or another durable backend. Its snapshot identity and integrity metadata must be bindable to the global checkpoint.

Whitebox agent server

The whitebox agent participant must:

• track one active /run execution per rollout attempt;
• request parking without cancelling an operation midway;
• accept a versioned AgentBoundaryRecord after a completed nonterminal turn;
• commit output items, usage, cookies, the last committed model-call ID, resources revisions, and agent-specific serializable state;
• restore that boundary as attempt N+1;
• omit terminal executions from the resumable set.

The in-memory task, coroutine, and HTTP connection are runtime handles and are never serialized as an agent boundary.

Stateful resources server

The resources participant must:

• key sessions by (rollout_id, attempt_index);
• serialize state mutations per session or declare another proven concurrency contract;
• return a monotonic state revision after successful mutations;
• close new mutation admission during prepare;
• wait for accepted mutations using the session lock, then export a versioned ResourceSnapshot;
• validate the complete replacement set before activating it under attempt N+1;
• retire successful terminal sessions while retaining state after failed verification.

Each environment supplies codecs for the state it owns. Environments that cannot export a safe state declare restart-only behavior instead of inheriting a false continuation guarantee.

Artifacts and publication invariants

A component writes payloads to temporary names, flushes and renames them, and writes its digest-bearing manifest last. The framework integration writes the aggregate Gym manifest only after every required component manifest validates. The external coordinator publishes the global checkpoint only after Gym state, framework state, and the referenced token-store snapshot all validate under the same checkpoint ID.

The initial prototype uses per-boundary, per-session, and per-rollout files. Production scale may require archive shards plus a member index to avoid tens of thousands of shared-filesystem metadata operations. That storage optimization must not change the participant lifecycle or restore semantics.

Recovery classification

For each execution, the integrating framework selects exactly one disposition:

• completed: the framework durably accepted the result before the checkpoint;
• continue: every required agent, resources, model-lineage, and token-store reference validates;
• fresh dispatch: partial state can be discarded safely and the original input is retried as attempt N+1;
• reject checkpoint: an unsafe external effect or required state cannot be reconciled.

A rollout cannot be both completed and resumable. A continuation cannot become runnable until model lineage and resources state have been restored and reconciled.

Delivery milestones

1. Stateless token-capture recovery: preserve completed token rows and groups; restart every unfinished rollout from input.
2. Whitebox agent continuation: resume supported Gym agents at completed operation boundaries with stateless resources.
3. Stateful resources continuation: restore versioned logical environment sessions before the agent resumes.
4. Blackbox runtime continuation: freeze and restore the harness process through a runtime provider and checkpoint-aware egress relay.

Each milestone provides a useful recovery guarantee without claiming support from environments or runtimes that cannot serialize their state safely.

Generation-safe cuts for active policy calls

The experimental policy-model participant cuts accepted generations at a durable token boundary instead of waiting for every HTTP request to return. Closing admission freezes the accepted-call membership. Each call then reaches exactly one checkpoint disposition: it never started generation, completed with durable custody, produced a durable active prefix, failed with a durable terminal record, or was explicitly excluded.

Each admitted request keeps two independent lifetimes. Response lifetime ends only after ASGI egress completes. Generation-cut safety can be reached earlier, when the model call has a durable completed record or the generation backend atomically swaps its active token buffer and makes the frozen pre-cut buffer durable. The checkpoint therefore does not wait only for the caller to consume a completed response.

A generation cut acknowledgement identifies the model call, rollout and attempt, immutable pre-cut buffer, token boundary, digest, policy metadata, and external token-store receipt. Every managed model-server worker must acknowledge the same checkpoint ID. A missing worker, partial acknowledgement, receipt mismatch, or expired deadline fails preparation.

Gym stores token-free lineage and cut coordinates. Prefix custody remains in the external token store. Delivery differs by inference backend: a worker with token-store access can receive a staged lineage chain, while a backend that needs token IDs receives an inline prefix materialized and validated by the framework integration. Inline and staged representations are explicit and mutually exclusive.

The Gym stack supplies these contracts and deterministic fake-backend tests. Real vLLM, Megatron inference, and TransferQueue buffer-swap implementations remain dependent work under #3224.

Sub-issues

The complete resources-server and agent-server audit is tracked through sub-issues rather than a flat checklist:

• #3218: generation-safe policy-model cuts
• #3219: stateful resources adapters (19 server children)
• #3220: whitebox agent continuation (13 agent children)
• #3221: opaque runtime providers (13 harness children)
• #3222: sandbox snapshot and reattachment (18 provider/server children)
• #3223: grouped, global, drain-only, and restart-only services (4 children)
• #3224: vLLM and Megatron inference integrations
• #3225: training-framework integrations, beginning with the NeMo-RL NemoGym actor

Stateless servers, servers already covered by the draft stack, and deliberately restart-only atomic agents are recorded in the engineering inventory without duplicate implementation issues.

First prototype (closed)

The first prototype stack is closed; the environment server design in the current status above replaces it:

• feat(checkpoint): propagate stable execution identity #2939: stable execution identity; replaced by the episode identity in feat(checkpoint): add the participant control plane, episode steps, and coordination #3882 and feat(checkpoint): continue environment server episodes from their boundaries #3885
• feat(checkpoint): add fenced control contracts #2940: fenced control and admission contracts; replaced by feat(checkpoint): add the participant control plane, episode steps, and coordination #3882
• feat(checkpoint): add policy model participant #2941: policy model participant and lineage checkpoint; replaced by feat(checkpoint): make policy model servers checkpoint participants with generation cuts #3883 and feat(token-capture): add worker staging helpers for generation cuts #3884
• feat(checkpoint): add whitebox agent participant #2942: whitebox agent participant; replaced by feat(checkpoint): add the agent session participant and Simple Agent continuation #3888
• feat(checkpoint): resume simple agent turns #2943: SimpleAgent continuation; replaced by feat(checkpoint): add the agent session participant and Simple Agent continuation #3888
• feat(checkpoint): resume Gymnasium agent turns #2944: Gymnasium agent continuation; replaced by fix(gymnasium): build Gymnasium servers on the shared resources server setup #3895 and feat(checkpoint): export Blackjack game state for partial-rollout checkpoints #3896
• feat(checkpoint): add resources state participant #2945: resources state participant; replaced by feat(checkpoint): add the resources server participant with asynchronous session hooks #3886
• feat(checkpoint): adapt initial stateful environments #2946: initial stateful environment adapters; replaced by feat(checkpoint): declare twenty training verifiers stateless with replayable verification #3887 and the resources server ports feat(checkpoint): checkpoint Workplace Assistant sessions #3894 to feat(checkpoint): continue proof refinement agent turns from their boundaries #3898

Out of scope

• Owning the global training-framework checkpoint transaction inside Gym.
• Defining how a framework schedules, batches, groups, or consumes completed rollouts.
• Requiring a particular actor system, controller implementation, or external token store.
• Transparent recovery for an environment that cannot serialize or reconstruct its state.
• Claiming deterministic replay for external tools, judges, simulators, clocks, random state, or sandboxes without an explicit idempotency contract.
• Restoring arbitrary live sockets, threads, GPU KV caches, or process memory through the application-level whitebox APIs.
• Treating a filesystem-only sandbox snapshot as a complete opaque-agent runtime snapshot.
• Requiring all environments to support the same recovery level.
• Moving token arrays into Gym's public rollout results or checkpoint manifests.

Acceptance criteria

• rollout_id and _ng_attempt_index propagate through agent, model, and resources calls independently of observability and capture settings.
• Gym servers reject conflicting identity and stale source attempts.
• Every checkpoint-capable server reports authenticated capabilities and supports fenced, idempotent control operations.
• Policy-model admission closes across every configured worker process, reports missing workers as an error, and uses bounded waits.
• Active policy calls reach a durable completed, prefix-cut, failure, or excluded disposition without waiting only for HTTP response delivery.
• Generation-cut acknowledgements validate immutable buffer identity and external token-store receipts across every managed worker.
• Model checkpoint commit and restore preserve lineage references that reconcile with the selected external token-store snapshot.
• Supported whitebox agents emit versioned boundaries and resume without repeating completed model or tool operations.
• Supported stateful resources servers export and atomically restore versioned session state under a replacement attempt.
• Terminal sessions, failed verification, cancellation, timeout, lost acknowledgements, corrupt files, and stale requests have fault-injection coverage.
• A documented framework-integration contract covers the live-execution registry, result-handoff acknowledgement, recovery classification, and ordered prepare, commit, restore, and resume orchestration.
• At least one framework adapter validates the complete protocol end to end without adding framework-specific types to Gym's public checkpoint APIs.
• The global checkpoint is published only after every required participant manifest and external token-store snapshot validate under the same checkpoint ID.
• Unsupported environments and runtimes advertise restart-only behavior or fail checkpoint preparation clearly.
• Production-scale storage avoids unbounded per-rollout files and unbounded in-memory execution payload retention.