Where #4508's 1,758 added lines go, and the ~210 that can go without changing behavior

NVIDIA-NeMo/RL #4508, part 9 of 11 in the partial-checkpointing stack. Line counts are the PR against its real base f47c6b4 (#4612). Permalinks pin the PR head 6c72a0f. The review comments are in a pending review; their links work once it is published.

Before
Round 1 cut the PR from +2,802 to +1,758 lines: benchmark scripts out, one save loop, no repeated defaults. All 7 round-1 comments are fixed.
This PR
About 767 product lines (164 of them moved, not new) and about 991 test lines.
Status
Proposed, not in the PR. The cuts below are review asks C8–C11. Only C8 changes behavior, and only after an interruption, where it matches what the code did before this PR. Not run, except where a row says so.

Four pages on this PR. The round-1 pages describe an older head. This one: what is still bigger than it needs to be.

Most of the source change is reuse: the row encoding is the old single-row body moved into _encode_record. What is left over comes from earlier versions of this PR. The batch used to succeed or fail as a whole. Since commit 6c72a0f each call fails on its own, but the code that only the whole-batch version needed is still there. On the read side, a later commit added a second retry and kept the first.

Lines added per file, and what each comment removes

gray = added by this PR · dark gray = moved from existing code · orange = what C8–C11 remove test_vllm_token_capture_hosting.py +510 −30 (C11) tq_token_sink.py 164 moved, not new +268 −21 (C9) test_prefix_read_batcher.py +257 −70 (C10, C11) generation_cut_capture.py +254 −28 (C8) prefix_read_batcher.py +175 −18 (C10) test_tq_token_sink.py +172 −45 (C9, C11) vllm_worker_async.py +38 config plumbing, other tests +84 Total about −67 product and −145 test lines. C10 also adds about 35 test lines, for two cases no test checks today.
CommentWhat goesWhy behavior stays
C8The second yield in _generation_cut_transaction, the second loop in flush(), and a copy of the one-row write in stage()Each call already fails on its own. No test reaches the two branches that go: both were made to raise, and 144 tests still pass.
C9_stage_at_key (now one caller), a one-line _discard_failed_write wrapper, and an attachments argument no caller passesLeftovers from deleting the one-row save. Nothing in the product code reaches them.
C10The key-by-key retry in PrefixReadBatcher._read, and max_pending, which no caller setsThe request already re-reads its own keys when the shared read fails, so a bad row still fails only its owner.
C11Repeated test helpers and checks in three test filesTest-only. The trimmed batcher file passes 12/12; the rollback helper passes 144.

C8 — generation_cut_capture.py:1299 C9 — tq_token_sink.py:531 C10 — prefix_read_batcher.py:95 C11 — test_prefix_read_batcher.py:23

The biggest source cut: one step per call, one loop (C8) — proposed, not in the PR

Where the fix lands: RL only, one file. _generation_cut_transaction and flush() in generation_cut_capture.py. No Gym change: Gym already takes one acknowledgement per call.

@@ generation_cut_capture.py L1296-L1304 · _generation_cut_transaction
                         terminal_stop_reason=terminal_stop_reason,
                     )
-                if chunk_token_ids:
-                    # Validate every acknowledgement in a batch before sealing
-                    # any frozen cut into its live request state. A failed
-                    # acknowledgement leaves the cut available for rollback.
-                    yield acknowledgement
                 with state.lock:
                     state.seal_frozen_cut(checkpoint_id)
@@ generation_cut_capture.py L1421-L1447 · flush
                 keys = stage(batch)
-                ready = []
                 for (transaction, _, _, prefix), key in zip(batch, keys, strict=True):
                     handed += 1
                     if key is None:
                         # Closing rolls the frozen cut back.
                         transaction.close()
                         continue
                     try:
-                        # A failure here rolls the call back and clears its row.
-                        ready.append((transaction, prefix, transaction.send(key)))
-                    except Exception:
-                        fail(prefix)
-                # Every acknowledgement in the batch is built before any call
-                # seals its frozen cut into its live request state.
-                for transaction, prefix, acknowledgement in ready:
-                    try:
-                        transaction.send(None)
-                    except StopIteration:
-                        acknowledgements.append(acknowledgement)
-                    except Exception:
-                        fail(prefix)
-                    else:
-                        transaction.close()
-                        acknowledgements.append(
-                            self._failed_generation_cut_ack(prefix, reason="cut_error")
-                        )
+                        # Builds the ack and seals; a failure rolls back and clears the row.
+                        transaction.send(key)
+                    except StopIteration as done:
+                        acknowledgements.append(done.value)
+                    except Exception:
+                        fail(prefix)

What each piece does: the generator now builds the acknowledgement and seals in one step, as the code before this PR did for every call (base _checkpoint_active_generation_cut). send(key) then ends with StopIteration, which carries the acknowledgement. A failure still rolls the call back, through the generator's own except BaseException. The one change: if the batch is interrupted partway (for example by KeyboardInterrupt), calls earlier in that batch stay saved instead of rolling back. Batch size 1 already behaves this way today, and so did the code before this PR. One test, test_worker_batch_interruption_clears_rows_no_call_received, needs its expectation updated. Not run.

C8 — generation_cut_capture.py:1299

So what. Nothing here is a bug in normal use, and the design is sound. But this is part 9 of a stack past 20k lines. About 67 source lines and 145 test lines can go without changing what the code does. Two smaller items are in the same review: aclose() can wait forever on a stopped event loop (C12, a one-line fix). And before merge the PR still needs CI on this head, the end-to-end save time, and the size of the largest put with routes on.

C12 — prefix_read_batcher.py:156