feat(checkpoint): add durable turn-level rollout recovery by macandro96 · Pull Request #3349 · NVIDIA-NeMo/Gym

macandro96 · 2026-09-13T20:11:11Z

What does this PR do?

Summary

Adds durable turn-level rollout recovery to NeMo Gym on top of #2946.

This change allows an accepted multi-turn rollout to stop at a committed agent/environment boundary, persist coordinated Gym state, and continue under a replacement rollout attempt after process or job restart.

What this adds above #2946

#2946 introduces checkpoint adapters for the initial stateful environments.

This PR integrates those adapters into the complete Gym rollout lifecycle:

durably acknowledge completed agent executions
preserve agent continuations and model-call lineage across attempts
restore conversation cookies and resource-state revisions
make resource mutations idempotent across recovery
coordinate model, agent, and resources admission during checkpointing
classify environments as export/restore, restart-only, or stateless
restart only executions that depend on non-restorable resources
use indexed checkpoint artifacts instead of directory scans
archive model lineage and agent state in bounded, checksummed shards

Recovery flow

The checkpoint coordinator closes model admission.
Accepted model calls finish and their capture lineage becomes durable.
The agent parks at its latest committed turn boundary.
Resource servers export the corresponding environment state and revision.
Model, agent, and resource manifests are committed under one checkpoint.
After restart, source attempt N remains fenced.
The checkpoint is installed for replacement attempt N+1 while admission
remains paused.
The participants resume together and the rollout continues from the saved
turn without replaying completed resource mutations.

sequenceDiagram
    participant C as Checkpoint coordinator
    participant M as Model server
    participant A as Agent server
    participant R as Resource servers
    participant S as Checkpoint store

    Note over C,S: Checkpoint

    C->>M: Close admission and prepare
    Note over M: Accepted calls finish<br/>No mid-generation recovery
    M-->>C: Model-lineage manifest

    C->>A: Prepare checkpoint
    A-->>C: Parked turn and continuation manifest

    C->>R: Export environment state
    R-->>C: State and revision<br/>or restart-only classification

    C->>S: Commit coordinated checkpoint
    S-->>C: Checkpoint committed

    Note over C,S: Restore after restart

    C->>S: Load committed checkpoint
    S-->>C: Return participant manifests

    C->>M: Restore lineage and fence attempt N
    C->>A: Install continuation as attempt N+1
    C->>R: Restore state or restart dependent execution

    C->>M: Reopen admission
    C->>A: Resume execution
    C->>R: Resume resources

Correctness guarantees

completed /run results remain available until acknowledged
source attempts cannot mutate restored successor state
resource mutations are replay-safe and revision checked
the first restored model call links to the saved parent exactly once
failures during prepare or restore leave participants paused or abortable
environments without export/restore support restart only their dependent unfinished executions

Scalability

Checkpoint artifacts use:

explicit continuation and storage-reference indexes
scoped artifact retrieval
bounded model-lineage tar shards
bounded agent-state tar shards
per-file size and SHA-256 validation

This avoids scanning or materializing one loose checkpoint file per active
rollout during restore.

Testing

unit coverage for model, agent and resources checkpoint participants
recovery across replacement attempts
completed-result acknowledgement and replay
exactly-once resource mutation behavior
restart-only resource fallback
SimpleAgent turn recovery
deterministic Workplace Assistant state recovery
archive integrity, corruption and unsafe-member validation

Scope

Included:

turn-boundary recovery
agent conversation/continuation recovery
supported environment-state recovery
restart fallback for unsupported stateful environments

Not included:

token-prefix recovery inside an unfinished model call
vLLM KV-cache persistence
TransferQueue checkpoint orchestration
sandbox process or in-memory snapshotting

Stack

Depends on feat(checkpoint): adapt initial stateful environments #2946
NeMo-RL controller integration: feat(checkpoint): Support turn-level checkpoint RL#4117

Checklist

I have read the contributing guidelines.
The change is focused; unrelated "drive-by" edits are tracked as separate issues/PRs.
Tests added or updated and pass locally, or N/A for docs-only / non-code changes (so CI unit/server checks pass when applicable).
Pre-commit checks pass locally (pre-commit run --all-files) (so CI lint/format/copyright pass).
All commits have DCO sign-off (git commit -s) (so the DCO check passes).

copy-pr-bot · 2026-09-13T20:11:14Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

