The background delete can reuse the checkpoint gate that already exists

Size review of NVIDIA-NeMo/RL #4655, stacked on #4508. All permalinks pin the PR head 495ce13. Review comments are in a pending review; their links work once it is published.

Before
When a call that was cut at a checkpoint finishes, its reply waits for a synchronous delete (a TQ CLEAR) of the old chunk rows.
This PR
Behind an opt-in flag, moves that delete to a background thread. To do it, it adds a 253-line module with its own slot reservation and its own checkpoint pause/resume, wired into the worker at 4 places.
Status
Proposed — not implemented. The smaller version below is a review suggestion. A prototype passed 5 checks against the real gate class; it has not been run inside NeMo-RL.

Two pages on this PR — this one: what the new module repeats, and what can go.

Background. During a rollout checkpoint, a call that is still generating is cut: its tokens so far are saved as chunk rows in TQ. When the call later finishes, its final row is written and the old chunk rows are deleted. A checkpoint must not save TQ while a final-row write is half done, so the worker has a _TokenCaptureSnapshotGate: writers enter() and exit() it, and a checkpoint closes it and waits until no one is inside. Dotted words have a short explanation: hover or tap them.

The PR's background thread needs one rule: no delete may run while a checkpoint saves TQ. The gate already enforces exactly that for anything that holds a permit. Instead of taking one, the new class (BackgroundPrefixCleanup) keeps a second copy of the gate's epoch bookkeeping and adds pause, wait_paused and resume, which the checkpoint code must now call next to the gate's own calls.

1 · Today, flag off (the #4508 path) The reply waits for the CLEAR. The checkpoint already waits for anything inside the gate. finishing call enter gate PUT final row CLEAR old chunks exit gate reply checkpoint close gate wait until no one is inside save TQ reopen 2 · This PR, flag on The reply no longer waits — but the cleanup thread gets its own pause/resume, and every call takes a slot. finishing call reserve slot enter gate PUT final row queue keys exit gate reply cleanup thread collect ≤32 calls own epoch check CLEAR checkpoint close gate pause(epoch) wait gate wait_paused save TQ reopen resume() 3 · Proposed, flag on Same reply time. The cleanup thread takes a gate permit, so the checkpoint is exactly today's. finishing call enter gate PUT final row exit gate queue.put(keys) reply cleanup thread collect ≤32 calls enter gate CLEAR exit gate checkpoint close gate wait until no one is inside save TQ reopen reply waits on this new fence code in this PR reuses the existing gate existing gate (unchanged)

What goes away

Piece in this PRIn the proposed version
_epoch/_released_epoch, pause, wait_paused, resume (L157-L174)Gone: the gate's own epochs already cover them.
4 hooks in the worker: pause, wait, resume, shutdown resume (L3388, L3421, L3483, L3497)Gone: checkpoint code stays as it is today.
CleanupReservation and a slot taken by every finishing call (L1508)A stdlib queue.Queue(maxsize=max_pending), filled after the call leaves the gate, only when it has keys.
close(), its paused check, stats() and 5 counters only tests readGone. The thread is a daemon; the PR already accepts that lost deletes only leave extra rows.
253-line module; 19 unit tests, mostly for the fence aboveAbout 25 lines; tests for the removed fence go with it.

review comment — background_prefix_cleanup.py:157 review comment — background_prefix_cleanup.py:221

The change — proposed, not in the PR

Nothing below is implemented in NeMo-RL. The same design ran as a stand-alone prototype against the real gate class.

Where the fix lands: RL only, two files. background_prefix_cleanup.py shrinks to a queue and a thread; vllm_worker_async.py loses the reservation and the 4 hooks. No Gym change.

@@ background_prefix_cleanup.py · the cleanup thread, whole loop
+    def _drain(self):
+        while True:
+            batch = [self._queue.get()]
+            # ... collect up to batch_size calls within wait_seconds ...
+            keys = list(dict.fromkeys(k for rows in batch for k in rows))
+            self._gate.enter()  # a checkpoint waits for this CLEAR, and blocks the next one
+            try:
+                self._clear(keys)
+            except Exception:  # a lost delete only leaves extra rows
+                LOGGER.exception("background prefix cleanup failed for %d keys", len(keys))
+            finally:
+                self._gate.exit()

@@ vllm_worker_async.py L1505-L1519 · _finish_request_capture
-        # Reserve outside the terminal gate: a paused, full cleanup queue must
-        # never hold a permit that the snapshot fence is trying to drain.
-        cleanup = self._prefix_cleanup
-        reservation = cleanup.reserve() if cleanup is not None else None
-        try:
-            self._token_capture_snapshot_gate.enter()
-            try:
-                return self._finish_request_capture_after_snapshot_fence(
-                    request, content, cleanup_reservation=reservation
-                )
-            finally:
-                self._token_capture_snapshot_gate.exit()
-        finally:
-            if reservation is not None:
-                reservation.close()
+        self._token_capture_snapshot_gate.enter()
+        try:
+            content, keys = self._finish_request_capture_after_snapshot_fence(
+                request, content
+            )
+        finally:
+            self._token_capture_snapshot_gate.exit()
+        # After exit, so a full queue never holds a permit a checkpoint is draining.
+        if keys and self._prefix_cleanup is not None:
+            self._prefix_cleanup.submit(keys)
+        return content

@@ vllm_worker_async.py L3388-L3397 · begin_token_capture_snapshot_fence_async
-        cleanup = self._prefix_cleanup
-        if cleanup is not None:
-            cleanup.pause(epoch)
         await asyncio.get_running_loop().run_in_executor(
             self._token_capture_fence_executor, gate.close_and_wait, epoch
         )
-        if cleanup is not None:
-            await asyncio.get_running_loop().run_in_executor(
-                self._token_capture_fence_executor, cleanup.wait_paused, epoch
-            )
   … same for resume() at L3421-L3422 and the shutdown hooks at L3483-L3485, L3497-L3498

What each piece does, in order: (1) the thread takes a gate permit around each CLEAR, so close_and_wait waits for a running delete and enter() holds the next one until the checkpoint reopens the gate; (2) the finishing call queues its keys only after it leaves the gate, the same reason the PR reserves outside it today; (3) the checkpoint code goes back to what #4508 has. The prototype checked: a reply returns while a CLEAR is still running; a checkpoint waits for that CLEAR; no CLEAR starts while the gate is closed; 1,024 calls blocked on a full queue plus a second checkpoint do not deadlock; 2,048 keys clear in 34–36 calls at batch size 32 (the PR's test gets 32). Do not reuse PrefixReadBatcher from #4508 instead: it returns awaited results on one event loop, while this delete is fire-and-forget from a worker thread.

review comment — background_prefix_cleanup.py:157

So what. The feature stack is already about 20k lines, and this PR adds a second checkpoint fence for one thread. Taking a permit on the gate that exists keeps the reply-time gain and drops the module to about 25 lines, with no checkpoint changes. It also removes the bug on the next page. The PR still shows no latency number for the gain; the review asks for one.