Skip to content

feat(checkpoint): persist recoverable Gym rollout state - #4265

Open
macandro96 wants to merge 1 commit into
amahishi/gym-checkpoint-participationfrom
amahishi/gym-durable-rollout-state
Open

macandro96 wants to merge 1 commit into
amahishi/gym-checkpoint-participationfrom
amahishi/gym-durable-rollout-state

Conversation

@macandro96

@macandro96 macandro96 commented Sep 25, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Defines and persists the rollout-level state required to recover Gym executions: stable execution ownership, completed-result receipts and ACK obligations, agent continuations, resource dependencies, policy lineage/TQ references, and source-to-replacement attempt mappings.

This is 2/3 in the stacked decomposition of #4117 and depends on #4264.

Stack PR Scope
1/3 #4264 Gym participant protocol
2/3 This PR Durable rollout state and validation
3/3 #4266 Single Controller orchestration

This layer defines what a recoverable rollout snapshot contains. The next PR schedules, publishes, and restores those snapshots from Single Controller.

Why

Participant files alone are insufficient to resume a rollout. Recovery must also prove that:

  • a completed Gym result is durably owned exactly once;
  • Gym completion acknowledgements are not silently lost;
  • each saved agent continuation refers to valid policy lineage;
  • every external token reference exists in the Transfer Queue;
  • resource revisions match the continuation that depends on them;
  • an unfinished source attempt maps to one replacement attempt;
  • executions using restart_only resources restart without discarding unrelated continuable work.

Without these relationships, recovery can duplicate a rollout, acknowledge the wrong result, or restore agent state against missing model/resource data.

Durable-state flow

flowchart LR
    A[Gym execution completes] --> B[Completion receipt]
    B --> C[Rollout manager seals result]
    C --> D[Recovery ledger records ACK obligation]
    D --> E[Gym validates and accepts ACK]
    E --> F[Recovery ledger removes obligation]

    C --> G[Agent continuation index]
    G --> H[Policy lineage and TQ reference index]
    G --> I[Resource revision dependencies]
    H --> J[Cross-artifact validation]
    I --> J
    F --> J
    J --> K[Recoverable rollout state]

    K --> L[Source rollout attempt N]
    L --> M[Replacement attempt N plus 1]
    M --> N{Resource mode}
    N -->|export_restore| O[Continue from saved boundary]
    N -->|restart_only| P[Restart only dependent execution]
Loading

Main changes

  • Add a Gym execution registry with atomic registration/freeze boundaries.
  • Carry receipt-bound completion acknowledgements instead of positional or response-only identity.
  • Add a persistent RolloutManager acknowledgement sink and durable pending-ACK records.
  • Extend the rollout recovery ledger with attempt mappings, stable ownership, sealed sibling/group state, and ACK obligations.
  • Persist agent continuation roots and per-continuation resource dependencies.
  • Collect policy lineage and external TQ storage references.
  • Validate participant manifests, continuation relationships, digests, capture keys, and TQ references before restore.
  • Reject unsupported checkpoint schema versions with an actionable message; this stack writes schema v3.
  • Support selective restart of continuations that depend on restart_only resources.
  • Add APIs to discard restored continuations that are intentionally restarted.

Correctness rules

  • A receipt identifies the exact durable Gym result being acknowledged; stale or mismatched receipts are rejected.
  • Pending ACK obligations remain explicit until Gym accepts them and RL records their removal.
  • PENDING_MODEL continuations must point to real captured model lineage.
  • Continuation coordinates cannot silently move backward or disappear.
  • External TQ keys are deduplicated only when their metadata is identical; conflicting duplicates fail closed.
  • Restored source attempt N is never reused as the new physical owner; recovery creates a tracked replacement attempt.
  • Restart-only resource dependencies affect only the continuations that reference those resources.

Scope

  • Builds on the participant protocol in feat(checkpoint): add Gym participant coordination #4264.
  • Does not yet start a periodic checkpoint pump or publish snapshots from Single Controller.
  • Does not add in-generation token-prefix restoration.
  • The checkpoint schema is experimental; backward restore from earlier private schema revisions is intentionally unsupported.

Test plan

Targeted coverage includes:

  • tests/unit/environments/test_gym_checkpoint.py
  • tests/unit/environments/test_nemo_gym.py
  • tests/unit/environments/test_nemo_gym_checkpoint.py
  • tests/unit/environments/test_nemo_gym_token_capture.py
  • tests/unit/experience/test_rollout_manager.py
  • tests/unit/experience/test_rollout_recovery.py
  • tests/unit/data_plane/test_rollout_reassembler.py
  • tests/unit/experience/test_rollout_generation_failures.py

Suggested command:

uv run pytest -q \
  tests/unit/environments/test_gym_checkpoint.py \
  tests/unit/environments/test_nemo_gym.py \
  tests/unit/environments/test_nemo_gym_checkpoint.py \
  tests/unit/environments/test_nemo_gym_token_capture.py \
  tests/unit/experience/test_rollout_manager.py \
  tests/unit/experience/test_rollout_recovery.py \
  tests/unit/data_plane/test_rollout_reassembler.py \
  tests/unit/experience/test_rollout_generation_failures.py

Before review

  • Contributor conventions followed.
  • Unit coverage added for receipts, continuations, validation, and recovery mappings.
  • Full target environment test run recorded in CI.

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

copy-pr-bot Bot commented Sep 25, 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.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant