[ananthsub]

What changed and why

An inference worker that serves generation cuts needs three helpers from RolloutTokenCapture. This PR adds them:

• begin_call(generation_cut=..., generation_cut_staging_keys=...) validates the durable prefix a replacement call continues against its admission: the staging keys, source call, digest, lineage, generated-token count, and that the prefix's policy version is not newer than the current one. It also checks that the prefix extends the replacement's exact parent tokens: it recomputes the chain hash from the admission's parent chain hash and the prefix's tokens, and the cumulative hash from the resolved prefix plus those tokens. The snapshot's digest authenticates only its own fields, so without this a prefix built on other parent tokens, with the same parent call ID and length, would be spliced in silently. This costs one hash of the prefix per recovered call.
• build_prefix_record snapshots a live call without completing it, so a cut can stage the prefix while the request keeps decoding. For a continued call it produces the old prefix plus the new tail, stamped with the oldest contributing policy version.
• build_generation_chunk_record stages only newly generated tokens, for workers that flush a long generation in chunks.

complete_call now builds its record through build_prefix_record and claims completion only after the record is valid. Both builders use StagedCallRecord.from_components from #4213, so every locally built record computes its digest once.

Why: the policy model PR sends cut requests and attaches restored cuts, but a real worker (for example NeMo RL's vLLM worker) could not serve them, because these helpers existed only on the old prefix-recovery branch.

Policy workers are logged as they come and go. The policy checkpoint coordinator in the main process now logs each worker by process ID, so its lines match uvicorn's own Child process [<pid>] died:

• an info line when a worker joins, with how many of the expected workers are connected;
• an info line when a worker stops cleanly: a worker that shuts down tells the coordinator first, with a short timeout;
• a warning when a worker is lost without saying so (a crash, or uvicorn killing a hung worker), with how many workers remain and whether an open checkpoint can no longer commit.

Before, a lost worker showed up only as a readiness blocker or an export_failed commit during a checkpoint, and outside one only in the status route's workers count. uvicorn restarts a lost worker on its own, and the replacement registers with the coordinator and adopts its state.

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,<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 #3938, #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 W 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
  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

Loading

This PR

What an inference worker does with these helpers when the policy model asks for a cut, and when a replacement call continues it.

sequenceDiagram
  participant P as Policy model server
  participant W as Inference worker
  participant S as Token staging
  P->>W: POST /ng-control/v1/generation-cut with the in-flight calls
  W->>S: build_prefix_record stages the tokens so far, the call keeps decoding
  W-->>P: receipt - staging keys and digest per call
  Note over P,W: after a crash and restore
  P->>W: replacement call carrying the continuation
  W->>S: begin_call checks keys, source call, digest, lineage, parent tokens, token count, policy version
  W->>W: decode only the rest of the generation

Loading

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 #3938 (capture ledger retire and delete) and #3939 (complete-record retire and delete), which have merged. #3882 and #3883 have merged too, so #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 (#4268) and the striped lock files (#3937) are independent of the stack.

1. feat(checkpoint): add the participant control plane, episode steps, and coordination #3882 (ananthsub/partial-ckpt-core): feat(checkpoint): add the participant control plane, episode steps, and coordination (merged)
2. feat(checkpoint): make policy model servers checkpoint participants with generation cuts #3883 (ananthsub/partial-ckpt-policy-model): feat(checkpoint): make policy model servers checkpoint participants with generation cuts (merged)
3. feat(token-capture): add worker staging helpers for generation cuts #3884 (ananthsub/partial-ckpt-model-worker-cuts): feat(token-capture): add worker staging helpers for generation cuts (this PR)
4. feat(checkpoint): continue environment server episodes from their boundaries #3885 (ananthsub/partial-ckpt-environment): feat(checkpoint): continue environment server episodes from their boundaries
5. feat(checkpoint): add the resources server participant with asynchronous session hooks #3886 (ananthsub/partial-ckpt-resources): feat(checkpoint): add the resources server participant with asynchronous session hooks
6. feat(checkpoint): declare twenty training verifiers stateless with replayable verification #3887 (ananthsub/partial-ckpt-verifier-declarations): feat(checkpoint): declare five training verifiers stateless with replayable verification
7. feat(checkpoint): add the agent session participant and Simple Agent continuation #3888 (ananthsub/partial-ckpt-agent): feat(checkpoint): add the agent session participant and Simple Agent continuation
8. test(checkpoint): add a process-level end-to-end suite driven by coordination #3889 (ananthsub/partial-ckpt-e2e): test(checkpoint): add a process-level end-to-end suite driven by coordination
9. feat(checkpoint): checkpoint evaluation runs from rollout collection #3893 (ananthsub/partial-ckpt-rollout-collection): feat(checkpoint): checkpoint evaluation runs from rollout collection
10. feat(checkpoint): spans and metrics for partial-rollout checkpoints #3903 (ananthsub/partial-ckpt-telemetry): feat(checkpoint): spans and metrics for partial-rollout checkpoints
11. feat(checkpoint): resources, agent, and environment servers with several workers #3909 (ananthsub/partial-ckpt-multi-worker): feat(checkpoint): resources, agent, and environment servers with several workers
12. test(checkpoint): check checkpoint invariants over random sequences of operations #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 (#2939 to #2946)

None of #2939 to #2946 had these helpers. They are ported from #3412 (prefix token-level recovery), which was built on the old stack; the wire contract is unchanged.

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 (#2939 to #2946) carry the original discussion.

Validation

Run on this branch, on the current main, which includes #4213:

• test_workers_are_logged_by_pid_as_they_join_and_a_lost_one_as_a_warning, test_a_worker_lost_during_a_checkpoint_is_logged_as_blocking_its_commit and test_a_worker_that_stops_cleanly_is_not_reported_as_lost: a worker connection dropped without notice is a warning, a clean stop is not; all three fail without the change, which logged nothing.
• test_generation_chunk_record_computes_staging_digest_once and perf(token-capture): bulk-pack columns and compute staging digests once #4213's test_complete_call_computes_staging_digest_once: each builder computes its digest once; all three cases fail with this PR's previous builders, which hashed before building the record.
• 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 tests/unit_tests/test_token_capture_*.py tests/unit_tests/test_token_id_capture.py tests/unit_tests/test_base_responses_api_model.py: passed. test_generation_cut_resume_rejects_a_prefix_built_on_other_parent_tokens fails without the parent-token check.
• 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

• The end-to-end suite in the top PR uses these helpers in its fake inference worker, so the cut record Gym restores is validated exactly as a real worker validates it (generation-cut scenario).
• A real vLLM worker serving generation cuts: pending.

Compatibility

Additive keyword arguments and new methods on RolloutTokenCapture; existing callers are unaffected. complete_call produces the same records as before.

[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.

[macandro96]

LGTM! let me know if its worth merging #4213 before this?

[ananthsub]

thanks, approved and merged #4213 and rebased the stack on this.

Both builders build_prefix_record and build_generation_chunk_record now use StagedCallRecord.from_components.

complete_call goes through build_prefix_record, so every locally built record is hashed once

[macandro96]

thanks, lgtm!

[github-actions]

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/38488609063