[ananthsub]

What changed and why

A process-level end-to-end suite for partial-rollout checkpointing, in tests/e2e/checkpoint. It starts real Gym server processes (environment server, Simple Agent, resources server, policy model server) against a fake inference backend and token store that survive a Gym crash. The fake worker stages generation cuts with the real staging helpers and continues them through begin_call, so the cut record Gym restores is validated exactly as a real worker validates it.

Each scenario checkpoints mid-episode through the coordination functions, usually kills every Gym process group, restores into fresh processes, and runs the replacement attempt, asserting on the exact model calls made:

• episodes on the typed single_agent_turn request path and legacy /run episodes continue from their last boundary;
• restored resources state (the counter example) is required for the right reward, with a negative control that skips the resources restore;
• token lineage continues across the crash;
• an in-flight generation is cut and its prefix continued;
• a checkpoint without a crash parks and then releases the episode;
• verification in flight is replayed or waited for, according to checkpoint_verify;
• an episode that finishes during prepare is neither lost nor redone;
• a second crash, after a checkpoint taken before the replacement attempt starts, still continues the episode;
• the same checkpoint restores again after its replacement attempt made model calls and Gym crashed before the next checkpoint. The second run continues from the checkpoint's boundary, and none of the dead run's calls are in its token lineage;
• in a deployment that mixes a saved environment with one whose resources server is restart_only, the restart never holds up a checkpoint and starts over after a crash, while the saved episode continues; without a crash, the restart keeps running through the checkpoint and finishes;
• an episode on the typed single_agent_turn request path retired mid-flight, with token capture on, stops everywhere: afterwards no server tracks anything for it except one entry refusing the retired attempt, the model makes no further call for it, and its capture ledger is retired. A duplicate /run of the retired attempt is refused with stale_attempt, and coordination.forget clears the entry on every server, including the resources server's retired sessions.

Every server refuses a retired attempt until the controller forgets its rollout, so a stale /run for a retired attempt, such as a client retry, is refused with stale_attempt rather than running a wasted episode. A duplicate /run of an attempt that is still running is separate work, idempotent /run.

The scenarios on the typed single_agent_turn request path, and the double-crash, token-lineage, generation-cut, and park scenarios, run with the policy model in one process and with two uvicorn workers.

Three opt-in scenarios are skipped by default:

• two real-model tests (NEMO_GYM_CHECKPOINT_VLLM_URL). One sends rollouts 20 ms apart until the first finishes, checkpoints, and checks the controller contract: every rollout either replied before the checkpoint or is exported, none replies while Gym is prepared, and every continued rollout completes after a crash and restore. The other mixes saved rollouts with rollouts on a restart_only resources server, and checks that prepare reports every restart in flight and no saved rollout, and that after a crash the restarts start over while the saved rollouts continue;
• a scale test (NEMO_GYM_CHECKPOINT_SCALE=<rollouts>) that holds each rollout at its second model call, checkpoints thousands of in-flight rollouts with token capture and generation cuts, crashes, restores, checks the same contract, and reports per-stage prepare, commit, and restore times.

The suite sends every asynchronous request through Gym's global aiohttp client, once and without retries, and polls the fake backend through GET /_ctl/call_count instead of fetching its call list. The fake backend's hold reuses the active event when it is placed twice, and a per-rollout hold ({"after_calls_per_rollout": 1}) puts every rollout at the same boundary. The scale case asserts every stage's prepare result and that every rollout was exported; the real-model test takes which rollouts finished before the checkpoint from the environment server's commit, not from client task timing.

The whole suite is skipped unless NEMO_GYM_CHECKPOINT_E2E=1 is set, so it does not run in the default unit test job. It takes about two minutes on this branch.

How it works

Where this PR sits in the overall flow

The highlighted part is what this PR adds.

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
  T["Process-level e2e suite<br/>plays the controller"]
  T --> CO
  classDef this fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937
  class T 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

How the suite runs real server processes and crashes them.

flowchart LR
  T["pytest process<br/>drives coordination"] -->|"start, then crash by killing every process group"| G["Gym servers, one process group each<br/>environment, policy model, agent, resources"]
  G --> B["Fake inference backend<br/>holds calls, serves generation cuts"]
  T -->|"/_ctl/hold and /_ctl/release"| B

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 to #3887 have merged too, so #3888's base is main. The checkpoint stack uses those operations to free the ledgers of retired attempts and to clear a dead execution's capture files before a restore. The capture-ledger cache bound (#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 (merged)
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 (merged)
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 (merged)
4. feat(checkpoint): continue environment server episodes from their boundaries #3885 (ananthsub/partial-ckpt-environment): feat(checkpoint): continue environment server episodes from their boundaries (merged)
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 (merged)
6. feat(checkpoint): declare twenty training verifiers stateless with replayable verification #3887 (ananthsub/partial-ckpt-verifier-declarations): feat(checkpoint): declare twenty training verifiers stateless with replayable verification (merged)
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 (this PR)
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 the in-process end-to-end test from feat(checkpoint): adapt initial stateful environments #2946, replaced by real processes and a real crash (every Gym process group is killed).
• The Blackjack and Workplace Assistant adapters from feat(checkpoint): adapt initial stateful environments #2946 are not carried (see the resources PR).

