Repository navigation
feat(checkpoint): trace rollout events to check what a restore cost - #5048
Merged
Merged
Conversation
github-actions
Bot
added
the
Documentation
terrykong
reviewed
Nov 5, 2026
terrykong
left a comment
terrykong
left a comment
Collaborator
There was a problem hiding this comment.
Thanks for this. A per-rollout record is the right tool for checking what a restore costs: the existing totals can't say which rollouts were thrown away. The ledger hook is small and well placed, and the feature is opt-in with no checkpoint format change.
On size: this is [2/11] of a stack that is already past 20k lines, so this review looked for smaller ways to get the same result. This PR is close to its needed size, and half of it is tests. The asks below save about 55–65 lines (about 9%):
- (C9) Merge the four ledger tests into one table test and drop a test that can't fail (about −40).
- (C6, C7) Replace the design doc's event table with a short usage section, and keep the event list only in the trace.py docstring (about −22).
- (C5, C11) Two small cleanups.
Correctness of the trace:
4. (C3, C4) A restart reports a try that was already abandoned as abandoned again, so the doc's thrown-away time counts it twice.
5. (C2) A restore copies the whole ledger even with tracing off.
6. (C8, C10) Tests: the "failed" label is untested, and the weight-sync check is too loose.
The queued writer thread is worth keeping: it keeps file writes off the event loop, which matters on a shared filesystem at large rollout counts. No existing writer in the repo fits better (logging.handlers.QueueListener prints a traceback on every failed write; the logger's JSONL helper opens the file on every call).
CI has not run the lint and unit-test workflows on this head yet.
Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-5048-rollout-trace.html
Generated by Claude Code
terrykong
added
the
CI:Lfast
macandro96
force-pushed
the
amahishi/rollout-trace
branch
from
November 6, 2026 00:39
8c75341 to
2202eff
Compare
Contributor
Author
|
/ok to test 2202eff |
The run logs only averages and counts for rollout recovery, and the recovery ledger drops each group's history when the group finishes, so a restored run cannot be checked rollout by rollout. rollout_checkpointing.trace (default false) appends one JSON line per event to rollout_trace.jsonl in the log directory. Every line has the wall-clock time, a run ID that is new per process, and the event: - the recovery ledger reports every transition through an optional event sink: group reserved, admitted, sealed and discarded; attempt dispatched, sealed, refused, resumed, and abandoned with its reason (dropped, failed, cancelled or restart, and whole_group for a prompt-group restart). The rollout manager now tells a cancelled rollout, such as one the run's end stops, from a failed one; - the controller reports each trained step with its group IDs, and each weight sync; - the snapshotter reports each published trainer checkpoint and rollout snapshot, with when the snapshot's state was captured; - the restorer reports each restore and the unfinished groups it loaded. Events are queued and a background thread appends them, so tracing never blocks the controller's event loop, where many of them are raised under a data-plane mutation cut; a killed run loses only events not yet written. A failed write warns once and stops tracing without failing the run. A disabled trace holds no queue or thread, so the controller class that defaults to one still pickles for Ray. Rows match across runs by prompt_id (the dataset index) and generation_index. Restoring the ledger reports nothing; the restart after it reports which interrupted attempts resume and which start over. The switch is under rollout_checkpointing, which the bootstrap fingerprint already excludes. Nothing a checkpoint saves changes. The design doc lists the events and how a stop's thrown-away time and lost progress follow from them. Source: +311 / -15 (trace.py 126 lines). Tests: +350 (13 new tests). Unit: 2109 passed (default lane), 121 passed (Gym lane). ruff and pyrefly clean. The Gym turn-recovery functional test passes with the trace on. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
restore_finished listed the unfinished groups through ledger.groups(), which deep-copies every group, before emit checked whether tracing was on, so a restore paid for that copy with tracing off. It is now built only when tracing is enabled. The weights_synced emit drops its guard instead: its only argument is an int and emit returns at once when tracing is off. Two bare test controllers that drive _sync_weights now set _trainer_version, as the actor always does. Review comments C2 and C5. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
A try that failed before the stop stays ABANDONED in its group until the next dispatch, so a snapshot in that window saves it. On restore, prepare_for_restart skipped only SEALED tries and reported it again as attempt_abandoned with reason restart, and _abandon_entire_group did the same for every sibling of a prompt group. The trace then counted the failed try as thrown away by the stop. Both paths now skip tries already abandoned or failed; retry treats ABANDONED and FAILED alike, so redispatch is unchanged. Review comments C3 and C4. Adds a regression test for both paths. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
The config docstring and both exemplar YAML comments restated the event list and had already drifted from it. They now point to the trace.py docstring, which names group_discarded's step as discard. Review comment C7. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
- Check the failed abandon reason end to end in the rollout manager's failed dispatch test, and drop the test that only passed reason=cancelled through. - Restore a recording ledger with _recording_ledger(from_ledger=...) instead of repeating the steps, and drop a disabled-trace test that could not fail. - Assert the exact weight_version of each weights_synced event. Review comments C8, C9 and C10. C11's line in the test helper is gone: the helper now calls the actor's _build_rollout_checkpointing. Unit, on #5011 rebased onto main: 3452 passed (default lane), 95 passed (Gym lane, on the files this PR touches). ruff and pyrefly clean. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
macandro96
force-pushed
the
amahishi/rollout-trace
branch
from
November 6, 2026 15:51
2202eff to
1ea04dd
Compare
The design doc's event table copied each event's fields from the code and had drifted from it: group_admitted, group_reserved and trace_opened were listed wrongly or not at all. Every line in the trace already names its event and fields, so the section now says how to turn tracing on, what each line carries, where the events are listed (trace.py), how to match rows across runs and how to work out a stop's thrown-away time. Review comment C6. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Contributor
Author
|
/ok to test 1ea04dd |
@macandro96, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/ |
Contributor
Author
|
/ok to test 76c81b1 |
terrykong
enabled stacked auto-merge (squash)
November 6, 2026 18:29
terrykong
approved these changes
Nov 6, 2026
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?
Adds an opt-in, per-rollout event log for checking what a rollout checkpoint stop and restore actually cost.
Stacked on #5011. Follow-up: the batch-age PR on top of this one.
Why
Rollout recovery only logs averages and counts, and the recovery ledger drops each group's history when the group finishes, so a restored run cannot be checked rollout by rollout. The validation plan for partial checkpointing needs per-rollout records: when each try started and ended, why it ended, which tries a restore resumed or restarted, and when each step and checkpoint happened.
What it records
With
rollout_checkpointing.trace: true(defaultfalse), the controller appends one JSON line per event torollout_trace.jsonlin the log directory:group_reserved,group_admitted,group_sealed,group_discardedattempt_dispatched,attempt_sealed,attempt_refused,attempt_resumed,attempt_abandoned(reason:dropped,failed,cancelled,restart)step_trained,weights_syncedtrainer_checkpoint_published,snapshot_published(with when its state was captured)restore_started,restore_finished(with the unfinished groups loaded)Rows match across runs by dataset prompt index and sibling index. The design doc lists every field and how a stop's thrown-away time and lost progress follow from the events.
Design
emitserializes and queues; a background thread appends in batches. A killed run loses only events not yet written. A failed write warns once and stops tracing; it never fails the run.rollout_checkpointing, which the bootstrap fingerprint already excludes.RolloutRecoveryLedgerreports every transition, so no call site needs its own logging. The rollout manager now tells a cancelled rollout (for example, one the run's end stops) from a failed one.Testing
tests/unit/single_controller/test_rollout_trace.py): file format, ledger event sequences (dispatch to seal, drop and retry, turn-level resume versus restart, prompt-group restart, cancel versus fail), the controller's step, weight-sync, checkpoint, snapshot and restore events, write failure, and pickling the controller class.🤖 Generated with Claude Code