Repository navigation
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
ananthsub
changed the title
This was referenced Sep 1, 2026
ananthsub
force-pushed
the
ananthsub/checkpoint-agent-base
branch
from
September 3, 2026 09:12
1f9f857 to
399842a
Compare
ananthsub
force-pushed
the
ananthsub/checkpoint-agent-base
branch
from
September 3, 2026 11:12
399842a to
2dec5fc
Compare
ananthsub
force-pushed
the
ananthsub/checkpoint-agent-base
branch
from
September 10, 2026 20:48
2dec5fc to
6604c8e
Compare
ananthsub
force-pushed
the
ananthsub/checkpoint-agent-base
branch
from
September 22, 2026 17:28
6604c8e to
15a158e
Compare
Track one active run per rollout attempt, park it at typed turn boundaries, and commit or restore continuation records only through checkpoint control operations. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Block publication on boundary-less parks or unacknowledged results, separate retirement from acknowledgement, and cancel retired execution tasks. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
ananthsub
force-pushed
the
ananthsub/checkpoint-agent-base
branch
from
September 22, 2026 17:37
15a158e to
2a345d5
Compare
2 tasks done
macandro96
reviewed
Sep 29, 2026
| app = FastAPI() | ||
|
|
||
| self.setup_session_middleware(app) | ||
| app.add_middleware(RolloutContextMiddleware) |
Contributor
There was a problem hiding this comment.
Adding this middleware breaks the existing rollout-prefixed agent routes. The middleware removes /ng-rollout/ from scope["path"] before FastAPI routes the request. As a result, training-capture calls become /training-token-capture/v1/responses, which is not registered and returns 404. Normal prefixed calls reach /v1/responses without request.path_params["rollout_id"], so agents that read that value lose rollout correlation. Could we register the stripped training route, make handlers fall back to current_rollout_id(), and add tests for both prefixed paths? This is partially addressed by commit eab140c, which registers the stripped training route and updates the base helper. Direct request.path_params readers still need the context fallback.
| allowed_phases=frozenset( | ||
| { | ||
| CheckpointPhase.IDLE, | ||
| CheckpointPhase.PREPARING, |
Contributor
There was a problem hiding this comment.
resume can retire the current checkpoint while prepare is running. We may have to serialize the operations. Had implemented implemented in d7b4963.
| if execution.parked_task is not None: | ||
| execution.parked_task.cancel() | ||
| else: | ||
| execution.resume_event.set() |
Contributor
There was a problem hiding this comment.
This wakes the parked task but leaves its state as PARKED until the task runs. A new prepare can count it as safely parked, after which the old wakeup changes it to RUNNING. Could resume update the state immediately and use a checkpoint generation so an old wakeup cannot cross a newer checkpoint? This is fixed later in 1518b4e.
| require_control_auth(authorization, auth_token) | ||
|
|
||
| async def run() -> dict[str, Any]: | ||
| return await asyncio.to_thread(restore_agent_state, participant, Path(body.checkpoint_dir)) |
Contributor
There was a problem hiding this comment.
Restore can run while old executions still exist. install_restored() tombstones their identities but does not remove the live execution objects, so those old tasks can still commit or finish. It also modifies participant state from the worker thread. Could we reject restore while executions are live, load files in the worker, and install the result on the event loop? This is fixed later in 9a1c3b4.
| return response | ||
| if execution is None or self._checkpoint_participant is None: | ||
| return response | ||
| await self._checkpoint_participant.park(execution) |
Contributor
There was a problem hiding this comment.
This can park after the checkpoint has already resumed. Reading the 409 response body is asynchronous, so resume may finish during that read. This call then marks the execution PARKED, clears the wake-up event, and waits forever. It can also park before agent prepare was attempted, so partial-prepare rollback may never resume it. Could we only park when an agent prepare is currently active, and otherwise retry with bounded backoff?
|
|
||
| def install_restored(self, records: list[AgentBoundaryRecord]) -> None: | ||
| for record in records: | ||
| self._tombstones.add((record.rollout_id, record.attempt_index)) |
Contributor
There was a problem hiding this comment.
Do we need to reject requests with attempt_id < record.attempt_index too? can probably create a watermark and raise if attempt_id < latest_attempt_idx
| raise AgentCheckpointError("cannot commit while agent executions still block durable publication") | ||
| return [ | ||
| execution.boundary | ||
| for execution in self._executions.values() |
Contributor
There was a problem hiding this comment.
A restored continuation remains in _restored until its replacement /run begins, but this method commits only live _executions. A new checkpoint taken before dispatch therefore drops valid recovery state. Could we include pending _restored records in the committed snapshot and add a restore → resume → checkpoint-before-run → restore test?
This was referenced Oct 1, 2026
dims
pushed a commit
to dims/Nemo-Gym
that referenced
this pull request
Oct 14, 2026
…ndaries (NVIDIA-NeMo#3885) ## What changed and why Every environment server takes part in checkpoints once checkpointing is on. The environment server owns the episode protocol, so it records where each episode can continue. - An episode records a boundary after each completed step and runs each step in `wait` or `replay` mode. Prepare closes episode admission and is ready once every live episode is parked at a boundary or inside a replay step. - A restored boundary is handed to the replacement attempt's `/run`, which continues from it. The task must match the one recorded. - `single_agent_turn` runs seed, invoke agent, close agent, verify, and return as explicit stages with typed session handles. Seeding and the agent invocation replay; closing the agent waits; verification uses the mode the resources server reported on its seed reply, kept in the session handles so a restored episode still knows it. Once the agent's activation returns, the episode records a `close_agent` boundary carrying the agent's response, so a checkpoint that closes then parks the episode there instead of waiting on the close step, and a restore at that boundary closes the agent and verifies without invoking the agent again. - **Every stage starts with its boundary, including `close_agent` and `return`.** After a restore and resume, a later checkpoint can prepare before the replacement `/run` arrives, and the replacement is then admitted while admission is closed. A replacement restored at `close_agent` used to start closing the agent before it recorded a boundary, which made a prepared participant not ready; one restored at `return` ran final cleanup and called `end()`, so an episode could finish while Gym was prepared. Both boundaries now sit at the start of their stage, as the others do, so a replacement at any stage parks before doing any work and a commit exports the boundary it was restored from. For a first attempt each boundary is recorded at the same point, with the same state, as before. - **A failed agent activation closes the agent session at once** to keep its observations, as `main` does since NVIDIA-NeMo#4562. The close runs after the failed replay step, which parks with no boundary while a checkpoint is open, so it waits for resume and never runs while Gym is prepared. - A boundary's state can be given as a function, which is called only when checkpointing is on, so `single_agent_turn` builds no boundary payloads otherwise. - The `single_agent_turn_legacy` adapter checkpoints too, since its rows run the same protocol. - An unclaimed restored episode, one whose replacement has not started yet, is exported again by the next checkpoint while the controller's commit scope still continues it, so a second crash still continues it. Once the scope leaves it out, it is not exported, and the restored record is retired. - A replacement `/run` is refused with `admission_closed` until the controller resumes after the restore, and the restored record stays until a replacement actually starts. A retire of a replacement that had already started would wait for its final cleanup, whose calls wait for that same resume. - A retired attempt's late `/run`, such as a client retry, is refused with `stale_attempt` until the controller forgets the rollout. An episode with no boundary yet is exported with a null boundary and starts over from its input after a crash. - A `/run` of an attempt that is already running is refused, and the refused request owns nothing: it neither ends nor cleans up the live episode. An episode is untracked even if its final cleanup is cancelled. - A replay step that fails while a checkpoint is open clears the episode's boundary instead of blocking the checkpoint, so the checkpoint commits and the episode starts over from its input after a crash. The episode then parks at that cleared boundary until resume before the failure propagates. It stays out of the blockers and is exported with no boundary, but its final cleanup, its `end()`, and its `/run` reply wait for resume; otherwise it could finish and drop its record before the commit, and a commit that names it would export nothing for it. A retire still stops a parked episode. With no checkpoint open, the episode goes back to running as before. - Final cleanup, which closes the episode's agent and resources sessions, runs before `/run` replies on every path, and a retire never interrupts it. A retire cancels an episode only before cleanup starts, then waits until the episode has ended, so its reply means the sessions are closed. This matters because anyio's shielded scope does not stop a direct `task.cancel()`. A retire cut short by its deadline keeps the episode and its task digest until the episode has ended, so the next prepare reports it as a blocker instead of exporting it. - **Restarts.** An environment server that does not record boundaries (`checkpoint_boundaries = False`, the default) cannot continue its episodes, so each one is a restart: it never blocks prepare, is never exported, and is reported in prepare's `restarts`. `single_agent_turn` sets `checkpoint_boundaries = True`, and marks an episode as a restart when its resources or agent seed reply carries `x-ng-checkpoint-restart`, because a server it depends on cannot capture its part. - A commit exports only the live episodes inside its scope, as it already did for restored episodes. The controller leaves every participant's restarts out of the scope, so a live episode that another server reports as a restart must not be exported here either. - The participant answers the readiness wait's `ready()` from its running blocker count, so a prepare does not list every episode after each change. - Checkpointing refuses `num_workers > 1` on environment servers, since episodes live in one process. ## How it works ### Where this PR sits in the overall flow The highlighted part is what this PR adds. ```mermaid flowchart LR C["Controller<br/>NeMo RL, or rollout collection"] CO["Coordination<br/>prepare, commit, restore, resume,<br/>retire, forget"] K["Control plane on every server<br/>phases, lease, storage,<br/>retire: refuse, stop, free;<br/>forget"] subgraph G["One participant per Gym server"] E["Environment server<br/>episode steps"] M["Policy model<br/>held responses, generation cuts"] A["Agent<br/>sessions parked at boundaries"] R["Resources server<br/>session state"] end W["Inference worker<br/>stages cut prefixes"] D[("Checkpoint directory<br/>records, then manifest")] L[("Capture ledger<br/>retire and delete from NVIDIA-NeMo#3938, NVIDIA-NeMo#3939")] C --> CO --> K K --> E & M & A & R M --> W M --> L G --> D classDef this fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937 class E this ``` One checkpoint, a crash, and the restore, end to end: ```mermaid sequenceDiagram participant C as Controller participant G as Gym participants participant D as Checkpoint directory C->>G: prepare, in order environment, model, agent, resources Note over G: admission closes, in-flight work parks at a boundary,<br/>undelivered model responses are held G-->>C: prepared, or blockers at the deadline C->>G: commit with the episodes the controller continues G->>D: each participant writes its records, then its manifest Note over G: restored state the commit's scope leaves out is deleted C->>C: publish the checkpoint with the controller's own state C->>G: resume, in order resources, agent, model, environment Note over C,G: crash - every Gym process dies C->>G: restore the checkpoint in fresh processes, all or nothing D-->>G: records installed under attempt N+1, attempt N's capture ledger retired C->>G: resume C->>G: /run as attempt N+1 continues each episode from its boundary Note over C,G: retiring an episode, only while no checkpoint is open C->>G: retire - environment, then agent, then model and resources Note over G: each server marks the attempt, stops its work, waits, frees its state, then replies C->>G: forget, once the rollout is done for good - same order Note over G: the mark refusing the retired attempt goes away ``` ### This PR `single_agent_turn` records a boundary at the start of every stage. A checkpoint exports the latest boundary, and a restored attempt starts at its `next` step, parking there first if a checkpoint is open. ```mermaid flowchart LR B1(["next: seed"]) --> S["seed the resources session"] S --> B2(["next: invoke_agent"]) --> I["invoke the agent<br/>which parks its own turn loop"] I --> BC(["next: close_agent"]) --> CA["close the agent session"] CA --> B3(["next: verify"]) --> V["verify<br/>wait or replay, per the resources server"] V --> B4(["next: return"]) --> R["reply to /run"] ``` Final cleanup closes the episode's sessions before `/run` replies. A retire never interrupts it: it cancels an episode only before cleanup starts, and then waits for the episode to end. ```mermaid flowchart TD X{"retire arrives"} -->|"episode still in its protocol"| C["cancel the episode"] --> F X -->|"final cleanup already started"| F["wait"] F --> CL["final cleanup closes the agent and resources sessions"] --> D["retire replies"] ``` ## Where this sits in the stack This is one PR in a stack of draft PRs that re-cut partial-rollout checkpointing onto environment servers. Each PR is one commit, and its base is the branch of the PR before it, so each diff shows only that PR's change. The stack builds on the token-capture cleanup PRs NVIDIA-NeMo#3938 (capture ledger `retire` and `delete`) and NVIDIA-NeMo#3939 (complete-record `retire` and `delete`), which have merged. NVIDIA-NeMo#3882 and NVIDIA-NeMo#3883 have merged too, so NVIDIA-NeMo#3884's base is `main`. The checkpoint stack uses those operations to free the ledgers of retired attempts and to clear a dead execution's capture files before a restore. The capture-ledger cache bound (NVIDIA-NeMo#4268) and the striped lock files (NVIDIA-NeMo#3937) are independent of the stack. 1. NVIDIA-NeMo#3882 (`ananthsub/partial-ckpt-core`): feat(checkpoint): add the participant control plane, episode steps, and coordination (merged) 2. NVIDIA-NeMo#3883 (`ananthsub/partial-ckpt-policy-model`): feat(checkpoint): make policy model servers checkpoint participants with generation cuts (merged) 3. NVIDIA-NeMo#3884 (`ananthsub/partial-ckpt-model-worker-cuts`): feat(token-capture): add worker staging helpers for generation cuts 4. NVIDIA-NeMo#3885 (`ananthsub/partial-ckpt-environment`): feat(checkpoint): continue environment server episodes from their boundaries **(this PR)** 5. NVIDIA-NeMo#3886 (`ananthsub/partial-ckpt-resources`): feat(checkpoint): add the resources server participant with asynchronous session hooks 6. NVIDIA-NeMo#3887 (`ananthsub/partial-ckpt-verifier-declarations`): feat(checkpoint): declare five training verifiers stateless with replayable verification 7. NVIDIA-NeMo#3888 (`ananthsub/partial-ckpt-agent`): feat(checkpoint): add the agent session participant and Simple Agent continuation 8. NVIDIA-NeMo#3889 (`ananthsub/partial-ckpt-e2e`): test(checkpoint): add a process-level end-to-end suite driven by coordination 9. NVIDIA-NeMo#3893 (`ananthsub/partial-ckpt-rollout-collection`): feat(checkpoint): checkpoint evaluation runs from rollout collection 10. NVIDIA-NeMo#3903 (`ananthsub/partial-ckpt-telemetry`): feat(checkpoint): spans and metrics for partial-rollout checkpoints 11. NVIDIA-NeMo#3909 (`ananthsub/partial-ckpt-multi-worker`): feat(checkpoint): resources, agent, and environment servers with several workers 12. NVIDIA-NeMo#4197 (`ananthsub/partial-ckpt-fault-tests`): test(checkpoint): check checkpoint invariants over random sequences of operations Partial-rollout checkpointing lets a training controller, such as NeMo RL, checkpoint Gym while rollouts are in flight and, after a crash, continue those rollouts from their last safe point instead of starting them over. Checkpointing is off by default; with the `checkpoint:` block unset, no server installs any checkpoint routes or behavior. ## Relationship to the old stack (NVIDIA-NeMo#2939 to NVIDIA-NeMo#2946) New compared with the old stack, which had no environment servers. It takes over the protocol side of NVIDIA-NeMo#2942 (where an episode can continue), which the old stack reconstructed from agent and resources state after the fact. ## Issue No tracking issue exists for this re-cut. The design and the mapping from the old stack are described in this PR series, and the old stack's PRs (NVIDIA-NeMo#2939 to NVIDIA-NeMo#2946) carry the original discussion. ## Validation Run on this branch, on the current `main`: - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_checkpoint_*.py`: passed, including, each failing without its change: - a retire during an episode's final cleanup does not interrupt it; - a retire replies only after the episode and its cleanup have ended; - a commit whose scope leaves out a restored episode retires it; - a refused duplicate `/run` leaves the live episode tracked, and an episode whose cleanup is cancelled is still untracked; - a replay step that fails during a checkpoint lets the checkpoint commit with the episode exported with no boundary, the episode has not finished or replied before resume even with its cleanup running at once, and it starts over after a crash; - a retire stops an episode parked after its replay step raised; - a commit does not export live episodes outside its scope; - an agent that returns during a checkpoint parks at `close_agent` and stays ready, and a restore at `close_agent` closes and verifies without invoking the agent; - boundary payloads are not built with checkpointing off; - a replacement restored at `close_agent` or at `return` and admitted during a second checkpoint parks before its stage: readiness stays true, no close request and no `end()` happen before resume, and the commit exports the same boundary (without the fix, the agent close or the final resources close started and readiness was false); - a retire stops an episode parked after its replay step raised, waiting on the step and on the `[None]` boundary rather than on the task ending (the previous version passed without reaching the retire, and the new one fails if the failed step does not park or the retire skips the episode); - an agent failure during a checkpoint closes the agent session only after resume, and the failure carries the agent's observations; It also checks that an episode whose retire was cut short blocks the next prepare, is not committed, and is freed by a retried retire after resume, and that readiness is answered from the blocker count. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_environment_server.py`: passed. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider environment_servers/single_agent_turn/tests`: passed. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider environment_servers/single_agent_turn_legacy/tests`: passed. - `responses_api_agents/codex_agent/tests`, `responses_api_agents/openclaw_agent/tests`, `responses_api_agents/opencode_agent/tests`, and `responses_api_agents/pi_agent/tests`, each run in its own server environment, created from that server's `requirements.txt` as `gym env test` does: passed, including the integration tests that collect rows through an environment server's `/run`. Their fake responses gain an empty `headers` mapping, because the environment server now reads the seed and agent response headers. One Pi outcome test, `test_outcome_extension_uses_pinned_pi_context_classifier`, needs a newer Node.js than the development machine's Node.js 20, and fails the same way on `main`. - The full core unit suite (`tests/unit_tests`, eight processes, without `test_opensandbox_compose.py`, which needs the optional `opensandbox` package): passed except a sandbox retry test, which fails the same way on `main` in this development environment. - `pre-commit run --files <files changed by this PR>`: all hooks passed, and no hook modified a file. - Every commit carries a `Signed-off-by` line. ## Rollout evidence - Process-level evidence is in the top PR (an episode on the typed `single_agent_turn` request path continues from its last boundary; second crash before the replacement starts). - **Real model:** see the rollout evidence in NVIDIA-NeMo#3889, whose opt-in tests drove this PR's code with `Qwen/Qwen3-4B-Instruct-2507` on vLLM: saved rollouts continued after a kill and restore, restarts started over, and every rollout finished with reward 1.0. ## Compatibility - Nothing changes unless the global `checkpoint:` block sets `enabled: true`. - With checkpointing on, an environment server with `num_workers > 1` fails at startup. - A commit exports only the live episodes inside its `episode_ids` scope. A commit request without `episode_ids` still exports every live episode. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
OliviaViessmann
pushed a commit
to OliviaViessmann/Gym
that referenced
this pull request
Oct 16, 2026
…continuation (NVIDIA-NeMo#3888) ## What changed and why This PR builds on the agent session lifecycle that NVIDIA-NeMo#3784 added to `main`. The base server's seed and close also open and close the session in the checkpoint participant, and Simple Agent exports, restores and retires the sessions the base server holds. Every agent takes part in checkpoints once checkpointing is on. - Agents that set `checkpoint_sessions_supported = True` and implement the session hooks get an `AgentSessionParticipant`. A session is idle, running, or at a boundary. A running activation parks at its next boundary when a checkpoint closes admission, and one waiting on a model call counts as parked at its last boundary. If that call's reply arrives while the checkpoint is open, the session stays parked at the boundary before the call and holds the reply until resume, so its readiness never regresses between prepare and commit. The model may already have recorded the call, because it delivered the reply before its own prepare. After a crash the restored agent calls again from that boundary, and the token-capture builder treats the two rows like a client retry: it keeps the call the agent continued from. Boundaries are lazy snapshots, so a long turn loop does not serialize its history at every step. Restored sessions are installed under the replacement attempt. - The hooks are coroutines, and export takes every session at once: `export_agent_sessions(session_keys)`, `restore_agent_sessions(sessions)`, and `retire_agent_session(session_key)`. An agent whose state lives outside the process, such as a sandbox it runs tools or a harness in, can checkpoint all of it concurrently within the commit. - A legacy `/run` is an episode of its own: seed and the turn loop replay, and verification uses the mode the resources server reported on its seed reply. - Agents without the hooks get a restart-only participant. The server framework installs it on whatever app `setup_webserver` returns, so agents that build their own app, such as `harbor_agent` and the `mini_swe_agent` agents, take part too; it is installed once per app. Their rollouts are restarts: they never block prepare, keep running through the checkpoint, and are reported in prepare's `restarts`, so one opaque harness no longer stalls checkpointing for every environment. The seed reply carries `x-ng-checkpoint-restart`, and only a new seed waits while a checkpoint is open. After a crash, the controller retires the rollout and starts it over from its input. - With checkpointing on, agent work without a rollout ID is refused with `rollout_id_required`, and every agent refuses `num_workers > 1`. - Simple Agent implements the hooks for its environment server sessions and for its legacy `/run` loop. A session's export carries how many activations it has started, because each activation reports its observations under its own invocation ID (`root`, then `activation-2`, and so on). A restored session's next activation therefore gets the next ID, and an activation continued from its restored boundary keeps the ID it already had. - Simple Agent's session export carries its resources cookie jar as it is. Since NVIDIA-NeMo#4449 on `main`, a session without a direct tool grant holds no resources cookie jar (`resources_cookies` is `None`), so the environment server keeps the session cookie from its own seed for `/verify` and `/close_session`. Export and restore keep that `None` instead of turning it into an empty jar, so a restored session's close still reports no jar and the environment server keeps the cookie from its resources seed. - A commit exports only the live sessions whose episode is inside its scope, as it already did for restored sessions. The controller leaves every participant's restarts out of the scope, so a session another server reports as a restart is not exported here either. - An unclaimed restored session, one whose replacement has not started yet, keeps its restored boundary in the next checkpoint while the controller's commit scope still continues it, so a second crash still continues it. Once the scope leaves it out, it is not exported, and it is retired. - **Activations are idempotent.** The environment server sends each activation's index, its 1-based count within the session, in the `x-ng-agent-activation` header. With checkpointing on, the agent exports the index and the reply of its last completed activation with the session. An invocation of an activation that already completed gets the stored reply back without running the loop: no model or tool call, and no change to observations or the activation count. This keeps the environment server and the agent consistent when the environment server's checkpoint still has the invocation ahead of it while the agent's checkpoint has the activation done. An index out of order is refused with 409 `activation_out_of_order`. Without the header, every invocation is a new activation, as before. - A request for an activation that is still running waits for it to end and then gets its reply, instead of being refused. - With checkpointing on, every activation of an agent session must carry the `/ng-rollout/<id>` prefix. One without it is refused with 400 `rollout_id_required` before any replay, claim or activation, since a restore rewrites the session's episode to the replacement attempt and the session can no longer stand in for a missing path. This holds for restored sessions and for sessions seeded in the running process; every caller that activates an agent session is a Gym environment server, and each sends the prefix. - Before it replays a reply or starts an activation, the agent checks the attempt the request's `/ng-rollout/<id>` path names against the session's, and refuses a mismatch with `stale_attempt`, so a crashed attempt's harness cannot use a restored session. - A replayed reply claims a restored session, like a new activation does, so a later commit that leaves the episode out does not free a session the live episode still uses. - The reply is kept only when the request carried an activation index. - **A legacy `/run` sends activation index 1 on its self-dispatch.** Its turn loop records its reply like an environment server session's activation, and the reply is exported with the legacy session. A checkpoint taken after the loop replied but before `/run` recorded its `verify` boundary therefore restores the reply instead of running the loop again; the reply leaves when the `/run` ends or a retire frees it. - **Direct `/v1/responses` calls are restarts.** With checkpointing on, a direct call with a rollout path is reported as a restart on every agent, and one without a rollout path is refused with `rollout_id_required`, as restart-only agents already did. - **Restart-only seeds respect retires.** A restart-only seed checks the retired mark after it waits for resume, a retire waits for seeds in progress, and the restart-only retire frees seeded sessions through `retire_agent_session`. The base `retire_agent_session` now releases a session through the agent's close hook, so a restart-only agent stops a sandbox it started. - **Restarts from earlier stages do not hold up agent prepare.** The agent participant leaves the restarts carried from earlier prepare stages (NVIDIA-NeMo#3886) out of `readiness`, `ready` and its export. An activation of such an episode still parks at its next boundary until resume, but it no longer holds up prepare, so an episode `single_agent_turn` marked as a restart, from a restart-only resources seed or an MCP grant, does not stop the checkpoint while it is inside a slow tool call. - **Retire cost is per rollout.** Sessions, seeds in progress and direct restart calls are indexed by rollout, so a retire touches only its own rollout's, and each index entry leaves with the last item of its rollout. - **Simple Agent boundaries leave out each turn's `question`**, which repeats the request's input and the loop's earlier outputs. They store `question_outputs`, how many outputs each question held, and a restore rebuilds each question from the request's input and the restored outputs. The boundary snapshot is built directly, without a validate-and-dump round trip. - **A completed activation's reply is kept once.** A session-scoped reply to a request that carries an activation index leaves out `_ng_trajectory`, in the original and in the copy kept to replay it, since `single_agent_turn` takes the session's evidence from the `AgentObservationBundle` at close; callers that send no index, such as the UserSim environment server, still get it. A legacy `/run` drops its turn loop's kept reply in the same step that records its `verify` or `return` boundary, which holds the reply. Kept replies stay unserialized until a replay or an export needs them. - **An episode that grants MCP access is a restart.** `single_agent_turn` marks an episode as a restart when it grants an `MCPToolAccess`, so no checkpoint continues an agent that holds an MCP token past its seed; with NVIDIA-NeMo#3886's token fence, a replaced attempt's token stops working after a restore. - **Direct HTTP tool calls name their attempt.** With checkpointing on, `single_agent_turn`'s direct HTTP grant carries the seeding attempt's `/ng-rollout/<capture key>` prefix in its base URL, since NVIDIA-NeMo#3886 refuses a call on a restored session that names no attempt. Simple Agent sends each direct HTTP tool call under the rollout prefix of the activation it is serving, in place of the prefix in the grant, so after a restore the continued session's calls name the replacement attempt; its legacy `/run` tool calls carry the same prefix. No other agent with checkpoint support uses direct HTTP grants, and restart-only agents seed again after a crash and get a fresh grant. - **A legacy `/run` closes its resources session after a replayable verification.** It sends `x-ng-checkpoint-close-follows: 1` on `/verify`, so NVIDIA-NeMo#3886 keeps the session tracked and exported until the close, and then closes the session, also when a replacement is restored at its `return` step. A failed close is logged and keeps the run's verdict. - **A `skip_verification` legacy `/run` records its `return` boundary.** It records the same boundary as a verified run, with the configured reward as its result, so a checkpoint that finishes while the turn loop's reply is on its way parks the run at `return` until resume instead of letting it finish uncaptured. A run restored at `return` returns its recorded result. - **Closing a session stops its activation before cleanup.** A close first stops an activation still running for it, so an activation the environment server gave up on stops calling the model and tools, and only then runs the agent's close hook. A close cancelled while it waits for the activation has not run the hook, so a retried close runs it exactly once. The session stays tracked, and holds up prepare, until the activation has stopped. A close cancels only an activation nothing has cancelled yet, so it never interrupts a retire's wait for the activation's cleanup. - **A close waits for resume while agent admission is closed**, as a seed does, and checks admission before it marks the session as closing. So a commit that has checked readiness always exports an open session, never a closing one. This cannot stall prepare through the environment server's final cleanup: coordination prepares environment servers before agents, and an episode in final cleanup keeps its environment server from being ready, so its agent close finishes before agent admission closes. A close that a retire took over is refused with `stale_attempt`, and checks after every await that its session record is still current, so it never writes a receipt for a session a retire freed. - A retire cancels the attempt's activations, or a restart-only agent's `/run` calls, and waits until they have stopped before the retire hook runs. The hook is where a sandboxed agent stops its sandbox. A retired session stays tracked until its activation has stopped and the legacy episode's retire and the retire hook have finished, so a retire cut short by its deadline at any of those points leaves it for a retry, which waits again without cancelling a second time. Meanwhile the session blocks the next checkpoint and is not exported, and a legacy `/run` blocks whether or not its session is still tracked. - A retired attempt's late activation is refused with `stale_attempt` until the controller forgets the rollout. - A session seed that arrives while a checkpoint is open waits for resume, and one already in progress holds up prepare and is exported. The environment server seeds inside a replay step that a checkpoint does not wait for, so the agent's session must match whichever boundary the environment server records. Refusing the seed failed the episode, and admitting it but leaving it out of the checkpoint let a replacement invoke a session nothing restored; a random checkpoint sequence against a real `single_agent_turn` deployment found both. A new session that is not a seed, such as a legacy `/run`, is still refused. Seeds are counted per session key, so a retried seed overlapping the first is tracked until both finish. A retire also waits for seeds of its attempt that are in progress, and a seed refused because its attempt was retired frees the state it built before the retire replies. A restore is refused while a seed is in progress. - `SimpleResponsesAPIAgent._free_session` frees a retired session that has no close receipt; it was `_discard_session`. - The participant answers the readiness wait's `ready()` without building its blocker list, and notifies only while a checkpoint is open. ## How it works ### Where this PR sits in the overall flow The highlighted part is what this PR adds. ```mermaid flowchart LR C["Controller<br/>NeMo RL, or rollout collection"] CO["Coordination<br/>prepare, commit, restore, resume,<br/>retire, forget"] K["Control plane on every server<br/>phases, lease, storage,<br/>retire: refuse, stop, free;<br/>forget"] subgraph G["One participant per Gym server"] E["Environment server<br/>episode steps"] M["Policy model<br/>held responses, generation cuts"] A["Agent<br/>sessions parked at boundaries"] R["Resources server<br/>session state"] end W["Inference worker<br/>stages cut prefixes"] D[("Checkpoint directory<br/>records, then manifest")] L[("Capture ledger<br/>retire and delete from NVIDIA-NeMo#3938, NVIDIA-NeMo#3939")] C --> CO --> K K --> E & M & A & R M --> W M --> L G --> D classDef this fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937 class A this ``` One checkpoint, a crash, and the restore, end to end: ```mermaid sequenceDiagram participant C as Controller participant G as Gym participants participant D as Checkpoint directory C->>G: prepare, in order environment, model, agent, resources Note over G: admission closes, in-flight work parks at a boundary,<br/>undelivered model responses are held G-->>C: prepared, or blockers at the deadline C->>G: commit with the episodes the controller continues G->>D: each participant writes its records, then its manifest Note over G: restored state the commit's scope leaves out is deleted C->>C: publish the checkpoint with the controller's own state C->>G: resume, in order resources, agent, model, environment Note over C,G: crash - every Gym process dies C->>G: restore the checkpoint in fresh processes, all or nothing D-->>G: records installed under attempt N+1, attempt N's capture ledger retired C->>G: resume C->>G: /run as attempt N+1 continues each episode from its boundary Note over C,G: retiring an episode, only while no checkpoint is open C->>G: retire - environment, then agent, then model and resources Note over G: each server marks the attempt, stops its work, waits, frees its state, then replies C->>G: forget, once the rollout is done for good - same order Note over G: the mark refusing the retired attempt goes away ``` ### This PR An agent session's states. Prepare is ready once no activation is between boundaries. ```mermaid flowchart LR I(["idle"]) -->|an activation starts| R(["running"]) R -->|before a model call, or after a tool round| B(["at boundary"]) B -->|next step, when no checkpoint is open| R B -->|parked while a checkpoint is open| B R -->|the activation ends| I ``` Which participant an agent gets. ```mermaid flowchart TD Q{"checkpoint_sessions_supported"} -->|yes| P["AgentSessionParticipant<br/>export_agent_sessions, restore_agent_sessions,<br/>retire_agent_session"] Q -->|no| O["Restart-only: reported in prepare, never blocks,<br/>retired and started over after a crash"] ``` Closing or retiring a session stops its activation first: close cancels it, and retire also waits for it before the retire hook, which is where a sandboxed agent stops its sandbox. ```mermaid flowchart LR E["close from the environment server,<br/>or retire"] --> X["cancel the running activation"] --> W["retire waits for it"] --> H["retire_agent_session hook<br/>stops a sandbox the agent owns"] ``` ## Where this sits in the stack This is one PR in a stack of draft PRs that re-cut partial-rollout checkpointing onto environment servers. Each PR is one commit, and its base is the branch of the PR before it, so each diff shows only that PR's change. The stack builds on the token-capture cleanup PRs NVIDIA-NeMo#3938 (capture ledger `retire` and `delete`) and NVIDIA-NeMo#3939 (complete-record `retire` and `delete`), which have merged. NVIDIA-NeMo#3882 to NVIDIA-NeMo#3887 have merged too, so NVIDIA-NeMo#3888's base is `main`. The checkpoint stack uses those operations to free the ledgers of retired attempts and to clear a dead execution's capture files before a restore. The capture-ledger cache bound (NVIDIA-NeMo#4268) and the striped lock files (NVIDIA-NeMo#3937) are independent of the stack. 1. NVIDIA-NeMo#3882 (`ananthsub/partial-ckpt-core`): feat(checkpoint): add the participant control plane, episode steps, and coordination (merged) 2. NVIDIA-NeMo#3883 (`ananthsub/partial-ckpt-policy-model`): feat(checkpoint): make policy model servers checkpoint participants with generation cuts (merged) 3. NVIDIA-NeMo#3884 (`ananthsub/partial-ckpt-model-worker-cuts`): feat(token-capture): add worker staging helpers for generation cuts (merged) 4. NVIDIA-NeMo#3885 (`ananthsub/partial-ckpt-environment`): feat(checkpoint): continue environment server episodes from their boundaries (merged) 5. NVIDIA-NeMo#3886 (`ananthsub/partial-ckpt-resources`): feat(checkpoint): add the resources server participant with asynchronous session hooks (merged) 6. NVIDIA-NeMo#3887 (`ananthsub/partial-ckpt-verifier-declarations`): feat(checkpoint): declare twenty training verifiers stateless with replayable verification (merged) 7. NVIDIA-NeMo#3888 (`ananthsub/partial-ckpt-agent`): feat(checkpoint): add the agent session participant and Simple Agent continuation **(this PR)** 8. NVIDIA-NeMo#3889 (`ananthsub/partial-ckpt-e2e`): test(checkpoint): add a process-level end-to-end suite driven by coordination 9. NVIDIA-NeMo#3893 (`ananthsub/partial-ckpt-rollout-collection`): feat(checkpoint): checkpoint evaluation runs from rollout collection 10. NVIDIA-NeMo#3903 (`ananthsub/partial-ckpt-telemetry`): feat(checkpoint): spans and metrics for partial-rollout checkpoints 11. NVIDIA-NeMo#3909 (`ananthsub/partial-ckpt-multi-worker`): feat(checkpoint): resources, agent, and environment servers with several workers 12. NVIDIA-NeMo#4197 (`ananthsub/partial-ckpt-fault-tests`): test(checkpoint): check checkpoint invariants over random sequences of operations Partial-rollout checkpointing lets a training controller, such as NeMo RL, checkpoint Gym while rollouts are in flight and, after a crash, continue those rollouts from their last safe point instead of starting them over. Checkpointing is off by default; with the `checkpoint:` block unset, no server installs any checkpoint routes or behavior. ## Relationship to the old stack (NVIDIA-NeMo#2939 to NVIDIA-NeMo#2946) - Supersedes NVIDIA-NeMo#2942 (agent participant) and NVIDIA-NeMo#2943 (Simple Agent turns). Completed-result acknowledgements, prepare polling, and the refusal retry loop are gone. - NVIDIA-NeMo#2944 (Gymnasium agent turns) is not carried; Gymnasium continuation moves to a separate PR, and until then Gymnasium agents are restart-only. ## Issue No tracking issue exists for this re-cut. The design and the mapping from the old stack are described in this PR series, and the old stack's PRs (NVIDIA-NeMo#2939 to NVIDIA-NeMo#2946) carry the original discussion. ## Validation Run on this branch, on the current `main`: - Each of these behaviors has a test that fails without its change: - an agent session activation without the rollout prefix is refused, for a restored session and for one seeded in the running process, with no model call, replay or claim, while checkpointing off still accepts it; - a legacy `/run` that skips verification parks at its return boundary until resume; - an episode the environment server reports as a restart, from a restart-only resources seed or an MCP grant, holds up neither the agent stage nor the resources stage of prepare, checked through coordination with a `single_agent_turn` server and real environment, agent and resources participants, and the agent neither waits for nor exports such an episode's activation until resume; - an indexed session reply and its replay carry no trajectory, nor does the export, and a kept reply stays unserialized until a replay or an export needs it; - a checkpoint of a legacy `/run` parked at verify holds its reply once, and a restore still replies; - a restore after the legacy turn loop replied returns that reply without running the loop again; - an episode that grants MCP access is a restart; - a close that arrives after prepare waits for resume and the commit exports its session, a cleanup close of a parked episode after the agent stage waits without holding up prepare, and a retire that frees a session while its close waits leaves no receipt behind; - a close stops the activation before its cleanup starts, and a close cancelled while it stops the activation runs the cleanup once when retried; - a request for an activation still running waits for it and gets its reply; - a request on the replaced attempt's path is refused for the restored session; - a restart-only seed that arrives after its rollout was retired is refused, and a restart-only retire frees its rollout's sessions and waits for a seed in progress; - a direct `/v1/responses` call is a restart, and one without a rollout path is refused; - a boundary leaves out the turn questions, and a resumed loop rebuilds them; - an activation without an index keeps no reply; - a restored session answered with its stored reply survives a commit that leaves its episode out; - with checkpointing on, a direct HTTP grant names the seeding attempt, and a restored session's direct HTTP tool call reaches the restored resources session under the replacement's prefix, through Simple Agent and a real resources server; - a legacy `/run` closes its resources session after a verification a checkpoint does not wait for. - A separate ordering check confirms that an environment episode's cleanup close finishes before the agent stage of prepare; it passes with or without the change. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_checkpoint_*.py`: passed, including: - closing a session stops its running activation; - a retire stops the activation before it replies, and a straggler is retired after resume; - a restart-only agent's `/run` has stopped when its retire replies; - a commit whose scope leaves out a restored session retires it; - a session whose retire hook was cut short stays blocking, is not exported, and is freed by a retry; - a seed arriving during an open checkpoint waits for resume, and one in progress when admission closes holds up prepare and is exported; - overlapping seeds of one session each hold up prepare, a retire waits for a seed in progress and frees what it built, and a restore is refused while a seed is in progress; - a closed session holds up prepare until its activation has stopped, and neither a close nor a retire cancels an activation twice; - a restored legacy `/run` claimed at verify is no longer pending and survives a commit that leaves it out; - an agent that builds its own app still reports its runs as restarts, and the participant is installed once per app; - `ready()` does not build the readiness report, and nothing notifies outside a checkpoint; - a commit does not export live sessions outside its scope. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider responses_api_agents/simple_agent/tests`: passed, including that a restored session's next activation does not reuse an invocation ID and that an activation continued from its restored boundary keeps its own; that a restored activation keeps its boundary through a second checkpoint; that a model reply arriving after prepare is held at the boundary before the call; that an activation already completed is answered with its stored reply without running again; that a legacy `/run` restored at verify or return still reports its trajectory; and that a session without a resources cookie jar restores without one. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider environment_servers/single_agent_turn/tests`: passed, including the timeline where the agent's reply to an invocation is lost to a restore: the environment server invokes again and the restored agent answers with the stored reply while the model sees no new call. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_terminal_attribution.py`: passed, including that a call delivered before a checkpoint and issued again after a restore keeps only the re-issued call. - The full core unit suite (`tests/unit_tests`, eight processes, without `test_opensandbox_compose.py`, which needs the optional `opensandbox` package): passed except a sandbox retry test, which fails the same way on `main` in this development environment. - `pre-commit run --files <files changed by this PR>`: all hooks passed, and no hook modified a file. - Every commit carries a `Signed-off-by` line. ## Rollout evidence - Process-level evidence is in the top PR (episodes on the typed `single_agent_turn` request path and legacy `/run` episodes continue from their last boundary). - **Real model:** see the rollout evidence in NVIDIA-NeMo#3889, whose opt-in tests drove this PR's code with `Qwen/Qwen3-4B-Instruct-2507` on vLLM: saved rollouts continued after a kill and restore, restarts started over, and every rollout finished with reward 1.0. Every rollout ran Simple Agent. ## Compatibility - New agent hooks and class attribute: `checkpoint_sessions_supported`, `export_agent_sessions`, `restore_agent_sessions`, `retire_agent_session`. Existing agents need no change and are restart-only under checkpointing. - With checkpointing on, an agent session activation without the `/ng-rollout/<id>` prefix gets 400 `rollout_id_required`. A caller outside Gym that activates Simple Agent sessions with only the session cookie must add the prefix; no caller in the repository is affected. - With checkpointing on, agent calls without a rollout ID get 400 `rollout_id_required`, and agents with `num_workers > 1` fail at startup. - With checkpointing on, a direct HTTP grant's base URL ends in the seeding attempt's rollout prefix. An agent that continues a session after a restore must send its tool calls under its own activation's prefix; Simple Agent does. - With `skip_verification` and checkpointing on, a legacy `/run` caught by a checkpoint waits at its `return` boundary until resume. Nothing changes with checkpointing off. - New request header `x-ng-agent-activation` on agent session invocations, and a 409 `activation_out_of_order` refusal. An environment server that can invoke an activation again after a restore must send the header. - `AgentSessionParticipant.close_session` is now a coroutine, and the agent implements `SimpleServer.setup_checkpoint_participant(app)`, the hook NVIDIA-NeMo#3886 adds that the framework calls after `setup_webserver`. - A commit exports only the live sessions inside its `episode_ids` scope. A commit request without `episode_ids` still exports every live session. - Nothing changes unless the global `checkpoint:` block sets `enabled: true`. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Contributor
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note
This checkpoint implementation is experimental. Its Python modules live under
nemo_gym._checkpoint; they are not a supported public import surface yet.Summary
/runexecutions and committed turn boundaries/runtasks and active inner handlers when an attempt is retiredWhat an agent boundary means
A whitebox agent can resume only at a point it explicitly declares safe.
AgentBoundaryRecordcaptures the completed conversation output, usage, last committed model-call ID, resources revisions, and agent-specific serializable state after an operation finishes. The in-memory Python task is never serialized.Agent execution state
stateDiagram-v2 [*] --> Running: begin rollout attempt Running --> BoundaryCommitted: agent completes a nonterminal turn BoundaryCommitted --> Running: no checkpoint requested Running --> ParkRequested: prepare arrives between boundaries BoundaryCommitted --> ParkedWithBoundary: prepare observes a safe boundary ParkRequested --> ParkedWithBoundary: next safe boundary is committed Running --> ParkedWithoutBoundary: downstream checkpoint refusal before a boundary Running --> CompletedUnacknowledged: /run result awaits framework durability ParkedWithBoundary --> Committed: write boundary and manifest Committed --> RestoredPaused: restore as attempt N+1 RestoredPaused --> Running: global resume Running --> Retired: controller abandons attempt ParkedWithBoundary --> Retired: controller abandons attemptPrepare closes new
/runadmission and reports every condition that prevents publication. Only executions parked with committed boundaries are checkpointable. A parked execution without a boundary must make progress or be retired. A completed result remains cached and blocks prepare until a later framework-owned acknowledgement layer confirms that the result is durable.Retirement is not result acknowledgement. It tombstones an abandoned active attempt and cancels both its outer request task and any inner self-dispatched handler. This prevents an old handler from mutating state after a replacement attempt starts.
Commit writes immutable boundary payloads before its instance-namespaced manifest. Restore registers each continuation under a replacement attempt and tombstones the source attempt. A repeated commit validates and returns the existing manifest.
Stack
Test plan
Generation-safe checkpoint update
Layer validation: 24 tests passed.