feat(checkpoint): continue proof refinement agent turns from their boundaries

[ananthsub]

• Branch: ananthsub/partial-ckpt-proof-refinement
• Base branch: ananthsub/partial-ckpt-e2e
• Suggested labels: feature, area:agent, area:env-infra (draft, so no state label yet)

---

What changed and why

The proof refinement agent runs its own seed, generate, verify, and correct loop inside a legacy /run. Before this PR it did not opt into checkpointing, so it was restart-only: every checkpoint had to retire its in-flight rollouts and restart them from the task input, which throws away proofs already generated and compiled.

This PR:

• Runs the /run loop under the agent participant's legacy_run and records a boundary before each step:

  • seed, before /seed_session;
  • generate, before each policy call, holding the current input and the completed attempts;
  • verify, before each /verify, holding the generated proof;
  • return, once the final result is known.

  A replacement attempt continues from that boundary. It skips the seed, keeps the completed attempts, and never regenerates a proof that was already generated.

• Runs the seed and each generation as replay steps. A checkpoint does not wait for them. The policy model holds an undelivered response until resume, and after a crash the step runs again from its boundary. Verification uses the mode the resources server reports in its seed reply (x-ng-checkpoint-verify), as Simple Agent does.

• Makes the agent's /v1/responses proxy pass the policy model's 409 checkpoint_parked through to /run instead of raising a 500. /run backs off for 0.1 seconds, returns to the boundary before the generation, and calls again. Once the agent's own admission closes, that boundary parks the run until resume. The backoff covers the short window when the policy model has closed but the agent has not yet been prepared.

• Sets checkpoint_sessions_supported = True. The three session hooks are no-ops, because the boundary holds the whole episode.

• Declares math_formal_lean as checkpoint_mode = "stateless" and checkpoint_verify = "replay". Its /verify builds the proof from the request, compiles it in the sandbox's stateless /execute route, and returns a result that depends only on the request and the compiler output. It keeps no session state, so running it again after a crash is safe and changes nothing a checkpoint holds.

Without checkpointing, /run behaves as before. The only visible difference is that the first attempt's input record is dumped in JSON mode, which serializes to the same bytes.

How it works

Where this PR sits in the overall flow

The highlighted part is what this PR adds.

flowchart LR
  C["Controller<br/>NeMo RL, or rollout collection"]
  CO["Coordination<br/>prepare, commit, restore, resume, retire"]
  K["Control plane on every server<br/>phases, fencing, lease, storage"]
  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")]
  C --> CO --> K
  K --> E & M & A & R
  M --> W
  G --> D
  classDef this fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937
  class A,R this

Loading

One checkpoint, a crash, and the restore, end to end:

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
  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 + 1
  C->>G: resume
  C->>G: /run as attempt + 1 continues each episode from its boundary
  G-->>C: a late call from attempt 0 gets 409 stale_attempt

Loading

This PR

The proof refinement agent's legacy /run loop, with a boundary after every step. A restored attempt starts after its last boundary and never repeats a finished generation.

flowchart LR
  B1(["boundary: seed"]) --> S["seed the Lean session"]
  S --> B2(["boundary: generate"]) --> G["generate a proof<br/>replay step"]
  G --> B3(["boundary: verify"]) --> V["verify with Lean<br/>replay, the server is stateless"]
  V -->|"proof fails, turns left"| B2
  V -->|"proof passes, or no turns left"| B4(["boundary: return"]) --> R["reply to /run"]

Loading

Where this sits in the stack

This PR builds on the partial-rollout checkpointing stack: #3882 (core) through #3889 (end-to-end suite), plus #3893 (rollout collection). It is one of five PRs that port the environments the old stack checkpointed onto the new hooks. Each one is based on #3889 and can be reviewed on its own, except Blackjack, which builds on the Gymnasium fix:

• feat(checkpoint): checkpoint Workplace Assistant sessions #3894 (ananthsub/partial-ckpt-workplace-assistant): feat(checkpoint): checkpoint Workplace Assistant sessions
• fix(gymnasium): build Gymnasium servers on the shared resources server setup #3895 (ananthsub/partial-ckpt-gymnasium): fix(gymnasium): build Gymnasium servers on the shared resources server setup
• feat(checkpoint): export Blackjack game state for partial-rollout checkpoints #3896 (ananthsub/partial-ckpt-blackjack): feat(checkpoint): export Blackjack game state for partial-rollout checkpoints (on the Gymnasium PR)
• feat(checkpoint): export indirect prompt injection session state #3897 (ananthsub/partial-ckpt-indirect-prompt-injection): feat(checkpoint): export indirect prompt injection session state
• feat(checkpoint): continue proof refinement agent turns from their boundaries #3898 (ananthsub/partial-ckpt-proof-refinement): feat(checkpoint): continue proof refinement agent turns from their boundaries (this PR)

