[Bugfix][KV Offloading] Offload last block at request finish and prevent reuse race

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

@@ -10,7 +10,9 @@
     generate_store_output,
     to_keys,
 )
-from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID
+from tests.v1.kv_connector.unit.utils import (
+    EOS_TOKEN_ID,
+)
 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
     OffloadingConnectorStats,
     _ConnectorMetricName,
@@ -61,7 +63,7 @@ def test_scheduler_reports_allocation_failure(request_runner):
     runner.run(decoded_tokens=[EOS_TOKEN_ID])
 
     reduced = _reduce_kv_connector_stats(runner)
-    assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 1
+    assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 2
 
 
 def test_scheduler_reports_lookup_sync_delay(request_runner):
@@ -2757,3 +2759,148 @@ def capture_fence():
     # Verify fence is empty after full lifecycle (cleanup happened).
     assert runner.connector_scheduler._block_id_to_pending_jobs == {}
     assert len(runner.connector_scheduler._jobs) == 0
+
+
+@pytest.mark.parametrize("async_scheduling", [True, False])
+@pytest.mark.parametrize("prompt_offset", [-1, -2])
+def test_last_block_offload_at_request_finish(
+    request_runner, async_scheduling: bool, prompt_offset: int
+):
+    """The last block that fills at request finish time is offloaded.
+
+    Regression test for orozery's TODO: when the finishing token (EOS)
+    completes a block, that block's KV must still be offloaded. Before
+    schedule time the block wasn't full, so _build_store_jobs skipped it;
+    only request_finished has the complete token count.
+
+    prompt_offset=-1: prompt has block_size-1 tokens, decode EOS fills
+    the block → store job created → offloaded.
+    prompt_offset=-2: prompt has block_size-2 tokens, decode EOS leaves
+    the block partial → no store job → not offloaded.
+    """
+    block_size = 4
+    block_size_factor = 1
+    should_offload = prompt_offset == -1
+
+    runner = request_runner(
+        block_size=block_size,
+        num_gpu_blocks=10,
+        async_scheduling=async_scheduling,
+        block_size_factor=block_size_factor,
+    )
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+
+    runner.new_request(token_ids=[0] * (block_size + prompt_offset))
+
+    cs = runner.connector_scheduler
+    finished_req_store_jobs_seen = []
+
+    original_update = runner.scheduler.update_from_output
+
+    def patched_update(*args, **kwargs):
+        result = original_update(*args, **kwargs)
+        finished_req_store_jobs_seen.append(dict(cs._finished_req_store_jobs))
+        return result
+
+    runner.scheduler.update_from_output = patched_update
+
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID],
+        expected_stored=(0,) if should_offload else (),
+        drain_finished_req_stores=should_offload,
+    )
+
+    if should_offload:
+        non_empty_states = [s for s in finished_req_store_jobs_seen if s]
+        assert non_empty_states, (
+            "_finished_req_store_jobs was never non-empty. "
+            "Expected request_finished to create a store job for the last block. "
+            f"All states seen: {finished_req_store_jobs_seen}"
+        )
+        assert runner.manager.complete_store.call_count == 1
+    else:
+        assert runner.manager.complete_store.call_count == 0
+
+
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_eos_block_reuse_race(request_runner, async_scheduling: bool):
+    """Verify EOS block store is safe when block is reused.
+
+    When request_finished creates a store job in _finished_req_store_jobs,
+    and a new request reuses the same block, handle_preemptions now submits
+    the store job BEFORE calling wait() — so wait() is effective and blocks
+    are properly fenced until the store completes.
+
+    This test uses pre_schedule_fn to inject Request B at the exact moment
+    when _finished_req_store_jobs is non-empty (after request_finished but
+    before build_connector_meta merges the jobs).
+
+    Scenario (block_size=4, num_gpu_blocks=2):
+    - Request A: prompt=3 tokens, decode EOS → 4 tokens = 1 full block
+    - pre_schedule_fn: when _finished_req_store_jobs is non-empty, add Request B
+    - Request B: prompt=3 tokens (same prefix) → reuses Request A's block
+    - Fixed: wait() called AFTER submit_store → effective fence
+    """
+    block_size = 4
+
+    runner = request_runner(
+        block_size=block_size,
+        num_gpu_blocks=2,
+        async_scheduling=async_scheduling,
+        block_size_factor=1,
+    )
+    runner.manager.prepare_store.side_effect = lambda keys, req_context: (
+        generate_store_output(keys)
+    )
+
+    # Request A: prompt=3 tokens, will decode EOS
+    runner.new_request(token_ids=[0] * (block_size - 1))
+
+    handler = runner.offloading_spec.handler
+    handler.wait_before_submit.clear()
+
+    request_b_added = [False]
+
+    def inject_request_b():
+        """Inject Request B when _finished_req_store_jobs is non-empty."""
+        if request_b_added[0]:
+            return
+        cs = runner.connector_scheduler
+        if cs._finished_req_store_jobs:
+            # Reset prefix cache to allow block reuse
+            runner.scheduler.reset_prefix_cache()
+            # Request B: same prefix as Request A → reuses block 0
+            runner.new_request(token_ids=[0] * (block_size - 1))
+            request_b_added[0] = True
+
+    # Run Request A with pre_schedule_fn to inject Request B at the right time
+    # Use complete_transfers=True to allow jobs to complete and loop to exit
+    # decoded_tokens has three EOS tokens: one for Request A, two for Request B
+    # (async mode may need extra steps)
+    runner.run(
+        decoded_tokens=[EOS_TOKEN_ID, EOS_TOKEN_ID, EOS_TOKEN_ID],
+        complete_transfers=True,
+        drain_finished_req_stores=True,
+        pre_schedule_fn=inject_request_b,
+        expected_stored=(0,),  # Request A stores block 0
+        expected_flushed=(0,),  # Block 0 is flushed after wait() fences it
+        max_steps=30,
+    )
+
+    # Verify Request B was injected
+    assert request_b_added[0], (
+        "Request B was not injected. "
+        "_finished_req_store_jobs was never non-empty at pre_schedule_fn time."
+    )
+
+    # Race verification: wait_before_submit should be empty (no race)
+    # handle_preemptions now submits store_jobs BEFORE wait(), so
+    # wait() always operates on submitted jobs.
+    assert not handler.wait_before_submit, (
+        "Race detected: wait() was called on an unsubmitted job. "
+        "This means block could be reused before store reads it "
+        "→ silent data corruption. "
+        f"wait_before_submit={handler.wait_before_submit}"
+    )