…playable verification (NVIDIA-NeMo#3887) ## What changed and why A resources server that declares nothing is restart-only, so the rollouts that use it start over from their input after a crash. These twenty training verifiers keep no session state and their verification can run again after a crash, so they declare `checkpoint_mode = "stateless"` and `checkpoint_verify = "replay"`: - `code_gen`, `competitive_coding_challenges`, `equivalence_llm_judge`, and `math_with_judge` verify each request on its own. - `abstention`, `calendar`, `equivalence_rule`, `ether0`, `format_verification`, `instruction_following`, `inverse_if`, `jailbreak_detection`, `mcqa`, `multichallenge`, `nvarc`, `reasoning_gym`, `string_match`, `structured_outputs`, and `terminus_judge` also verify each request on its own, from the request body, with a judge model call, or with a subprocess that ends with the call (from NVIDIA-NeMo#4540, by macandro96). A replayed judge call can return a different verdict; the declaration says that repeating it is safe, not that the result is identical. - `genrm_compare` keys its cohorts by prompt, not by session. Its verification must replay: a member waits for its siblings, so waiting on it would deadlock a checkpoint while siblings are still generating. After a crash, members that had not recorded their reward re-verify and rebuild the cohort. A crash between a cohort's result and every member recording it leaves the remaining members waiting; they fail after `cohort_collection_timeout_s` and are retried from input. The declarations are class attributes, not config, because they describe what each server's code can do. The test reads them from source, because each server's dependencies live in its own environment. ## How it works ### Where this PR sits in the overall flow The highlighted part is what this PR adds. ```mermaid flowchart LR C["Controller NeMo RL, or rollout collection"] CO["Coordination prepare, commit, restore, resume, retire, forget"] K["Control plane on every server phases, lease, storage, retire: refuse, stop, free; forget"] subgraph G["One participant per Gym server"] E["Environment server episode steps"] M["Policy model held responses, generation cuts"] A["Agent sessions parked at boundaries"] R["Resources server session state"] end W["Inference worker stages cut prefixes"] D[("Checkpoint directory records, then manifest")] L[("Capture ledger retire and delete from NVIDIA-NeMo#3938, NVIDIA-NeMo#3939")] C --> CO --> K K --> E & M & A & R M --> W M --> L G --> D V["Five training verifiers stateless, replayable verify"] R --- V classDef this fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937 class V 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, 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 These five servers declare that verification can run again after a crash, so a checkpoint never waits for a slow judge. ```mermaid flowchart LR E["Environment server"] -->|"/seed_session"| R["code_gen, competitive_coding_challenges, equivalence_llm_judge, math_with_judge, genrm_compare"] R -->|"x-ng-checkpoint-verify: replay"| E E -->|"verify in flight at a checkpoint"| N["The checkpoint does not wait. Verify runs again after a crash."] ``` ## 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#3885 have merged too, so NVIDIA-NeMo#3886'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 6. NVIDIA-NeMo#3887 (`ananthsub/partial-ckpt-verifier-declarations`): feat(checkpoint): declare twenty training verifiers stateless with replayable verification **(this PR)** 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) None of NVIDIA-NeMo#2939 to NVIDIA-NeMo#2946 declared these servers. The declarations come from NVIDIA-NeMo#3349 (durable turn-level recovery), where they were config on the agent; here they are class attributes on the resources servers. ## 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`: - `tests/unit_tests/test_checkpoint_server_declarations.py`: each of the twenty servers declares `stateless` and `replay`, read from its source. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_checkpoint_*.py`: passed. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider resources_servers/code_gen/tests`: passed. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider resources_servers/competitive_coding_challenges/tests`: passed. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider resources_servers/equivalence_llm_judge/tests`: passed. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider resources_servers/genrm_compare/tests`: passed. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider resources_servers/math_with_judge/tests`: not run in this environment. It needs the server's `math-verify==0.8.0` requirement, which this environment does not install. It passed on an earlier version of this branch with that requirement installed. - 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 - These are declarations only: with checkpointing off, nothing changes for these servers. With it on, they no longer force their rollouts to be retired at each checkpoint. - A checkpointed training run on these environments with a real model: **pending**. ## Compatibility - Nothing changes unless the global `checkpoint:` block sets `enabled: true`. With it on, these five servers no longer block checkpoints, and their `/verify` may run again after a crash. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com> Co-authored-by: Anish Mahishi <amahishi@nvidia.com>

…o#3894) - **Branch:** `ananthsub/partial-ckpt-workplace-assistant` - **Base branch:** `main` - **Suggested labels:** `feature`, `area:env-infra` --- ## What changed and why Workplace Assistant used the default `restart_only` mode, so every rollout with a live Workplace session was a restart and started over from scratch after a crash. This PR lets those sessions be checkpointed and continued. - The server declares `checkpoint_mode = "exported"` and implements the three session hooks over its per-session state. That state is the tool environment created at `/seed_session`: a set of in-memory pandas tables for email, calendar, analytics, project management, the customer relationship manager, and the company directory. - `export_session_states` encodes every table column by column. It keeps column dtypes, such as the boolean `user_engaged` column, and index labels, which matter because the analytics tool appends rows by label. It also keeps the difference between empty CSV cells, which load as NaN, and values a tool stored as `None`. JSON writes both as `null`, so NaN positions are listed separately. Encoding runs in a worker thread. Admission is closed during a commit, so no tool call changes the tables while they are encoded. - `restore_session_states` rebuilds every session's tool environment first and installs them only after all of them are valid. A state with missing or extra containers or tables, or with mismatched column lengths, raises `ValueError` and installs nothing. - Export leaves out a session the server no longer holds. `/verify` discards the session's tool environment in a `finally` block. If verify raises, the session is gone but the participant still tracks it, so the commit must not fail on it. - `retire_session_state` drops the session's tool environment. - `checkpoint_verify` stays at the default `wait`. Verify computes the reward only from the response and the ground truth, but it deletes the session's state, so it is not replay-safe. - The toolkit list is now one module constant, `_TOOLKITS`, instead of being repeated in `seed_session` and `mcp_tools`. ## How it works ### Where this PR sits in the overall flow The highlighted part is what this PR adds. ```mermaid flowchart LR C["Controller NeMo RL, or rollout collection"] CO["Coordination prepare, commit, restore, resume, retire, forget"] K["Control plane on every server phases, lease, storage, retire: refuse, stop, free; forget"] subgraph G["One participant per Gym server"] E["Environment server episode steps"] M["Policy model held responses, generation cuts"] A["Agent sessions parked at boundaries"] R["Resources server session state"] end W["Inference worker stages cut prefixes"] D[("Checkpoint directory records, then manifest")] L[("Capture ledger 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 R 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, 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 Each session's tool environment is a set of tables. A checkpoint exports them whole, and a restore rebuilds them before any of them is installed. ```mermaid sequenceDiagram participant A as Agent participant W as Workplace Assistant participant D as Checkpoint directory A->>W: /seed_session creates the session's tables A->>W: tool calls read and change the tables Note over W: commit - export_session_states writes each table column by column, with column types, index labels, and where cells are empty W->>D: one record per session Note over A,D: crash, then restore into a fresh server D-->>W: restore_session_states rebuilds every session, then installs them all A->>W: tool calls continue on the same cookie session A->>W: /verify scores and drops the session ``` ## Where this sits in the stack This PR uses the resources server checkpoint participant from NVIDIA-NeMo#3886, which has merged, and needs nothing else from the open part of the partial-rollout checkpointing stack. It is one of five PRs that port the environments the old stack checkpointed onto the new hooks. Workplace Assistant and the Gymnasium fix target `main`, and Blackjack builds on the Gymnasium fix. Indirect prompt injection and proof refinement are still based on the end-to-end suite (NVIDIA-NeMo#3889): - NVIDIA-NeMo#3894 (`ananthsub/partial-ckpt-workplace-assistant`): feat(checkpoint): checkpoint Workplace Assistant sessions **(this PR)** - NVIDIA-NeMo#3895 (`ananthsub/partial-ckpt-gymnasium`): fix(gymnasium): build Gymnasium servers on the shared resources server setup - NVIDIA-NeMo#3896 (`ananthsub/partial-ckpt-blackjack`): feat(checkpoint): export Blackjack game state for partial-rollout checkpoints (on the Gymnasium PR) - NVIDIA-NeMo#3897 (`ananthsub/partial-ckpt-indirect-prompt-injection`): feat(checkpoint): export indirect prompt injection session state - NVIDIA-NeMo#3898 (`ananthsub/partial-ckpt-proof-refinement`): feat(checkpoint): continue proof refinement agent turns from their boundaries ## Relationship to the old stack - Ported from NVIDIA-NeMo#2946, which carried the Workplace Assistant adapter on the old hooks. It covers gap 7 of the re-cut, where this adapter was not carried. - Dropped from NVIDIA-NeMo#2946: - the `execution_to_session` map and the identity lookups in the tool route and verify. The v2 hooks are keyed by the cookie session ID, so neither is needed. - route classification of tool calls as mutations. In v2, admission refuses every data route while a checkpoint is open. - synthetic `checkpoint:<rollout>:a<attempt>` session IDs on restore. The restored session keeps its original cookie session ID. - The old encoding used pandas `to_json(orient="split")` and repaired types after reading the JSON back. That turned NaN into `None` and lost the difference between them. The new encoding keeps it. - Not ported from NVIDIA-NeMo#3349: the audit-trail test fixture (`checkpoint_test_app.py`) and its test. Both depend on old-stack rollout identities and on the test agent from that PR. ## Issue No tracking issue exists for this re-cut. The re-cut series and the old stack's PRs carry the discussion. ## Validation Run on this branch, on the current `main`: - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q resources_servers/workplace_assistant/tests`: passed, including the new `tests/test_checkpoint.py`. Its tests drive the real control routes over ASGI: - A session is seeded, mutated by three tool calls, prepared, and committed. It is restored into a fresh server instance with the same cookies and continues with four more tool calls and `/verify`. Tool outputs, every table, and the reward all match an uncheckpointed run of the same steps. - When one of two sessions has an invalid state, restore raises and installs neither session. - A retired session's tool environment is gone and the next commit exports nothing. A retire stops and frees the session instead of fencing it, so a later tool call is refused and does not recreate the environment. - A session whose `/verify` raised is left out of the export. The commit succeeds with only the live session, which then continues after restore. - A mutation check confirmed the equivalence test is not vacuous. Without restoring NaN cells, the test fails. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q tests/unit_tests/test_checkpoint_*.py`: passed. - `ruff check` and `ruff format --check` on the changed files: clean. - `pre-commit run --files <changed files>`: all hooks passed. - The commit carries a `Signed-off-by` line. ## Rollout evidence ### Setup The model was `Qwen/Qwen3-4B-Instruct-2507`, served by vLLM 0.24.0 with the `hermes` tool-call parser on one NVIDIA RTX 6000 Ada Generation (48 GB). Gym's `vllm_model` server ran with `uses_reasoning_parser=false`, because the served model is an instruct model without a reasoning parser. Every rollout sampled at temperature 1.0, used the five tasks in `resources_servers/workplace_assistant/data/example.jsonl`, and ran Simple Agent with `max_steps: 6` behind the flat-row wrapper. Each tree ran in its own worktree, with server venvs created by `gym env start`. ### Uninterrupted baseline This change on `main` ran with checkpointing off, eight rollouts per task: ```bash gym env start --config resources_servers/workplace_assistant/configs/workplace_assistant.yaml --model-type vllm_model \ --model Qwen/Qwen3-4B-Instruct-2507 --model-url http://127.0.0.1:8100/v1 --model-api-key dummy \ ++policy_model.responses_api_models.vllm_model.uses_reasoning_parser=false gym eval run --no-serve --config resources_servers/workplace_assistant/configs/workplace_assistant.yaml --model-type vllm_model \ --model Qwen/Qwen3-4B-Instruct-2507 --model-url http://127.0.0.1:8100/v1 --model-api-key dummy \ ++policy_model.responses_api_models.vllm_model.uses_reasoning_parser=false \ --agent workplace_assistant_simple_agent --input resources_servers/workplace_assistant/data/example.jsonl \ --output rollouts.jsonl --num-repeats 8 --concurrency 16 --temperature 1.0 ``` The 40 rollouts had a mean reward of 0.25. By task, the email reply succeeded 7 of 8 times, the CRM status move 3 of 8 times, and the calendar rename and the two other CRM tasks never succeeded. The same command on the checkpoint stack with this change, still with checkpointing off, also gave 0.25 over 40 rollouts. Neither run had failed or masked rollouts or tracebacks in the server logs. ### Checkpoint, crash and restore The tree was the checkpoint stack (`ananthsub/partial-ckpt-e2e`) with this change cherry-picked on top; it applied without conflicts. A driver built on the process-level checkpoint suite (`tests/e2e/checkpoint/checkpoint_deployment.py`) started the policy model, the Workplace Assistant resources server, Simple Agent and the flat-row wrapper, each as its own process from its own venv. It sent the rollouts to the wrapper's `/run` and drove the checkpoint through `nemo_gym._checkpoint.coordination`, as a training controller would: 1. Prepare and commit a checkpoint while rollouts are mid-turn. 2. Kill every Gym server process group with `SIGKILL`, start fresh processes from the same configs, restore the committed episodes and resume. 3. Take a second checkpoint of the restored state before the replacements start, then resume again. 4. Run every unfinished rollout again as attempt 1, which continues from its checkpointed boundary. Four runs cut the checkpoint at different points, so that later cuts catch episodes after their state-changing tool calls: | Run | Rollouts | Finished before the checkpoint | Continued | Continued with changed Workplace state | Mean reward of the run | |---|---|---|---|---|---| | 1 | 40 | 2 | 38 | 0 | 0.275 | | 2 | 40 | 13 | 27 | 0 | 0.300 | | 3 | 60 | 28 | 32 | 7 | 0.317 | | 4 | 80 | 46 | 34 | 10 | 0.288 | Every prepare returned in 8 ms or less and reported no restarts. Simple Agent reported every in-flight episode at a boundary, and the Workplace server exported one session per continued episode. The 17 episodes with changed state had already made one or two successful tool calls that changed their session's tables. Twelve had sent the email reply and five had moved one or two CRM customers to "Lost" before the checkpoint. Each export differs from a freshly seeded session exactly where those calls changed it. An email episode's export has one added row, the reply "RE: Task Update on Develop prototype for report generation", and every other email row equal to the seeded data. A CRM episode that had updated customers `00000037` and `00000195` exports both with status "Lost" where the seeded data has "Proposal", and no other cell differs. After the restore, all 131 continued episodes resumed rather than restarted: - All 131 replacements returned HTTP 200 with a reward. - Each final trajectory begins with the tool calls made before the crash, with the same call IDs and arguments, and the model then continued from there. - The second checkpoint, taken in the fresh processes before any replacement ran, exported the same tables as the first for all 131 sessions. - The five CRM episodes kept updating customers on their restored sessions: three or four more updates each, every one answered "Customer updated successfully." - None of the 276 tool results returned after the restore reported a missing session or a stale attempt. At the end, the Workplace server and Simple Agent held no sessions. The rewards are comparable to the uninterrupted run. Over all 220 rollouts the mean was 0.295 (65 of 220), against 0.25 (10 of 40) uninterrupted. By task, the email reply succeeded 44 of 44 times and the CRM status move 20 of 44 times, against 7 of 8 and 3 of 8. The calendar rename succeeded once in 44 and the other two CRM tasks never, as in the baseline. Of the 17 episodes caught after changing their state, 16 earned a reward of 1. ### What is not covered - `/verify` scores a trajectory by replaying its tool calls on a fresh environment, so the reward does not depend on the restored session. The restored state is shown instead by the second export and by the tool calls that ran on the restored sessions. - Only the flat-row wrapper ran. The typed `single_agent_turn` request path did not run with Workplace Assistant. - Only the five example tasks ran; the train and validation splits are downloaded from Hugging Face and were not used. The model rarely solves three of the five tasks. - MCP tool exposure, token capture and generation cuts stayed off. ## Compatibility - Without the global `checkpoint:` block set to `enabled: true`, behavior is unchanged. The hooks are only called by the checkpoint control routes. - With checkpointing on, Workplace Assistant sessions no longer block a checkpoint. They are exported, at about 350 KB of JSON per session. - Tool behavior, verify, and rewards are unchanged. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

yaoyu-33 added area:core Shared APIs, servers, telemetry, health, and registries feature New capabilities, enhancements, or enablement work needs-review PR is ready for code review and waiting on a reviewer labels Sep 13, 2026

macandro96 force-pushed the amahishi/gym-turn-level-recovery branch from 6473f49 to 6bba7b8 Compare September 14, 2026 01:40

macandro96 marked this pull request as draft September 14, 2026 03:47

macandro96 changed the title ~~feat(checkpoint): Gym turn level recovery~~ feat(checkpoint): add durable turn-level rollout recovery Sep 14, 2026

macandro96 force-pushed the amahishi/gym-turn-level-recovery branch from dbb8b76 to 80bd2a9 Compare September 15, 2026 16:33

zyzhou5 mentioned this pull request Sep 19, 2026

feat(checkpoint): add indirect prompt injection state recovery #3548

Open

macandro96 added 21 commits September 22, 2026 17:13

test(checkpoint): exercise SimpleAgent turn recovery

184a38c

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): isolate agent snapshot state