Relationship to the old stack

This PR supersedes #3634, which added the same continuation on the old stack. The loop shape, the stop rules, and the attempt records carry over, and so does the replay-safe Lean verification test. These parts are not carried:

• Capture-outcome handling. feat(checkpoint): resume proof refinement agent turns #3634 raised on a capture_failed model call and turned a no_generation call into a masked zero-reward result with failure_kind = agent_no_generation. Both read the model-call capture outcome header and the failure kind added by the token-capture fixes in feat(checkpoint): add durable turn-level rollout recovery #3349, which are not on this stack. The agent does not inspect capture outcomes here.
• Model-call coordinates. PendingModelPayload, the committed model call ID and capture key, and the forwarding of the model-call ID and outcome headers through /v1/responses. On v2 an undelivered generation is left out of the checkpoint and re-issued from the generate boundary, and the model ledger restores rows under the next attempt's key.
• Fresh-parent calls for correction turns. feat(checkpoint): resume proof refinement agent turns #3634 called take_checkpoint_parent() before each correction so its capture started a new root. v2 has no parent relay. A correction prompt replaces the conversation, so it does not extend the earlier proof's capture.
• Resource request IDs and revisions, retry_checkpoint_refusal, and checkpointable_external_wait. Ordering replaces receipts. The 409 handling above replaces the refusal retry, and the seed reply's verify mode replaces the agent-side checkpoint_replayable_verify config flag.

Issue

No separate issue. This ports #3634 onto the v2 checkpoint hooks.

Validation

• RAY_TMPDIR=/tmp python -m pytest -q responses_api_agents/proof_refinement_agent/tests: 9 passed. The new checkpoint tests drive the agent's real ASGI app and control routes. The agent's self-call to /v1/responses also goes through the app, so only the policy model and the Lean server are faked.
  • A /run is checkpointed in five positions: generation in flight, replayable verification in flight, a waited-on verification, the correction generation in flight, and the final verification finishing during prepare. The checkpoint is committed, the original attempt is retired, and the record is restored into a fresh agent instance. The replacement attempt returns exactly the uncheckpointed result. It never re-seeds, and it generates and verifies only the steps after its boundary. For example, with a replayable verification in flight it verifies turns 0, 1, and 2 but generates only turns 1 and 2.
  • The retired attempt is gone: a later prepare reports no sessions.
  • The policy model's 409 checkpoint_parked reaches /run through the proxy. The run stops calling the model once the agent is prepared, and it completes with the uncheckpointed result after resume.
  • A /run without a rollout ID is refused with 400 rollout_id_required.
  • Two mutation checks failed as expected: ignoring the restored boundary failed 5 tests, and removing the 409 pass-through failed the refusal test.
• RAY_TMPDIR=/tmp python -m pytest -q resources_servers/math_formal_lean/tests: 45 passed, 2 skipped. New tests:
  • The same verification gives an identical result on the original server, after an unrelated verification, and on a fresh server.
  • With checkpointing on, the seed reply carries x-ng-checkpoint-verify: replay, status reports stateless/replay, /verify is served while the server is prepared, and a commit exports no records.
  • A real-compilation replay test is skipped unless a Lean sandbox is reachable at NEMO_SKILLS_SANDBOX_HOST:NEMO_SKILLS_SANDBOX_PORT. No sandbox was available, so these 2 tests were skipped.
• RAY_TMPDIR=/tmp python -m pytest -q tests/unit_tests/test_checkpoint_*.py: 116 passed, including the math_formal_lean entry added to the server declaration test.
• pre-commit run --files <changed files>: passed.

Rollout evidence

Pending. #3634 ran a training and recovery experiment on the old stack, with a checkpoint taken after a failed first proof and a resume in fresh processes. That needs to be rerun on this stack with a real policy model and a Lean sandbox. The process-level e2e suite does not cover this agent yet.

Compatibility

• Without checkpointing, nothing changes for /run callers or datasets.
• With checkpointing on, this agent is no longer restart-only. Its /run now needs a rollout ID (_ng_rollout_id, or task and rollout indices) and is refused with 400 rollout_id_required without one, as Simple Agent's is.
• The capture-outcome behavior feat(checkpoint): resume proof refinement agent turns #3634 proposed (masking no_generation turns) is not part of this PR. It should land with the token-capture fixes it depends on.

[copy-pr-bot]

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.