Skip to content

feat(checkpoint): add the participant control plane, episode steps, and coordination - #3882

Merged
ananthsub merged 1 commit into
mainfrom
ananthsub/partial-ckpt-core
Oct 13, 2026
Merged

ananthsub merged 1 commit into
mainfrom
ananthsub/partial-ckpt-core

Conversation

@ananthsub

@ananthsub ananthsub commented Oct 1, 2026 •

Copy link
Copy Markdown
Contributor

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 (feat(checkpoint): add Gym turn-level rollout recovery RL#4500), got a server error.
    • The suffix is reserved for attempt-qualified rollout IDs. An explicit "<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_id ends in -a<N>. It raises a ConfigError when 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_index on every retried row whether or not checkpointing is on, made its /run fail mid-run. Rows passed straight to run_examples are not checked. A request body that sets both an -a<N> ID and _ng_attempt_index is still refused, because the two contradict each other. The external agent harnesses tutorial, the comment on _ng_rollout_id and the docstring of maybe_rollout_id_from_run_body now say the suffix is reserved.
  • 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.
    • A /run reply 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 /run calls whose episode is neither in the committed episode_ids nor 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.
  • 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.
    • A commit can publish its manifest and then fail to sync the participant directory. Its retry finds the identical manifest already published, and syncs the participant directory before it reports success, so a commit that succeeds is always durable.
    • 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.

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 this
Loading

One checkpoint, a crash, and the restore, end to end:

sequenceDiagram
  participant C as Controller
  participant G as Gym participants
  participant D as Checkpoint directory
  C->>G: prepare, in order environment, model, agent, resources
  Note over G: admission closes, in-flight work parks at a boundary,<br/>undelivered model responses are held
  G-->>C: prepared, or blockers at the deadline
  C->>G: commit with the episodes the controller continues
  G->>D: each participant writes its records, then its manifest
  Note over G: restored state the commit's scope leaves out is deleted
  C->>C: publish the checkpoint with the controller's own state
  C->>G: resume, in order resources, agent, model, environment
  Note over C,G: crash - every Gym process dies
  C->>G: restore the checkpoint in fresh processes, all or nothing
  D-->>G: records installed under attempt N+1, attempt N's capture ledger retired
  C->>G: resume
  C->>G: /run as attempt N+1 continues each episode from its boundary
  Note over C,G: retiring an episode, only while no checkpoint is open
  C->>G: retire - environment, then agent, then model and resources
  Note over G: each server marks the attempt, stops its work, waits, frees its state, then replies
  C->>G: forget, once the rollout is done for good - same order
  Note over G: the mark refusing the retired attempt goes away
Loading

This PR

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.

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
Loading

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."]
Loading

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.

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
Loading

Where this sits in the stack

This is one PR in a stack of draft PRs that re-cut partial-rollout checkpointing onto environment servers. Each PR is one commit, and its base is the branch of the PR before it, so each diff shows only that PR's change.

The stack builds on the token-capture cleanup PRs #3938 (capture ledger retire and delete) and #3939 (complete-record retire and delete), which have merged: #3882's base is main. The checkpoint stack uses those operations to free the ledgers of retired attempts and to clear a dead execution's capture files before a restore. The capture-ledger cache bound (#4268) and the striped lock files (#3937) are independent of the stack.

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

Partial-rollout checkpointing lets a training controller, such as NeMo RL, checkpoint Gym while rollouts are in flight and, after a crash, continue those rollouts from their last safe point instead of starting them over. Checkpointing is off by default; with the checkpoint: block unset, no server installs any checkpoint routes or behavior.

Relationship to the old stack (#2939 to #2946)

  • Supersedes feat(checkpoint): propagate stable execution identity #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 feat(checkpoint): add fenced control contracts #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

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:
    • 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;
    • a commit whose first sync of the participant directory fails is retried, and the retry, which finds the manifest the failed write published, syncs that directory before it reports the manifest.
  • 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:
    • a dataset row whose explicit rollout ID is "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;
    • a resumed run whose materialized row has the ID "job-a2" and failed once fails when it loads that row, instead of failing mid-run when its retry stamps attempt 1;
    • a resume retry of a row with the plain explicit ID "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 /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.
  • The rollout ID suffix -a<N> is reserved for attempt-qualified rollout IDs: an explicit "<id>-a<N>" always means attempt N of <id>. Rollout collection now raises a ConfigError for a dataset row whose explicit _ng_rollout_id ends 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. On main such a row runs, and its resume retry runs as "<id>-a<N>-a1". Rows with such IDs must be renamed. Rows passed straight to run_examples are not checked.
  • Nothing changes unless the global checkpoint: block sets enabled: true. With it on, model and resources calls always carry the rollout path prefix.

@ananthsub ananthsub added feature New capabilities, enhancements, or enablement work area:core Shared APIs, servers, telemetry, health, and registries labels Oct 1, 2026
@copy-pr-bot

copy-pr-bot Bot commented Oct 1, 2026

Copy link
Copy Markdown

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
ananthsub requested a review from zyzhou5 October 1, 2026 19:04
@ananthsub
ananthsub force-pushed the ananthsub/partial-ckpt-core branch from 0a24eb7 to c2b5103 Compare October 1, 2026 19:38
@ananthsub
ananthsub force-pushed the ananthsub/partial-ckpt-core branch from c2b5103 to 9040654 Compare October 1, 2026 21:53
@ananthsub
ananthsub force-pushed the ananthsub/partial-ckpt-core branch from 9040654 to 846c721 Compare October 1, 2026 22:05
@ananthsub
ananthsub force-pushed the ananthsub/partial-ckpt-core branch from 846c721 to f4092dc Compare October 2, 2026 19:43
@ananthsub
ananthsub changed the base branch from main to ananthsub/token-capture-store-retire October 2, 2026 19:44
@pthombre
pthombre added this pull request to stack #3962 October 2, 2026 22:50
This was referenced Oct 13, 2026

This branch was previously deployed

1 inactive deployment
public — 082b250c Deployed Oct 12, 2026 by copy-pr-bot[bot] via release / finalize / notify #5120
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core Shared APIs, servers, telemetry, health, and registries complexity:high Cross-component or shared-contract change with a large review surface feature New capabilities, enhancements, or enablement work needs-review PR is ready for code review and waiting on a reviewer sla:review-overdue Review response is over the one-business-day SLA

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants