Repository navigation
Conversation
|
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. |
2 tasks done
ananthsub
changed the title
This was referenced Sep 1, 2026
ananthsub
force-pushed
the
ananthsub/checkpoint-control-contracts
branch
from
September 22, 2026 17:28
5439535 to
046d0c5
Compare
Expose versioned checkpoint capabilities on every Gym server and provide a shared phase fence for idempotent control operations. Reject stale non-final operations after resume so an old coordinator cannot mistake an accepting server for a prepared one. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Close admission atomically in each server process and aggregate worker acknowledgements and in-flight counts through a service-level Unix socket coordinator. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Keep a failed restore paused and observable so the coordinator can retry or abort it safely. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Move the checkpoint implementation under nemo_gym._checkpoint until its contracts are ready for public support. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
ananthsub
force-pushed
the
ananthsub/checkpoint-control-contracts
branch
from
September 22, 2026 17:37
046d0c5 to
5b9f21a
Compare
macandro96
reviewed
Sep 25, 2026
|
|
||
| self._validate(checkpoint_id, allowed_phases) | ||
|
|
||
| entry_phase = self.phase |
Contributor
There was a problem hiding this comment.
There could be race conditions here we might need to have a lock here - had added a fix in my stack, logically it can live here
| if message.get("type") != "state": | ||
| return | ||
| state = AdmissionState(message["state"]) | ||
| self._coordinator_sequence = int(message["seq"]) |
Contributor
There was a problem hiding this comment.
older message can arrive after newer ones - we should probably exit if that is the case; had a fix here
| body: ModelAdmissionResumeRequest, | ||
| authorization: Optional[str] = Header(default=None), | ||
| ) -> dict[str, Any]: | ||
| require_control_auth(authorization, auth_token) |
Contributor
There was a problem hiding this comment.
I think we need to incorporate deadline here too? like it is pause()
| f"operation after the checkpoint completes" | ||
| ) | ||
| if rollout_id is not None and attempt_index is not None: | ||
| self._seen_attempts.add((rollout_id, attempt_index)) |
Contributor
There was a problem hiding this comment.
_seen_attempts are retained throughout the entire run and is part of the checkpoint. Although it just stores a tuple - do we need to retire it at some point or is it preserve it?
| task=task, | ||
| ) | ||
| # Kept private and intentionally absent from the public inventory. | ||
| ticket._limiter = self |
Contributor
There was a problem hiding this comment.
nit: instead of explicitly setting in admit - do you think its cleaner to add it in AdmissionTicket.__init__?
This was referenced Oct 1, 2026
pachmu
pushed a commit
to pachmu/Gym
that referenced
this pull request
Oct 13, 2026
…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<br/>NeMo RL, or rollout collection"] CO["Coordination<br/>prepare, commit, restore, resume,<br/>retire, forget"] K["Control plane on every server<br/>phases, lease, storage,<br/>retire: refuse, stop, free;<br/>forget"] subgraph G["One participant per Gym server"] E["Environment server<br/>episode steps"] M["Policy model<br/>held responses, generation cuts"] A["Agent<br/>sessions parked at boundaries"] R["Resources server<br/>session state"] end W["Inference worker<br/>stages cut prefixes"] D[("Checkpoint directory<br/>records, then manifest")] L[("Capture ledger<br/>retire and delete from 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,<br/>undelivered model responses are held G-->>C: prepared, or blockers at the deadline C->>G: commit with the episodes the controller continues G->>D: each participant writes its records, then its manifest Note over G: restored state the commit's scope leaves out is deleted C->>C: publish the checkpoint with the controller's own state C->>G: resume, in order resources, agent, model, environment Note over C,G: crash - every Gym process dies C->>G: restore the checkpoint in fresh processes, all or nothing D-->>G: records installed under attempt N+1, attempt N's capture ledger retired C->>G: resume C->>G: /run as attempt N+1 continues each episode from its boundary Note over C,G: retiring an episode, only while no checkpoint is open C->>G: retire - environment, then agent, then model and resources Note over G: each server marks the attempt, stops its work, waits, frees its state, then replies C->>G: forget, once the rollout is done for good - same order Note over G: the mark refusing the retired attempt goes away ``` ### 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.<br/>The episode parks at the boundary after it."] Q -->|replay| P["The checkpoint does not wait.<br/>The boundary before the step is exported,<br/>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>
hesam-shams
pushed a commit
to hesam-shams/Gym
that referenced
this pull request
Oct 13, 2026
…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<br/>NeMo RL, or rollout collection"] CO["Coordination<br/>prepare, commit, restore, resume,<br/>retire, forget"] K["Control plane on every server<br/>phases, lease, storage,<br/>retire: refuse, stop, free;<br/>forget"] subgraph G["One participant per Gym server"] E["Environment server<br/>episode steps"] M["Policy model<br/>held responses, generation cuts"] A["Agent<br/>sessions parked at boundaries"] R["Resources server<br/>session state"] end W["Inference worker<br/>stages cut prefixes"] D[("Checkpoint directory<br/>records, then manifest")] L[("Capture ledger<br/>retire and delete from 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,<br/>undelivered model responses are held G-->>C: prepared, or blockers at the deadline C->>G: commit with the episodes the controller continues G->>D: each participant writes its records, then its manifest Note over G: restored state the commit's scope leaves out is deleted C->>C: publish the checkpoint with the controller's own state C->>G: resume, in order resources, agent, model, environment Note over C,G: crash - every Gym process dies C->>G: restore the checkpoint in fresh processes, all or nothing D-->>G: records installed under attempt N+1, attempt N's capture ledger retired C->>G: resume C->>G: /run as attempt N+1 continues each episode from its boundary Note over C,G: retiring an episode, only while no checkpoint is open C->>G: retire - environment, then agent, then model and resources Note over G: each server marks the attempt, stops its work, waits, frees its state, then replies C->>G: forget, once the rollout is done for good - same order Note over G: the mark refusing the retired attempt goes away ``` ### 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,<br/>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,<br/>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<br/>gate: tickets, held responses"] CO -->|"control call on the shared port"| W2["uvicorn worker 2<br/>gate: tickets, held responses"] W1 <-->|"Unix socket: forward, close, reports, cut claims"| C["Coordinator in the main process<br/>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>
Contributor
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note
This checkpoint implementation is experimental. Its Python modules live under
nemo_gym._checkpoint; they are not a supported public import surface yet.Summary
Why this layer exists
A checkpoint cannot safely copy server state while new operations continue to enter. This PR supplies the shared control-plane vocabulary used by later model, agent, and resources participants. It does not save participant payloads itself.
Admission state machine
stateDiagram-v2 [*] --> Accepting Accepting --> Draining: prepare(checkpoint_id) Draining --> Paused: all accepted operations finish Draining --> Draining: new operation is parked or rejected Paused --> Accepting: resume(same checkpoint_id) Paused --> RestoreFailedPaused: restore fails RestoreFailedPaused --> Paused: retry restore RestoreFailedPaused --> Accepting: abort and resume Paused --> Paused: repeated prepare / status Accepting --> Accepting: stale checkpoint ID rejectedFor a multi-worker service, one pause must cover every Uvicorn process behind the port:
flowchart LR C[Checkpoint controller] --> K[Service coordinator] K --> W1[Worker 1 limiter] K --> W2[Worker 2 limiter] K --> WN[Worker N limiter] W1 --> A[Aggregate acknowledgements and in-flight counts] W2 --> A WN --> A A -->|every expected worker acknowledged| P[Service reports paused] A -->|worker missing or deadline exceeded| F[Prepare fails closed]The coordinator protocol and a real two-worker test are included. Production startup does not yet launch and route this coordinator, so
num_workers > 1advertisesunmanagedinstead of claiming checkpoint safety.Stack
Test plan
Generation-safe checkpoint update
Layer validation: 32 tests passed.