Skip to content

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

Draft
macandro96 wants to merge 43 commits into
mainfrom
amahishi/gym-turn-level-recovery
Draft

feat(checkpoint): add durable turn-level rollout recovery#3349
macandro96 wants to merge 43 commits into
mainfrom
amahishi/gym-turn-level-recovery

Conversation

@macandro96

@macandro96 macandro96 commented Sep 13, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

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

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

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

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

copy-pr-bot Bot commented Sep 13, 2026

Copy link
Copy Markdown

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.

@yaoyu-33 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
ananthsub and others added 26 commits September 13, 2026 20:36
@ananthsub @macandro96
feat(checkpoint): propagate stable execution identity
b692978 
Keep the logical rollout ID and attempt index independent from the attempt-qualified capture key so checkpoint participants can fence retries without depending on capture configuration. Expose the authoritative model call ID on captured responses for later boundary reconciliation.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
feat(checkpoint): add fenced control contracts
ac0a97e 
Expose versioned checkpoint capabilities on every Gym server and provide a shared phase fence for idempotent control operations. Reject stale non-final operations after resume so an old coordinator cannot mistake an accepting server for a prepared one.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
feat(checkpoint): coordinate admission across workers
95d1bcf 
Close admission atomically in each server process and aggregate worker acknowledgements and in-flight counts through a service-level Unix socket coordinator.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
fix(checkpoint): preserve failed restore state
9fa1768 
Keep a failed restore paused and observable so the coordinator can retry or abort it safely.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
refactor(checkpoint): keep experimental APIs internal
fe0a54f 
Move the checkpoint implementation under nemo_gym._checkpoint until its contracts are ready for public support.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
feat(checkpoint): add policy model participant
710574e 
Pause policy generation at a quiescent cut and checkpoint token-free capture custody while leaving framework-owned staged token arrays to TransferQueue.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
fix(checkpoint): preserve model lineage across restore
350daea 
Namespace model artifacts, reject changing commit exclusions, and link the first restored call to its verified parent exactly once.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
feat(checkpoint): add whitebox agent participant
38c4250 
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>
@ananthsub @macandro96
fix(checkpoint): make agent boundaries safe to publish
d5f4676 
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>
@ananthsub @macandro96
feat(checkpoint): resume simple agent turns
3140c1f 
Capture complete tool turns in the shared agent participant and reconstruct the next attempt from the last committed model call and conversation boundary.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
fix(checkpoint): park SimpleAgent at durable boundaries
d807bdb 
Retry checkpoint refusals without recording false tool failures and stop retired inner handlers from mutating successor state.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
feat(checkpoint): resume Gymnasium agent turns
d115856 
Preserve accumulated episode state at completed environment steps so replacement attempts skip reset and continue with consistent reward, usage, and model history.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
fix(checkpoint): preserve Gymnasium turn state
eb3ab07 
Retry refused model and resource calls from the committed boundary without adding an extra turn at the step limit.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
feat(checkpoint): add resources state participant
9750356 
Serialize per-attempt mutations, export all known sessions as one prepared cut, and require environments to atomically activate the complete restored state set.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
fix(checkpoint): fence resources state mutations
1632387 
Bind state to execution identity, serialize mutations with snapshots, and reject unsafe restore or untracked stateful traffic.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
feat(checkpoint): adapt initial stateful environments
b21cb72 
Restore Gymnasium, workplace, and counter sessions under replacement attempt identities while preserving resource revisions at agent boundaries.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
fix(checkpoint): harden stateful environment adapters
3673b0d 
Require explicit checkpoint opt-in, bind state only after successful session creation, and retire adapter-owned state safely.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
test(checkpoint): cover a complete recovery cycle
1726db9 
Exercise prepare, commit, restore, explicit resume, lineage handoff, resource revisions, and straggler retirement across participants.

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@macandro96
feat(checkpoint): acknowledge completed agent executions
d8ef489 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
test(checkpoint): exercise SimpleAgent turn recovery
23bad8f 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
fix(checkpoint): isolate agent snapshot state
7ac1d99 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
feat(checkpoint): add artifact reference indexes
0141a15 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
refactor(checkpoint): require artifact indexes
f82ea58 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
refactor(checkpoint): require scoped artifact requests
0af3414 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
perf(checkpoint): avoid scanning inactive lineage
7169bb0 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
fix(checkpoint): align restore with capture contracts
b9b46ba 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
macandro96 and others added 16 commits September 13, 2026 20:36
@macandro96
test(checkpoint): use current capture ledger records
68aa574 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
fix(checkpoint): support nested token-capture calls
213fa5b 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
fix(checkpoint): preserve receipt lineage across attempts
e0de784 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@ananthsub @macandro96
feat(checkpoint): make agent completion durable
ea63c55 
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
fix(checkpoint): preserve SimpleAgent continuation cookies
bd35bd4 
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
fix(checkpoint): preserve Gymnasium continuation cookies
301f18a 
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
feat(checkpoint): make resource mutations exactly once
8a274c6 
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
feat(checkpoint): make generation admission atomic
afcf886 
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
feat(checkpoint): bind model cuts to capture ledger
b8dbeba 
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@ananthsub @macandro96
test(checkpoint): validate generation-safe recovery
5c83688 
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
@macandro96
fix(checkpoint): integrate generation-safe recovery
b1fffd8 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
test(checkpoint): add workplace turn recovery fixture
bede4bc 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
perf(checkpoint): add exact completion receipt lookup
a11dce8 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
feat(checkpoint): classify restart-only resources
4d61c82 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
fix(checkpoint): scope restart-only recovery per execution
366a574 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
perf(checkpoint): archive rollout lineage in shards
6bba7b8 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
macandro96 force-pushed the amahishi/gym-turn-level-recovery branch from 6473f49 to 6bba7b8 Compare September 14, 2026 01:40
@macandro96
perf(checkpoint): archive agent state in shards
80bd2a9 
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
@macandro96
macandro96 marked this pull request as draft September 14, 2026 03:47
@macandro96 macandro96 changed the title feat(checkpoint): Gym turn level recovery feat(checkpoint): add durable turn-level rollout recovery Sep 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

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

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@macandro96 @yaoyu-33 @ananthsub