0d43301

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

feat(checkpoint): add artifact reference indexes

40778e0

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

refactor(checkpoint): require artifact indexes

8dde86d

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

refactor(checkpoint): require scoped artifact requests

d1e0db0

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

perf(checkpoint): avoid scanning inactive lineage

506939e

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): align restore with capture contracts

7c036a6

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

test(checkpoint): use current capture ledger records

d6e2ad1

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): support nested token-capture calls

eab140c

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): preserve receipt lineage across attempts

b9838b5

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): integrate generation-safe recovery

b47749a

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

test(checkpoint): add workplace turn recovery fixture

e9191ff

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

perf(checkpoint): add exact completion receipt lookup

f836bea

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

feat(checkpoint): classify restart-only resources

1302fda

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): scope restart-only recovery per execution

a60c048

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

perf(checkpoint): archive rollout lineage in shards

d00fd50

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

perf(checkpoint): archive agent state in shards

525494d

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

feat(checkpoint): freeze replayable verification waits

b45d211

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): harden turn recovery coordination

b1faf58

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): expose test agent to uvicorn workers

e1cc264

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): preserve model lineage across recovery

e113528

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

macandro96 added 10 commits September 22, 2026 17:46

fix(checkpoint): validate shared model lineage union

1fb63b9

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): harden model capture outcomes

