Repository navigation
fix(checkpoint): harden Gym turn recovery for failures and scale - #4501
Closed
macandro96 wants to merge 32 commits into
Closed
macandro96 wants to merge 32 commits into
macandro96 wants to merge 32 commits into
Conversation
This was referenced Oct 5, 2026
macandro96
force-pushed
the
amahishi/gym-turn-recovery-pr1
branch
from
October 5, 2026 21:54
fcff1d2 to
f881d6a
Compare
macandro96
force-pushed
the
amahishi/ckpt-v2-gym-integration
branch
from
October 5, 2026 21:54
26e5ca4 to
158bb81
Compare
1 similar comment
macandro96
force-pushed
the
amahishi/ckpt-v2-gym-integration
branch
from
October 5, 2026 22:52
158bb81 to
e1d1959
Compare
macandro96
force-pushed
the
amahishi/gym-turn-recovery-pr1
branch
from
October 5, 2026 22:52
f881d6a to
1ed92e7
Compare
1 similar comment
This was referenced Oct 6, 2026
macandro96
force-pushed
the
amahishi/ckpt-v2-gym-integration
branch
from
October 7, 2026 01:36
e1d1959 to
160e8ec
Compare
4 tasks done
macandro96
force-pushed
the
amahishi/gym-turn-recovery-pr1
branch
from
October 11, 2026 19:25
b11af3a to
341b9ad
Compare
macandro96
force-pushed
the
amahishi/ckpt-v2-gym-integration
branch
from
October 11, 2026 22:01
160e8ec to
cf124ff
Compare
macandro96
force-pushed
the
amahishi/ckpt-v2-gym-integration
branch
from
October 15, 2026 02:17
cf124ff to
4558c22
Compare
macandro96
force-pushed
the
amahishi/gym-turn-recovery-pr1
branch
from
October 15, 2026 02:17
341b9ad to
399aff4
Compare
- Snapshots: abort the Gym phase when a train step commits, warn when snapshots keep losing to train steps, report gym_parked_seconds, and require stall_timeout_s to outlast a Gym park. - Gym control: set lease_grace_seconds from control_timeout_s, retry a transient resume on the failed instances only, make the coordinator errors picklable, share one restore cleanup through coordinator.restored(), and route discover/restore errors through the same unavailable path. - Adapter: drop the unused participant summaries and reply-shape checks; initialize adapters inside the shared spinup deadline. - Dispatch: replace RolloutDispatchAdmissionGate with a second DataPlaneCheckpointBarrier; use Gym's no-result failure classes. - Restore: pass restored_from_rollout_snapshot from setup instead of the gym_restart_fallback marker file; restore the trainer save to main's shape; name the Gym instances whose manifest disagrees with the ledger; require the Gym fields in v4 ledger attempts. - Config: reject the removed *_granularity keys in check_outdated_config; derive target-level order from RecoveryTargetLevel.rank; document the turn requirements in the exemplar YAMLs and guide. - Tests: drive the real GymCheckpointCoordinator through a shared fake Gym actor, test the adapter against real in-process Gym participants, and assert Gym restore and pre-cut model calls in the functional test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…L needs Move the Gym pin from main's 6f295d7d to bf631d0a on NVIDIA-NeMo/Gym amahishi/rl-gym-turn-recovery-pin: - Gym main at 9ef5ed0f, which carries the merged checkpoint stack (#3882-#3889) and descends from main's pin 6f295d7d; - NVIDIA-NeMo/Gym#3894: checkpoint Workplace Assistant sessions; - NVIDIA-NeMo/Gym#4690: name the episodes an agent's own /run drives in the agent commit reply (run_episode_ids); - the checkpoint test agent the turn-recovery functional test runs, with its participant installed once per app now that Gym installs it twice. - NVIDIA-NeMo/Gym#3895-#3898: Gymnasium servers on the shared resources server setup, and checkpointing for Blackjack, indirect prompt injection and the proof refinement agent. Gym adds requests to its core dependencies, so the lock records it for nemo-gym; requests is already locked, and every resolved package keeps its version. Gym's coordination now sends each control body pre-encoded as data rather than json, so the adapter test's in-process client forwards both. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Address the review of turn-level recovery against the Gym main pin: - Restarts: carry PrepareResult.restarts() back from each Gym actor, leave them out of the commit scope (Gym refuses a commit that names one), skip them in the drain, accept them in the cut, and record them in gym_checkpoint.json (schema 3) so a restore checks saved episodes plus restarts against the ledger and starts the restarts over from their tasks. - Forget: retire leaves a mark per rollout on every Gym participant until the controller forgets it. Forget each rollout right after its retire, and each restored rollout once its continued attempt seals, before the next prepare; a deployment that refuses gets the retire again and then the forget, and one that still refuses keeps them queued. - Resume order: Gym resumes before the snapshot is published, so new rollouts wait only for the Gym commit and TQ capture. A resume that still fails after a successful cut is reported, not raised: Gym's lease reopens the deployment, and the snapshot is kept. - Ownership: read the committed episodes an agent's own /run drives from Gym's run_episode_ids commit reply instead of Gym's on-disk agent records, and name the two kinds of agent session in Gym's terms. Document the resume order, restarts, and retire/forget in the guide. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…up test setup gym_checkpoint.json is new in this PR and only RL writes it, so start its schema at version 1 and parse it plainly. The version check and the checks of what it means (non-empty IDs, no duplicates, matching topologies, no episode both saved and restarting) stay; a file of the wrong shape still fails to load, with a KeyError or TypeError instead of a tailored message. Build the admitted, dispatched turn-level group the checkpoint race tests use in one helper instead of seven copies, and merge the two equivalent ledger-test helpers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…t agent The Workplace turn-recovery test needs a published checkpoint holding a turn whose calendar write was already applied. It held that boundary with a Gym checkpoint test agent that subclassed private participant internals and was switched by environment variables. A test-only proxy between Gym's policy model server and vLLM now does it through production paths. It holds the first policy call that carries a tool result: Gym counts a session at its boundary while it waits on a policy call made right after one, and a held call never blocks a checkpoint prepare. The rows use the plain workplace_assistant_simple_agent, so the test needs no Gym-side agent. After the write, the proxy forces one calendar_search_events call and then sends tool_choice none, so no turn writes a second event. The restored episode's search must find the event written before the crash exactly once, which checks that the Workplace state itself was restored, not only the model calls; its reward must be 1.0. The selector accepts only the held episode, since another episode also counts as at its boundary while its own post-tool call is in flight. The proxy forwards every path, so token capture is unaffected, and listens on the node's address so replicas on other nodes reach it. Passed with one Gym replica on one node and with two replicas on two nodes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
The pin no longer needs commits that exist only on an unmerged branch. The turn-recovery test holds its boundary through a policy-call proxy, so it no longer uses the checkpoint test agent (NVIDIA-NeMo/Gym#4773), and the agent commit reply's run_episode_ids (NVIDIA-NeMo/Gym#4690) is on main. The other pin-only commits, the indirect prompt injection and proof refinement checkpoint exports (NVIDIA-NeMo/Gym#3897, #3898), are for environments RL's tests do not use. Gym's pyproject.toml is identical to the previous pin's, so uv.lock is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…unless allowed A restart-only Gym server reports its in-flight rollouts as restarts at every checkpoint, so under target_level='turn' they quietly start over from their tasks after a crash, and nothing told the user. Restart-only is Gym's default for resources servers that declare no checkpoint mode and for agents without checkpoint session support. Before the first rollout, SC reads every Gym participant's checkpoint status: a resources server reports its mode, and an agent reports restart_only (NVIDIA-NeMo/Gym#4957; an agent on an older Gym is not flagged). If any server is restart-only, setup fails and names it, unless the new rollout_recovery.allow_restart_only_gym_servers is true, which accepts it with a warning. Setting it when no server is restart-only warns that it has no effect, and setting it below target_level='turn' is rejected. Each saved rollout snapshot also logs gym_restarted_episodes, the episodes that would start over after a crash, on its save line and as a metric. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…ly reachable replies Only typed environment servers take part in Gym checkpoints, and turn recovery refuses those at setup; the legacy_agent relay it accepts is not a participant. The adapter's environment-ownership branch was therefore unreachable. It is removed, and discover() now fails, naming the server, if an environment participant ever appears, instead of attributing its episodes wrongly. The branch can return, tested against Gym's real participant, once typed environment servers are supported. The adapter tests no longer build replies for environment servers or the participant manifest fields the adapter never reads. The stateless-resources case runs on Gym's real participants in the save/restore test. Hand-written replies remain for states valid participants cannot produce, plus three that had no test: a commit from an unexpected participant set, an export outside the requested scope, and a /run episode two agents claim. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…s rollout metrics Main now passes full_results to _compute_rollout_metrics and computes per-group reward diagnostics from each result's full_result (#4435). This test predates that and still stubbed the old signature and built results without full_result, so it failed in CI's Nemo Gym unit lane after the rebase. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…rom discovery Gym main now has NVIDIA-NeMo/Gym#4957: every checkpoint participant reports its checkpoint mode, and discover() returns it. The restart-only check reads that mode instead of asking every participant for its status again, and no longer depends on an agent status field that never shipped. The pin also brings NVIDIA-NeMo/Gym#4997's renames: the policy model's checkpoint_policy moves into a checkpoint block, restored_pending is unclaimed, and control error codes carry a checkpoint_ prefix. RL reads the codes from Gym's error classes; only test fixtures named them. NVIDIA-NeMo/Gym#3909 lets checkpointing agents and resources servers run several workers, so the guide drops the num_workers: 1 requirement. Gym's dependencies are unchanged, so uv.lock is too. Signed-off-by: Anish Mahishi <amahishi@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…and 3 Turn-level recovery moved the rollout-recovery ledger to schema 4: a group records its restore level instead of its recovery granularity, and each attempt records its Gym instance and attempt. Schema 4 was the only one accepted, so a run resumed from a checkpoint main wrote before this change failed at startup. #4551 made main load schema 2 checkpoints again. Schemas 2 and 3 now load and are rewritten in the schema 4 shape. A group's recovery_granularity becomes its restore level; the values are the same. Each attempt starts with no Gym turn state, as a new attempt does, because those schemas predate it. An old checkpoint that already carries a restore level, Gym attempt state or the turn level is refused rather than guessed at. Signed-off-by: Anish Mahishi <amahishi@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Main's Lens export (#4730) sends each rollout save and restore metric through a fixed mapping, and a drift test fails when the controller logs a key the mapping does not name. Turn-level recovery logs four such keys: how long Gym held rollout dispatch for a save, the episodes that save could not capture, and the siblings a restore continued or restarted. gym_parked_seconds becomes a save duration phase, gym_parked. The other three are sibling counts, beside the reused and rerun siblings a restore already reports: continued, restarted and gym_restarted. Signed-off-by: Anish Mahishi <amahishi@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…out IDs Gym runs every row of a /run as its own task, and a stream that ends early cancels none of them. The sibling retry re-sent the missing rows under the same rollout ID while the old copies were still running in Gym. When a stream ends without some rows and attempts remain, the rows that never came back (and that Gym did not refuse at admission) are abandoned in the recovery ledger, queued for the next checkpoint to retire, and re-sent as fresh attempts under new rollout IDs; attempts that never reached Gym keep their identity. On the last attempt the rows are left to the group's abandon, which queues them the same way. A stream that ends cleanly but short spends its attempt like a broken one, so the admission re-send cannot send a lost row again under its old rollout ID. The streamed-completion identity check reads the current gate ID from the ledger, and the functional-test hook resolves its group by ID and forwards the new recorder hook. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…eckpoint saved gym_checkpoint.json recorded actor instances, episodes, staging keys and restarts, but not which Gym participants committed each instance. Gym restores only the participants that are live, so one missing from the restored deployment, for example a resources server dropped from the config, had its state silently lost while the restore succeeded. Each commit summary now carries the names of the participants that committed, and the manifest records them per instance. Before Gym installs anything, the adapter refuses a restore whose live deployment lacks any of them, naming the missing ones, so the mismatch leaves Gym untouched. A live participant the checkpoint does not name is allowed: it has nothing to restore. Gym already verifies each participant's records digest when it loads its state, and the adapter already refuses state from another checkpoint, so RL does not read Gym's state files. The manifest schema moves to version 2, and the loader checks the version before any field, so a manifest written before participants were recorded is refused by version rather than by a missing field. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…y rollout After a Gym commit, the checkpoint waits for the few episodes Gym did not export because their reply is already on the wire. It did so by rebuilding the full Gym inventory from the ledger every 50 ms. That scan is synchronous and scales with every live rollout -- about 0.2 s per pass at 200k rollouts -- so it blocked the event loop that delivers those very replies, while every rollout stayed parked behind the checkpoint. The drain still polls every 50 ms, but each check now looks up only the pending episodes, through a ledger lookup that applies the same ownership rule as the inventory. The drain's timeout and stranded-session errors now report a count and a short sample rather than every episode identity. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
The NemoGym actor raised a failed row's error, which ended the stream for every row of the dispatch while their /runs kept going in Gym. The retry then re-sent every row without a result, including siblings that were still running. The actor now yields a GymRowFailure naming the row, the way it yields GymAdmissionClosed, and the stream carries on. The single-controller consumer raises a data failure at once, since the prompt fails as a whole, and holds a transport failure until the stream ends, so only the failed rows are re-sent, under new rollout IDs. The sync path, which retries whole batches, still fails the batch on a failed row and stops that instance's stream. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
When checkpoint admission refused part of a prompt group, the collector re-sent those rows only after the original stream finished. A cohort verifier such as GenRM keeps the admitted siblings waiting for the refused ones, so that stream never finished and the refused rows were never sent. Each refused row is now re-sent on its own, from inside the stream, once the dispatch gate reopens, with a bounded backoff if Gym is still closed; the admitted rows keep streaming. A refusal keeps the row's identity and spends no row attempt. Every re-send runs as its own task, so a failed re-send cancels the stream still blocked on it, and a stream failure or cancellation cancels and awaits the re-sends in flight. A refused row that was already re-sent when its stream broke is in Gym again, so it is replaced with the rest. When no row of a cohort came back and the whole cohort is replaced, the rows also move to the group attempt the ledger now records, the coordinate a restart rebuilds: GenRM keeps failed group attempts, so new rollout IDs alone would send every row straight back to the same failure. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Keeping a group's stream open after one row failed left the failed attempt dispatched in the recovery ledger until every sibling returned. A checkpoint drain waits on every dispatched attempt Gym did not export, and Gym cannot export a row whose /run already failed, while the siblings that would end the stream stay parked until the checkpoint resumes. The drain could only time out, failing that snapshot. Under sibling recovery the failed row is now replaced as soon as its failure arrives: its attempt is abandoned and queued for retirement, and the fresh attempt waits for the re-send. The cleanup after the stream finds it already replaced. Under prompt-group recovery the whole cohort is replaced, so a failed row ends the stream at once, as before the stream kept going. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
macandro96
force-pushed
the
amahishi/gym-turn-recovery-pr1
branch
from
October 28, 2026 14:58
d1d83f1 to
70ba0c1
Compare
macandro96
force-pushed
the
amahishi/ckpt-v2-gym-integration
branch
from
October 28, 2026 14:58
0714200 to
3f24e0a
Compare
Contributor
Author
|
/ok to test 3f24e0a |
macandro96
force-pushed
the
amahishi/gym-turn-recovery-pr1
branch
4 times, most recently
from
November 2, 2026 16:23
789cedf to
3ca7028
Compare
Contributor
Author
|
Closing in favor of #4500, which now carries these six commits on top of the turn-level recovery stack, rebased onto
As agreed on #4500, |
An error occurred while trying to automatically change base from
amahishi/gym-turn-recovery-pr1
to
main
November 2, 2026 20:50
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 does this PR do ?
Hardens Gym turn-level rollout recovery (#4500) against broken or partly failed rollout streams, checkpoint-refused rows inside a cohort verifier, Gym participant drift on restore, and drain cost at large rollout counts.
Stacked on #4500. The base branch is
amahishi/gym-turn-recovery-pr1, so this diff shows only the 5 commits below. It will be retargeted tomainonce #4500 merges. #4710 is folded in as the last commit.Issues
How it works
Broken and partly failed streams. Gym runs every row of a
/runas its own task, and an ended stream cancels none of them. RL therefore never re-sends a row under a rollout ID Gym may still be running:sequenceDiagram autonumber participant RM as RolloutManager participant L as Recovery ledger participant G as Gym replica RM->>G: /run rows 0..2 (rollout IDs a0, a1, a2) G-->>RM: row 0 G-->>RM: row 1 failed (yielded, stream continues) G-->>RM: row 2 Note over RM: transport failure: raised once the stream ends RM->>L: abandon row 1 and queue a1 for the next checkpoint to retire RM->>L: mint a fresh attempt (new rollout ID b1) RM->>G: /run row 1 (rollout ID b1) G-->>RM: row 1Checkpoint-refused rows. A row Gym refuses at checkpoint admission never ran. It is re-sent on its own, from inside the stream, once the dispatch gate reopens, while the admitted rows keep streaming. A cohort verifier such as GenRM can hold the admitted rows until the refused one arrives, so waiting for the stream to end first would deadlock. When no row of a cohort came back and the whole cohort is replaced, the rows also move to the new group attempt, so GenRM does not send them straight back to the failed one.
Changes
fix(rollout): re-send rows a broken Gym stream dropped under new rollout IDsfix(checkpoint): refuse a Gym restore that lacks a participant the checkpoint savedgym_checkpoint.jsonrecords the names of the participants that committed each instance. Restore refuses a live deployment missing any of them, by name, before Gym installs anything. Extra live participants are allowed.perf(checkpoint): wait on Gym on-wire replies without rescanning every rolloutfix(rollout): keep streaming a group's rows when one Gym row failsGymRowFailureinstead of raising. The single-controller consumer raises a data failure at once and a transport failure once the stream ends, so only the failed rows are re-sent. The sync path still fails the batch.fix(rollout): redispatch checkpoint-refused siblings independently(was #4710)Lines: source +494 / −153, tests +999 / −28.
Usage
No new configuration.
Testing
bf631d0a6):tests/functional/grpo_async_gym_single_controller_turn_recovery.sh: passed;examples/nemo_gym/grpo_workplace_assistant_two_replicas.yaml) passes on Slurm on29dad27188. Restored episodes split 3/3 across replicas, and the selected turn resumed on its original replica.Before your PR is "Ready for review"
Pre checks:
Additional Information
main.🤖 Generated with Claude Code