tests/v1/kv_connector/unit/offloading_connector/utils.py

@@ -84,6 +84,11 @@ def __init__(self):
         self.waiting_jobs: set[int] = set()
         self.completed_jobs: list[int] = []
         self.flushed_jobs: set[int] = set()
+        # Tracks job_ids that were wait()ed before submit_store() — exposes
+        # the race where _finished_req_store_jobs jobs end up in
+        # jobs_to_flush but were never submitted, making wait() a no-op
+        # and allowing block reuse before the store reads the data.
+        self.wait_before_submit: set[int] = set()
 
     def get_finished(self) -> list[TransferResult]:
         finished = self.completed_transfers
@@ -118,6 +123,9 @@ def complete_jobs(self, job_ids: set[int]) -> None:
                 self.completed_transfers.append(result)
 
     def wait(self, job_ids: set[int]) -> None:
+        for job_id in job_ids:
+            if job_id not in self.transfer_specs:
+                self.wait_before_submit.add(job_id)
         self.flushed_jobs |= job_ids
         self.complete_jobs(job_ids)
 
@@ -459,11 +467,36 @@ def _update_gpu_blocks(self):
                 for block_idx, block in enumerate(blocks):
                     self.gpu_blocks[block.block_id] = GPUBlock(group_idx, block_idx)
 
+    def _discard_finished_req_store_jobs(self):
+        """Discard _finished_req_store_jobs and clean up _jobs/transfer_jobs.
+
+        In production these jobs would be delivered by the next request's
+        schedule(); in single-request tests with drain_finished_req_stores=False
+        there is no next request, so we must fully clean up to avoid leaving
+        stale jobs in _jobs (which would prevent the loop from exiting).
+        """
+        cs = self.connector_scheduler
+        for job_id in list(cs._finished_req_store_jobs.keys()):
+            job = cs._jobs.pop(job_id, None)
+            if job is None:
+                continue
+            cs._remove_pending_job(job_id, job.sliding_window_block_ids)
+            cs._remove_pending_job(job_id, job.non_sliding_window_block_ids)
+            req_status = cs._req_status.get(job.req_id)
+            if req_status is not None:
+                req_status.transfer_jobs.discard(job_id)
+                if not req_status.transfer_jobs and req_status.req.is_finished():
+                    del cs._req_status[job.req_id]
+        cs._finished_req_store_jobs.clear()
+
     def _run(
         self,
         decoded_tokens: list[int],
         complete_transfers: bool,
         post_step_fn: Callable[[], None] | None = None,
+        pre_schedule_fn: Callable[[], None] | None = None,
+        drain_finished_req_stores: bool = False,
+        max_steps: int = 100,
     ):
         """
         Runs multiple engine (scheduler + worker) steps.
@@ -474,20 +507,42 @@ def _run(
             complete_transfers: complete transfers immediately
             post_step_fn: optional callback invoked after each step's
                 update_from_output(), before the next schedule().
+            pre_schedule_fn: optional callback invoked before each step's
+                schedule(). Useful for injecting requests at specific points
+                in the scheduling cycle (e.g., after request_finished creates
+                _finished_req_store_jobs but before build_connector_meta
+                merges them).
+            drain_finished_req_stores: when True, keep looping after
+                token exhaustion to deliver _finished_req_store_jobs
+                created by request_finished via the next build_connector_meta.
+            max_steps: safety limit to prevent infinite loops in tests.
         """
 
         tokens_iter = iter(decoded_tokens)
         token_id = next(tokens_iter, None)
         prev_scheduler_output = None
         prev_model_runner_output = None