88fd612

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): preserve restored execution boundary

d29add9

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

feat(checkpoint): resume stateless verifier rollouts

59abc18

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

test(checkpoint): cover restored model transition

6a0bd63

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): harden external wait races

1518b4e

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): serialize control mutations

d7b4963

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): order worker state broadcasts

5c57bae

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): install restored agents on event loop

9a1c3b4

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

fix(checkpoint): reconcile turn recovery with adapter stack

f4fcf8c

macandro96 force-pushed the amahishi/gym-turn-level-recovery branch from 5b0b444 to f4fcf8c Compare September 23, 2026 00:36

macandro96 changed the base branch from main to ananthsub/checkpoint-environment-adapters September 23, 2026 00:39

terrykong mentioned this pull request Sep 26, 2026

feat(checkpoint): add Gym participant coordination NVIDIA-NeMo/RL#4264

Open

2 of 3 tasks

macandro96 added 2 commits September 28, 2026 16:20

fix(checkpoint): preserve restored parent across parked retries

366590f

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

feat(checkpoint): return receipts with agent runs

d1abebd

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

macandro96 marked this pull request as ready for review September 29, 2026 21:52

macandro96 mentioned this pull request Sep 29, 2026

feat(checkpoint): prefix token level recovery #3412

Open

5 tasks

