Repository navigation
feat(checkpoint): add the participant control plane, episode steps, and coordination - #3882
Merged
Merged
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. |
This was referenced Oct 1, 2026
Merged
Merged
ananthsub
force-pushed
the
ananthsub/partial-ckpt-core
branch
from
October 1, 2026 19:38
0a24eb7 to
c2b5103
Compare
This was referenced Oct 1, 2026
ananthsub
force-pushed
the
ananthsub/partial-ckpt-core
branch
from
October 1, 2026 21:53
c2b5103 to
9040654
Compare
This was referenced Oct 1, 2026
ananthsub
force-pushed
the
ananthsub/partial-ckpt-core
branch
from
October 1, 2026 22:05
9040654 to
846c721
Compare
ananthsub
force-pushed
the
ananthsub/partial-ckpt-core
branch
from
October 2, 2026 19:43
846c721 to
f4092dc
Compare
ananthsub
changed the base branch from
main
to
ananthsub/token-capture-store-retire
October 2, 2026 19:44
5 tasks done
This was referenced Oct 13, 2026
7 tasks done
This was referenced Oct 29, 2026
This branch was previously deployed
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.
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.
rfor the first attempt andr-aNafter a restore. When checkpointing is on,ServerClientalways attributes model and resources calls to their rollout through the rollout path prefix, not only when observability is on.mainalready reserves the-a<N>suffix in rollout IDs; this PR addsEpisodeId.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 typedsingle_agent_turnrequest path (_episode_request_bodyinnemo_gym/rollout_collection.py) and the legacy flat-row wrapper (_episode_request_from_rowinenvironment_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'starget_level: turndoes (feat(checkpoint): add Gym turn-level rollout recovery RL#4500), got a server error."<id>-a<N>"always means attempt N of<id>, which is what NeMo RL sends for restored turns. Rollout collection therefore rejects dataset rows whose explicit_ng_rollout_idends in-a<N>. It raises aConfigErrorwhen it loads the input file, and again when a resumed run loads the materialized rows of an earlier run. The error names the file, the line, the ID, and the reserved suffix. Without this check, such a row ran on its first attempt, but a resume retry, which stamps_ng_attempt_indexon every retried row whether or not checkpointing is on, made its/runfail mid-run. Rows passed straight torun_examplesare not checked. A request body that sets both an-a<N>ID and_ng_attempt_indexis still refused, because the two contradict each other. The external agent harnesses tutorial, the comment on_ng_rollout_idand the docstring ofmaybe_rollout_id_from_run_bodynow say the suffix is reserved.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).steps.py). An episode records a boundary after each completed step. A step in flight when a checkpoint closes runs inwaitmode (the checkpoint waits for it) orreplaymode (it runs again after a crash). Time parked for a checkpoint does not count against the episode's deadline.ready()), and lists blockers once, for the reply./verifymay be replayed in anx-ng-checkpoint-verifyheader on its/seed_sessionreply.restart_only, is a restart.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.PrepareResult.continued(...), which leaves the restarts out. A commit whose scope names a restart is refused with 409restart_in_scope.x-ng-checkpoint-restart. Marking one waits while a checkpoint is open, so the restarts a prepare reported cannot change before its commit./runresend. With checkpointing on,ServerClientnever sends a/runagain 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.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).renewextends every lease when publishing a checkpoint takes longer than expected./runreply can already be on the wire when the checkpoint is prepared. Before it builds the checkpoint's rollout set, the controller waits only for outstanding/runcalls whose episode is neither in the committedepisode_idsnor a reported restart. It does not wait for restarts: they keep running through the checkpoint, and after a crash a restart starts over from its input.stale_attempt.forgetoperation, 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'sforgetcalls bound it.forget, so a controller must forget each restored rollout once it is finished for good, not only the ones it retired.forgetfollows the same phase rule and runs in the same caller-first order.delete_restoredhook, 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_restoredmust therefore leave nothing that refuses the attempt later, such as a capture-ledger fence. It defaults toretire, which is enough for participants whose retire refuses nothing beyond itself; the policy model deletes the attempt's ledger instead.invalid_checkpoint_stateerror naming its episode.invalid_checkpoint_stateerrors, 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.discovercaps 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.retireandforgetreply with counts,retired_countandforgotten_count, not every ID; commit'sepisode_idsand restore'srestoredstay exact./runcalls parked during a checkpoint) cannot starve prepare, commit, or resume. The control pool is built like the data pool onmain, with the same keepalive socket options and, when metrics export is on, the same timed connector.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.checkpoint_conflictinstead of the earlier result.deadline_exceededand 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.JsonPayload) is checked once for JSON-ness, then written as is, rather than deep-validated and copied.orjsonis 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 valueorjsonaccepts butjsoncannot write, such as adatetime, 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.checkpoint:block withenabled,control_auth_token, andlease_grace_seconds. It is off by default, and a control token is required when it is on.checkpointis a reserved top-level key, and the block is validated once, when theServerClientis 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
mainand 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.
flowchart LR C["Controller<br/>NeMo RL, or rollout collection"] CO["Coordination<br/>prepare, commit, restore, resume,<br/>retire, forget"] K["Control plane on every server<br/>phases, lease, storage,<br/>retire: refuse, stop, free;<br/>forget"] subgraph G["One participant per Gym server"] E["Environment server<br/>episode steps"] M["Policy model<br/>held responses, generation cuts"] A["Agent<br/>sessions parked at boundaries"] R["Resources server<br/>session state"] end W["Inference worker<br/>stages cut prefixes"] D[("Checkpoint directory<br/>records, then manifest")] L[("Capture ledger<br/>retire and delete from #3938, #3939")] C --> CO --> K K --> E & M & A & R M --> W M --> L G --> D classDef this fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937 class CO,K thisOne checkpoint, a crash, and the restore, end to end:
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 inidleandrestored.Episode steps: what happens to a step that is running when a checkpoint closes.
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.
Where this sits in the stack
This is one PR in a stack of draft PRs that re-cut partial-rollout checkpointing onto environment servers. Each PR is one commit, and its base is the branch of the PR before it, so each diff shows only that PR's change.
The stack builds on the token-capture cleanup PRs #3938 (capture ledger
retireanddelete) and #3939 (complete-recordretireanddelete), which have merged: #3882's base ismain. The checkpoint stack uses those operations to free the ledgers of retired attempts and to clear a dead execution's capture files before a restore. The capture-ledger cache bound (#4268) and the striped lock files (#3937) are independent of the stack.ananthsub/partial-ckpt-core): feat(checkpoint): add the participant control plane, episode steps, and coordination (this PR)ananthsub/partial-ckpt-policy-model): feat(checkpoint): make policy model servers checkpoint participants with generation cutsananthsub/partial-ckpt-model-worker-cuts): feat(token-capture): add worker staging helpers for generation cutsananthsub/partial-ckpt-environment): feat(checkpoint): continue environment server episodes from their boundariesananthsub/partial-ckpt-resources): feat(checkpoint): add the resources server participant with asynchronous session hooksananthsub/partial-ckpt-verifier-declarations): feat(checkpoint): declare five training verifiers stateless with replayable verificationananthsub/partial-ckpt-agent): feat(checkpoint): add the agent session participant and Simple Agent continuationananthsub/partial-ckpt-e2e): test(checkpoint): add a process-level end-to-end suite driven by coordinationananthsub/partial-ckpt-rollout-collection): feat(checkpoint): checkpoint evaluation runs from rollout collectionananthsub/partial-ckpt-telemetry): feat(checkpoint): spans and metrics for partial-rollout checkpointsananthsub/partial-ckpt-multi-worker): feat(checkpoint): resources, agent, and environment servers with several workersananthsub/partial-ckpt-fault-tests): test(checkpoint): check checkpoint invariants over random sequences of operationsPartial-rollout checkpointing lets a training controller, such as NeMo RL, checkpoint Gym while rollouts are in flight and, after a crash, continue those rollouts from their last safe point instead of starting them over. Checkpointing is off by default; with the
checkpoint:block unset, no server installs any checkpoint routes or behavior.Relationship to the old stack (#2939 to #2946)
x-nemo-gym-rollout-id,x-nemo-gym-attempt-index, andx-nemo-gym-model-call-idheaders are gone.Issue
Part of #3024
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:forget, a failed retire keeps its attempts refused, and a retry stops them;forgetis refused while a checkpoint is open, and coordination runs it, like retire, environment servers first;discoverreports a server that is down, one that never replies, and one of an unknown kind; a badcheckpoint:block fails when theServerClientis built;single_agent_turnrequest path decodes an explicit"<id>-aN", a stamped attempt index, and a first attempt into the same rollout and attempt as the agent;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 tests/unit_tests/test_rollout_collection.py tests/unit_tests/test_rollout_correlation.py: passed. New:"job-a2"is rejected when the input file loads, with an error that names its line, the ID, the reserved suffix, and the attempt and rollout it would name;"job-a2"and failed once fails when it loads that row, instead of failing mid-run when its retry stamps attempt 1;"job"still runs as"job-a1".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/runwith the explicit rollout ID"restored-a1"seeds attempt 1 ofrestoredand sends its/v1/responsescall under the rollout path prefixrestored-a1.tests/unit_tests, eight processes, withouttest_opensandbox_compose.py, which needs the optionalopensandboxpackage): passed except a sandbox retry test, which fails the same way onmainin this development environment.pre-commit run --files <files changed by this PR>: all hooks passed, and no hook modified a file.Signed-off-byline.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
checkpoint:block withenabled(defaultfalse),control_auth_token(required when enabled), andlease_grace_seconds(default 600).global_aiohttp_control_connector_limitandglobal_aiohttp_control_connector_limit_per_host, both at least 0.export_failed: the checkpoint cannot be committed; resume, then prepare again under a new checkpoint ID.retireandforgetreply withretired_countandforgotten_countinstead of ID lists.forgetevery rollout it restored once that rollout is finished for good.single_agent_turnrequest path and the legacy flat-row wrapper accept an explicit rollout ID that ends in-aNand run it as attempt N of the rest of the ID; onmainthey refused it. This applies whether or not checkpointing is on.-a<N>is reserved for attempt-qualified rollout IDs: an explicit"<id>-a<N>"always means attempt N of<id>. Rollout collection now raises aConfigErrorfor a dataset row whose explicit_ng_rollout_idends in-a<N>, when it loads the input file and when a resumed run loads the materialized rows of an earlier run, whether or not checkpointing is on. Onmainsuch a row runs, and its resume retry runs as"<id>-a<N>-a1". Rows with such IDs must be renamed. Rows passed straight torun_examplesare not checked.checkpoint:block setsenabled: true. With it on, model and resources calls always carry the rollout path prefix.