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. |
ananthsub
changed the title
This was referenced Sep 1, 2026
ananthsub
force-pushed
the
ananthsub/checkpoint-model-participant
branch
2 times, most recently
from
September 3, 2026 11:12
c92b187 to
a6bb6c5
Compare
ananthsub
force-pushed
the
ananthsub/checkpoint-model-participant
branch
from
September 10, 2026 20:48
a6bb6c5 to
dfd685a
Compare
ananthsub
force-pushed
the
ananthsub/checkpoint-model-participant
branch
from
September 22, 2026 17:28
dfd685a to
a52154d
Compare
Pause policy generation at a quiescent cut and checkpoint token-free capture custody while leaving framework-owned staged token arrays to TransferQueue. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Namespace model artifacts, reject changing commit exclusions, and link the first restored call to its verified parent exactly once. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
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-model-participant
branch
from
September 22, 2026 17:37
a52154d to
e3590a2
Compare
macandro96
reviewed
Sep 29, 2026
| if manifest.get("server_name") != self.server_name: | ||
| raise LedgerMismatchError("ledger checkpoint belongs to a different model server") | ||
|
|
||
| expected_names = {name for metadata in manifest["rollouts"].values() for name in metadata["files"]} |
Contributor
There was a problem hiding this comment.
Can we identify failures with a union approach? for eg. policy model A and policy model B can restore into the same directory but when B sees A files it calls them unexpected. Instead restore everything and validate the combined inventory. Eg. here
| else: | ||
| request_messages = None | ||
| if "request" in inspect.signature(self.chat_completions).parameters: | ||
| mark_current_generation_started() |
Contributor
There was a problem hiding this comment.
A checkpoint can pause here after we recorded that this model call was starting. The code marks the request safe too early, so the checkpoint may save an unfinished call and incorrectly mark the rollout as failed. Could we first cancel the pending call, record that no generation happened, and only then mark the request safe?
| and get_first_server_config_dict(self.global_config_dict, server_name).get("instance_role", "policy") | ||
| == "policy" | ||
| ) | ||
| is_agent_parent_relay = ( |
Contributor
There was a problem hiding this comment.
We send the restored parent to the agent here, but the agent server does not install RolloutContextMiddleware, so it never reads these headers or forwards the parent to the policy model. The model can then treat the restored conversation as a new root. Could we add RolloutContextMiddleware in SimpleResponsesAPIAgent.setup_webserver() and test this using the real agent app instead of adding the middleware manually in the test? II believe this was fixed in later stack eb386809
| checkpoint_id=checkpoint_id, | ||
| server_name=server_name, | ||
| tombstones=tuple(limiter.checkpoint_exclusions()), | ||
| source_attempts=tuple(limiter.seen_attempts()), |
Contributor
There was a problem hiding this comment.
Restored attempts are blocked in memory, but they are not included in seen_attempts(). This means checkpoint B can forget the stale-attempt blocks restored from checkpoint A, allowing an old attempt after restoring B. Could we carry restored fences forward in source_attempts—separately from checkpoint_exclusions, so their lineage remains available—and extend the existing test through a second restore to verify attempt 0 is still rejected?
| and request_path.endswith("/v1/responses") | ||
| ) | ||
| source_capture_key, parent_model_call_id = ( | ||
| take_checkpoint_parent() if is_policy_generation or is_agent_parent_relay else (None, None) |
Contributor
There was a problem hiding this comment.
This marks the restored parent as used before the HTTP request is sent. If the request is refused before model admission or fails before reaching the server, a retry in the same rollout loses its parent and may be treated as a new root. Could we treat this as a claim: release it when we know the request was not admitted, and consume it only when the downstream server confirms admission? For an uncertain network failure, we should fail closed or use an idempotent request ID instead of silently retrying without the parent.
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
|
Closing: this first prototype is replaced by the environment server checkpoint stack tracked in #3024. |
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
What this participant checkpoints
Generation workers put token IDs, masks, and log probabilities in
TransferQueue. The Gym policy model server retains the smaller custody ledger that maps rollout and model-call identities to those staged rows. This PR saves and restores that ledger without moving token arrays through Gym checkpoint files.Prepare, commit, and continuation flow
sequenceDiagram participant C as Checkpoint controller participant M as Gym policy model server participant G as Generation worker participant TQ as TransferQueue participant D as Controller-provided directory C->>M: pause(checkpoint_id, deadline) M->>M: Close new policy admission M->>G: Previously accepted call continues G->>TQ: Stage token delta and await durability G-->>M: Response with token-free coordinates M->>M: Append custody row and finish HTTP response M-->>C: paused when in-flight count reaches zero C->>M: commit(checkpoint_id, directory) M->>D: Write model-specific lineage, then manifest last M-->>C: manifest path, digest, and counts C->>M: restore into attempt N+1 M->>M: Verify source call and complete staging chain M-->>G: First resumed request carries source parent coordinates M->>M: Consume restored-parent relay M-->>G: Later requests link within attempt N+1Restore validates the model-server namespace and installs source attempts as tombstones. Commit retries must use the same exclusions and source-attempt set as the published manifest. A changed set is rejected instead of silently returning an incompatible checkpoint.
The replacement state remains paused until the global controller restores
TransferQueueand every other required participant, then explicitly resumes Gym.This PR implements the drain-to-zero model participant used by the first integration. Continuing a generation through a live token-prefix cut would require a different readiness condition and additional generation-worker/TQ contracts; it is not claimed here.
Stack
Test plan
Generation-safe checkpoint update
Layer validation: 261 tests passed.