yaoyu-33 added the complexity:high Cross-component or shared-contract change with a large review surface label Sep 29, 2026

fix(checkpoint): harden turn recovery capabilities

023ae00

Signed-off-by: Anish Mahishi <amahishi@nvidia.com>

macandro96 force-pushed the amahishi/gym-turn-level-recovery branch from f47cdcf to 023ae00 Compare September 30, 2026 15:07

yuhezhang-ai mentioned this pull request Sep 30, 2026

feat(eval): persist failures and resume unfinished evaluations #3302

Merged

5 tasks done

github-actions Bot added the sla:triage-overdue Review assignment is over the one-business-day SLA label Sep 30, 2026

github-actions Bot mentioned this pull request Sep 30, 2026

[Tracker] Pull request handoff SLA #3402

Open

ananthsub deleted the branch ananthsub/checkpoint-environment-adapters October 19, 2026 20:49

ananthsub closed this Oct 19, 2026

Navigation Menu

Repository navigation

feat(checkpoint): add durable turn-level rollout recovery - #3349

feat(checkpoint): add durable turn-level rollout recovery#3349
macandro96 wants to merge 39 commits into
ananthsub/checkpoint-environment-adaptersfrom
amahishi/gym-turn-level-recovery

macandro96 commented Sep 13, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Sep 13, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

macandro96 commented Sep 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Summary

What this adds above #2946

Recovery flow

Correctness guarantees

Scalability

Testing

Scope

Stack

Checklist

Uh oh!

copy-pr-bot Bot commented Sep 13, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

macandro96 commented Sep 13, 2026 •

edited

Loading