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.
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.
| Comment | What goes | Why behavior stays |
|---|---|---|
| C8 | The 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 passes | Leftovers from deleting the one-row save. Nothing in the product code reaches them. |
| C10 | The key-by-key retry in PrefixReadBatcher._read, and max_pending, which no caller sets | The request already re-reads its own keys when the shared read fails, so a bad row still fails only its owner. |
| C11 | Repeated test helpers and checks in three test files | Test-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
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
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.