Size review, round 2, of NVIDIA-NeMo/RL #4506 (part 5 of 11 in the partial-checkpointing stack, base #5056). All permalinks pin the PR head b6aff18. Review comments are in a pending review; their links work once it is published.
Three pages on this PR, one question each — this one: what can go.
target_level: prefix, each rollout snapshot asks every vLLM worker to save the tokens its running calls have decoded so far, as chunk rows. When a call finishes, the worker writes its full terminal row and deletes its chunk rows. That delete is the only write that can break a snapshot in progress: Gym's snapshot may still point at those chunk rows. Dotted words have a short explanation: hover or tap them.To keep that delete out of the snapshot window, the PR closes a gate on every worker from the snapshot's start until the TQ save is done. Every terminal write waits, so a call that finishes inside the window cannot return its reply. The gate needs begin/end calls on the worker, VllmGeneration and the interface, plus GenerationCutFence. The PR's own notes say the terminal write itself is safe; only the delete is not. So queue the delete instead.
Why the proposed order is safe. The driver already waits for the fence to begin before Gym prepares its snapshot (async with order), so the list swap at the start of snapshot N+1 always runs after snapshot N's TQ save. Rows swapped out belong to calls that had already finished. Gym acknowledges a finished call through its terminal row only (terminal ack), and exports only this round's cuts (Gym export), so snapshot N+1 never names those chunk rows. A terminal row written mid-window is already accepted: validation runs with clear_unreferenced=False (snapshots.py:421) and restore clears extra rows (restore.py:466). One cost: if a worker restarts, its list is lost and those rows wait for the next restore to clear them.
C30 review comment — generation_cut_capture.py:247
Nothing below is implemented, and it has not been run.
Where the change lands: RL only. The worker keeps a list of chunk-row keys to delete; the three delete sites append to it; the begin/end pair becomes one call that GenerationCutFence.held awaits before Gym prepares. No Gym change: Gym already acknowledges finished calls through their terminal row. It also removes the gate that #4655 would otherwise build on before a later PR deletes it.
| Ask | What goes | Why it stays correct | lines |
|---|---|---|---|
| C28 | Land the prefix functional test once, in #4507, which rewrites about a third of it (L1 line) C28 review comment — L1_Functional_Tests_SingleController_3.sh:72 | At this head it would fail at Gym start-up (wrong env var name) and FAST lanes skip it, so nothing is tested here today | −615 test |
| C30 | Snapshot gate, fence executor, begin/end calls on three layers | Only the delete is unsafe; it moves after the TQ save (above) | ≈−100 src ≈−110 test |
| C31 | 11 per-reason failure counters and their plumbing to step metrics (reasons) | Gym already reports failed cuts per snapshot; one INFO log line per cut keeps the breakdown. The restore and progress-read counters stay. | ≈−60 src ≈−45 test |
| C33 | Five GenerationCutCaptureMixin._x(self, …) calls (first one) and a test fake that copies 25 worker fields by hand | No subclass overrides those methods; a real worker object in the tests makes plain self._x work | ≈−35 test |
| C34 | A copy of vLLM's greedy-temperature constant (seed check) | vLLM's own SamplingType.RANDOM_SEED is the same test | −3 src |
| C38 | A "prefix needs turn" check that cannot fire (nemo_gym.py:2824) | The prefix level always implies the turn checkpoint in the level enum | −4 src |
| Question | The private vLLM output wraps and kept finished-request state (wraps) | vLLM's supported streaming output could replace them, at a per-step cost on the engine loop. Not measured, so only asked. | ~150 src? |
| Follow-up | Turn the capture mixin into its own object (already in the PR's follow-up list) | Mostly a testing gain, about −60 source | later |