Address KV offload finalization review

Centralize finished-request signaling after final store construction and make the scheduler test runner drain finished request steps and transfer completions.

Co-authored-by: Or Ozeri <oro@il.ibm.com>
Signed-off-by: Rui Yin <2260891073@qq.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
 ):
-    """Prepare the final store before notifying the manager of request finish.
+    """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
@@ -94,85 +96,20 @@ def test_last_block_offloaded_at_request_finish(
     )
     # prompt = block_size + prompt_offset tokens
     runner.new_request(token_ids=[0] * (block_size + prompt_offset))
-    manager_events: list[str] = []
-
-    def prepare_store(keys, req_context):
-        manager_events.append("prepare_store")
-        return generate_store_output(list(keys))
-
-    runner.manager.prepare_store.side_effect = prepare_store
-    runner.manager.on_request_finished.side_effect = lambda req_context: (
-        manager_events.append("on_request_finished")
-    )
-
-    # Run with one step (EOS)
-    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."
-    )
-    assert manager_events == []
-
-    scheduler_output = runner.scheduler.schedule()
-    metadata = scheduler_output.kv_connector_metadata
-    assert isinstance(metadata, OffloadingConnectorMetadata)
-
-    if prompt_offset == -1:
-        assert metadata.store_jobs
-        assert manager_events == ["prepare_store", "on_request_finished"]
-        assert "0" in cs._req_status
-    else:
-        assert not metadata.store_jobs
-        assert manager_events == ["on_request_finished"]
-        assert "0" not in cs._req_status
-
-
-@pytest.mark.parametrize("async_scheduling", [True, False])
-def test_store_completion_at_finish_keeps_pending_final_store(
-    request_runner, async_scheduling: bool
-):
-    """A completed older store must not discard the pending final store."""
-    block_size = 4
-    blocks_per_chunk = 3
-    tokens_per_chunk = block_size * blocks_per_chunk
-    runner = request_runner(
-        block_size=block_size,
-        num_gpu_blocks=10,
-        async_scheduling=async_scheduling,
-        blocks_per_chunk=blocks_per_chunk,
-    )
-    runner.new_request(token_ids=[0] * (2 * tokens_per_chunk - 1))
     runner.manager.prepare_store.side_effect = lambda keys, req_context: (
         generate_store_output(list(keys))
     )
 
-    # The first chunk's store completes in the same output that finishes the
-    # request. The EOS-filled second chunk still needs the next schedule step.
-    runner.run(
-        decoded_tokens=[EOS_TOKEN_ID],
-        complete_transfers=True,
-        expected_stored=tuple(range(blocks_per_chunk)),
-    )
+    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
-    req_id = str(runner.req_id)
-    assert req_id in cs._req_status
-    assert not cs._req_status[req_id].transfer_jobs
-    runner.manager.on_request_finished.assert_not_called()
-
-    scheduler_output = runner.scheduler.schedule()
-    metadata = scheduler_output.kv_connector_metadata
-    assert isinstance(metadata, OffloadingConnectorMetadata)
-    assert metadata.store_jobs
-    runner.manager.on_request_finished.assert_called_once_with(
-        cs._req_status[req_id].req_context
-    )
+    # After the full run completes, req_status is cleaned up.
+    assert "0" not in cs._req_status
 
 
 @pytest.mark.parametrize("async_scheduling", [True, False])
@@ -632,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
@@ -676,9 +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 following
-    # scheduler step prepares any final store and fires the hook before any
-    # complete_store has arrived.
+    # 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,
@@ -688,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,
@@ -702,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
@@ -895,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()
 
@@ -908,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
@@ -1797,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
@@ -1835,8 +1805,7 @@ def test_reset_cache_finalizes_finished_request_with_pending_store(
     assert any(job.is_store for job in cs._jobs.values())
 
     # Finish the request while its store is still in flight. The manager hook
-    # is deferred until the final store decision, and the entry stays tracked
-    # so later completions can still call complete_store().
+    # 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 == []

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()