Issue

No tracking issue exists for this re-cut. The design and the mapping from the old stack are described in this PR series, and the old stack's PRs (#2939 to #2946) carry the original discussion.

Validation

Run on this branch, on the current main:

• NEMO_GYM_CHECKPOINT_E2E=1 RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/e2e/checkpoint: passed, except the opt-in real-model and scale tests, which were skipped.
• The restore-again scenario fails without the ledger fix in the policy model PR, at the second restore, with one and with two policy workers: capture ledger for again-1-a1 already holds rows from another execution.
• RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_checkpoint_*.py: passed.
• The full core unit suite (tests/unit_tests, eight processes, without test_opensandbox_compose.py, which needs the optional opensandbox package): passed except a sandbox retry test, which fails the same way on main in this development environment.
• pre-commit run --files <files changed by this PR>: all hooks passed, and no hook modified a file.
• Every commit carries a Signed-off-by line.

Rollout evidence

• Process-level crash and restore: the end-to-end suite above, with a fake inference backend: all scenarios pass.

• Scale, measured on one workstation with checkpoint_dir on its local disk, before the scale case moved to the per-rollout hold, so the in-flight column reflects the earlier global hold; a run on shared Lustre storage is tracked separately, with every Gym server and the fake inference backend (weather rollouts on the typed single_agent_turn request path, token capture and generation cuts on, each rollout's second model call held so most rollouts are in flight at the checkpoint), run with NEMO_GYM_CHECKPOINT_E2E=1 NEMO_GYM_CHECKPOINT_SCALE=<N> pytest tests/e2e/checkpoint -k scale -s:

  Rollouts	Policy workers	In flight at checkpoint (exported)	Prepare (s)	Commit (s)	Restore (s)	Replacement failures
  8,000	1	4,201	0.67	0.60	0.60	0
  8,000	2	4,083	0.56	0.69	0.56	0
  16,000	2	14,320	7.57 (model 7.38)	4.64	4.97	0
  16,000	2 (rerun)	12,228	3.60 (model 3.53)	1.83	3.23	0

  Every rollout either finished before the checkpoint or was exported, and every exported rollout completed after the restore. These numbers come from the lane branches this stack was cut from; the scale test was not rerun on this stack.

• Real model, on one workstation: Qwen/Qwen3-4B-Instruct-2507 served by vLLM 0.24.0 on one RTX 6000 Ada (--enable-auto-tool-choice --tool-call-parser hermes), driving real Gym server processes that the opt-in tests in test(checkpoint): add a process-level end-to-end suite driven by coordination #3889, feat(checkpoint): checkpoint evaluation runs from rollout collection #3893 and feat(checkpoint): resources, agent, and environment servers with several workers #3909 checkpoint mid-run, kill and restore:

  • Saved rollouts. Weather rollouts on the typed single_agent_turn request path, sent 20 ms apart until the first one finished. The checkpoint caught about 20 of 26 inside their agent turn. After every Gym process group was killed and the checkpoint restored into fresh processes, every continued rollout finished with reward 1.0, as did every rollout that finished before the checkpoint.
  • Restarts. Half the rollouts used a restart_only resources server. Prepare listed every restart still in flight and no saved rollout, and the commit left the restarts out. After the kill, the restarts were retired and started over from their input while the saved rollouts continued, and every rollout finished with reward 1.0.
  • Rollout collection. 48 rows, four at a time: a SIGUSR1 checkpoint, four more rows sent, then the collector killed with SIGKILL while the Gym servers kept running. The resumed run retired the four rows sent after the checkpoint and sent them again as attempt 1, and all 48 rows finished exactly once with reward 1.0.
  • The first two run with one and two policy workers and with one and two workers per server. Each scenario passed repeated runs.
  • NEMO_GYM_CHECKPOINT_E2E=1 NEMO_GYM_MULTIWORKER_E2E=1 NEMO_GYM_CHECKPOINT_VLLM_URL=http://127.0.0.1:8100/v1 NEMO_GYM_CHECKPOINT_VLLM_MODEL=Qwen/Qwen3-4B-Instruct-2507 RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/e2e/checkpoint -k real_model

• Not covered yet: a real vLLM worker serving generation cuts, and several hosts.

Compatibility

Tests only. The suite is skipped unless NEMO_GYM_CHECKPOINT_E2E=1 is set.

[copy-pr-bot]

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.

[macandro96]

Non-blocking experiment question: What filesystem backed checkpoint_dir for the commit timings reported in the PR description? The test passes tmp_path / "ckpt", so I would expect node-local temporary storage unless pytest's temp root was explicitly placed on shared storage.
If node local dir was used - do you think it makes sense to do a version where we use lustre as filesystem to observe the gap?

[macandro96]

[RISK] This negative branch can pass on any server failure or malformed response, so it does not prove that skipping resource restore is what changed the behavior. Please require status 200 and assert the expected negative result (reward == 0.0; ideally also the fresh counter value 2). That keeps this control from passing vacuously.

[ananthsub]

Agreed. The negative control now requires status 200 and a reward of exactly 0.0, so a server failure or a malformed reply fails it.
The counter server's verify reply does not include the counter value, so the test cannot assert the fresh count of 2 directly; the comment states it, and a reward of 0.0 on a successful verification is only possible with a counter other than 6.

[macandro96]

[BLOCKER] This scale case can issue 16k requests through httpx. The repository guidance explicitly forbids httpx.AsyncClient in async paths because its connection pooling becomes O(n^2) and hangs at this concurrency. Please route these requests through Gym's global aiohttp client (server_utils.request()); otherwise the test can hang before exercising checkpointing.

[ananthsub]

Agreed. Every asynchronous request in the suite now goes through Gym's global aiohttp client, through a small GymHTTP helper in checkpoint_deployment.py.
That covers the scale case's /run calls and the other tests' /run calls, status reads and capture manifest reads, which also used httpx.
The later layers' tests move the same way: the multi-worker session test, the MCP test, and the random-sequence suite.
The helper sends each request once and fails if the server cannot be reached, so a /run cut off by a crash is never sent again as a second episode attempt.
The scale case raises the test process's per-host connection limit first, since every /run holds a connection until its episode replies.
Run once with NEMO_GYM_CHECKPOINT_SCALE=1000 to confirm the path works.

[macandro96]

[RISK] This predicate calls deployment.backend_calls(), which performs a blocking requests.get() and returns/deserializes the entire call list on every 100 ms poll. At scale that blocks the same event loop driving the rollout tasks and makes the timing partly measure the poller. Please add a cheap count endpoint and poll it asynchronously through aiohttp, or otherwise avoid synchronous full-list polling.

[ananthsub]

Agreed. The fake backend now has GET /_ctl/call_count, which returns how many calls it received and how many returned, without the call list.
Every test that waited on the call list now polls that count asynchronously through the aiohttp client, the scale case included, so the poll no longer blocks the event loop that drives the rollouts.

[macandro96]

[RISK] Replacing the event here strands any requests already awaiting the previous event if /_ctl/hold is called again before /_ctl/release. The subsequent release sets only the new event, so those calls hang forever. Please reuse the active event, or reject/release an existing hold before installing another one.

[ananthsub]

Fixed: a hold placed before the previous one was released keeps the event its held calls already wait on, so the next release frees all of them.
A hold after a release starts a new event, as before.
Covered by test_a_second_hold_before_a_release_keeps_the_calls_the_first_one_held in tests/e2e/checkpoint/test_fake_backend.py, which serves the backend in process, holds twice and releases once.
It needs no Gym servers, so it always runs; on the previous code the first call never returned.

[macandro96]

[RISK] run.done() is a client-side scheduling observation, not the server's prepare classification. A response may already have completed server-side while its HTTP task becomes done immediately after this snapshot, making finished_during_checkpoint flaky. Please derive the partition from the prepare replies/server state, or synchronize completed tasks before snapshotting.

[ananthsub]

Agreed. The real-model test now takes the partition from the environment server.
It commits with every dispatched rollout in scope, and the environment server exports exactly the episodes it still holds, so the rollouts it leaves out are the ones that finished before prepare.
The test then asserts that no exported rollout's /run has replied while Gym is prepared, which is the property the old snapshot was after, without reading it from client task timing.

[macandro96]

[RISK] These per-kind prepare results are discarded, so a timed stage can report blockers or failures without failing here; only a later second prepare across all participants is asserted. Please retain and assert every stage result, or collect timings around one coordinated prepare. Also, the global hold_after=count above does not guarantee that every rollout's first call completed—fast rollouts can consume multiple call slots—so it does not reliably establish the per-rollout boundary described by the comment.

[ananthsub]

Both fixed.
The scale case keeps every stage's prepare result and asserts each one prepared, naming the stage and its blockers if not.
The fake backend now has a per-rollout hold, {"after_calls_per_rollout": 1}, which lets each rollout's first call return and holds its second, keyed by the rollout the capture admission names.
The test waits until exactly one call per rollout has returned and every second call has arrived, so every rollout is at the same boundary when the checkpoint starts.
With that, no rollout can finish first, and the test asserts the checkpoint exported every one of them.

[ananthsub]

On the filesystem question: node-local.
The scale test writes checkpoint_dir under pytest's tmp_path, so the commit and restore timings in the description were measured on the machine's local disk, not on shared storage.

Yes, a Lustre run is worth doing, because training runs write checkpoints to shared storage.
Each participant writes a records file and a manifest and syncs its directories, and a restore imports many capture ledger files, so those steps can cost noticeably more on Lustre.
I'll run the scale case on the cluster with checkpoint_dir on Lustre and on node-local storage at the same rollout counts, and compare commit and restore per participant kind.
That is tracked separately, so it does not hold up this PR.

[github-actions]

🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA-NeMo/Gym/actions/runs/38821201371