How the checkpoint became one atomic cut

NVIDIA-NeMo/RL PR #3480 — the fix for the snapshot-skew bug from the last round. Links pinned to bca9a721. Part 2 of 3 · the ledger · the proof

Last round's bug: the save read the sampler's dispatch cursor, then did ~15 awaits of model and optimizer writing, then snapshotted the replay buffer — and the rollout pump kept running in the gap, so the two reads could disagree by one step. On resume the anti-duplication guard then silently dropped a full batch of prompts. The fix is not a patch on that one pair — the author made every restart-facing read happen inside one exclusive window, and split sampler admission so its cursor advance takes part in the same locking. Nothing that moves a prompt can run while the checkpoint looks.

Every writer takes a ticket; the checkpoint closes the booth

The lock is a shared DataPlaneCheckpointBarrier (replay_buffer.py:196) with two sides. mutation() — the "ticket" — just counts you in as an in-progress writer: it blocks only while a checkpoint is open, and writers never block each other. checkpoint() is the exclusive side: it stops new writers from entering, waits for the counted ones to finish, and only then reads. Note the inversion — the code that writes state takes the shared side; the snapshot reader takes the exclusive one. Exclusivity is between the group of writers and the save, not among writers.

the shared barrier (DataPlaneCheckpointBarrier) mutation() = register as an in-progress writer: writers never block each other, they only wait while a checkpoint is open. Each box below is one such write: dataloader advance + ledger reserve_group one cut: cursor never saved past an unowned prompt sampler commit_admission + mark ADMITTED cursor advance and phase flip in one cut (see below) TQ commit (rollouts finished → replay buffer) groups join the replay buffer only between cuts drop / substitute-from-spare-pool decision old owner out, new owner in — atomically watchdog abort of a stale rollout the one intentional drop; also a single cut the exclusive cut — barrier.checkpoint(): new mutations wait, active ones drain first single_controller.py:2485-2554 — everything below reads one consistent instant 1 sampler dispatch cursor (the read that used to happen 15 awaits too early) 2 dataloader position 3 spare pool snapshot (replacement_reserve.pt contents) 4 replay buffer index (finished groups) 5 ledger snapshot − groups already in 4 ← the subtraction: never two owners on disk 6 TQ data-plane snapshot + sha256/count of 5 stamped into its metadata all blocked while the cut is open Consistency between the train pump's own counters and the cut is free: _save_checkpoint runs in the train pump task, so trainer_version / step counters cannot move while it executes.
Six restart-facing artifacts, one instant. Last round only item 4 was inside the window; items 1–3 were read early and item 5 did not exist.

The admission split — locking the cursor without deadlocking the gate

The subtle part. Sampler admission does two things: it waits at a gate until the trainer catches up, then it advances the dispatch cursor and stamps the batch's step. Wrapping both in a mutation slot would be wrong: the gate wait can only end when the trainer advances, and the trainer's own save is what waits for mutation slots to drain — lock held across the wait means the two sides wait for each other. So the PR splits the old admit() into a two-step protocol, TransactionalAdmissionSampler: 

# single_controller.py:816-822 — the dispatch path
await sampler.wait_until_admissible(...)      # the gate wait — OUTSIDE the barrier,
                                              # a checkpoint can run while this polls
async with barrier.mutation():                # then the short part takes a slot:
    target_step = sampler.commit_admission()  #   cursor += 1, stamp the step
    ledger.mark_group_admitted(...)           #   RESERVED → ADMITTED, same cut
    # already-buffered duplicate check + dispatch bookkeeping, same cut

Between the wait ending and the slot being taken, nothing can regress the gate: the rollout pump is the only admitter and runs these two lines back to back, and the trainer version only grows. All four built-in samplers implement the split. A custom sampler that does not gets its whole admit() wrapped in a slot as a fallback — correct, at the cost of the checkpoint waiting out that sampler's gate.

Same tiny numbers as last round's bug page — gate window 1, 4 prompts per step, trainer at step 7, cursor at 6:

last round (bug)this round (fixed)
cursor readbefore the save's ~15 awaitsinside the exclusive cut
pump admits during the save?yes — stamps step 7 mid-savecannot — commit_admission needs a slot, slots wait
what disk sayscursor 6, buffer already has step 7cursor and buffer agree by construction of the cut
resumeguard sees step 7 covered → silently drops a fresh batch of 4nothing to reconcile
Does it hold up? Two agents in this review independently attacked the cut and found no interleaving that breaks it. The one ledger write outside any slot — dropping a group's ledger entry right after its TQ commit — is safe in every position relative to the cut, because the save subtracts buffer-owned groups from the ledger snapshot and restore subtracts again — the buffer's copy always wins. And it is tested with the real code: the races suite blocks a real save mid-write, fires a real commit, and asserts exactly one owner lands on disk (test_checkpoint_dispatch_races.py).
The cost, honestly. The exclusive window got wider: checkpoint-directory setup and the sidecar write now happen inside it, and every commit, admission, and dataloader advance waits for the whole window. On a slow shared filesystem that is real time, and the PR posts no before/after number for it. The review asks for one — a per-checkpoint "exclusive window duration" with recovery on vs off would settle whether this matters.