+        step = 0
         while True:
-            # Strict-always-False frees the request immediately on EOS, but
-            # the worker may still have a deferred store queued. In production
-            # the next request's step drains it; in single-request tests we
-            # must keep stepping until the scheduler sees no in-flight jobs.
+            step += 1
+            assert step <= max_steps, (
+                f"Test exceeded max_steps={max_steps} — "
+                f"scheduler.requests={bool(self.scheduler.requests)}, "
+                f"connector_scheduler._jobs={bool(self.connector_scheduler._jobs)}, "
+                f"_finished_req_store_jobs="
+                f"{bool(self.connector_scheduler._finished_req_store_jobs)}"
+            )
+            # Exit when no requests and no pending connector work.
             if not self.scheduler.requests and not self.connector_scheduler._jobs:
                 break
 
+            # Invoke pre_schedule_fn before schedule() to allow injecting
+            # requests at specific points (e.g., after request_finished
+            # creates _finished_req_store_jobs but before build_connector_meta
+            # merges them into store_jobs).
+            if pre_schedule_fn is not None:
+                pre_schedule_fn()
+
             scheduler_output = self.scheduler.schedule()
             self._update_gpu_blocks()
 
@@ -549,6 +604,8 @@ def _run(
                 and (self.scheduler.requests or self.connector_scheduler._jobs)
             ):
                 # continue for one more step to allow offloading to kick off
+                if not drain_finished_req_stores:
+                    self._discard_finished_req_store_jobs()
                 continue
 
             if token_id is None:
@@ -558,6 +615,19 @@ def _run(
                         prev_scheduler_output, prev_model_runner_output
                     )
                     self._record_kv_connector_stats(engine_outputs)
+                    prev_scheduler_output = None
+                    prev_model_runner_output = None
+                # update_from_output (async above, or sync at loop bottom)
+                # may have triggered request_finished, creating
+                # _finished_req_store_jobs that need delivery via the next
+                # schedule() → build_connector_meta.
+                if drain_finished_req_stores and (
+                    self.connector_scheduler._finished_req_store_jobs
+                    or self.connector_scheduler._jobs
+                ):
+                    continue
+                # Discard undelivered _finished_req_store_jobs when not draining.
+                self._discard_finished_req_store_jobs()
                 break
 
         self._parse_transfers()
@@ -590,6 +660,9 @@ def run(
         expected_loaded: tuple[int | tuple[int, int], ...] = (),
         expected_flushed: tuple[int | tuple[int, int], ...] = (),
         post_step_fn: Callable[[], None] | None = None,
+        pre_schedule_fn: Callable[[], None] | None = None,
+        drain_finished_req_stores: bool = False,
+        max_steps: int = 100,
     ):
         """
         Runs multiple engine (scheduler + worker) steps.
@@ -608,14 +681,27 @@ def run(
             A GPU block is either a (group_idx: int, request_block_offset: int)
             or just request_block_offset: int.
             The latter case is a convenience for representing all groups.
+            post_step_fn: optional callback invoked after each step's
+                update_from_output(), before the next schedule().
+            pre_schedule_fn: optional callback invoked before each step's
+                schedule(). Useful for injecting requests at specific points
+                in the scheduling cycle.
+            max_steps: safety limit to prevent infinite loops (default 100).
         """
 
         expected_stored_gpu_blocks = self._to_gpu_blocks(expected_stored)
         expected_loaded_gpu_blocks = self._to_gpu_blocks(expected_loaded)
         expected_flushed_gpu_blocks = self._to_gpu_blocks(expected_flushed)
 
         self.manager.reset_mock()
-        self._run(decoded_tokens, complete_transfers, post_step_fn=post_step_fn)
+        self._run(
+            decoded_tokens,
+            complete_transfers,
+            post_step_fn=post_step_fn,
+            pre_schedule_fn=pre_schedule_fn,
+            drain_finished_req_stores=drain_finished_req_stores,
+            max_steps=max_steps,
+        )
 
         loaded_gpu_blocks: set[GPUBlock] = set()
         for transfer in self.completed_loads: