[Bugfix][KV Offloading] Defer request finalization until final store (#49671)

Signed-off-by: Rui Yin <2260891073@qq.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>

Author: Palaiologos1453

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

@@ -68,15 +68,17 @@ 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
+    # Two attempts: once while running (block becomes full during prefill),
+    # once from finished_req_ids on the next step.
+    assert reduced[_ConnectorMetricName.ALLOCATION_FAILURE] == 2
 
 
 @pytest.mark.parametrize("async_scheduling", [True, False])
 @pytest.mark.parametrize("prompt_offset", [-1, -2])
 def test_last_block_offloaded_at_request_finish(
     request_runner, async_scheduling: bool, prompt_offset: int
 ):
-    """EOS fills the last block at request finish — verify req_status is kept alive.
+    """EOS fills the last block at request finish - verify the final block is stored.
 
     prompt = block_size + prompt_offset tokens → not a full block at schedule time,
     so _build_store_jobs creates no store job. After EOS, request_finished
@@ -98,18 +100,16 @@ def test_last_block_offloaded_at_request_finish(
         generate_store_output(list(keys))
     )
 
-    # Run with one step (EOS)
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID],
-    )
+    if prompt_offset == -1:
+        # EOS fills the block, so a store job is created for block 0.
+        runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0,))
+    else:
+        # Block remains partial, so no store job is created.
+        runner.run(decoded_tokens=[EOS_TOKEN_ID])
 
     cs = runner.connector_scheduler
-    # Verify req_status is kept alive for _build_store_jobs to process
-    # regardless of whether there are storable blocks
-    assert "0" in cs._req_status, (
-        "req_status was deleted but should be kept alive "
-        "for _build_store_jobs to process finished_req_ids."
-    )
+    # After the full run completes, req_status is cleaned up.
+    assert "0" not in cs._req_status
 
 
 @pytest.mark.parametrize("async_scheduling", [True, False])
@@ -569,14 +569,14 @@ def test_request_preemption(request_runner, async_scheduling: bool):
 
 
 @pytest.mark.parametrize("async_scheduling", [True, False])
