Size review of NVIDIA-NeMo/RL #4508. Line counts are from git diff 2548c67 5498a0a (the PR against its real base, #4612). All permalinks pin the PR head 5498a0a. The review comments are in a pending review; their links work once it is published.
Four pages on this PR. This round-1 page describes an older head (5498a0a); see round 2 for the current one. This one: how big it was, and what could go.
The write side mostly reuses what exists. About 140 of the 274 lines added to tq_token_sink.py are an existing body moved into _encode_record, and batching goes through TransferQueue's own put_samples.
The lines that can go are the benchmark scripts, a second copy of the save loop, plumbing for two settings, and the shared read path. The shared read path has no numbers behind it yet.
| What the PR adds | Ask | lines |
|---|---|---|
| Two benchmark scripts (writes 781, reads 371) plus a test that checks only the script's own counters (81). Both call the sink directly, so they time TransferQueue, not the save path. | Delete them and their 2 pyrefly.toml lines; put a results table in the PR description | −1,235 |
The shared read path: PrefixReadBatcher 186, wiring ~40, tests 214. No restore timing exists, and it has one confirmed bug.
|
Move to its own PR that leads with restore timings | ~−440 |
| A second save loop kept for batch size 1 (L1399), its driver, and the single-row write it calls. No test reaches anything that only it does. | One loop: diff below | ~−110 |
| Default values 256 / 4,194,304 written into 5 function signatures, plus a limits class that repeats the config's own "must be positive" check (L38) | Keep the 2 config keys; drop the copies and the class | ~−30 |
| Per-row failure results in the batch write (L741) that its only caller never uses: it stops at the first failure | Stop at the first failure in the sink too | ~−25 |
Where the fix lands: RL only, two files. generation_cut_capture.py and tq_token_sink.py. No Gym change: Gym never calls the single-row write.
@@ generation_cut_capture.py L1396-L1433 · _checkpoint_generation_cut - acknowledgements = [] - single_prefixes = inventory.active_prefixes - batch_limits = self._require_generation_prefix_batch_limits() - if batch_limits.max_rows > 1: - acknowledgements = self._checkpoint_generation_cut_batched( - inventory, deadline=deadline - ) - single_prefixes = () - for index, prefix in enumerate(single_prefixes): … 28 more deleted lines: the per-call loop with its own deadline check … + # Batch size 1 is one row per write through the same loop. + acknowledgements = self._checkpoint_generation_cut_batched( + inventory, deadline=deadline + ) receipt = GenerationCutReceipt( @@ generation_cut_capture.py L977-L1003 · delete _checkpoint_active_generation_cut (its only caller is gone) @@ tq_token_sink.py L530-L549 · delete stage_generation_prefix (its only product caller was L990)
Why this is safe: with a batch limit of 1, _checkpoint_generation_cut_batched writes after every row, with the same key, deadline and rollback rules. The test worker already defaults to 256 rows (test L168). In review, routing batch size 1 through the batched loop left all 62 tests passing, and so did turning off the old loop's deadline check. To keep batch size 1 covered, add one assert to the [1] case of test_worker_batch_consecutive_cuts_reuse_then_append: assert sink.batches == [[3], [4]]. That assert was run: it passes after the change and fails today.