Repository navigation
perf(checkpoint): batch route snapshots within cut inventories - #4843
Open
macandro96 wants to merge 51 commits into
Open
macandro96 wants to merge 51 commits into
macandro96 wants to merge 51 commits into
Conversation
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…fast on stranded ones A legacy_agent environment server only relays /run, so its environment participant never holds the episode; the Simple Agent's legacy session does. The commit summary counted only environment exports as owned, so every in-flight Workplace episode looked unexported and the drain waited for replies that cannot arrive while Gym is prepared. Every periodic checkpoint timed out. When no environment participant exported an episode, a legacy agent session (a record with an episode step boundary) now owns it. The adapter reads the agent's committed records through Gym's digest-checked store. A parked agent session that nobody owns fails the checkpoint immediately instead of waiting out the drain deadline. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…oint refusals An attempt was marked DISPATCHED before it waited at the dispatch gate, so a checkpoint could count an attempt that cannot reach Gym until the checkpoint releases the gate; its drain could only time out, fatally at a trainer checkpoint. A /run that reached Gym after prepare closed admission was refused with 409 checkpoint_parked, charged as a data failure, and stayed DISPATCHED while it retried behind the closed gate. The ledger now marks an attempt DISPATCHED inside the gate admission, in the same mutation as the Ray submission. NemoGym takes /run failures as Gym sidecar rows so a checkpoint_parked refusal names its row without ending the stream. The refused attempt returns to RESERVED with its identity and is re-sent once admission reopens, spending no row attempt and no data budget. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Two-process crash/restore test for a Workplace Assistant turn. Phase one publishes a bootstrap rollout snapshot that holds one turn after its calendar mutation, then is killed; phase two restores it and must dispatch that episode once, as Gym attempt + 1 on its original replica, and train every step exactly once with no invalid or poisoned rows. The turn is captured deterministically: the Gym checkpoint test agent holds the first post-mutation boundary until a real checkpoint parks it, and the rows force a single calendar_create_event whose sentinel the selector requires exactly once in the saved Workplace state. The selector copies the chosen bootstrap snapshot before step 1's trainer checkpoint prunes it, and phase two drops every trainer checkpoint so the restore starts from that anchor. The hook records dispatches once the ledger marks them, plus any checkpoint refusals, and patches run_rollout so RolloutManager's NeMo-Gym type checks still hold. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…kpoint Abandoning or discarding a group drops its attempts from the recovery ledger, but an episode Gym already admitted keeps running. Gym exports it at the next commit, outside the candidate set, and every checkpoint fails until the episode finishes on its own. The ledger now records each dropped attempt Gym may still hold (dispatched, or reserved with restored Gym state), in memory only. With dispatch closed and before preparing Gym, while Gym is still idle (it refuses a retire once a checkpoint is open), the controller retires them and forgets only what it retired, so a drop recorded meanwhile waits for the next checkpoint. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
…ing the run Three ways a Gym checkpoint could take the run down with it: - Control calls had no RL-side timeout. Gym enforces deadline_ts only once it runs a call, so an actor with a stuck event loop, or a call Ray never scheduled, left the coordinator waiting forever with dispatch admission closed. Each call is now bounded by its own deadline plus a short reply grace; an unanswered instance is reported as a TimeoutError and its Ray call is cancelled. Giving up is safe: Gym's lease reopens a parked deployment if the resume is lost too. - The periodic rollout checkpoint treated every Gym failure as a broken invariant and ended training. A checkpoint that could not park, went unanswered, hit a dead actor, or failed in Gym's control plane changed nothing RL holds, so it is now retried under a gym_unavailable reason, still bounded by max_consecutive_failures. Ownership, topology and reply mismatches still abort. Gym's CoordinationError cannot be pickled and reached RL as an untyped UnserializableException, so the adapter raises a picklable GymCheckpointUnavailable in its place for live checkpoint calls. - Rollout deadlines kept ticking while Gym held rollouts parked, so a long checkpoint could time out healthy rollouts as soon as it resumed. The request-deadline registry now tracks named holders (ported from #4266), and a Gym checkpoint holds its own pause from closing admission until Gym resumes, without releasing a colocated engine's pause. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Parking a turn checkpoint waits for every in-flight agent and environment step to reach a boundary, but the prepare deadline reused token_capture.control_timeout_s, which bounds each control-plane call. The only way to give long steps time to park was to loosen unrelated token-capture RPC deadlines. rollout_recovery.gym_prepare_timeout_s (default 300 s) now sets the prepare deadline alone; renew, commit, retire, restore and resume keep the control timeout. Gym extends its lease from each call's deadline, so a longer prepare cannot let the lease lapse. The value is operational, so it is left out of the bootstrap fingerprint. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
With turn-level recovery, a trainer checkpoint parked Gym and committed its state inside the trainer save. A Gym prepare that missed its deadline, or a transient Gym control error, therefore failed the save and ended the run; the periodic snapshot path already retries those. The trainer checkpoint now saves weights, dataloader, TQ and the recovery ledger without pausing Gym, and marks itself with gym_restart_fallback.json: it holds no Gym state. Gym state lives only in rollout snapshots. The trainer save wakes the snapshot pump, which waits for the checkpoint to be durable off the training loop and then captures Gym anchored to it, with the gym_unavailable retry and max_consecutive_failures budget it already has. The trainer cut no longer counts as a rollout snapshot, so that snapshot is not skipped for lack of data-plane mutations. Restore prefers the newest rollout snapshot. A trainer checkpoint restored on its own restarts its unfinished turn-level rollouts from their tasks; one with neither Gym state nor the marker is refused. A reserved attempt that continued earlier restored Gym state, and that this restart did not restore again, starts over as well. Ported from the trainer-boundary fallback in #4117, without its forced boundary snapshot, which blocked training on the asynchronous weight write. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…tart fallback Turn recovery now keeps Gym state only in rollout snapshots: the trainer checkpoint never pauses Gym, marks itself with gym_restart_fallback.json and wakes the snapshot loop. Document that, the restore of a trainer checkpoint without a later snapshot, and rollout_recovery.gym_prepare_timeout_s. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…de key Address review feedback on turn-level rollout recovery: - Derive a group's recovery granularity from its persisted restore level instead of storing both; recovery_granularity is no longer persisted. Schema 4 is new in this change, so no released checkpoint carries it. - GymCheckpointEpisode.capture_key is the single source of Gym's attempt-qualified key; the ledger and the controller use the wire type directly instead of converting tuples at each call. - Refuse turn recovery behind an environment server other than legacy_agent: typed servers build EpisodeId from the restored "<id>-a<N>" key, which Gym rejects, so every restored turn would restart until its row budget ran out. - Reject target-level overrides with target_level=turn: Gym exports every live episode, so a coarser route would fall outside the commit scope. - Mint the Gym checkpoint control token per run, classify a failed environment-server episode by Gym's terminal flag, and keep the /run timing on refused rows. - Abandon a reserved Gym attempt that this restart did not restore again. - Pass gym_instance_id and dispatch_recorder to run_rollout explicitly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Move the event writer and the ledger-group lookup that the turn, sibling and rollout recovery hooks each copied into _single_controller_hook_utils.py, drop the turn hook's unused delegate passthrough, and forward gym_instance_id and dispatch_recorder from the sibling hook. The sibling recovery script now checks the persisted restore_level, since recovery_granularity is no longer saved. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Note that only the current rollout_recovery.pt schema can be resumed and that target-level overrides are rejected with target_level=turn, since a checkpointing Gym deployment exports every live episode. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
- 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 120d516f, a merge of the pre-merge checkpoint stack that is on no Gym branch head and does not descend from main's pin, 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 da9f6714; - 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. Relock with uv 0.11.18: Gym moves ray and jinja2 out of its core dependencies and adds requests; 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>
…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. 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>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
NeMo Gym renamed two generation-cut wire fields so every model names the same thing the same way: - GenerationCutPrefix names its attempt attempt, as EpisodeId does (it was attempt_index). - GenerationCutContinuation names the staged prefix prefix_token_count and prefix_digest, as the worker's acknowledgement already does (they were generation_token_count and digest). The vLLM worker reads and checks the renamed fields. NeMo RL's own names for a restored prefix's length follow them: the output-limit helper's argument, _RestoredPrefixTerminal and its classifier, and the inherited count in the capture state. A call's total output keeps the name generation_token_count, because it is a different count. The prefix recovery snapshot reads the renamed keys from Gym's checkpoint records. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Correctness: - Replay a restored prefix through vLLM's own incremental detokenizer, primed on the original prompt, so SentencePiece keeps its leading space and a stopped call drops its stop token or stop string exactly as vLLM did. - Keep calls using stop strings, bad_words, or presence/frequency penalties out of generation cuts; they regenerate after a restore. - Apply a restored call's remaining budget even when the request set neither max_tokens nor max_completion_tokens. - Abort a failed request's capture off the event loop when a cut or terminal write holds its lifecycle lock. Simplification: - Request CUMULATIVE vLLM outputs and drop the two delta accumulators. - Replace the active/frozen/sealed capture buffers with one token list and a staged count. - Move the snapshot fence into a context manager. - Remove the no-op completed-ack fallback, the snapshot checks Gym repeats in begin_call, the unshipped staging-key compatibility branch, and the duplicate selected-token logprob loop. - Use Gym's GENERATION_CUT_ROUTE (pinned by a test) and EpisodeId.capture_key; pass generation_prefix_cuts_enabled as an explicit keyword. - Share the functional-test snapshot helpers between turn and prefix recovery. Observability: log the fence hold time and Gym's cut/failed/skipped counts per rollout checkpoint, and restored calls/prefix tokens per training step. Tests cover the replay, uncuttable calls, abort, cut route, fence error paths, setup guards, and metrics; docs describe the prefix target's requirements. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
- Keep a call with a per-request seed out of generation cuts when it samples: vLLM seeds its own generator at admission, so a restored call would reuse the draws that chose its prefix. - Clear the chunk rows of an aborted call, and of a call whose terminal staging fails, through the token-capture snapshot gate; the abort hands the delete to a cleanup thread so it never blocks the serving loop. - Acknowledge a failed cut per call instead of failing the receipt for every call on the worker, roll back its frozen buffer, and check the effective output limit before freezing. - Regenerate a restored call from its turn boundary when its prefix cannot be fetched or does not match its admission, instead of failing the rollout; its restored rows are retired through the same cleanup. - Report a resumed call's restored prefix as completion tokens, as the uninterrupted call would, not only for a terminal prefix. - Observe only the tokens vLLM has published with their logprobs. The HTTP thread reads the engine's live lists while its loop appends, so a token read before its logprob waits for the next output instead of making the call uncuttable. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Make the parity test exercise many restored prefixes and separate recovery faults from run-to-run sampling noise: - Workload: four Workplace prompts per step whose closing turn lists 1..N (N = 150/250/350/450) at natural length. Group lengths are far apart, so prompt-group readiness order is deterministic. One prompt expects a different duration so rewards are not constant. Runs Qwen3-1.7B for 3 steps. - Selection/verification: require a recoverable prefix in every Workplace prompt group at the crash (Gym refuses dispatches that land during a checkpoint, so siblings may start late), and verify each cut returns to its Gym owner under attempt+1, restores its exact prefix digest, and completes once with its trained generated-token count equal to the restored prefix plus the regenerated tail. - Parity: classify token divergences relative to the restored prefix. A divergence inside the prefix is excused as phase-one noise only when the trained prefix tokens hash-match what the worker restored, and excused rows still require equal rewards and sample masks; a prompt divergence always fails. Gate vLLM generation_logprobs in aggregate over token- matched rows (bit-identical fraction, Spearman, max abs diff) since they are not batch-invariant, and require exact prompt-group readiness order. Report Kendall tau for arrival and sealing order, per-rollout timelines, and logprob statistics in parity-report.json before the strict checks run. - Hook: timestamp every event and record completion arrival separately from the ledger seal. - Relies on the Gym checkpoint test agent's natural-length closing turn (NEMO_GYM_TEST_PREFIX_MIN_TOKENS=0), pinned by the base branch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…very parity The prefix-recovery parity harness reads each restored cut from Gym's checkpoint records and reports it. NeMo Gym now names the staged prefix prefix_token_count and prefix_digest, so the harness, its report, and its unit tests use those keys. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
… test The prefix-recovery snapshot helper now imports _gym_recovery_snapshot_helpers as a sibling module. Put tests/functional on sys.path before loading it by file path, as running it as a script does. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Parity data and gates: - Compare runs on the trainer's own async_rl.log_full_train_data dump (read from the run's exp_NNN directory) instead of a recorder in the hook; the hook now only probes live staging keys per train step. - Require exact parity for every row. The one exception is a restored row whose phase-1 prefix already differed from the reference before the cut (vLLM sampling is not batch-invariant): it must start with exactly the restored tokens (prefix_ids_sha256), match up to its closing call, and keep its reward and admission; at least one restored row must match outright. - Gate generation logprobs on bit-identical fraction and max |diff|. - Default to one train step and say the test checks data, not learning (top_p=1e-6 zeroes the policy gradient). Removed: the ordering/timeline/Kendall report, the hand-rolled Spearman/Pearson gate and parity-report.json, the metric comparison, the prefix+tail==total tautology, the duplicate trained-once check, the unreachable successor Gym-manifest check, the single-cut selection fallback, and the swallowed Workplace layout errors. The JSON readers and the restored-prefix log parser are shared in _gym_recovery_snapshot_helpers. The L1 runner comment now says the test performs three runs. The Workplace calendar sentinel count and the Workplace test-row jq builder are shared with the turn-recovery test (_workplace_calendar_rows.jq). Gym: bump to 4f996d10, which drops the checkpoint test agent's min_tokens knob; the script no longer sets it. Also fix AdvantageStageConfig.from_master_config, which indexed LoggerConfig as a dict and crashed every single-controller run with async_rl.log_full_train_data=true. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…ixel copy A call's first generation-prefix row stored a full copy of its media tensors, but restore only used it to check that the reprocessed pixels were byte-identical; the terminal row stages the pixels again for training. Every media call cut by a checkpoint paid an extra pixel write under the cut's deadline, for a check a digest can do. Fingerprint each media tensor (dtype, shape, sha256 of its bytes) once, when the call is admitted and its media is captured off the event loop. The prefix row's digest-covered extras now carry those digests next to the media spans, with no attachments. On restore, the reprocessed media is checked against the staged items and digests by a pure helper in captured_media.py that names the mismatch (items, tensor set, dtype or shape, bytes), and the reprocessed tensors go on to the terminal row. This removes the media half of TQTokenSource.fetch_generation_cut (the cut now reads chunks with fetch_for_finalization), FetchedStagedCall.media, the hand-built imgs/imgs_sizes/num_frames rebuild, and the first-row attachments plumbing. Also from review: - Drop the getattr fallbacks that only served a test fake: the cut calls fetch_for_finalization, reads chunk.snapshot/extras, and calls _generation_cut_record_payload, which the mixin now declares. - Parse a row's media spans in one helper for both the retained-parent and the prefix-restore paths, so both reject malformed spans. - Extract the resumed media prompt check so it can be unit-tested. - Test the prefix round trip for image and video calls, every restore rejection, the extras-after-first-chunk rule, and a recovery request whose source media disappeared failing before any TQ read. - Restore the TURN periodic-checkpointing setup test and cover multimodal prefix recovery in its own setup test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…s end The prefix-recovery parity checks located a restored call at the start of its row's last unmasked segment. When top-k/top-p filtering is on, the trainer drops positions whose recomputed logprob is -inf from token_mask (#4154), and the rows these checks read are dumped after that pass, so a masked token inside the call made the check start the call too late, miscount its tokens, and fail a correct restore. A restored call closes its row, so the total_generation_tokens it logs on completion end the row. _assert_prefix_boundary now counts back from the row's end, requires the token before the call to be untrained, checks the completion's prefix + tail = total accounting, and requires the restored prefix at that position; a dropped or duplicated token moves the position. It returns where each call starts, which _rows_diverged_before_cut reuses to compare the reference row at the same position. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…onal test A Gym-v2 checkpoint is taken while vLLM is mid-decode on image prompts; the run is hard-killed and restored, and each restored call must reprocess its image, verify it against the staged media digests, resume from its exact durable prefix, and finish and train. The Reasoning model's free-form thinking makes token parity with an uninterrupted run flaky, so this test checks each restore without a reference run. - select --stateless-resources picks cuts from a stateless resources server (CLEVR's string_match), which exports no state to check; a Workplace caller must still pass its calendar sentinel. - verify-completion reads the run's log_full_train_data dump and applies the shared checks: each row trained once, each restored call completes once and starts with its restored prefix (_assert_prefix_boundary), and no restored prefix row is live when its step trains. - stop_phase1 and the artifact cleanup move into _gym_prefix_recovery_phases.sh, sourced by this and the Workplace test. - Registered in the GB200 Omni single-controller L1 list (full mode). Requires a Gym that includes NVIDIA-NeMo/Gym#3887. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
The cut writes a call's media spans and tensor digests into its first prefix row's extras, and restore recomputes the cumulative digest from the staged chunks; Gym's begin_call rejects a restore whose digest differs from the cut's. Nothing tied the two sides together, so a change on either side would first surface as a failed recovery on GPU. Cut a text call once and an image call once and twice, restore each through _resolve_generation_cut, and require the rebuilt digest to equal the one in the cut's receipt; a changed media tensor digest must change it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
… capture This PR removes the setup guard that rejected multimodal capture under prefix recovery, so the multimodal case of the parent stack's rejection test now fails on a later, unrelated processor check. Keep that test for router replay, which is still unsupported here; test_prefix_recovery_accepts_multimodal_capture covers the accepted case. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…orts
media_tensor_digests indexed its tensor map after an `or {}` guard, which
pyrefly cannot narrow, so it reported indexing a possibly-None mapping.
Iterating the guarded map's items avoids the second lookup.
Also sort two test files' imports, which ruff's isort hook flagged.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…t be read After rebasing onto #4507's fail-soft prefix restore, a call's media is restored from the same chunk fetch as its prefix, inside the same fallback. A prefix row whose media metadata cannot be read therefore regenerates the call from its turn boundary, like any other unusable prefix. Reprocessed media that does not match the staged digests is still rejected with retained_media_changed. Cover the first case and say both in the guide. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
A generation cut now stages the routes of the tokens it saves, and a resumed call restores them. vLLM keeps an active request's routes in its output state, so the HTTP client copies each step's new route chunks on the engine loop. The capture state keeps them beside its generated tokens. A cut lines them up with the token count it froze: the first chunk carries the full delta-aligned routes, and each later chunk carries one route per new token. On restore, the route fragments are checked against their digests and joined. The resumed call's response keeps the restored prefix's routes. The client also records the token count when it copies the chunks. vLLM keeps appending to its live token list, so only the tokens the chunks cover are observed. The rest come with the next output, whose chunks cover them. Routes that still do not line up fail only that call's cut; the call restarts from its last saved turn. Gym's generation-chunk records now carry extras, so the pin moves to 25062eb18, one commit past 4f996d10. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…dary routes NeMo RL now stages a multi-turn call's routes with one leading row for its parent's last token (#4610): that token only gets a route once the next call prefills it, and the reassembler writes that row back over the parent's placeholder. The sink requires the extra row whenever prev_len > 0. Generation-cut chunks keep their own convention, one route per staged token, because only the resume path reads them back. The sink rejected every such chunk of a multi-turn router-replay call, so the call could not be cut. Chunks are now checked against their token delta alone. The resumed call's canonical row is an ordinary capture row and carries the boundary row its own prefill computed, as an uncut call would. The merge that restores the prefix's original routes now starts after that row; it used to expect one row fewer and failed the call. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Routes are read once per cut instead of on every decode step. vLLM 0.29 only appends to an active request's route chunks, so a cut freezes its token count, then reads the chunks on the engine loop and keeps the rows its frozen tokens need. This drops the per-step copies, the chunk cursor and the side-channel attributes on vLLM's RequestOutput. Setup refuses to run when the installed vLLM is not the release whose private route state the read was verified on. A call whose routes cannot be read, are malformed, or do not cover every frozen token is no longer cut. Its cut is rolled back with durable_failure and it restarts from its last saved turn; the rest of the checkpoint goes on. Missing route rows are no longer padded, which shifted later rows onto the wrong tokens and broke later cuts' digests. A restored prefix that is already finished now answers with its restored routes, so router replay no longer fails that call. A resumed call whose routes do not fit its restored prefix is logged and staged without routes instead of raising after the handler's abort path closed. The final placeholder route row comes from one helper on both the save and the restore path, and holds ROUTE_MISSING_SENTINEL like an uncut call's final row. decode_routed_experts keeps the wire dtype when none is given, replacing three copies of the dtype table. The capture mixin declares its router-replay members instead of reading them with getattr defaults. Restore and completion now log the hash of the prefix routes a resumed call keeps. A tiny-MoE functional test kills and restores a router-replay run mid-generation and requires the routes each restored call staged for training to hash like the ones it restored. It also reports gen_kl_error over restored prefixes, resumed tails and the other rows. The guides say which routes a resumed row trains with. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Register the tiny-MoE router-replay prefix-recovery test in the fast tier of L1_Functional_Tests_SingleController_3, next to the text prefix-recovery test. It takes two short runs on two GPUs. The test rebuilds its model on every run, but Megatron caches the HF conversion by model path under HF_HOME. On a runner that keeps HF_HOME, a later run could load an earlier run's conversion. If those weights differ (another transformers version), vLLM and Megatron would serve different models. Convert into the run's own directory instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
…refix I/O tests NeMo Gym's GenerationCutPrefix names its attempt attempt, as EpisodeId does; it was attempt_index. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
- PrefixReadBatcher: a failed multi-key GET re-reads each key alone, once, for any error, so one bad row fails only its own callers and a TQ outage costs N+1 GETs instead of 2N-1. - Reads are capped by rows only. Every key was charged max_model_len tokens, so the token budget silently acted as a ~32-key cap tied to the write knob. - One cut path: every cut goes through the batched transaction loop (batch size 1 writes one row per put). Remove _checkpoint_active_generation_cut and TQTokenSink.stage_generation_prefix. - stage_generation_prefix_batch raises on the first failure and returns the written keys. A failed put discards every key of the batch; flush() clears written rows whose transaction never sealed. The transaction receives the staged key instead of a StageResult. - Config defaults live only on RolloutRecoveryConfig: setup_token_capture and PrefixReadBatcher take the limits without defaults, and GenerationPrefixBatchLimits with its helpers is gone (Field(gt=0) already validates). - Drop the benchmark tools and their test; results go in the PR description and the tools live on amahishi/ckpt-v2-prefix-io-benchmarks. Signed-off-by: Anish Mahishi <amahishi@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The class-qualified calls only existed so a SimpleNamespace test fake could run them. #4506 replaced that fake with a real worker; the two setup tests that build the prefix reader now use one too (_real_worker), so production code calls the method on self. Signed-off-by: Anish Mahishi <amahishi@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The base now acknowledges a call that cannot be cut as a durable_failure and lets it regenerate after a restore, instead of failing the receipt for every call on the worker. Do the same in the batched path: - A failed freeze, acknowledgement or seal fails only its own call. A cut frozen before the transaction's rollback scope begins is released by the caller, as the base's per-call loop did. - A failed batch put is retried one row per put (the sink has already discarded the batch), so one bad row fails only its own call. - Rows written but never handed to a call, which only an interruption can leave, are still cleared. Restores take the base's regenerate-on-failure path: the worker's batched reader fetches the chunks and _resolve_usable_generation_cut uses them, reading the rows alone if the shared read failed. Signed-off-by: Anish Mahishi <amahishi@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Contributor
Author
|
/ok to test 94ec17f |
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
macandro96
force-pushed
the
amahishi/batch-prefix-route-snapshots
branch
from
October 21, 2026 21:12
94ec17f to
f049700
Compare
Contributor
Author
|
/ok to test f049700 |
macandro96
force-pushed
the
amahishi/ckpt-v2-prefix-batched-io
branch
from
November 5, 2026 00:02
c3eb2e4 to
6c72a0f
Compare
terrykong
added
CI:Lfast
This branch has not been 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.
Summary
Stacked on #4508; base is
amahishi/ckpt-v2-prefix-batched-io, not main. This adds within-inventory engine route-read batching to that PR's existing TQ I/O batching.min(generation_prefix_batch_size, 64)calls and the existing token-position budget. Tensor assembly remains on the cut thread and happens one call at a time; the engine callback only shallow-copies route lists.No new configuration or checkpoint-format change. Dense/replay-off calls skip engine lookups. The single-call host adapter is retained.
Scope
Only route-fetch batching and its regressions. This deliberately excludes the integration checkout's incremental route cache, version-2 digest layout, raw-route tensor reuse, TQ CPU-topology cache, route-padding fix, and diagnostic instrumentation. It does not coalesce separate Gym inventories.
Review fixes and remaining stack integration
state.absorb_engine_route_pieces(...), and remove the temporary_materialize_engine_routed_expertshooks. Do not resolve that rebase by retaining full-prefix per-piece materialization. This amendment does not rewrite either parent PR branch or change the live training checkout.Validation
git diff --checkpass.nemo_rl/models/generation/generation_cut_capture.pywith Torch available. Full project type checking, normal-import worker tests and GPU functional parity have not been run locally on this isolated branch; CI requested.Production motivation (not a controlled end-to-end speedup claim)
In the existing 10-node integration experiments, the first snapshots had 430 calls before grouping versus 428 afterward: engine handoffs fell from 212 to 8, and summed route-fetch wait from 12.70 s to 0.45 s. Loads/timing differ between runs, and later batched cuts still exhibit engine scheduling stalls. This PR reduces handoff count; it does not eliminate a busy engine loop.