What the new rollout trace records across a stop and restore, and how it could be smaller

NVIDIA-NeMo/RL #5048 (part 2 of a stack that starts at #5011; #5049 builds on it). All permalinks pin the PR head 8c75341b.

Before
A restore logged only totals and timings: how many rollouts restarted, how long the save and the restore took. No record said which rollouts were thrown away, or how long each had run.
This PR
Adds rollout_checkpointing.trace (off by default). When on, the controller writes one JSON line per rollout event to rollout_trace.jsonl in the log folder.
Status
New feature, in review. The design is sound. The review asks for a smaller test file and design doc section (about 9% of the PR), plus one fix so a restart does not report the same lost rollout twice. The cuts below are proposed, not in the PR.
Background. In the single-controller loop, each prompt becomes a group of rollouts (one per sibling). The recovery ledger tracks every try at every rollout. A rollout snapshot saves the ledger at one moment, the cut. After a stop, the next run restores the snapshot and restarts what was still running. Dotted words have a short explanation: hover or tap them.

The PR adds a small writer, RolloutTrace, and one hook on the ledger, set_event_sink. Every time the ledger changes a try's state, it calls the hook, so the ledger's 15 change points need no logging code of their own. The train pump, weight sync, snapshotter and restorer add six more events. Each line has the wall-clock time t and a run id that is new per process, so two runs can share one file.

One prompt through a stop and a restore

Prompt 7, siblings 0 and 2 (sibling 1 acts like sibling 0). Seconds since start. 040 6090 run a1 sibling 0 sibling 2 thrown away: 40 − 12 = 28 s snapshot cut, cut_time = 40 process killed (t = 52) run b2 sibling 2, try 1 1 2 3 4 5 6 7 finished and saved ran before the cut, then thrown away ran after the cut, never saved rollout_trace.jsonl (fields cut down to the ones that matter here) 1 {"t":12,"run":"a1","event":"attempt_dispatched","prompt_id":"7","generation_index":2,"attempt":0} 2 {"t":25,"run":"a1","event":"attempt_sealed","prompt_id":"7","generation_index":0} 3 {"t":43,"run":"a1","event":"snapshot_published","step":4,"cut_time":40} 4 {"t":60,"run":"b2","event":"restore_started","path":"…/snapshot_000004"} 5 {"t":60,"run":"b2","event":"attempt_abandoned","generation_index":2,"attempt":0,"reason":"restart"} 6 {"t":63,"run":"b2","event":"attempt_dispatched","generation_index":2,"attempt":1} 7 {"t":90,"run":"b2","event":"group_sealed","prompt_id":"7"} line 5 comes from prepare_for_restart (rollout_recovery.py:836-847)

The numbers are made up but every event and field name is real. To get a stop's thrown-away time, take each attempt_abandoned with reason: "restart" (line 5) and find that try's last attempt_dispatched (line 1). The thrown-away time runs from that dispatch to the snapshot's cut_time (line 3), here 40 − 12 = 28 s. The design doc gives this rule (design doc L230-238). The totals that existed before (restarted_rollouts, restore timings) would say only "1 rollout restarted".

What could be smaller — proposed, not in the PR

Nothing below is implemented. The PR is 694 added lines; 348 of them are one test file and 126 are the writer. No existing writer in the repo fits better. The review checked these:

So the savings come from writing less, not from reusing more.

The writer's thread stays. The PR writes each line from a background thread, so the event loop never waits on the disk. Timed with 2 million events on a laptop SSD, the event loop spends 2.5 µs per event with the thread and 3.6 µs writing each line directly. At the 1B example's 1,150 events per step that is nothing. At a million events per step (hundreds of thousands of rollouts) it is about a second per step, and every direct write is a system call on the event loop, which on a shared filesystem can stall far longer. Removing the thread would save about 40 lines, but it is not worth that risk.

The cuts the review does ask for, about 55–65 lines:

AskLinesComment
Test file: one table test for the ledger event orders, drop a check that overlaps the off-by-default test, move the cancelled-vs-failed check next to the rollout-manager test that owns it≈ −40 to −50C9, C8
Design doc: replace the event table (already wrong in three places) with a short section on how to turn tracing on and read the file≈ −14C6
Event list written out in five places: keep one list, in the trace.py docstring, and point to it from the rest≈ −8C7
Small cleanups: a redundant enabled check, a getattr fallback that never runs≈ −3C5, C11
Additions: skip the ledger copy when tracing is off; skip re-reporting a try abandoned before the stop; a test for the "failed" label≈ +10C2, C3, C8
So what. This PR is not much larger than it needs to be. Half of it is tests, and the feature is a thin hook plus a file writer whose thread earns its keep at large scale. Trimming the tests and the design doc saves about 9% of the PR. The one change that affects the numbers the trace gives you: after a restart, a try that had already been abandoned before the stop is reported as abandoned again, so the thrown-away time above counts it twice. The review asks for a short skip in two places.

💬 C3 — rollout_recovery.py:841-847C4 — rollout_recovery.py:863-872

The review comments above are pending until the review is published; until then the links open only for the reviewer.