A training step that finishes during the Gym pause throws the snapshot away

Explainer for NVIDIA-NeMo/RL #4266. Permalinks pin the PR head 4db8f66; "before" links pin the base ed14402.

Before
A periodic rollout snapshot took the data-plane barrier and checked that no training step had finished since it started. If one had, it dropped the attempt. That cost almost nothing, because nothing had been paused yet.
This PR
With Gym checkpointing on, each attempt first pauses Gym and waits for it to drain and write its state, and only then runs the same check. A step that finishes during that wait still drops the attempt, but now after rollouts were stopped for the whole pause.
Status
Open — not fixed in this PR. No data is lost; the cost is wasted pauses and snapshots that never land. Raised in the review; a fix is proposed below.

Ten pages on this stack — this one: #4266: a training step that lands while Gym is paused.

Background. Single Controller trains and generates at the same time. Every few minutes it tries a rollout snapshot. The snapshot is filed under the newest trainer checkpoint, say step_10, and must describe the rollout data exactly as it was when those weights were saved. With this PR, Gym's paused agents are saved in the same snapshot, so Gym has to stop first. Dotted words have a short explanation: hover or tap them.

An attempt records the current step (10) and prepares a folder under step_10 (single_controller.py:4626). It then closes rollout admission and asks Gym to drain and write its state (:4648). The training loop is not paused, so it keeps going. Only after Gym is done does the attempt take the data-plane barrier and compare the step again (:4664). If step 11 finished in between, the attempt deletes the folder, aborts Gym, and logs skipped (:4911). A skipped attempt never counts as a failure.

One attempt, second by second

save_period 5 → last trainer checkpoint is step_10. Step 11 is training. Times in seconds. 0 25 40 95 training step 11 (40 s) step 11 commits at 40 s: its trained groups are deleted, step count 10 → 11 single_controller.py:3412-3459 Gym snapshot drain + write (70 s), began at step 10 _prepare_and_commit_gym_checkpoint new rollouts stopped for 70 s ✗ step is 11, not 10 → throw it away step check :4664 training Gym paused for the snapshot no new rollouts (wasted) rollouts flowing

Why it can't just finish into step_10 anyway. When step 11 commits, it deletes the rollout groups it trained on (_cleanup_consumed_metas_unlocked). A step_10 snapshot taken after that is missing those groups. But the step_10 weights never learned from them, so a restart from it would silently drop them. So the check is right to throw the attempt away. The problem is that the throw-away now happens after a 70 s Gym pause, not before one.

How often. An attempt is wasted whenever a step finishes during the pause. In this example the pause (70 s) is longer than a step (40 s), so every attempt overlaps the end of some step, as long as the buffer holds enough finished rollouts to keep training going while new rollouts are stopped. The check before the pause only looks at whether a trainer checkpoint exists for the current step (:4598), not at whether a step is about to finish.

What would fix it — proposed, not in the PR

Nothing below is implemented.

Where the fix lands: RL only, one file. single_controller.py: the snapshot attempt opens a hold before Gym prepare and releases it after its cut; the training loop waits on that hold just before its optimizer step. No Gym change: Gym only sees a normal prepare, commit and resume.

# snapshot attempt (_save_rollout_checkpoint_locked) self._snapshot_cut_pending.clear() # t = 25 s: "don't commit a step yet" try: await prepare_and_commit_gym(...) # 25 s → 95 s, step count still 10 async with barrier.checkpoint() as cut: # 95 s: step 10 == 10 ✓ → cut saved ... finally: self._snapshot_cut_pending.set() # 95 s: training may commit again # training loop, just before the optimizer step await self._snapshot_cut_pending.wait() # step 11 ready at 40 s, waits 55 s policy_result = finish_train_step() # 95 s: step 11 commits, count → 11

The two places the loop starts an optimizer step: GRPO finish_train_step :3399 PPO critic epochs :3272

Holding before the optimizer step keeps the model and the data plane at step 10 until the cut is taken. The code already treats that point as safe: rows picked for the running step are kept in the snapshot and offered again after a restart (comment at :656). The cost moves from rollouts to training: step 11 waits 55 s, but only when a step happens to finish during a pause. Rollouts are stopped for the same 70 s either way, and now the snapshot lands. The hold must be released on every exit path, including a Gym abort, or training stops for good.

So what. Nothing is lost or corrupted: throwing the attempt away is the correct response. But when a training step is shorter than Gym's drain plus write, the periodic Gym snapshots between trainer checkpoints can nearly all be wasted. Each one stops new rollouts for the full pause, and each is logged only as skipped with reason trainer_state_changed. The review asks the author to measure that rate first; the hold above is the fix if it is high.

🔴 review comment — single_controller.py:4626