-def test_on_request_finished_is_not_deferred_until_store_completion(
+def test_on_request_finished_not_deferred_until_store_completion(
     request_runner, async_scheduling: bool
 ):
-    """on_request_finished fires when no more stores will be submitted.
+    """on_request_finished fires after the last prepare_store is submitted.
 
-    A request can finish while its GPU->primary store is still in flight. The
-    manager-level hook should not wait for that completion; complete_store may
-    still arrive afterward for already-submitted transfer jobs.
+    The manager contract guarantees no more submit-side calls (prepare_store)
+    after on_request_finished. However, complete_store callbacks for
+    already-submitted transfers may still arrive afterward.
     """
     block_size = 4
     blocks_per_chunk = 3
@@ -613,8 +613,9 @@ def test_on_request_finished_is_not_deferred_until_store_completion(
         complete_transfers=False,
     )
 
-    # Finish the request while its stores are still in flight. The hook should
-    # fire immediately even though no complete_store has arrived yet.
+    # Finish the request while its stores are still in flight. The hook fires
+    # once the last prepare_store is issued (on the next schedule step), even
+    # though complete_store has not yet been called.
     runner.run(
         decoded_tokens=[EOS_TOKEN_ID],
         complete_transfers=False,
@@ -624,8 +625,7 @@ def test_on_request_finished_is_not_deferred_until_store_completion(
 
     assert calls == [("on_request_finished", req_id)], calls
 
-    # Drain the stores afterward. The already-submitted complete_store calls
-    # are allowed to arrive after on_request_finished.
+    # Drain the stores afterward. complete_store is allowed after the hook.
     runner.run(
         decoded_tokens=[],
         complete_transfers=True,
@@ -638,11 +638,50 @@ def test_on_request_finished_is_not_deferred_until_store_completion(
     finished_idx = calls.index(("on_request_finished", req_id))
     store_indices = [i for i, c in enumerate(calls) if c == ("complete_store", req_id)]
 
-    # The request-level hook no longer waits for already-submitted transfers.
+    # complete_store arrives after on_request_finished, as allowed by the contract.
     assert store_indices, calls
     assert finished_idx < min(store_indices), calls
 
 
+@pytest.mark.parametrize("async_scheduling", [True, False])
+def test_on_request_finished_fires_after_final_block_store(
+    request_runner, async_scheduling: bool
+):
+    """on_request_finished fires after the final-block prepare_store at EOS.
+
+    When EOS fills a partial block, request_finished() keeps req_status alive
+    so _build_store_jobs can create a store job for it on the next step.
+    """
+    block_size = 4
+    runner = request_runner(
+        block_size=block_size,
+        num_gpu_blocks=10,
+        async_scheduling=async_scheduling,
+    )
+
+    calls: list[tuple[str, str]] = []
+    runner.manager.on_request_finished.side_effect = lambda req_context: calls.append(
+        ("on_request_finished", req_context.req_id)
+    )
+
+    def prepare_store(keys, req_context):
+        calls.append(("prepare_store", req_context.req_id))
+        return generate_store_output(keys)
+
+    runner.manager.prepare_store.side_effect = prepare_store
+
+    runner.new_request(token_ids=[0] * (block_size - 1))
+    runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(0,))
+
+    req_id = str(runner.req_id)
+    assert calls.count(("on_request_finished", req_id)) == 1, calls
+
+    finished_idx = calls.index(("on_request_finished", req_id))
+    prepare_indices = [i for i, c in enumerate(calls) if c == ("prepare_store", req_id)]
+    assert prepare_indices, calls
+    assert finished_idx > max(prepare_indices), calls
+
+
 @pytest.mark.parametrize("async_scheduling", [True, False])
 def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling: bool):
     block_size = 4
@@ -831,7 +870,10 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
     touch_calls = runner.manager.touch.call_args_list
     assert len(touch_calls) == 6
 
-    runner.run(decoded_tokens=[EOS_TOKEN_ID])
+    # EOS fills the 7th block (offset 6). The extra schedule step processes
+    # finished_req_ids and stores block 6 for both groups before the request's
+    # GPU blocks are freed.
+    runner.run(decoded_tokens=[EOS_TOKEN_ID], expected_stored=(6,))
 
     runner.scheduler.reset_prefix_cache()
 
@@ -844,19 +886,11 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo
         # Group 1 (sliding window, window=2): only the last 2 blocks
         #   are within the window → loads blocks 1,2
         expected_loaded=((0, 0), (0, 1), (0, 2), (1, 1), (1, 2)),
-        # The deferred store from the previous request's last block
-        # completes during this step, and its blocks are flushed because
-        # they were reallocated to the new request.
-        # Only block 1 (sliding window group) is stored — block 0's
-        # deferred store is flushed because it was reallocated.
-        expected_stored=((0, 1),),
-        expected_flushed=((0, 1),),
     )
 
-    # 4 touch calls: 2 from get_num_new_matched_tokens (2 groups)
-    # + 2 from _get_reqs_to_store (2 groups)
+    # 2 touch calls from get_num_new_matched_tokens (2 groups)
     touch_calls = runner.manager.touch.call_args_list
-    assert len(touch_calls) == 4
+    assert len(touch_calls) == 2
     # full attention group touched all 3 blocks
     assert len(touch_calls[0].args[0]) == 3
     # sliding window group touched just the last 2 blocks
@@ -1733,8 +1767,8 @@ def test_reset_cache(request_runner, async_scheduling: bool):
 def test_reset_cache_finalizes_finished_request_with_pending_store(
     request_runner, async_scheduling: bool
 ):
-    """reset_cache drops a finished request whose in-flight stores it discards
-    without calling on_request_finished twice.
+    """reset_cache fires on_request_finished for a finished request whose
+    in-flight stores it discards, exactly once.
     """
     block_size = 4
     blocks_per_chunk = 3
@@ -1770,16 +1804,15 @@ def test_reset_cache_finalizes_finished_request_with_pending_store(
     assert req_status.transfer_jobs, "expected an in-flight store before finish"
     assert any(job.is_store for job in cs._jobs.values())
 
-    # Finish the request while its store is still in flight. request_finished
-    # fires the hook eagerly, but the entry stays tracked so later completions
-    # can still call complete_store().
+    # Finish the request while its store is still in flight. The manager hook
+    # is deferred because the final store decision has not happened yet.
     req_status.req.status = RequestStatus.FINISHED_STOPPED
     cs.request_finished(req_status.req)
-    assert finalized == [req_id]
+    assert finalized == []
     assert req_id in cs._req_status
 
-    # reset_cache discards the in-flight store and drops the state without a
-    # duplicate on_request_finished call.
+    # reset_cache discards both the in-flight and not-yet-prepared final stores,
+    # so it issues the deferred notification before dropping the state.
     cs.reset_cache()
     assert finalized == [req_id]
     assert req_id not in cs._req_status

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

@@ -482,8 +482,13 @@ def _run(
             # 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.
-            if not self.scheduler.requests and not self.connector_scheduler._jobs:
+            # must keep stepping until the scheduler sees no in-flight jobs
+            # and no pending finished_req_ids awaiting build_connector_meta.
+            if (
+                not self.scheduler.requests
+                and not self.connector_scheduler._jobs
+                and not self.scheduler.finished_req_ids
+            ):
                 break
 
             scheduler_output = self.scheduler.schedule()
@@ -544,19 +549,30 @@ def _run(
             if (
                 prev_token_id == EOS_TOKEN_ID
                 and prev_token_id != token_id
-                and (self.scheduler.requests or self.connector_scheduler._jobs)
+                and (
+                    self.scheduler.requests
+                    or self.connector_scheduler._jobs
+                    or self.scheduler.finished_req_ids
+                )
             ):
                 # continue for one more step to allow offloading to kick off
                 continue
 
             if token_id is None:
                 if self.async_scheduling:
-                    # sample last token
+                    # Flush the previous step's output.
                     engine_outputs = self.scheduler.update_from_output(
                         prev_scheduler_output, prev_model_runner_output
                     )
                     self._record_kv_connector_stats(engine_outputs)
-                break
+                    prev_model_runner_output = None
+                if self.scheduler.requests:
+                    # Request still running, just exhausted decoded_tokens.
+                    break
+                if not self.scheduler.finished_req_ids and (
+                    not complete_transfers or not self.connector_scheduler._jobs
+                ):
+                    break
 
         self._parse_transfers()
 

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

@@ -278,6 +278,8 @@ class RequestOffloadState:
     # time.monotonic() of this request's first deferred offload lookup;
     # None once consumed (observed) or while no lookup is pending.
     deferred_lookup_start_time: float | None = None
+    # True once on_request_finished has been signaled to the manager.
+    finished_signaled: bool = False
 
     def __post_init__(self) -> None:
         self.group_states = tuple(
@@ -481,13 +483,6 @@ def _calc_num_offloadable_tokens(
             num = min(num, req_status.req.num_prompt_tokens)
         return num
 
-    def _maybe_cleanup_finished_req(
-        self, req_id: str, req_status: RequestOffloadState
-    ) -> None:
-        """Clean up req_status if finished and no in-flight jobs."""
-        if req_status.req.is_finished() and not req_status.transfer_jobs:
-            del self._req_status[req_id]
-
     def _maximal_prefix_lookup(
         self,
         keys: Iterable[OffloadKey],
@@ -1039,7 +1034,6 @@ def _build_store_jobs(
 
             if not new_offload_keys:
                 req_status.advance_stored_idx(num_offloadable_tokens)
-                self._maybe_cleanup_finished_req(req_id, req_status)
                 continue
 
             store_output = self.manager.prepare_store(
@@ -1050,12 +1044,10 @@ def _build_store_jobs(
                     _ConnectorMetricName.ALLOCATION_FAILURE
                 )
                 logger.warning("Request %s: cannot store chunks", req_id)
-                self._maybe_cleanup_finished_req(req_id, req_status)
                 continue
 
             if not store_output.keys_to_store:
                 req_status.advance_stored_idx(num_offloadable_tokens)
-                self._maybe_cleanup_finished_req(req_id, req_status)
                 continue
 
             self._touch(req_status)
@@ -1199,6 +1191,17 @@ def build_connector_meta(
             store_jobs=self._build_store_jobs(scheduler_output),
             jobs_to_flush=self._current_batch_jobs_to_flush,
         )
+
+        # All prepare_store calls for finished requests have been issued.
+        # Signal on_request_finished and clean up state where possible.
+        for req_id in scheduler_output.finished_req_ids or ():
+            req_status = self._req_status.get(req_id)
+            if req_status is None:
+                continue
+            req_status.finished_signaled = True
+            self.manager.on_request_finished(req_status.req_context)
+            if not req_status.transfer_jobs:
+                del self._req_status[req_id]
         self._current_batch_load_jobs = {}
         self._current_batch_jobs_to_flush = set()
         self._current_batch_allocated_block_ids = set()
@@ -1289,7 +1292,7 @@ def update_connector_output(self, connector_output: KVConnectorOutput):
 
             del self._jobs[job_id]
             req_status.transfer_jobs.remove(job_id)
-            if not req_status.transfer_jobs and req_status.req.is_finished():
+            if req_status.finished_signaled and not req_status.transfer_jobs:
                 del self._req_status[job_status.req_id]
 
     def get_stats(self) -> OffloadingConnectorStats | None:
@@ -1331,7 +1334,6 @@ def request_finished(
             self.manager.on_request_finished(req_context)
             return False, None
 
-        self.manager.on_request_finished(req_status.req_context)
         self._maybe_observe_lookup_async_delay(req_status)
 
         # Update offload keys with final block hash so _build_store_jobs can
@@ -1375,6 +1377,8 @@ def reset_cache(self) -> None:
 
         for req_id, status in list(self._req_status.items()):
             if status.req.is_finished():
+                if not status.finished_signaled:
+                    self.manager.on_request_finished(status.req_context)
                 del self._req_status[req_id]
 
         # Reset offloading manager cache