feat(checkpoint): propagate stable execution identity by ananthsub · Pull Request #2939 · NVIDIA-NeMo/Gym

ananthsub · 2026-09-01T22:03:49Z

Summary

propagate a logical rollout ID and attempt index through agent, model, and resources requests
make execution identity independent of observability and token capture
reject conflicting caller-supplied identity instead of silently changing attribution

Why this layer exists

Checkpoint state is distributed across several Gym services. Those services need one stable key before later layers can join an agent boundary, resources session, model-call lineage, retry, and restored execution. rollout_id names the logical rollout. _ng_attempt_index distinguishes the original execution from replacement attempts after retry or restore.

Execution identity flow

sequenceDiagram
    participant RL as NeMo-RL / caller
    participant Agent as Agent server
    participant Model as Model server
    participant Resources as Resources server

    RL->>Agent: /run with rollout_id and attempt index
    Agent->>Agent: Validate or establish request context
    Agent->>Model: Model call with identity headers
    Agent->>Resources: Tool or verifier call with identity headers
    Model->>Model: Attribute model-call ID to the same execution
    Resources->>Resources: Associate session state with the same execution
    Note over Agent,Resources: One logical rollout can later continue as attempt N+1 without reusing attempt N's mutable state.

If a caller omits a rollout ID, Gym can mint one at ingress. If nested headers conflict with the active context, the request fails rather than producing records under two identities. This PR establishes correlation only; admission gates and stale-attempt tombstones are added by later stack layers.

Stack

base: #2872
next: #2940

Test plan

execution identity unit tests
scoped pre-commit checks

copy-pr-bot · 2026-09-01T22:03:52Z

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.

copy-pr-bot · 2026-09-22T17:28:53Z

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.

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>

macandro96 · 2026-09-23T03:12:21Z

            if body is None:
                body = next((arg for arg in args if isinstance(arg, BaseRunRequest)), None)
-            with rollout_context(self.rollout_id_from_run(body)):
+            logical_rollout_id, attempt_index = execution_identity_from_run_body(body)


We need to install middleware for agents also right for eg. to cover cases such as agent calling itself, we need to extract identity from headers, otherwise it may not be reachable downstream?

I got confused about this when i first read this part. I asked my agent to trace it and it seems like it was covered in #2942. Haven't got a chance to review that PR yet.

Oh i see. In that case - okay with merging it there or in this PR.

macandro96 · 2026-09-23T03:23:51Z

+        async def _send_with_model_call_id(message: dict[str, Any]) -> None:
+            if message.get("type") == "http.response.start":
+                headers = list(message.get("headers") or ())
+                headers.append((MODEL_CALL_ID_HEADER.encode("ascii"), model_call_id.encode("ascii")))


I think we may need to skip attaching a model call id when a capture failed? I had to add this later in my branch to prevent fallbacks on failed / no generation.

zyzhou5 · 2026-09-23T21:52:50Z

One thing that cause confusion when I'm reading the PR is we use both URL prefix and http header mechanism for different things. Is there a reason why certain mechanism is chosen at certain places or is it just different choice at different stage of development. We might consider at least documenting this to avoid confusion.

…nd coordination (NVIDIA-NeMo#3882) ## What changed and why This PR adds the core that every checkpoint participant builds on. On its own it installs nothing on any server; the later PRs in the stack each add one participant on top of it. - **Identity.** An episode attempt is keyed by its capture key: `r` for the first attempt and `r-aN` after a restore. When checkpointing is on, `ServerClient` always attributes model and resources calls to their rollout through the rollout path prefix, not only when observability is on. `main` already reserves the `-a<N>` suffix in rollout IDs; this PR adds `EpisodeId.from_capture_key`, which parses a key back to its rollout and attempt. An explicit rollout ID may end with its own `-a<N>`, which is how a controller names a replacement attempt, but not as well as an attempt index: that key would decode to no rollout and attempt, so it is refused. The typed `single_agent_turn` request path (`_episode_request_body` in `nemo_gym/rollout_collection.py`) and the legacy flat-row wrapper (`_episode_request_from_row` in `environment_servers/single_agent_turn_legacy/app.py`) decode the rollout ID the same way the agent does: an explicit `"<id>-aN"` runs as attempt N of `<id>`, a row without one runs as `"<task>-<rollout>"` at its stamped attempt index, and a first attempt is unchanged. Before this, both paths took an explicit ID as the base rollout ID and failed the reserved-suffix check, so a restored turn sent through a typed environment server, as NeMo RL's `target_level: turn` does (NVIDIA-NeMo/RL#4500), got a server error. - **Control plane** (`nemo_gym/_checkpoint/control.py`). Each server has one participant behind bearer-protected routes under `/ng-control/v1/checkpoint/` (`status`, `prepare`, `renew`, `retire`, `forget`, `commit`, `restore`, `resume`). One control plane per server owns the phase machine, checkpoint ID fencing, the readiness wait, deadline-bounded file I/O, and manifest-last storage. A lease resumes a participant whose controller stopped calling. Resumed checkpoint IDs are remembered for the last 64 checkpoints, enough to refuse a stale controller's delayed request. Status reports how many rollouts the participant still refuses retired attempts of (`retired_rollouts`). - **Episode steps** (`steps.py`). An episode records a boundary after each completed step. A step in flight when a checkpoint closes runs in `wait` mode (the checkpoint waits for it) or `replay` mode (it runs again after a crash). Time parked for a checkpoint does not count against the episode's deadline. - Every live episode is checkpointable. One that has not recorded a boundary yet, such as an episode inside its first replay step, is exported with no boundary and starts over from its input after a crash. A restored episode starts at the boundary it was restored from. - A replay step that ends while a checkpoint is open leaves the episode at the boundary before the step, not a blocker, until its next boundary parks it with the step's result. - An episode that is being retired blocks a checkpoint whatever its state, and is never exported, until its task has ended. A retire cut short by its deadline therefore cannot carry the episode into the next checkpoint. - The steps keep a running count of blocking episodes, list at most the reported number of them, and notify only while a checkpoint is open, so readiness costs nothing per step outside a checkpoint. The readiness wait asks a participant only whether it is ready (`ready()`), and lists blockers once, for the reply. - A resources server reports whether its `/verify` may be replayed in an `x-ng-checkpoint-verify` header on its `/seed_session` reply. - **Rollouts that cannot be saved never block a checkpoint.** A rollout that some server cannot capture, because its agent has no session hooks or its resources server is `restart_only`, is a *restart*. - Each participant lists its restarts in prepare's report (`PrepareReport.restarts`, never capped, since the controller must see every one). A restart never blocks prepare and is never exported; it keeps running through the checkpoint. - The controller builds its commit scope with `PrepareResult.continued(...)`, which leaves the restarts out. A commit whose scope names a restart is refused with 409 `restart_in_scope`. - After a crash, the controller retires the restart's attempt everywhere and sends the rollout again from its input as the next attempt, so a harness that outlived the crash, for example in a remote sandbox, cannot write into the rerun. - An episode becomes a restart when a seed reply carries `x-ng-checkpoint-restart`. Marking one waits while a checkpoint is open, so the restarts a prepare reported cannot change before its commit. - **No `/run` resend.** With checkpointing on, `ServerClient` never sends a `/run` again after a dropped connection. The first request may already be running, and a second execution of the same attempt would orphan the first one's agent session. - **Coordination** (`coordination.py`). Gym owns the participant order, so a training controller calls these functions instead of sequencing servers itself: prepare environment servers, then the policy model, then agents, then resources servers; resume in reverse; commit everywhere; restore all or nothing; retire callers before the servers they call (environment servers, then agents, then the policy model and resources servers). `renew` extends every lease when publishing a checkpoint takes longer than expected. - **Retire refuses, stops, and frees; forget ends the refusal.** - A participant's retire first marks the attempt's rollout as retired up to that attempt, before any work its deadline limits, so even a retire that arrives late leaves the mark. Callers are retired before callees, but a request a caller sent just before its own retire travels on another connection and can reach a server after that server's retire finished; the mark refuses it with 409 `stale_attempt`. - The retire then stops the running work of the attempt and its earlier attempts, waits for it, frees its state, and only then replies. An episode being retired stays tracked until its task has ended, so a retire cut short by its deadline leaves it for a retry, which waits again without cancelling a second time. - The mark never goes down, and it stays even if the retire fails. It goes away only through the new `forget` operation, which the controller calls once a rollout is done for good: nothing of any of its attempts can still run, and it will not be dispatched again. Each participant keeps one entry per rollout retired and not yet forgotten, so the controller's `forget` calls bound it. - Retire is refused while a checkpoint is preparing, prepared or committed: stopping an episode waits for its final cleanup, whose calls wait for resume. A controller with a straggler resumes, retires, then checkpoints again. - **Restore marks the attempts it replaces.** Retire marks live in memory, so fresh processes after a crash would not refuse the crashed attempt of a continued rollout, and a harness that outlived the crash, for example in a remote sandbox, could still reach the restored servers. Restore therefore marks every episode in its scope as retired up to the attempt before its replacement, on every participant and every worker, before any work its deadline limits. These marks also end only through `forget`, so a controller must forget each restored rollout once it is finished for good, not only the ones it retired. - A retire stops its episodes concurrently, up to a fixed bound, and each rollout once, through its latest attempt; the cleanup of unscoped restored state does the same. - Retire stays allowed right after a restore, where nothing is running yet; a failed restore uses it to free what it installed. `forget` follows the same phase rule and runs in the same caller-first order. - The controller keeps its own side paused from prepare to resume. For NeMo RL that is rollout dispatch, the trainer actors, and the trajectory assembler actors. - **A commit deletes restored state the controller no longer continues.** A commit's scope names every episode the controller continues, including an unclaimed restored episode, named by its replacement attempt. Participants export only restored state inside the scope. After writing its records, within the commit's deadline, a participant deletes the restored episodes the scope leaves out through the `delete_restored` hook, so unclaimed restored state is never carried from checkpoint to checkpoint, and a commit retried after a failure there returns the same manifest. These episodes are not marked retired: their replacement never ran here, and the controller may still start the rollout over as that attempt. `delete_restored` must therefore leave nothing that refuses the attempt later, such as a capture-ledger fence. It defaults to `retire`, which is enough for participants whose retire refuses nothing beyond itself; the policy model deletes the attempt's ledger instead. - **Restore and discovery fail cleanly.** - Restore reads, verifies, and validates in one streaming pass in a thread, keeping only the records in its scope, so it never holds the whole checkpoint. A record that fails validation is a typed `invalid_checkpoint_state` error naming its episode. - A manifest that is not JSON or not an object, is missing fields, or names a records file outside the participant directory, a missing records file, and a directory that cannot hold state are all typed `invalid_checkpoint_state` errors, never server errors. - `install(records, scope)` receives the restore scope: every episode the restore continues, including ones a participant has no record of, so a participant that keeps files per episode can clear them too. - An install can outlive its restore's deadline, as the model ledger's import thread does. The control plane keeps the install, and a retire or the next restore waits for it, so the retire that cleans up a failed restore cannot run before the state it frees is installed. Admission reopens only once that install has finished, unless a later restore has taken over by then. - A restore whose close of admission fails part way, for example on one of several workers, reopens admission, as prepare does. - If a restore fails, coordination retires the replacement attempts and always resumes every participant, then raises the restore error, not a cleanup error. - `discover` caps its connection retries and bounds each probe with a timeout, so a server that is down, or accepts a connection but never replies, is reported instead of waited on forever. A server reporting an unknown participant kind is reported the same way. - **Control requests are cheap to fan out.** A control route validates its request once, and coordination encodes each fan-out body once and sends the same bytes to every participant. `retire` and `forget` reply with counts, `retired_count` and `forgotten_count`, not every ID; commit's `episode_ids` and restore's `restored` stay exact. - **Reserved control connections.** Checkpoint and other control calls use a separate aiohttp connection pool, so data calls that hold every connection to a host (for example, `/run` calls parked during a checkpoint) cannot starve prepare, commit, or resume. The control pool is built like the data pool on `main`, with the same keepalive socket options and, when metrics export is on, the same timed connector. - **One export, one write, and one cleanup per checkpoint.** The records writer runs in a thread, which a deadline cannot stop. The control plane therefore keeps each checkpoint's export, the task storing it, and the task deleting unscoped restored state as shielded tasks. - A commit retried while one of them runs, or after it finished, awaits it instead of starting another, and one retried after storing failed stores the same records again. A checkpoint never exports twice, so its cut stays the one taken at the first commit. An export that fails, or a participant whose readiness regressed after prepare before its export started, ends that checkpoint for good: every commit for its ID gets 409 `export_failed`, and the controller resumes and prepares again under a new ID, because a later export could take a later cut than the other participants committed. - A completed commit or restore is remembered with the arguments it ran with: the directory and the scope, compared in sorted order. A retry with different ones gets `checkpoint_conflict` instead of the earlier result. - Resume, including a lease expiry, cancels an export still running and waits, within its deadline, for a cleanup still deleting restored state, because a replacement may start as that attempt. A resume that runs out of time raises `deadline_exceeded` and keeps the phase, so the controller resumes again; an expired lease retries on its own every 30 seconds, after any error, not only a passed deadline. A reopen that fails during a restore's rollback stays pending, and the next resume retries it instead of answering as a no-op. Resume also stops the write: the writer checks before each record, and its last check, the rename of its records file, and the manifest's creation run under one lock that resume takes, so no manifest appears once resume has returned. The records file is named after its digest (`records-<sha256>.jsonl`) and the manifest is only created, never replaced, so two writers cannot leave a records file and a manifest from different writes. - **Faithful, cheap storage.** - Records are encoded, hashed, and written to the records file one at a time, and read back the same way, so a commit or restore never holds the whole file in memory. - Each participant's records file keeps every dict's key order, NaN and infinities exactly as exported, so a restored episode sees byte-for-byte the state it had. Tool outputs often serialize stored dicts. A test checks the round trip. - Each commit syncs the parent of every directory from the participant directory up to the controller's checkpoint directory before it publishes the manifest, so a host crash cannot lose a directory entry under a committed manifest. - Opaque server state in a record (`JsonPayload`) is checked once for JSON-ness, then written as is, rather than deep-validated and copied. `orjson` is the fast check. When it rejects a value, such as an integer beyond 64 bits or nesting deeper than it supports, the value is accepted only if it is plain JSON with string keys, so it reads back unchanged. A value `orjson` accepts but `json` cannot write, such as a `datetime`, fails the commit with an error naming the episode. Records read back at restore, or passed between a server's workers, are not checked again. For 1,000 sessions of 350 KB, commit drops from 11.5 s to 3.6 s and restore from 6.0 s to 3.4 s. - **Settings.** A global `checkpoint:` block with `enabled`, `control_auth_token`, and `lease_grace_seconds`. It is off by default, and a control token is required when it is on. `checkpoint` is a reserved top-level key, and the block is validated once, when the `ServerClient` is built, so a bad block fails at startup. Why: the old stack spread identity across headers and path prefixes, gave each participant kind its own control contracts, and left the participant order to the training controller. Environment servers now exist on `main` and own the episode protocol, which makes a smaller and fail-closed design possible. ## 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 CO,K 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 participant's phases, enforced by the control plane. A second checkpoint ID while one is open gets `checkpoint_conflict`. Retire is accepted only in `idle` and `restored`. ```mermaid flowchart LR I(["idle"]) -->|prepare closes admission| PG(["preparing"]) PG -->|every live execution is at a boundary| PD(["prepared"]) PD -->|commit writes records, then the manifest| CM(["committed"]) CM -->|resume reopens admission| I PG -->|resume aborts, or the lease expires| I PD -->|resume aborts, or the lease expires| I I -->|restore, in a fresh process only| RS(["restored"]) RS -->|resume| I ``` Episode steps: what happens to a step that is running when a checkpoint closes. ```mermaid flowchart TD S["A step is in flight when the checkpoint closes"] --> Q{"Step mode"} Q -->|wait| W["The checkpoint waits for the step. The episode parks at the boundary after it."] Q -->|replay| P["The checkpoint does not wait. The boundary before the step is exported, and the step runs again after a crash."] ``` Retire refuses, stops, and frees. Each server marks the attempt before it stops anything, so a late request from a caller is refused, and keeps the mark until the controller forgets the rollout. ```mermaid sequenceDiagram participant C as Controller participant E as Environment servers participant A as Agents participant MR as Model and resources servers C->>E: retire(r, N) Note over E: mark r up to N, cancel the episode, wait for it and its final cleanup E-->>C: stopped and freed C->>A: retire(r, N) Note over A: cancel activations and wait, then the retire hook A-->>C: stopped and freed C->>MR: retire(r, N) Note over MR: cancel calls and wait, drain session requests, retire ledgers and hooks MR-->>C: stopped and freed Note over C,MR: every server refuses r up to N from here on C->>E: forget(r), once r is done for good, then agents, then model and resources ``` ## 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'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 **(this PR)** 2. NVIDIA-NeMo#3883 (`ananthsub/partial-ckpt-policy-model`): feat(checkpoint): make policy model servers checkpoint participants with generation cuts 3. NVIDIA-NeMo#3884 (`ananthsub/partial-ckpt-model-worker-cuts`): feat(token-capture): add worker staging helpers for generation cuts 4. NVIDIA-NeMo#3885 (`ananthsub/partial-ckpt-environment`): feat(checkpoint): continue environment server episodes from their boundaries 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 five training verifiers stateless with replayable verification 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) - Supersedes NVIDIA-NeMo#2939 (execution identity). Only the rollout path prefix remains; the `x-nemo-gym-rollout-id`, `x-nemo-gym-attempt-index`, and `x-nemo-gym-model-call-id` headers are gone. - Supersedes the control-plane part of NVIDIA-NeMo#2940 (control contracts). The per-kind request models and capability endpoint become one route set and one request shape for every participant kind. The generation-cut models move to the policy model PR. The multi-worker admission coordinator is not carried; multi-worker policy model servers return in a different, smaller form in the policy model PR. - New compared with the old stack: episode steps and Gym-owned coordination. The completed-result acknowledgements are gone: no episode completes while Gym is prepared, so there is nothing to acknowledge. ## 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`: - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_checkpoint_*.py`: passed. New since the first version of this PR: - a retire replies only after the participant stopped the attempt; - a request that arrives after a retire finished is refused until `forget`, a failed retire keeps its attempts refused, and a retry stops them; - a failed export is final for its checkpoint and a new checkpoint commits the new state, also when the export outlived its commit; a commit refused because readiness regressed is final for its checkpoint; - resume retries a reopen that a failed restore could not finish, and waits for an install that outlived its restore; the lease recovery retries a reopen that failed; - a retire past its deadline still refuses its attempts; a restore refuses the attempts it replaces until the controller forgets them, also past its deadline; - participant directories are durable before the manifest is published; - a retire stops its episodes concurrently and a cut-short one cancels each once; a commit deletes unscoped restored state concurrently; - restore does not check payloads read back from storage again; a control route validates its request once; a fan-out encodes its body once; - `forget` is refused while a checkpoint is open, and coordination runs it, like retire, environment servers first; - a straggler is retired after resume, not while a checkpoint is open; - a retire cut short by its deadline leaves the episode tracked, and a retry waits for it without cancelling again; - an episode with no boundary yet is exported to start over, a restored one at its restored boundary, and a replay step that ends during a checkpoint does not block it again; - a commit deletes restored episodes its scope leaves out without retiring them, and a commit that failed while doing so can be retried; - restore validates only the records in scope; a failed restore resumes every participant and raises the restore error; - `discover` reports a server that is down, one that never replies, and one of an unknown kind; a bad `checkpoint:` block fails when the `ServerClient` is built; - records stream to disk, and a retried commit leaves no temporary file; - a commit retried while its write runs awaits that write; resume stops a write that outlived its commit before it publishes; a losing writer neither replaces the published records nor leaves its own; - payloads are checked by what the writer can write; - a restore whose close fails part way reopens admission, and a restored record that fails validation is a typed error naming its episode; - an episode whose retire was cut short blocks the next checkpoint and is not exported; - an explicit rollout ID with an attempt suffix cannot also take an attempt index; - the typed `single_agent_turn` request path decodes an explicit `"<id>-aN"`, a stamped attempt index, and a first attempt into the same rollout and attempt as the agent; - a commit retried after its export finished reuses that export; resume waits for a cleanup that outlived its commit before it reopens admission, and a resume or lease expiry that runs out of time keeps the phase and succeeds on retry; a completed commit or restore retried with another directory or scope conflicts, and one retried with the same scope in another order does not; the control pool refuses negative limits; - a commit retried after its write failed stores the same records; resume waits for a publication under way, so none appears after it; a restore that ran out of time has its install waited for by the next retire, and reopens admission only once that install has finished; install is given every episode the restore continues; a read keeps only selected records but verifies the whole file; prepare lists blockers only for its reply. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_checkpoint_contracts.py`: passed. Contract tests over many inputs: every payload the check accepts is restored unchanged or fails the commit by name; every capture key built from a run body decodes back to itself; every control route refuses invalid requests with a client error; every kind of corrupt manifest is a typed error. Before the fixes in this PR they failed on integer dict keys, very deep payloads, six kinds of corrupt manifest, and an unwritable directory. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider environment_servers/single_agent_turn_legacy/tests`: passed. New: the legacy flat-row wrapper decodes the same three cases as the agent, and a `/run` with the explicit rollout ID `"restored-a1"` seeds attempt 1 of `restored` and sends its `/v1/responses` call under the rollout path prefix `restored-a1`. - 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 N/A for this PR on its own: it adds no participant, so no server behaves differently. The process-level end-to-end suite and the scale results in the top PR of the stack exercise this core through every participant. ## Compatibility - New global config: a `checkpoint:` block with `enabled` (default `false`), `control_auth_token` (required when enabled), and `lease_grace_seconds` (default 600). - New config for the reserved control connection pool: `global_aiohttp_control_connector_limit` and `global_aiohttp_control_connector_limit_per_host`, both at least 0. - New 409 error `export_failed`: the checkpoint cannot be committed; resume, then prepare again under a new checkpoint ID. - `retire` and `forget` reply with `retired_count` and `forgotten_count` instead of ID lists. - A controller must also `forget` every rollout it restored once that rollout is finished for good. - The typed `single_agent_turn` request path and the legacy flat-row wrapper accept an explicit rollout ID that ends in `-aN` and run it as attempt N of the rest of the ID; on `main` they refused it. This applies whether or not checkpointing is on. - Nothing changes unless the global `checkpoint:` block sets `enabled: true`. With it on, model and resources calls always carry the rollout path prefix. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

…ith generation cuts (NVIDIA-NeMo#3883) ## What changed and why A policy model server with `checkpoint_policy: true` becomes a checkpoint participant. The checkpoint never waits for generations. - **Held responses, and new calls that wait.** Closing admission holds every response that has not been delivered yet. Only a response that is already streaming blocks readiness. - A new policy call that arrives while a checkpoint is open waits until resume before it is admitted, as a resources server's requests do. It used to get 409 `admission_closed`. A rollout that keeps running through the checkpoint, such as a restart, would have shown that refusal to the model or failed. - A waiting call is not admitted, so prepare never counts it, and its agent is still at the boundary it recorded before the call. If its attempt is retired while it waits, it is refused with `stale_attempt` at once instead of running after resume. - The gate tracks the calls waiting at closed admission by capture key. When an attempt is marked retired, each waiting call of that attempt or an earlier one is flagged as refused for good, and woken. The flag stays set after the controller forgets the rollout, and it is set by the mark, which comes before any work the retire's deadline limits, so a retire cut short by its deadline still refuses the call. Without it, a call that had not run again by the time the controller forgot the rollout found no mark and was served once admission reopened. - A retire returns only after those waiting calls have left the gate with `stale_attempt`. A waiting call leaves the tracking when it leaves the gate, so the tracking holds only calls still waiting. - **Undelivered calls are left out, because they have no ledger row.** An undelivered call is excluded from the checkpoint, so after a restore the agent simply makes the call again. - A call's capture-ledger row is published when its response starts, at its first answer chunk, not when the worker stages its tokens. Headers and keep-alive comments publish nothing. Behind the gate, the capture handler leaves the row on the call's capture context (`deferred_commit`), and the gate writes it after `deliver_response` and before the chunk is sent. - Rows used to be written at staging. Each policy server left out its own undelivered calls by `model_call_id`, but a second policy server sharing the ledger could not see them and exported their rows; and a call whose capture context attached after the commit's snapshot could land its row before the export read the ledger. Either way a checkpoint could hold the answer of a call the caller never received. - The write cannot be cut short by a cancellation, and the call stays tracked until the write or its failure handling is done, so a retire or a checkpoint waits for it. A failed write poisons the call and no answer is sent. A call that leaves before its response starts has no row and is not poisoned; its staged tokens are left for the training framework's cleanup, as before. - The export still leaves out undelivered calls' rows, now only as a guard. - With checkpointing on, every model server named in the config of an agent with token capture must be a checkpoint policy server, since the capture ledger is run-wide and a server outside the gate would write rows at staging. Otherwise the model servers refuse to start, with an error naming the server, the agent and the policy server. - **Generation cuts** (`checkpoint_generation_cuts: true`). Prepare asks each inference worker to cut its in-flight calls at a durable token prefix. The replacement call carries the cut in its capture admission, and the worker continues the prefix instead of regenerating it. The wire models (`GenerationCutInventory`, `GenerationCutReceipt`, `GenerationCutPrefixAck`) keep the field layout NeMo RL already uses. A call with constrained decoding (`tool_choice` other than `auto` or `none`, a non-text `response_format`, guided or structured outputs) is never cut and never continues a cut, because a token prefix cannot restore its structured decoder. A cut gets half of the prepare time remaining, so a worker that hangs on its cut still leaves time for the stages after the model. - **Token lineage.** Ledger rows are restored under the replacement attempt's capture key, so parent resolution continues the lineage by normal fingerprint matching, with no parent header relay. - **Several uvicorn workers.** A policy model server may run `num_workers > 1`. It is still one participant to the controller: a coordinator in the main process owns the phases, lease, storage, and restored cuts, and each worker runs a gate (admission, held responses, its own cut round) that reports to the coordinator over a Unix socket. A worker that has not reported, or that disconnects during a checkpoint, blocks the checkpoint instead of weakening it. A worker can also disconnect after the commit checked readiness but before the coordinator collects the workers' snapshots; the rows of that worker's undelivered calls would then be exported as delivered. The commit therefore fails with `export_failed` if a worker was lost or fewer workers answered than expected, and that failure is final for the checkpoint, because a later export could take a later cut. With one worker nothing changes. - **Generation cuts at a per-backend address** (from NVIDIA-NeMo#3971, by Teodor-Dumitru Ene). `generation_cut_control_url` lists one control URL per `base_url`, for backends that serve cuts somewhere other than their OpenAI URL. It is refused alongside `endpoint_file`, and a backend without an address fails its cut, which only costs regenerating the call. - **One runtime directory per server process.** The coordinator socket of a server with several workers lives in `/tmp/ng-<pid>-<random>/` (`nemo_gym/runtime_dir.py`), under `/tmp` because AF_UNIX paths are limited to about 100 bytes. It used to be a socket file in the temporary directory that nothing ever removed. - The directory is removed when the process exits. - Creating one first removes directories that killed servers left behind, but only a directory whose owning process no longer exists **and** none of whose sockets accepts a connection. The second check keeps a live server in another container that shares `/tmp` but not its process namespace, and a connection that times out counts as alive. Nothing is removed by age. - NVIDIA-NeMo#3905 adds the same file for session routing's sockets, so either PR can merge first. - **Scale fixes.** Commit and restore read and write ledgers on a worker thread. A restore installs ledgers without syncing each file (the checkpoint is the durable copy, and importing again is a no-op), with one directory listing and one directory sync: importing 15,686 episodes drops from 3.2 s to 1.5 s. Readiness is cheap to recompute. An episode keeps only its latest cut. The commit reply lists `staging_keys`, the staged token-store keys the checkpoint still refers to, so the training framework knows which rows to keep. Cut requests use the reserved control connection pool, and cut failures are counted in readiness and logged. - **Restored attempts the controller no longer continues.** The model tracks every attempt a restore continues. An attempt stops being pending at its first admitted call, on any worker. With several workers, only the first reopen after an install gives the workers its attempts, so a later reopen cannot make an attempt a worker has started pending again. - An unclaimed restored cut, one no re-issued call has claimed yet, is exported again by the next checkpoint while the controller's commit scope still continues its episode, so a second crash still continues the prefix. - Once the scope leaves out an attempt that never started, the commit drops its cut and deletes its ledger without a fence (`delete_restored`), because the controller may still start the rollout over as that attempt. A started replacement is never touched, so its training rows are kept. - Several policy servers can share one run-wide capture ledger, as `policy_model` and `policy_model_reasoning_off` do, but each tracks its restored attempts on its own. When the first call of a restored attempt reaches a server, before the call can write a row or be held, that server creates a shared "started" marker next to the attempt's ledger, `<capture key>.lineage.started` (`mark_started()`). The marker is created with `O_CREAT|O_EXCL` under the rollout's lock. It is not synced to the directory: other servers only need to see it, and after a crash the next step is a restore, which deletes started markers with the other stale capture files. The commit deletes an unscoped restored attempt through `delete_unstarted()`, which checks for the marker under the same lock and keeps the ledger if any server sharing it has started the attempt. Without the marker, a server that never saw the attempt start would delete the rows another server wrote for it, including the rows of a call that is held and has no row yet. - The marker ends with the attempt's ledger: a ledger retire or delete removes it, and so does a restore that clears what a dead execution left for the attempt. A retired attempt is never marked. - The commit deletes the ledger first and only then drops the attempt from its restored bookkeeping and its saved cut, in one process and with several workers. A delete that fails therefore leaves the attempt pending, and the retried commit deletes it again. - In one process, a restored cut is kept until the call that uses it delivers its response, so a call still undelivered at the next checkpoint continues from it again if its own cut fails. - After each cut round, the newest cut replaces the restored cut a retry would use. For each capture key, the cut of the latest undelivered call, by admission time, as export chooses it, takes the restored cut's place, and that call claims it. The checkpoint exports the newer cut and keeps only its staged rows, so a call that ends before its response starts leaves the newer cut for its retry, not a restored cut whose staged rows the training framework may already have deleted. Delivering the response frees it as before. - With several workers, a worker claims a restored cut from the coordinator before the call runs, and a call that used it never gives it back. A retry of that call therefore regenerates instead of continuing a stale cut. - **Restoring the same checkpoint again.** The capture directory survives a Gym restart. If the replacement attempt makes model calls and Gym crashes before the next checkpoint, restoring the same checkpoint again finds that attempt's files. - Install receives the restore scope and clears the target attempt of every episode in it, including an episode the checkpoint holds no model record for. The import first deletes what dead executions left for the target attempt and for every later attempt of the same rollout. It uses the `delete` operations of NVIDIA-NeMo#3938 and NVIDIA-NeMo#3939: ledgers, token records, intents, incomplete markers, capture state, and **retire fences**. - A fence left in place would silently discard the restored episode's rows once the episode reached that attempt again. Attempt numbers after a restore are deterministic, and the training framework now retires every finished rollout, so fences are common. - The restore then retires the ledgers of the attempts it continues, and of earlier attempts, in one batch. Attempt N continues as attempt N+1, so attempt N's ledger would otherwise stay on disk. - The cleanup lists the ledger directory once and indexes it by rollout. Rescanning it per rollout took about 19 minutes for 16k rollouts over 96k files; the index takes about 0.23 seconds. - A model server that has served any call refuses to restore, because a ledger it wrote may belong to a live episode. With several workers, each worker reports this to the coordinator, and once any worker has left, restore is refused until the server restarts, since a fresh replacement cannot report what the old worker served. This replaces tracking every ledger the process appended to, which grew without bound. - **Worker messages carry what the checkpoint writer carries.** Messages between the coordinator and its workers are framed with the standard `json` encoder, as the checkpoint writer is. `orjson` refused integers beyond 64 bits and silently wrote NaN and infinities, such as a `-inf` logprob, as `null`. A reply that still cannot be framed is answered with a typed error, so the caller does not wait until its deadline. A message's send and reply share one timeout, and a timeout or transport error is `CoordinatorUnavailableError`, so a cut claim against a stalled coordinator falls back to regenerating. A ticket's exit reports to the coordinator even if giving back its cut fails. - **Capture-ledger cleanup stays consistent with checkpoints.** The training framework retires and deletes ledgers through the model server's control routes (NVIDIA-NeMo#3938). - While a checkpoint is open on the model server, those routes refuse with 409, and the caller retries after resume. A commit reads the ledger of every episode it continues, so a retire landing between prepare and commit would leave the checkpoint holding an episode with no lineage, with no error. - A checkpoint `retire` of attempt N cancels the calls of attempt N and every earlier attempt, each only once, and waits for them, then retires those attempts' ledgers, so their files are freed and no late row recreates them. The policy gate refuses the retired attempt's late calls until the controller forgets the rollout. A retire skips calls without a rollout prefix, such as evaluation, judge, or health-probe calls, and a malformed rollout prefix is a 400 `invalid_rollout_id`. With several workers, the coordinator marks the attempt on every worker before any of them stops its calls, relays `forget`, and gives a worker that registers later the current marks. - **Calls a restore carried over are verified against their own capture key.** A restored episode continues as attempt N+1, but the calls it made before the checkpoint stay staged under the earlier attempt's capture key. Rebuilding the episode's rollout rejected them as `wrong_rollout`, so the trainer dropped the row. Restore now stamps each carried-over ledger row with the capture key it was staged under (`CallRecord.capture_key`), and rebuild accepts the same or an earlier attempt of the rollout; anything else fails as `foreign_capture_key`. A chain of restores keeps each row's original stamp. - **Staged keys per episode.** The commit reply also returns `staging_keys_by_episode`: each committed episode's staged keys under its capture key, built from the same records as the flat `staging_keys` list. The controller keeps only the rows of episodes it continues, and a key's name does not always say whose it is: generation-cut keys name no rollout. - **`PolicyGate.deliver_response`** lets a response start, holding it while a checkpoint is open; it was `release_response`. - **A response starts with its first answer data, not its headers.** `StreamingResponse` sends `http.response.start` before it reads its body, and the buffered chat-streaming path sends SSE keep-alive comments while the backend is still generating. - The admission middleware used to call `deliver_response()` on the headers, so prepare treated a call that had produced nothing as streaming and waited for it instead of cutting it. If the snapshot fence held the call's terminal token write, the call could not finish and prepare stalled. - It now calls `deliver_response()` on the first body chunk that carries data, skipping empty chunks and, for an event stream, `SSE_KEEPALIVE`. `SSE_KEEPALIVE` moved to `nemo_gym/chat_streaming.py`, so the middleware and the streaming path share one definition of a liveness frame. - So headers and keep-alives leave the call cuttable, its first answer chunk is held while admission is closed, and an answer that already started still drains before prepare completes. This is the change proposed in NVIDIA-NeMo#4555. - **Generation cut wire names match.** The cut inventory and acknowledgement name the attempt `attempt`, as `EpisodeId` does, and the continuation a re-issued call carries names its prefix `prefix_token_count` and `prefix_digest`, as the worker's acknowledgement does. Conflict note: `CaptureContext` on `main` gained `external_worker_response_seen` (NVIDIA-NeMo#3670); this PR adds `admission_hook` next to it, and both fields are kept. ## 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 M 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 A model call across a checkpoint and a crash. ```mermaid sequenceDiagram participant A as Agent participant P as Policy model server participant W as Inference worker A->>P: model call for rollout r, attempt 0 P->>W: generate Note over P: prepare closes admission - new calls park, responses not yet delivered are held P->>W: generation cut, if enabled - stage the prefix so far W-->>P: receipt with staging keys per call Note over P: commit exports ledger rows minus undelivered calls, plus the latest cut of each episode Note over A,W: crash, restore, resume A->>P: the same call again, as attempt 1 P->>W: continue from the staged prefix, or generate again ``` With several uvicorn workers, one coordinator in the main process is the participant; every worker keeps its own gate and forwards control calls to the coordinator. ```mermaid flowchart LR CO["Coordination"] -->|"control call on the shared port"| W1["uvicorn worker 1 gate: tickets, held responses"] CO -->|"control call on the shared port"| W2["uvicorn worker 2 gate: tickets, held responses"] W1 <-->|"Unix socket: forward, close, reports, cut claims"| C["Coordinator in the main process phases, lease, ledger, restored cuts"] W2 <-->|"Unix socket"| C ``` The capture ledger across checkpoints, restores and the training framework's cleanup. ```mermaid sequenceDiagram participant T as Training framework participant P as Policy model server participant L as Capture ledger Note over P: checkpoint open - ledger retire and delete routes answer 409 T->>P: retry retire and delete after resume Note over P,L: restore - delete what dead executions left for attempt N+1 and later, fences included P->>L: import attempt N's rows as attempt N+1, then retire attempts 0 to N Note over P,L: checkpoint retire of attempt N - stop its calls, then retire the ledgers of attempts 0 to N ``` ## 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 has merged too, so NVIDIA-NeMo#3883'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 **(this PR)** 3. NVIDIA-NeMo#3884 (`ananthsub/partial-ckpt-model-worker-cuts`): feat(token-capture): add worker staging helpers for generation cuts 4. NVIDIA-NeMo#3885 (`ananthsub/partial-ckpt-environment`): feat(checkpoint): continue environment server episodes from their boundaries 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 five training verifiers stateless with replayable verification 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) - Supersedes NVIDIA-NeMo#2941 (policy model participant). The checkpointable ledger wrapper and parent relay are replaced by `main`'s capture ledger plus rows re-imported under the next attempt; admission no longer waits for generations. - Takes the generation-cut models from NVIDIA-NeMo#2940, under `nemo_gym._checkpoint.generation_cut` instead of `model_control_contracts`. - Multi-worker support follows the shape of NVIDIA-NeMo#3563 (a companion in the main process, a Unix socket, and a restored cut with a single owner) with much less machinery: no proof files, evidence files, or worker journals. The scale fixes cover the remaining needs listed in NVIDIA-NeMo#3755 (reserved control pool, restore I/O off the event loop). Constrained-decoding restart comes from NVIDIA-NeMo#3412. ## 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`, which includes NVIDIA-NeMo#3882: - `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. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider responses_api_models/vllm_model/tests`: passed. - Among the checkpoint tests, each failing without its change: - restoring the same checkpoint again after the replacement made a call (without the fix: `capture ledger for roll-1-a1 already holds rows from another execution`), with the files and fences of the target and later attempts deleted and look-alike rollouts untouched; - a restore retires the ledgers of the attempts it continues; - a model server that has served calls refuses to restore, with one process and with two workers; - ledger retire and delete return 409 while a checkpoint is open and succeed after resume; - a retire stops a running generation before it replies, and retires the ledgers of attempts 0 to N but not N+1; - a retire leaves no fence on any worker, including one that joins later; - worker messages carry NaN, infinities, and integers beyond 64 bits unchanged; - a commit that no longer names a started replacement leaves its ledger alone, and deletes a restored attempt nothing started without fencing it; with two workers, an attempt any worker started is no longer pending, and a reopen does not make it pending again; - restoring again clears a dead execution of an in-scope episode without a model record; - untagged calls neither break a retire nor surface as untyped errors; - a hung cut leaves the rest of the prepare deadline for the stages after it; - a message that cannot be sent or is never answered is a typed unavailable error, and a ticket leaving is reported even if giving back its cut fails; - restore is refused after a worker left, even if its replacement is fresh; - a call that arrives during a checkpoint waits for resume without being admitted, and is refused at once if its attempt is retired while it waits; - two policy servers sharing a ledger, one holding a completed response: neither snapshot, nor the restored ledger, holds its row (without the fix, the second server exported it); - a call whose capture context attaches after the commit's snapshot cannot land its row in the checkpoint (without the fix, it did); - a row is published at the first answer chunk, not at headers or keep-alives; a call that leaves first has no row and is not poisoned; a failed write poisons the call and sends no answer; a retire or a checkpoint during the write waits for it; each worker of a server with several workers follows the same rule; - a model server outside the checkpoint gate that an agent with token capture names is refused at startup; - a retried retire waits for a call an earlier retire cancelled without cancelling it again (without the fix, its cleanup was cancelled a second time); - marking a restored attempt started does not sync the ledger directory, and a restore after a crash removes the markers the dead execution left; - headers, empty bodies and keep-alives do not start a response, for JSON and event streams; a stream sending only keep-alives during a checkpoint is cut and its output is held until resume (without the fix, both time out); an event stream that already started finishes while admission is closed; - a worker that disconnects after the commit checked readiness and before the snapshot fails the commit, and a retried commit fails the same way; - a restored attempt started on one policy server keeps its ledger when a commit on another server sharing the ledger leaves it out of scope, both once the started call's row is recorded and while the call is admitted but not yet recorded; with several workers, a worker marks the attempt started in the shared ledger; - a restored attempt no server started is deleted, and no marker is left behind; - the started marker is removed by a ledger retire, a ledger delete, and a restore that clears the attempt; - a failed delete of a restored ledger keeps the attempt pending, and the next commit attempt deletes the remaining rows, in one process and with several workers; - a call waits at closed admission, its attempt is marked and retired, the controller forgets the rollout, and admission reopens; the call is refused with `stale_attempt` and nothing is served; - the same with a retire cut short by its deadline, in one process and on a worker of a server with several workers, where the waiting call also leaves the gate's tracking; - a resumed call that continues a restored cut is cut again by the next checkpoint, which exports the newer cut, and then ends before its response starts; its retry continues the newer cut, not the restored one; - with several workers, a retry of a call that used its restored cut gets no cut after a newer checkpoint, and the coordinator no longer holds the restored cut. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_runtime_dir.py`: passed. A dead server's directory with closed sockets is removed; a live process's directory is kept; a directory whose socket accepts is kept even when its process looks dead; only runtime directory names are considered; and a process creates its directory once, with mode 0700, and removes it at exit. - 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 - **Real model:** see the rollout evidence in NVIDIA-NeMo#3889, whose opt-in tests drove this PR's code with `Qwen/Qwen3-4B-Instruct-2507` on vLLM: saved rollouts continued after a kill and restore, restarts started over, and every rollout finished with reward 1.0. A real vLLM worker serving generation cuts is not covered: stock vLLM does not serve the cut route. - Scale, on one workstation with every Gym server and a fake inference backend (weather rollouts through the typed `single_agent_turn` request path, token capture and generation cuts on, each rollout's second model call held so most rollouts are in flight at the checkpoint; checkpoint, kill every Gym process group, restart, restore, run the replacement attempts). After the scale fixes in this PR, measured on the branch this stack was cut from, with the same commits: | Rollouts | Policy workers | In flight at checkpoint (exported) | Prepare (s) | Commit (s) | Restore (s) | Replacement failures | |---|---|---|---|---|---|---| | 8,000 | 1 | 4,201 | 0.67 | 0.60 | 0.60 | 0 | | 8,000 | 2 | 4,083 | 0.56 | 0.69 | 0.56 | 0 | | 16,000 | 2 | 14,320 | 7.57 (model 7.38) | 4.64 | 4.97 | 0 | | 16,000 | 2 (rerun) | 12,228 | 3.60 (model 3.53) | 1.83 | 3.23 | 0 | In the first 16,000-rollout run, the policy model logged 8 uvloop "File descriptor N is used by transport" errors and 16 requests failed with "ledger lineage resolution failed". The cause is a libuv accept race while clients disconnect, not checkpoint code; the ledger failed closed, clients retried, and no replacement failed. The rerun and every 8,000-rollout run had none. ## Compatibility - New model server config: `checkpoint_policy` (default `false`) makes the server a participant, and `checkpoint_generation_cuts` (default `false`) turns on generation cuts. Both only take effect with the global `checkpoint:` block enabled. - Users of the old stack's `token_id_capture.generation_prefix_cuts_enabled` key (never on `main`) switch to `checkpoint_generation_cuts`. - Generation-cut models are imported from `nemo_gym._checkpoint.generation_cut`. - The `CheckpointableLedger` protocol gains `mark_started()` and `delete_unstarted()`, so a custom ledger used with a checkpoint participant must implement both. `PolicyWorkerLink` takes a `ledger=` argument, which a worker needs to mark restored attempts started. - A restored attempt's first call writes a `<capture key>.lineage.started` file in the capture ledger directory. It is removed with the attempt's ledger. - With checkpointing off, model servers behave as before. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com> Co-authored-by: Anish Mahishi <amahishi@nvidia.com> Co-authored-by: Teodor-Dumitru Ene <teodord.ene@gmail.com>

…VIDIA-NeMo#3884) ## 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 NVIDIA-NeMo#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. ```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 W 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 What an inference worker does with these helpers when the policy model asks for a cut, and when a replacement call continues it. ```mermaid 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 ``` ## 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 and NVIDIA-NeMo#3883 have merged too, so NVIDIA-NeMo#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 (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 **(this PR)** 4. NVIDIA-NeMo#3885 (`ananthsub/partial-ckpt-environment`): feat(checkpoint): continue environment server episodes from their boundaries 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 five training verifiers stateless with replayable verification 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 had these helpers. They are ported from NVIDIA-NeMo#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 (NVIDIA-NeMo#2939 to NVIDIA-NeMo#2946) carry the original discussion. ## Validation Run on this branch, on the current `main`, which includes NVIDIA-NeMo#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 NVIDIA-NeMo#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. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

…ndaries (NVIDIA-NeMo#3885) ## What changed and why Every environment server takes part in checkpoints once checkpointing is on. The environment server owns the episode protocol, so it records where each episode can continue. - An episode records a boundary after each completed step and runs each step in `wait` or `replay` mode. Prepare closes episode admission and is ready once every live episode is parked at a boundary or inside a replay step. - A restored boundary is handed to the replacement attempt's `/run`, which continues from it. The task must match the one recorded. - `single_agent_turn` runs seed, invoke agent, close agent, verify, and return as explicit stages with typed session handles. Seeding and the agent invocation replay; closing the agent waits; verification uses the mode the resources server reported on its seed reply, kept in the session handles so a restored episode still knows it. Once the agent's activation returns, the episode records a `close_agent` boundary carrying the agent's response, so a checkpoint that closes then parks the episode there instead of waiting on the close step, and a restore at that boundary closes the agent and verifies without invoking the agent again. - **Every stage starts with its boundary, including `close_agent` and `return`.** After a restore and resume, a later checkpoint can prepare before the replacement `/run` arrives, and the replacement is then admitted while admission is closed. A replacement restored at `close_agent` used to start closing the agent before it recorded a boundary, which made a prepared participant not ready; one restored at `return` ran final cleanup and called `end()`, so an episode could finish while Gym was prepared. Both boundaries now sit at the start of their stage, as the others do, so a replacement at any stage parks before doing any work and a commit exports the boundary it was restored from. For a first attempt each boundary is recorded at the same point, with the same state, as before. - **A failed agent activation closes the agent session at once** to keep its observations, as `main` does since NVIDIA-NeMo#4562. The close runs after the failed replay step, which parks with no boundary while a checkpoint is open, so it waits for resume and never runs while Gym is prepared. - A boundary's state can be given as a function, which is called only when checkpointing is on, so `single_agent_turn` builds no boundary payloads otherwise. - The `single_agent_turn_legacy` adapter checkpoints too, since its rows run the same protocol. - An unclaimed restored episode, one whose replacement has not started yet, is exported again by the next checkpoint while the controller's commit scope still continues it, so a second crash still continues it. Once the scope leaves it out, it is not exported, and the restored record is retired. - A replacement `/run` is refused with `admission_closed` until the controller resumes after the restore, and the restored record stays until a replacement actually starts. A retire of a replacement that had already started would wait for its final cleanup, whose calls wait for that same resume. - A retired attempt's late `/run`, such as a client retry, is refused with `stale_attempt` until the controller forgets the rollout. An episode with no boundary yet is exported with a null boundary and starts over from its input after a crash. - A `/run` of an attempt that is already running is refused, and the refused request owns nothing: it neither ends nor cleans up the live episode. An episode is untracked even if its final cleanup is cancelled. - A replay step that fails while a checkpoint is open clears the episode's boundary instead of blocking the checkpoint, so the checkpoint commits and the episode starts over from its input after a crash. The episode then parks at that cleared boundary until resume before the failure propagates. It stays out of the blockers and is exported with no boundary, but its final cleanup, its `end()`, and its `/run` reply wait for resume; otherwise it could finish and drop its record before the commit, and a commit that names it would export nothing for it. A retire still stops a parked episode. With no checkpoint open, the episode goes back to running as before. - Final cleanup, which closes the episode's agent and resources sessions, runs before `/run` replies on every path, and a retire never interrupts it. A retire cancels an episode only before cleanup starts, then waits until the episode has ended, so its reply means the sessions are closed. This matters because anyio's shielded scope does not stop a direct `task.cancel()`. A retire cut short by its deadline keeps the episode and its task digest until the episode has ended, so the next prepare reports it as a blocker instead of exporting it. - **Restarts.** An environment server that does not record boundaries (`checkpoint_boundaries = False`, the default) cannot continue its episodes, so each one is a restart: it never blocks prepare, is never exported, and is reported in prepare's `restarts`. `single_agent_turn` sets `checkpoint_boundaries = True`, and marks an episode as a restart when its resources or agent seed reply carries `x-ng-checkpoint-restart`, because a server it depends on cannot capture its part. - A commit exports only the live episodes inside its scope, as it already did for restored episodes. The controller leaves every participant's restarts out of the scope, so a live episode that another server reports as a restart must not be exported here either. - The participant answers the readiness wait's `ready()` from its running blocker count, so a prepare does not list every episode after each change. - Checkpointing refuses `num_workers > 1` on environment servers, since episodes live in one process. ## 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 E 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 `single_agent_turn` records a boundary at the start of every stage. A checkpoint exports the latest boundary, and a restored attempt starts at its `next` step, parking there first if a checkpoint is open. ```mermaid flowchart LR B1(["next: seed"]) --> S["seed the resources session"] S --> B2(["next: invoke_agent"]) --> I["invoke the agent which parks its own turn loop"] I --> BC(["next: close_agent"]) --> CA["close the agent session"] CA --> B3(["next: verify"]) --> V["verify wait or replay, per the resources server"] V --> B4(["next: return"]) --> R["reply to /run"] ``` Final cleanup closes the episode's sessions before `/run` replies. A retire never interrupts it: it cancels an episode only before cleanup starts, and then waits for the episode to end. ```mermaid flowchart TD X{"retire arrives"} -->|"episode still in its protocol"| C["cancel the episode"] --> F X -->|"final cleanup already started"| F["wait"] F --> CL["final cleanup closes the agent and resources sessions"] --> D["retire replies"] ``` ## 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 and NVIDIA-NeMo#3883 have merged too, so NVIDIA-NeMo#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 (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 4. NVIDIA-NeMo#3885 (`ananthsub/partial-ckpt-environment`): feat(checkpoint): continue environment server episodes from their boundaries **(this PR)** 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 five training verifiers stateless with replayable verification 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) New compared with the old stack, which had no environment servers. It takes over the protocol side of NVIDIA-NeMo#2942 (where an episode can continue), which the old stack reconstructed from agent and resources state after the fact. ## 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`: - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_checkpoint_*.py`: passed, including, each failing without its change: - a retire during an episode's final cleanup does not interrupt it; - a retire replies only after the episode and its cleanup have ended; - a commit whose scope leaves out a restored episode retires it; - a refused duplicate `/run` leaves the live episode tracked, and an episode whose cleanup is cancelled is still untracked; - a replay step that fails during a checkpoint lets the checkpoint commit with the episode exported with no boundary, the episode has not finished or replied before resume even with its cleanup running at once, and it starts over after a crash; - a retire stops an episode parked after its replay step raised; - a commit does not export live episodes outside its scope; - an agent that returns during a checkpoint parks at `close_agent` and stays ready, and a restore at `close_agent` closes and verifies without invoking the agent; - boundary payloads are not built with checkpointing off; - a replacement restored at `close_agent` or at `return` and admitted during a second checkpoint parks before its stage: readiness stays true, no close request and no `end()` happen before resume, and the commit exports the same boundary (without the fix, the agent close or the final resources close started and readiness was false); - a retire stops an episode parked after its replay step raised, waiting on the step and on the `[None]` boundary rather than on the task ending (the previous version passed without reaching the retire, and the new one fails if the failed step does not park or the retire skips the episode); - an agent failure during a checkpoint closes the agent session only after resume, and the failure carries the agent's observations; It also checks that an episode whose retire was cut short blocks the next prepare, is not committed, and is freed by a retried retire after resume, and that readiness is answered from the blocker count. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_environment_server.py`: passed. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider environment_servers/single_agent_turn/tests`: passed. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider environment_servers/single_agent_turn_legacy/tests`: passed. - `responses_api_agents/codex_agent/tests`, `responses_api_agents/openclaw_agent/tests`, `responses_api_agents/opencode_agent/tests`, and `responses_api_agents/pi_agent/tests`, each run in its own server environment, created from that server's `requirements.txt` as `gym env test` does: passed, including the integration tests that collect rows through an environment server's `/run`. Their fake responses gain an empty `headers` mapping, because the environment server now reads the seed and agent response headers. One Pi outcome test, `test_outcome_extension_uses_pinned_pi_context_classifier`, needs a newer Node.js than the development machine's Node.js 20, and fails the same way on `main`. - 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 - Process-level evidence is in the top PR (an episode on the typed `single_agent_turn` request path continues from its last boundary; second crash before the replacement starts). - **Real model:** see the rollout evidence in NVIDIA-NeMo#3889, whose opt-in tests drove this PR's code with `Qwen/Qwen3-4B-Instruct-2507` on vLLM: saved rollouts continued after a kill and restore, restarts started over, and every rollout finished with reward 1.0. ## Compatibility - Nothing changes unless the global `checkpoint:` block sets `enabled: true`. - With checkpointing on, an environment server with `num_workers > 1` fails at startup. - A commit exports only the live episodes inside its `episode_ids` scope. A commit request without `episode_ids` still exports every live episode. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

…ous session hooks (NVIDIA-NeMo#3886) ## What changed and why Every resources server takes part in checkpoints once checkpointing is on, in one of three modes, declared with the `checkpoint_mode` class attribute: - `restart_only` (the default): nothing of a session is saved, so its rollout is a restart. It never blocks prepare, keeps running through the checkpoint, and is reported in prepare's `restarts`; the seed reply carries `x-ng-checkpoint-restart` so the environment server knows. After a crash the controller retires the rollout and starts it over from its input. Only a new session start waits while a checkpoint is open, so the restarts a prepare reported cannot change before its commit. A seed already in flight when admission closes is reported as a restart too, without waiting for it, because its reply will say restart; otherwise it would register during the checkpoint and the commit would fail with `restart_in_scope`. A waiting seed checks again each time it wakes, so a new checkpoint that closes admission before it runs keeps it waiting. - `stateless`: tool results depend only on the request, so nothing is exported. The weather example server uses it. - `exported`: the server implements the session hooks. The counter example server (`example_session_state_mgmt`) uses it. It exports only the sessions it still holds, so a dropped session is left out and the participant stops tracking it, instead of being exported as a new counter at zero. The session hooks are coroutines, and export takes every session at once: `export_session_states(session_ids)`, `restore_session_states(states)`, and `retire_session_state(session_id)`. A server whose session state lives outside the process, such as a sandbox per session, can snapshot all of it concurrently within the commit. A session missing from the export result is treated as already dropped (for example, after a failed verification) and is no longer tracked, instead of failing the commit. With `checkpoint_mode = "exported"`, a server that leaves any of these hooks as the base stub fails at startup; hooks from an intermediate base class or mixin count. `retire_session_state` is called in every mode, and its base version does nothing, so a `restart_only` or `stateless` server that keeps per-session state can free it on retire by overriding it. `checkpoint_mode` and `checkpoint_verify` are checked against their allowed values at startup, whether or not checkpointing is on. A typo such as `"exportd"` raises a `ValueError` naming the attribute, the value, and the allowed values, instead of acting as a fourth mode that exports nothing and reports no restarts. The session lifecycle comes from the server's session-start routes, `/close_session`, and `/verify`, keyed by the rollout attempt, so no request body is inspected. - **Declared session starts.** A server lists its session-start routes in `checkpoint_session_starts`, which defaults to `("/seed_session",)`; a Gymnasium-style server declares `/reset`. The checkpoint middleware treats every declared route as it treats `/seed_session`, before the handler runs: it waits out an open checkpoint in every mode, `restart_only` included, counts as a seed from admission, carries the seed reply headers, and binds the session when it succeeds. There is no other way to start a tracked session, so no start can slip in after a commit has fixed its scope. Sessions end through `/close_session`, or `checkpoint_session_ended` for a server that ends them elsewhere. - **A session start needs a rollout ID.** With checkpointing on, a session start without the `/ng-rollout/<id>` prefix is refused with 400 `rollout_id_required` before its handler runs, so no session can exist that the participant does not track. - **A resources server that bypasses the shared setup fails closed.** `SimpleServer.setup_checkpoint_participant(app)` runs for every server after `setup_webserver`. With checkpointing on, a resources server whose app did not install the checkpoint participant refuses to start, naming its class, instead of being silently left out of every checkpoint. Gymnasium servers on `main` build their own app, so they refuse to start with checkpointing on until NVIDIA-NeMo#3895 moves them onto the shared setup; with checkpointing off nothing changes. - **A session belongs to its attempt.** Any request on an existing session whose rollout prefix names an episode other than the session's is refused with `stale_attempt`, seeds excepted, so after a restore a replaced attempt's `/close_session` or `/verify` cannot end the replacement's session. A close from the owning attempt, including a retired attempt freeing its own session, still goes through. With checkpointing on, `/close_session` and `/verify` need the `/ng-rollout/<id>` prefix, as session starts do. - **A restored session needs its attempt named.** A request on a session a restore installed that names no attempt, with no rollout prefix and no MCP token, is refused with 400 `rollout_id_required`. The cookie alone cannot tell the replacement from the attempt the restore replaced. The refusal comes before the participant takes the session out of the restored sessions no request has used yet, so a refused call never claims one. Sessions seeded in the running process are not affected, and neither are restart-only servers, which never restore a session. - **Marking an attempt retired fences its sessions at once.** A retire or a restore that marks an attempt refuses the attempt's existing sessions before anything awaits, so a retire cut short by its deadline, or one that arrives after it and only marks, still leaves them refused, holding up readiness, and out of every export until a retried retire frees them. `forget` is refused while any tracked session of the rollout is at or below the retired mark. - **MCP tokens name their attempt.** With checkpointing on, the session token minted at seed carries an `episode` claim with the seeding attempt's capture key; with checkpointing off the token is unchanged. A call whose token names a fenced attempt, an attempt other than the session's, or no attempt gets 409 `stale_attempt` before its tool runs, and never claims a restored session. A restore therefore no longer lets the replaced attempt's harness keep using its token: the replacement gets a token by seeding under its own attempt with the session's cookie. `ResourcesParticipant.is_open(session_id)` reports whether a session is tracked and not fenced, for revocation at session close. - **Calls on closed or unknown sessions are out of scope.** A tool call that creates or changes state on a session that is closed or was never seeded is a session-lifecycle bug with checkpointing off as well, so it will be fixed on `main` for MCP and direct HTTP alike. - **Restarts reported by earlier prepare stages carry forward.** Coordination prepares stages in order (environment, model, agent, resources) and sends each later stage the restarts every earlier stage reported, as capture keys in a new optional `restarts` field of the prepare request. The commit leaves those episodes out, so a participant keeps them in `carried_restarts` until the checkpoint resumes, aborts or expires, and leaves them out of its blockers and its export. The resources participant no longer waits for requests in flight on their sessions, so an episode the environment server reported as a restart, for example from an MCP grant, keeps running through the checkpoint. Coordination sends the field only when an earlier stage reported a restart, and a prepare without it behaves as before. - **Commit cleanup after a restore costs as much as each rollout's own sessions**, through the rollout index, instead of scanning every restored session per episode. The `checkpoint_verify` class attribute declares whether `/verify` may run again after a crash (`replay`) or must finish before a checkpoint (`wait`, the default). The server reports it in the `x-ng-checkpoint-verify` header on every `/seed_session` reply and in its checkpoint status. A replayable `/verify` is admitted even after admission closes, because the checkpoint does not wait for the episode step that sends it. A `/seed_session` that arrives while a checkpoint is open waits for resume instead, and a seed already in flight holds up prepare and is exported. If the seed handler replaces the session ID, the session is tracked under the new one. The participant applies a seed, close, or verify to its session tracking as soon as the response starts, before the rest of the response is sent. A client that disconnects after the handler ran cancels the request, as an environment server's retire does to its calls in flight, and that cancellation can no longer leave a created session untracked or an ended one tracked. The environment server seeds inside a replay step that can finish while the checkpoint is open, so the session must match whichever boundary the environment server records. **Tools called over MCP** (`expose_tools_over_mcp: true`). A CLI agent harness calls tools over MCP with a signed session token and no session cookie. Before this change, such a server failed at startup with checkpointing on: MCP auto-exposure refuses middleware that direct MCP dispatch would skip, and the checkpoint middleware was one of them. - The checkpoint middleware works on the whole MCP request, so MCP auto-exposure accepts it. Admission and in-flight counting apply to the one POST that carries each tool call. - The middleware reads the session from the MCP session token, so a retire drains and refuses MCP calls like any other request on the session. The token key is deterministic, and the token's `episode` claim decides whether its attempt may still use it after a restore. - A request made while a checkpoint is open, over HTTP or MCP, waits for resume instead of getting a 409. It comes from a restart, which keeps running through a checkpoint, often a third-party harness that would pass the refusal to the model as a tool error. A waiting request is not in flight, so it does not hold up prepare. A replayable `/verify` is still admitted at once. - Checkpoint control routes are no longer harvested as MCP tools. **Ending sessions.** - `/close_session` is never refused. Closing is how a server frees a session's state. A close that arrives while a checkpoint is open, from an episode that ended during it, waits for resume as an MCP call does; a refusal would leave the session unfreed, since the environment server closes once. - A retire refuses the attempt's sessions, waits until no request is in flight on them (over HTTP or MCP), then frees them through `retire_session_state`, in every mode. A session stays tracked until its state is freed: a retire cut short by its deadline leaves it for a retry, and meanwhile it blocks the next checkpoint and is not exported. `retire_session_state` can therefore run again for the same session, and must tolerate state it already freed. The refusal lasts until the controller forgets the rollout: - a request that carries the rollout prefix is checked against the retired attempts, so a late `/seed_session` cannot create state nothing frees; - a tool call over MCP names only its session, so the retired session IDs are kept by rollout until `forget`, and counted in status as `retired_sessions`. - Every mode, `stateless` and `restart_only` included, records which attempt each session belongs to, so a retire drains and refuses their sessions too. Export and readiness still depend on the mode. - A retire waits for any `/seed_session` of the same rollout, up to the retired attempt, that is already admitted, then frees the session that seed created. A seed is counted at its admission, with no await in between, so a seed still waiting for resume cannot hold up the retire, and the retire's mark refuses it once it is admitted. Retire visits only the rollout's own sessions, through an index from rollout to sessions. - A seed can finish after its attempt was marked retired, for example when a retire's deadline passed while it waited for that seed. The session the seed created is then registered as retired at once instead of live: requests on it are refused, including MCP calls that carry only the session token, it is never exported, and it blocks checkpoints until a retried retire frees it. - `forget` is refused with `retire_incomplete`, changing nothing, while a retired session it targets is still tracked, or while a seed of a retired attempt of that rollout is still in flight; the controller retires again first. Forgetting during such a seed would clear the mark before the seed finished, and its session would then be bound as live. - A replayable `/verify` whose caller closes the session afterwards does not end it. The caller says so with the `x-ng-checkpoint-close-follows: 1` header, and the session stays tracked and exported until the `/close_session` that follows, a declared end, or a retire. A checkpoint can keep the environment server's boundary from before the verification even when the verification finishes before the resources server prepares, so the session must still be exported, and a restore runs the verification again against the restored state. `single_agent_turn`, the UserSim environment server, and the sandboxed Hermes agent send the header, since each closes the session in its final cleanup; Simple Agent's legacy `/run` sends it in NVIDIA-NeMo#3888. A `/verify` without the header, and a wait-mode `/verify`, end the session at once, so a caller that never closes cannot leave it tracked. - Unclaimed restored sessions, ones no request has used yet, are reported as pending, so a commit whose scope no longer continues their episode does not export them. It deletes them through `retire_session_state` without refusing anything later, since the controller may still start the rollout over as that attempt. A restored session that the server drops before any request uses it does not break the commit. - A commit exports only the live sessions whose episode is inside its scope, as it already did for restored sessions. The controller leaves every participant's restarts out of the scope, so a session another server reports as a restart is not exported here either. - The participant answers the readiness wait's `ready()` without building its blocker list, and admission refusals are the shared `AdmissionClosedError`. The MCP session token helpers live in `base_resources_server.py` and are shared with MCP auto-exposure. ## 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 How a resources server takes part, chosen by a class attribute. ```mermaid flowchart TD M{"checkpoint_mode"} -->|stateless| S["Nothing to export. Never blocks a checkpoint."] M -->|exported| X["export_session_states at commit restore_session_states after a crash retire_session_state on retire"] M -->|"restart_only, the default"| O["Its rollout is a restart: reported in prepare, never blocks, starts over after a crash"] ``` A session's lifecycle, as the participant learns it from the routes that already define it. ```mermaid sequenceDiagram participant E as Environment server or agent participant R as Resources server E->>R: /seed_session for rollout r R-->>E: session cookie, plus x-ng-checkpoint-verify - wait or replay E->>R: tool calls - wait while a checkpoint is open E->>R: /verify, with x-ng-checkpoint-close-follows if replayable E->>R: /close_session from final cleanup, which waits for resume if a checkpoint is open ``` Retire drains a session before freeing it. ```mermaid sequenceDiagram participant C as Controller participant R as Resources server C->>R: retire(r, N) Note over R: refuse the attempt's sessions, except /close_session Note over R: wait until no request is in flight on them, over HTTP or MCP Note over R: retire_session_state for exported servers, stop tracking R-->>C: stopped and freed ``` ## 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 **(this PR)** 6. NVIDIA-NeMo#3887 (`ananthsub/partial-ckpt-verifier-declarations`): feat(checkpoint): declare twenty training verifiers stateless with replayable verification 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) - Supersedes NVIDIA-NeMo#2945 (resources state participant). Route classification, mutation receipts, revision reconciliation, and the terminal context variable are gone; the prepare order (agents park before resources servers close) makes them unnecessary. - Takes the counter adapter from NVIDIA-NeMo#2946. The Blackjack and Workplace Assistant adapters are not carried: Blackjack moves to a separate Gymnasium PR, and Workplace Assistant can be ported onto these hooks later. ## 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`: - `test_each_prepare_stage_carries_the_restarts_every_earlier_stage_reported` and `test_a_request_in_flight_for_a_restart_an_earlier_stage_reported_does_not_hold_up_prepare`, each failing without its change: later stages get the earlier stages' restarts, the restarts are cleared at resume, and a resources request in flight for such an episode no longer holds up prepare. - This round's tests, each failing without its change: `test_a_retire_cut_short_while_a_seed_is_in_flight_still_fences_the_attempt_s_existing_session`, `test_a_retire_that_arrives_after_its_deadline_still_fences_the_attempt_s_sessions`, `test_after_a_restore_a_close_from_the_replaced_attempt_is_refused_and_keeps_the_session`, `test_a_request_on_a_session_from_another_rollout_is_refused`, `test_a_session_end_without_a_rollout_prefix_is_refused_before_its_handler_runs`, `test_deleting_one_rollout_s_restored_sessions_does_not_scan_every_restored_session`, `test_a_restore_refuses_the_replaced_attempt_s_token_and_a_seed_of_the_replacement_gets_one_that_works`, `test_a_stale_harness_cannot_create_state_for_a_session_the_fresh_process_never_saw`, `test_a_token_without_an_attempt_is_refused_when_checkpointing_is_on`, `test_a_second_restore_refuses_the_first_replacement_s_token`, `test_a_refused_stale_call_does_not_claim_a_restored_session`, `test_a_token_minted_with_checkpointing_names_the_attempt_that_seeded_it`, and the `is_open` tests; `test_a_token_minted_without_checkpointing_names_only_its_session_and_tools` guards the unchanged token with checkpointing off. - `test_a_declared_session_start_during_a_commit_s_export_waits_for_resume`, in all three modes: a declared `/reset` arriving while the commit is in export waits for resume, and the commit has no record for it (without the fix, the `restart_only` reset finished during the checkpoint); `test_a_restart_only_declared_session_start_in_flight_is_reported_as_a_restart` (without the fix, prepare reported no restart). - `test_a_session_start_without_a_rollout_prefix_is_refused_before_its_handler_runs`, for `/seed_session` and `/reset` (without the fix, a 200 that created untracked state). - `test_a_resources_server_whose_app_left_out_the_checkpoint_participant_refuses_to_start` and `test_a_resources_server_whose_app_installs_the_participant_passes_the_startup_check`, including a server with its own app starting with checkpointing off. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_checkpoint_*.py`: passed, including, each failing without its change: - a `/close_session` during a checkpoint waits for resume, then ends the session; - a retire waits for the session's request in flight, and refuses a new one meanwhile; - a commit whose scope leaves out a restored session retires it through the server's hook; - a session whose retire was cut short in the server's hook stays until a retry frees it, and blocks prepare meanwhile; - a seed arriving during an open checkpoint waits for resume, and a seed in flight when admission closes holds up prepare and is exported; - a retire waits for an admitted seed of its attempt and frees the session it creates, and a seed still waiting out a restore does not hold up the retire; - `forget` is refused while a retire has not freed its sessions; - a replayable `/verify` whose caller closes afterwards keeps the session tracked and exported until the close, including one that finishes between environment and resources prepare, after which a restore verifies again and closes it; - a replayable `/verify` without the close-follows header, and a wait-mode `/verify` with or without it, end the session at once; - after a restore, a call under the replaced attempt's prefix is refused with `stale_attempt` and a cookie-only call with `rollout_id_required`, across two restores, and neither changes the restored state or claims the session; - a stateless server drains and refuses a retired session; - deleting an unclaimed restored session refuses nothing later, and a restored session the server drops does not break the commit; - a request cancelled while notifying is no longer counted; - a session whose ID the seed replaces is tracked under its new ID; - an exported server without its hooks fails at startup; - retiring one rollout does not scan every session, and `ready()` does not list blockers; - in `restart_only` mode, a seed in flight when admission closes is reported as a restart by prepare, its reply says restart, and the commit of the remaining scope succeeds; - in `restart_only` mode, a seed waiting for resume waits again when a new checkpoint closes admission before it runs; - a seed, close, or verify whose reply has started still updates the session tracking when the request is cancelled before the reply finishes: the seeded session is tracked, and the closed or verified one is not; - an unknown `checkpoint_mode` or `checkpoint_verify` fails at startup; - a retire frees the sessions of `restart_only` and `stateless` servers through `retire_session_state`; - a commit does not export live sessions outside its scope; - a seed that finishes after its retire was cut short while waiting for it leaves a retired session: a request with only the session cookie is refused with `stale_attempt`, a checkpoint does not export it, `forget` is refused with `retire_incomplete`, and a retried retire frees it, after which `forget` succeeds; - `forget` is refused with `retire_incomplete` while a seed of a retired attempt is in flight, and the session that seed creates is refused afterwards. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_mcp_auto_exposure.py tests/unit_tests/test_checkpoint_resources_mcp.py`: passed. The MCP checkpoint tests go through the real control routes and MCP JSON-RPC: - a tool call during a checkpoint waits and runs after resume, and the checkpoint holds the state from before it; - a tool call in flight holds up prepare; - a retire waits for a tool call in flight over MCP; - a restored server refuses the replaced attempt's token, and a seed under the replacement gets one that continues the session; - control routes are not tools. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider resources_servers/example_session_state_mgmt/tests`: passed, including a new test that the counter example leaves a session it no longer holds out of its export. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider resources_servers/example_single_tool_call/tests`: passed. - `environment_servers/single_agent_turn/tests`, `environment_servers/nemo_user_sim/tests` (in its own environment, with the UserSim package) and `responses_api_agents/hermes_sandboxed_agent/tests`: passed, including that `/verify` says the session is closed afterwards. - The Codex, OpenClaw, OpenCode, and Pi agent tests, each in its own environment: passed, except one Pi test that needs a newer Node.js than this machine has and fails the same way on `main`. - 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 - Process-level evidence is in the top PR: restored counter state is required for the right reward, with a negative control that skips the resources restore and gets reward 0.0. - **Real model:** see the rollout evidence in NVIDIA-NeMo#3889, whose opt-in tests drove this PR's code with `Qwen/Qwen3-4B-Instruct-2507` on vLLM: saved rollouts continued after a kill and restore, restarts started over, and every rollout finished with reward 1.0. Its restart test uses a `restart_only` resources server. ## Compatibility - New resources server hooks and class attributes: `checkpoint_mode`, `checkpoint_verify`, `export_session_states`, `restore_session_states`, `retire_session_state`, `checkpoint_session_starts`, `checkpoint_session_ended`, and the `setup_checkpoint_participant` startup check. All are new; existing servers need no change and default to `restart_only` with `wait` verification. - An unknown `checkpoint_mode` or `checkpoint_verify` value raises a `ValueError` at startup, even with checkpointing off. - With checkpointing on, a retire calls `retire_session_state` in every mode, not only `exported`, and the base hook does nothing. A `restart_only` or `stateless` server that already overrides it must expect the call. - The prepare control request has a new optional `restarts` field. A participant from an older Gym build would refuse a prepare that carries it, which coordination sends only when an earlier stage reported a restart; within one deployment every server runs the same build. - A commit exports only the live sessions inside its `episode_ids` scope. A commit request without `episode_ids` still exports every live session. - `/seed_session` replies carry a new `x-ng-checkpoint-verify` header when checkpointing is on. - With checkpointing on, a caller that runs a replayable `/verify` and then closes the session should send `x-ng-checkpoint-close-follows: 1` on `/verify`; without it the session ends at the verification, as before. - With checkpointing on, a request on a restored session must carry its attempt's rollout prefix or MCP token; a cookie-only request is refused with `rollout_id_required`. - Nothing changes unless the global `checkpoint:` block sets `enabled: true`. Signed-off-by: Ananth Subramaniam <ansubramania@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>

…continuation (NVIDIA-NeMo#3888) ## What changed and why This PR builds on the agent session lifecycle that NVIDIA-NeMo#3784 added to `main`. The base server's seed and close also open and close the session in the checkpoint participant, and Simple Agent exports, restores and retires the sessions the base server holds. Every agent takes part in checkpoints once checkpointing is on. - Agents that set `checkpoint_sessions_supported = True` and implement the session hooks get an `AgentSessionParticipant`. A session is idle, running, or at a boundary. A running activation parks at its next boundary when a checkpoint closes admission, and one waiting on a model call counts as parked at its last boundary. If that call's reply arrives while the checkpoint is open, the session stays parked at the boundary before the call and holds the reply until resume, so its readiness never regresses between prepare and commit. The model may already have recorded the call, because it delivered the reply before its own prepare. After a crash the restored agent calls again from that boundary, and the token-capture builder treats the two rows like a client retry: it keeps the call the agent continued from. Boundaries are lazy snapshots, so a long turn loop does not serialize its history at every step. Restored sessions are installed under the replacement attempt. - The hooks are coroutines, and export takes every session at once: `export_agent_sessions(session_keys)`, `restore_agent_sessions(sessions)`, and `retire_agent_session(session_key)`. An agent whose state lives outside the process, such as a sandbox it runs tools or a harness in, can checkpoint all of it concurrently within the commit. - A legacy `/run` is an episode of its own: seed and the turn loop replay, and verification uses the mode the resources server reported on its seed reply. - Agents without the hooks get a restart-only participant. The server framework installs it on whatever app `setup_webserver` returns, so agents that build their own app, such as `harbor_agent` and the `mini_swe_agent` agents, take part too; it is installed once per app. Their rollouts are restarts: they never block prepare, keep running through the checkpoint, and are reported in prepare's `restarts`, so one opaque harness no longer stalls checkpointing for every environment. The seed reply carries `x-ng-checkpoint-restart`, and only a new seed waits while a checkpoint is open. After a crash, the controller retires the rollout and starts it over from its input. - With checkpointing on, agent work without a rollout ID is refused with `rollout_id_required`, and every agent refuses `num_workers > 1`. - Simple Agent implements the hooks for its environment server sessions and for its legacy `/run` loop. A session's export carries how many activations it has started, because each activation reports its observations under its own invocation ID (`root`, then `activation-2`, and so on). A restored session's next activation therefore gets the next ID, and an activation continued from its restored boundary keeps the ID it already had. - Simple Agent's session export carries its resources cookie jar as it is. Since NVIDIA-NeMo#4449 on `main`, a session without a direct tool grant holds no resources cookie jar (`resources_cookies` is `None`), so the environment server keeps the session cookie from its own seed for `/verify` and `/close_session`. Export and restore keep that `None` instead of turning it into an empty jar, so a restored session's close still reports no jar and the environment server keeps the cookie from its resources seed. - A commit exports only the live sessions whose episode is inside its scope, as it already did for restored sessions. The controller leaves every participant's restarts out of the scope, so a session another server reports as a restart is not exported here either. - An unclaimed restored session, one whose replacement has not started yet, keeps its restored boundary in the next checkpoint while the controller's commit scope still continues it, so a second crash still continues it. Once the scope leaves it out, it is not exported, and it is retired. - **Activations are idempotent.** The environment server sends each activation's index, its 1-based count within the session, in the `x-ng-agent-activation` header. With checkpointing on, the agent exports the index and the reply of its last completed activation with the session. An invocation of an activation that already completed gets the stored reply back without running the loop: no model or tool call, and no change to observations or the activation count. This keeps the environment server and the agent consistent when the environment server's checkpoint still has the invocation ahead of it while the agent's checkpoint has the activation done. An index out of order is refused with 409 `activation_out_of_order`. Without the header, every invocation is a new activation, as before. - A request for an activation that is still running waits for it to end and then gets its reply, instead of being refused. - With checkpointing on, every activation of an agent session must carry the `/ng-rollout/<id>` prefix. One without it is refused with 400 `rollout_id_required` before any replay, claim or activation, since a restore rewrites the session's episode to the replacement attempt and the session can no longer stand in for a missing path. This holds for restored sessions and for sessions seeded in the running process; every caller that activates an agent session is a Gym environment server, and each sends the prefix. - Before it replays a reply or starts an activation, the agent checks the attempt the request's `/ng-rollout/<id>` path names against the session's, and refuses a mismatch with `stale_attempt`, so a crashed attempt's harness cannot use a restored session. - A replayed reply claims a restored session, like a new activation does, so a later commit that leaves the episode out does not free a session the live episode still uses. - The reply is kept only when the request carried an activation index. - **A legacy `/run` sends activation index 1 on its self-dispatch.** Its turn loop records its reply like an environment server session's activation, and the reply is exported with the legacy session. A checkpoint taken after the loop replied but before `/run` recorded its `verify` boundary therefore restores the reply instead of running the loop again; the reply leaves when the `/run` ends or a retire frees it. - **Direct `/v1/responses` calls are restarts.** With checkpointing on, a direct call with a rollout path is reported as a restart on every agent, and one without a rollout path is refused with `rollout_id_required`, as restart-only agents already did. - **Restart-only seeds respect retires.** A restart-only seed checks the retired mark after it waits for resume, a retire waits for seeds in progress, and the restart-only retire frees seeded sessions through `retire_agent_session`. The base `retire_agent_session` now releases a session through the agent's close hook, so a restart-only agent stops a sandbox it started. - **Restarts from earlier stages do not hold up agent prepare.** The agent participant leaves the restarts carried from earlier prepare stages (NVIDIA-NeMo#3886) out of `readiness`, `ready` and its export. An activation of such an episode still parks at its next boundary until resume, but it no longer holds up prepare, so an episode `single_agent_turn` marked as a restart, from a restart-only resources seed or an MCP grant, does not stop the checkpoint while it is inside a slow tool call. - **Retire cost is per rollout.** Sessions, seeds in progress and direct restart calls are indexed by rollout, so a retire touches only its own rollout's, and each index entry leaves with the last item of its rollout. - **Simple Agent boundaries leave out each turn's `question`**, which repeats the request's input and the loop's earlier outputs. They store `question_outputs`, how many outputs each question held, and a restore rebuilds each question from the request's input and the restored outputs. The boundary snapshot is built directly, without a validate-and-dump round trip. - **A completed activation's reply is kept once.** A session-scoped reply to a request that carries an activation index leaves out `_ng_trajectory`, in the original and in the copy kept to replay it, since `single_agent_turn` takes the session's evidence from the `AgentObservationBundle` at close; callers that send no index, such as the UserSim environment server, still get it. A legacy `/run` drops its turn loop's kept reply in the same step that records its `verify` or `return` boundary, which holds the reply. Kept replies stay unserialized until a replay or an export needs them. - **An episode that grants MCP access is a restart.** `single_agent_turn` marks an episode as a restart when it grants an `MCPToolAccess`, so no checkpoint continues an agent that holds an MCP token past its seed; with NVIDIA-NeMo#3886's token fence, a replaced attempt's token stops working after a restore. - **Direct HTTP tool calls name their attempt.** With checkpointing on, `single_agent_turn`'s direct HTTP grant carries the seeding attempt's `/ng-rollout/<capture key>` prefix in its base URL, since NVIDIA-NeMo#3886 refuses a call on a restored session that names no attempt. Simple Agent sends each direct HTTP tool call under the rollout prefix of the activation it is serving, in place of the prefix in the grant, so after a restore the continued session's calls name the replacement attempt; its legacy `/run` tool calls carry the same prefix. No other agent with checkpoint support uses direct HTTP grants, and restart-only agents seed again after a crash and get a fresh grant. - **A legacy `/run` closes its resources session after a replayable verification.** It sends `x-ng-checkpoint-close-follows: 1` on `/verify`, so NVIDIA-NeMo#3886 keeps the session tracked and exported until the close, and then closes the session, also when a replacement is restored at its `return` step. A failed close is logged and keeps the run's verdict. - **A `skip_verification` legacy `/run` records its `return` boundary.** It records the same boundary as a verified run, with the configured reward as its result, so a checkpoint that finishes while the turn loop's reply is on its way parks the run at `return` until resume instead of letting it finish uncaptured. A run restored at `return` returns its recorded result. - **Closing a session stops its activation before cleanup.** A close first stops an activation still running for it, so an activation the environment server gave up on stops calling the model and tools, and only then runs the agent's close hook. A close cancelled while it waits for the activation has not run the hook, so a retried close runs it exactly once. The session stays tracked, and holds up prepare, until the activation has stopped. A close cancels only an activation nothing has cancelled yet, so it never interrupts a retire's wait for the activation's cleanup. - **A close waits for resume while agent admission is closed**, as a seed does, and checks admission before it marks the session as closing. So a commit that has checked readiness always exports an open session, never a closing one. This cannot stall prepare through the environment server's final cleanup: coordination prepares environment servers before agents, and an episode in final cleanup keeps its environment server from being ready, so its agent close finishes before agent admission closes. A close that a retire took over is refused with `stale_attempt`, and checks after every await that its session record is still current, so it never writes a receipt for a session a retire freed. - A retire cancels the attempt's activations, or a restart-only agent's `/run` calls, and waits until they have stopped before the retire hook runs. The hook is where a sandboxed agent stops its sandbox. A retired session stays tracked until its activation has stopped and the legacy episode's retire and the retire hook have finished, so a retire cut short by its deadline at any of those points leaves it for a retry, which waits again without cancelling a second time. Meanwhile the session blocks the next checkpoint and is not exported, and a legacy `/run` blocks whether or not its session is still tracked. - A retired attempt's late activation is refused with `stale_attempt` until the controller forgets the rollout. - A session seed that arrives while a checkpoint is open waits for resume, and one already in progress holds up prepare and is exported. The environment server seeds inside a replay step that a checkpoint does not wait for, so the agent's session must match whichever boundary the environment server records. Refusing the seed failed the episode, and admitting it but leaving it out of the checkpoint let a replacement invoke a session nothing restored; a random checkpoint sequence against a real `single_agent_turn` deployment found both. A new session that is not a seed, such as a legacy `/run`, is still refused. Seeds are counted per session key, so a retried seed overlapping the first is tracked until both finish. A retire also waits for seeds of its attempt that are in progress, and a seed refused because its attempt was retired frees the state it built before the retire replies. A restore is refused while a seed is in progress. - `SimpleResponsesAPIAgent._free_session` frees a retired session that has no close receipt; it was `_discard_session`. - The participant answers the readiness wait's `ready()` without building its blocker list, and notifies only while a checkpoint is open. ## 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 A 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 An agent session's states. Prepare is ready once no activation is between boundaries. ```mermaid flowchart LR I(["idle"]) -->|an activation starts| R(["running"]) R -->|before a model call, or after a tool round| B(["at boundary"]) B -->|next step, when no checkpoint is open| R B -->|parked while a checkpoint is open| B R -->|the activation ends| I ``` Which participant an agent gets. ```mermaid flowchart TD Q{"checkpoint_sessions_supported"} -->|yes| P["AgentSessionParticipant export_agent_sessions, restore_agent_sessions, retire_agent_session"] Q -->|no| O["Restart-only: reported in prepare, never blocks, retired and started over after a crash"] ``` Closing or retiring a session stops its activation first: close cancels it, and retire also waits for it before the retire hook, which is where a sandboxed agent stops its sandbox. ```mermaid flowchart LR E["close from the environment server, or retire"] --> X["cancel the running activation"] --> W["retire waits for it"] --> H["retire_agent_session hook stops a sandbox the agent owns"] ``` ## 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#3887 have merged too, so NVIDIA-NeMo#3888'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 (merged) 6. NVIDIA-NeMo#3887 (`ananthsub/partial-ckpt-verifier-declarations`): feat(checkpoint): declare twenty training verifiers stateless with replayable verification (merged) 7. NVIDIA-NeMo#3888 (`ananthsub/partial-ckpt-agent`): feat(checkpoint): add the agent session participant and Simple Agent continuation **(this PR)** 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) - Supersedes NVIDIA-NeMo#2942 (agent participant) and NVIDIA-NeMo#2943 (Simple Agent turns). Completed-result acknowledgements, prepare polling, and the refusal retry loop are gone. - NVIDIA-NeMo#2944 (Gymnasium agent turns) is not carried; Gymnasium continuation moves to a separate PR, and until then Gymnasium agents are restart-only. ## 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`: - Each of these behaviors has a test that fails without its change: - an agent session activation without the rollout prefix is refused, for a restored session and for one seeded in the running process, with no model call, replay or claim, while checkpointing off still accepts it; - a legacy `/run` that skips verification parks at its return boundary until resume; - an episode the environment server reports as a restart, from a restart-only resources seed or an MCP grant, holds up neither the agent stage nor the resources stage of prepare, checked through coordination with a `single_agent_turn` server and real environment, agent and resources participants, and the agent neither waits for nor exports such an episode's activation until resume; - an indexed session reply and its replay carry no trajectory, nor does the export, and a kept reply stays unserialized until a replay or an export needs it; - a checkpoint of a legacy `/run` parked at verify holds its reply once, and a restore still replies; - a restore after the legacy turn loop replied returns that reply without running the loop again; - an episode that grants MCP access is a restart; - a close that arrives after prepare waits for resume and the commit exports its session, a cleanup close of a parked episode after the agent stage waits without holding up prepare, and a retire that frees a session while its close waits leaves no receipt behind; - a close stops the activation before its cleanup starts, and a close cancelled while it stops the activation runs the cleanup once when retried; - a request for an activation still running waits for it and gets its reply; - a request on the replaced attempt's path is refused for the restored session; - a restart-only seed that arrives after its rollout was retired is refused, and a restart-only retire frees its rollout's sessions and waits for a seed in progress; - a direct `/v1/responses` call is a restart, and one without a rollout path is refused; - a boundary leaves out the turn questions, and a resumed loop rebuilds them; - an activation without an index keeps no reply; - a restored session answered with its stored reply survives a commit that leaves its episode out; - with checkpointing on, a direct HTTP grant names the seeding attempt, and a restored session's direct HTTP tool call reaches the restored resources session under the replacement's prefix, through Simple Agent and a real resources server; - a legacy `/run` closes its resources session after a verification a checkpoint does not wait for. - A separate ordering check confirms that an environment episode's cleanup close finishes before the agent stage of prepare; it passes with or without the change. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_checkpoint_*.py`: passed, including: - closing a session stops its running activation; - a retire stops the activation before it replies, and a straggler is retired after resume; - a restart-only agent's `/run` has stopped when its retire replies; - a commit whose scope leaves out a restored session retires it; - a session whose retire hook was cut short stays blocking, is not exported, and is freed by a retry; - a seed arriving during an open checkpoint waits for resume, and one in progress when admission closes holds up prepare and is exported; - overlapping seeds of one session each hold up prepare, a retire waits for a seed in progress and frees what it built, and a restore is refused while a seed is in progress; - a closed session holds up prepare until its activation has stopped, and neither a close nor a retire cancels an activation twice; - a restored legacy `/run` claimed at verify is no longer pending and survives a commit that leaves it out; - an agent that builds its own app still reports its runs as restarts, and the participant is installed once per app; - `ready()` does not build the readiness report, and nothing notifies outside a checkpoint; - a commit does not export live sessions outside its scope. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider responses_api_agents/simple_agent/tests`: passed, including that a restored session's next activation does not reuse an invocation ID and that an activation continued from its restored boundary keeps its own; that a restored activation keeps its boundary through a second checkpoint; that a model reply arriving after prepare is held at the boundary before the call; that an activation already completed is answered with its stored reply without running again; that a legacy `/run` restored at verify or return still reports its trajectory; and that a session without a resources cookie jar restores without one. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider environment_servers/single_agent_turn/tests`: passed, including the timeline where the agent's reply to an invocation is lost to a restore: the environment server invokes again and the restored agent answers with the stored reply while the model sees no new call. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_terminal_attribution.py`: passed, including that a call delivered before a checkpoint and issued again after a restore keeps only the re-issued call. - 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 - Process-level evidence is in the top PR (episodes on the typed `single_agent_turn` request path and legacy `/run` episodes continue from their last boundary). - **Real model:** see the rollout evidence in NVIDIA-NeMo#3889, whose opt-in tests drove this PR's code with `Qwen/Qwen3-4B-Instruct-2507` on vLLM: saved rollouts continued after a kill and restore, restarts started over, and every rollout finished with reward 1.0. Every rollout ran Simple Agent. ## Compatibility - New agent hooks and class attribute: `checkpoint_sessions_supported`, `export_agent_sessions`, `restore_agent_sessions`, `retire_agent_session`. Existing agents need no change and are restart-only under checkpointing. - With checkpointing on, an agent session activation without the `/ng-rollout/<id>` prefix gets 400 `rollout_id_required`. A caller outside Gym that activates Simple Agent sessions with only the session cookie must add the prefix; no caller in the repository is affected. - With checkpointing on, agent calls without a rollout ID get 400 `rollout_id_required`, and agents with `num_workers > 1` fail at startup. - With checkpointing on, a direct HTTP grant's base URL ends in the seeding attempt's rollout prefix. An agent that continues a session after a restore must send its tool calls under its own activation's prefix; Simple Agent does. - With `skip_verification` and checkpointing on, a legacy `/run` caught by a checkpoint waits at its `return` boundary until resume. Nothing changes with checkpointing off. - New request header `x-ng-agent-activation` on agent session invocations, and a 409 `activation_out_of_order` refusal. An environment server that can invoke an activation again after a restore must send the header. - `AgentSessionParticipant.close_session` is now a coroutine, and the agent implements `SimpleServer.setup_checkpoint_participant(app)`, the hook NVIDIA-NeMo#3886 adds that the framework calls after `setup_webserver`. - A commit exports only the live sessions inside its `episode_ids` scope. A commit request without `episode_ids` still exports every live session. - Nothing changes unless the global `checkpoint:` block sets `enabled: true`. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

…dination (NVIDIA-NeMo#3889) ## What changed and why A process-level end-to-end suite for partial-rollout checkpointing, in `tests/e2e/checkpoint`. It starts real Gym server processes (environment server, Simple Agent, resources server, policy model server) against a fake inference backend and token store that survive a Gym crash. The fake worker stages generation cuts with the real staging helpers and continues them through `begin_call`, so the cut record Gym restores is validated exactly as a real worker validates it. Each scenario checkpoints mid-episode through the coordination functions, usually kills every Gym process group, restores into fresh processes, and runs the replacement attempt, asserting on the exact model calls made: - episodes on the typed `single_agent_turn` request path and legacy `/run` episodes continue from their last boundary; - restored resources state (the counter example) is required for the right reward, with a negative control that skips the resources restore; - token lineage continues across the crash; - an in-flight generation is cut and its prefix continued; - a checkpoint without a crash parks and then releases the episode; - verification in flight is replayed or waited for, according to `checkpoint_verify`; - an episode that finishes during prepare is neither lost nor redone; - a second crash, after a checkpoint taken before the replacement attempt starts, still continues the episode; - the same checkpoint restores again after its replacement attempt made model calls and Gym crashed before the next checkpoint. The second run continues from the checkpoint's boundary, and none of the dead run's calls are in its token lineage; - in a deployment that mixes a saved environment with one whose resources server is `restart_only`, the restart never holds up a checkpoint and starts over after a crash, while the saved episode continues; without a crash, the restart keeps running through the checkpoint and finishes; - an episode on the typed `single_agent_turn` request path retired mid-flight, with token capture on, stops everywhere: afterwards no server tracks anything for it except one entry refusing the retired attempt, the model makes no further call for it, and its capture ledger is retired. A duplicate `/run` of the retired attempt is refused with `stale_attempt`, and `coordination.forget` clears the entry on every server, including the resources server's retired sessions. Every server refuses a retired attempt until the controller forgets its rollout, so a stale `/run` for a retired attempt, such as a client retry, is refused with `stale_attempt` rather than running a wasted episode. A duplicate `/run` of an attempt that is still running is separate work, idempotent `/run`. The scenarios on the typed `single_agent_turn` request path, and the double-crash, token-lineage, generation-cut, and park scenarios, run with the policy model in one process and with two uvicorn workers. Three opt-in scenarios are skipped by default: - two real-model tests (`NEMO_GYM_CHECKPOINT_VLLM_URL`). One sends rollouts 20 ms apart until the first finishes, checkpoints, and checks the controller contract: every rollout either replied before the checkpoint or is exported, none replies while Gym is prepared, and every continued rollout completes after a crash and restore. The other mixes saved rollouts with rollouts on a `restart_only` resources server, and checks that prepare reports every restart in flight and no saved rollout, and that after a crash the restarts start over while the saved rollouts continue; - a scale test (`NEMO_GYM_CHECKPOINT_SCALE=<rollouts>`) that holds each rollout at its second model call, checkpoints thousands of in-flight rollouts with token capture and generation cuts, crashes, restores, checks the same contract, and reports per-stage prepare, commit, and restore times. The suite sends every asynchronous request through Gym's global aiohttp client, once and without retries, and polls the fake backend through `GET /_ctl/call_count` instead of fetching its call list. The fake backend's hold reuses the active event when it is placed twice, and a per-rollout hold (`{"after_calls_per_rollout": 1}`) puts every rollout at the same boundary. The scale case asserts every stage's prepare result and that every rollout was exported; the real-model test takes which rollouts finished before the checkpoint from the environment server's commit, not from client task timing. The whole suite is skipped unless `NEMO_GYM_CHECKPOINT_E2E=1` is set, so it does not run in the default unit test job. It takes about two minutes on this branch. ## 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 T["Process-level e2e suite plays the controller"] T --> CO classDef this fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937 class T 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 How the suite runs real server processes and crashes them. ```mermaid flowchart LR T["pytest process drives coordination"] -->|"start, then crash by killing every process group"| G["Gym servers, one process group each environment, policy model, agent, resources"] G --> B["Fake inference backend holds calls, serves generation cuts"] T -->|"/_ctl/hold and /_ctl/release"| B ``` ## 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#3887 have merged too, so NVIDIA-NeMo#3888'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 (merged) 6. NVIDIA-NeMo#3887 (`ananthsub/partial-ckpt-verifier-declarations`): feat(checkpoint): declare twenty training verifiers stateless with replayable verification (merged) 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 **(this PR)** 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) - Supersedes the in-process end-to-end test from NVIDIA-NeMo#2946, replaced by real processes and a real crash (every Gym process group is killed). - The Blackjack and Workplace Assistant adapters from NVIDIA-NeMo#2946 are not carried (see the resources PR). ## 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`: - `NEMO_GYM_CHECKPOINT_E2E=1 RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/e2e/checkpoint`: passed, except the opt-in real-model and scale tests, which were skipped. - The restore-again scenario fails without the ledger fix in the policy model PR, at the second restore, with one and with two policy workers: `capture ledger for again-1-a1 already holds rows from another execution`. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_checkpoint_*.py`: passed. - 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 - **Process-level crash and restore:** the end-to-end suite above, with a fake inference backend: all scenarios pass. - **Scale**, measured on one workstation with `checkpoint_dir` on its local disk, before the scale case moved to the per-rollout hold, so the in-flight column reflects the earlier global hold; a run on shared Lustre storage is tracked separately, with every Gym server and the fake inference backend (weather rollouts on the typed `single_agent_turn` request path, token capture and generation cuts on, each rollout's second model call held so most rollouts are in flight at the checkpoint), run with `NEMO_GYM_CHECKPOINT_E2E=1 NEMO_GYM_CHECKPOINT_SCALE=<N> pytest tests/e2e/checkpoint -k scale -s`: | Rollouts | Policy workers | In flight at checkpoint (exported) | Prepare (s) | Commit (s) | Restore (s) | Replacement failures | |---|---|---|---|---|---|---| | 8,000 | 1 | 4,201 | 0.67 | 0.60 | 0.60 | 0 | | 8,000 | 2 | 4,083 | 0.56 | 0.69 | 0.56 | 0 | | 16,000 | 2 | 14,320 | 7.57 (model 7.38) | 4.64 | 4.97 | 0 | | 16,000 | 2 (rerun) | 12,228 | 3.60 (model 3.53) | 1.83 | 3.23 | 0 | Every rollout either finished before the checkpoint or was exported, and every exported rollout completed after the restore. These numbers come from the lane branches this stack was cut from; the scale test was not rerun on this stack. - **Real model**, on one workstation: `Qwen/Qwen3-4B-Instruct-2507` served by vLLM 0.24.0 on one RTX 6000 Ada (`--enable-auto-tool-choice --tool-call-parser hermes`), driving real Gym server processes that the opt-in tests in NVIDIA-NeMo#3889, NVIDIA-NeMo#3893 and NVIDIA-NeMo#3909 checkpoint mid-run, kill and restore: - **Saved rollouts.** Weather rollouts on the typed `single_agent_turn` request path, sent 20 ms apart until the first one finished. The checkpoint caught about 20 of 26 inside their agent turn. After every Gym process group was killed and the checkpoint restored into fresh processes, every continued rollout finished with reward 1.0, as did every rollout that finished before the checkpoint. - **Restarts.** Half the rollouts used a `restart_only` resources server. Prepare listed every restart still in flight and no saved rollout, and the commit left the restarts out. After the kill, the restarts were retired and started over from their input while the saved rollouts continued, and every rollout finished with reward 1.0. - **Rollout collection.** 48 rows, four at a time: a `SIGUSR1` checkpoint, four more rows sent, then the collector killed with `SIGKILL` while the Gym servers kept running. The resumed run retired the four rows sent after the checkpoint and sent them again as attempt 1, and all 48 rows finished exactly once with reward 1.0. - The first two run with one and two policy workers and with one and two workers per server. Each scenario passed repeated runs. - `NEMO_GYM_CHECKPOINT_E2E=1 NEMO_GYM_MULTIWORKER_E2E=1 NEMO_GYM_CHECKPOINT_VLLM_URL=http://127.0.0.1:8100/v1 NEMO_GYM_CHECKPOINT_VLLM_MODEL=Qwen/Qwen3-4B-Instruct-2507 RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/e2e/checkpoint -k real_model` - Not covered yet: a real vLLM worker serving generation cuts, and several hosts. ## Compatibility Tests only. The suite is skipped unless `NEMO_GYM_CHECKPOINT_E2E=1` is set. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

ananthsub · 2026-10-19T17:02:03Z

Closing: this first prototype is replaced by the environment server checkpoint stack tracked in #3024.
This PR's part is covered by the episode identity in #3882 and #3885: every request carries its rollout ID and attempt, and environment server episodes continue from their boundaries.
The status section of #3024 lists the whole stack and what is merged.

ananthsub changed the title ~~ananthsub/checkpoint execution identity~~ feat(checkpoint): propagate stable execution identity Sep 1, 2026

ananthsub mentioned this pull request Sep 1, 2026

feat(checkpoint): add fenced control contracts #2940

Closed

4 tasks done

ananthsub mentioned this pull request Sep 3, 2026

Partial rollout checkpointing for long-running Gym executions #3024

Open

15 tasks

An error occurred while trying to automatically change base from pthombre/tokcap/full-stack to main September 4, 2026 18:02

An error occurred while trying to automatically change base from pthombre/tokcap/full-stack to main September 4, 2026 19:53

zyzhou5 mentioned this pull request Sep 22, 2026

feat(checkpoint): propagate stable execution identity zyzhou5/Gym#1

Open

7 tasks

ananthsub force-pushed the ananthsub/checkpoint-execution-identity branch from cd3de54 to 9904983 Compare September 22, 2026 17:28

ananthsub removed this pull request from stack #2947 September 22, 2026 17:31

ananthsub changed the base branch from pthombre/tokcap/full-stack to main September 22, 2026 17:32

ananthsub added this pull request to stack #3623 September 22, 2026 17:32

ananthsub force-pushed the ananthsub/checkpoint-execution-identity branch from 9904983 to d5b9c87 Compare September 22, 2026 17:37

macandro96 reviewed Sep 23, 2026

View reviewed changes

zyzhou5 reviewed Sep 23, 2026

View reviewed changes

ananthsub closed this Oct 19, 2026

ananthsub deleted the ananthsub/checkpoint-execution-identity branch October 19, 2026 20:49

Navigation Menu

Repository navigation

feat(checkpoint): propagate stable execution identity - #2939

feat(checkpoint): propagate stable execution identity#2939
ananthsub wants to merge 1 commit into
mainfrom
ananthsub/checkpoint-execution-identity

ananthsub commented Sep 1, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Sep 1, 2026

Uh oh!

copy-pr-bot Bot commented Sep 22, 2026

Uh oh!

macandro96 Sep 23, 2026

Uh oh!

zyzhou5 Sep 23, 2026

Uh oh!

macandro96 Sep 28, 2026

Uh oh!

macandro96 Sep 23, 2026

Uh oh!

zyzhou5 Sep 23, 2026

Uh oh!

ananthsub commented Oct 19, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

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

Uh oh!

Summary

Why this layer exists

Execution identity flow

Stack

Test plan

Uh oh!

copy-pr-bot Bot commented Sep 1, 2026

Uh oh!

copy-pr-bot Bot commented Sep 22, 2026

Uh oh!

macandro96 Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

zyzhou5 Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

macandro96 Sep 28, 2026

Choose a reason for hiding this comment

Uh oh!

macandro96 Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

zyzhou5 Sep 23, 2026

Choose a reason for hiding this comment

Uh oh!

ananthsub commented Oct 19, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

ananthsub commented Sep 1, 2026 •

edited

Loading