diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 29d206b426..c65d02dda4 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -68,7 +68,9 @@ 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]) @@ -76,7 +78,7 @@ def test_scheduler_reports_allocation_failure(request_runner): 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 diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index b878e294a6..c1979a0b1f 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -482,8 +482,13 @@ class RequestRunner: # 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 @@ class RequestRunner: 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() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 2e1875ae73..b53a365abd 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/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 @@ class OffloadingConnectorScheduler: 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 @@ class OffloadingConnectorScheduler: 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 @@ class OffloadingConnectorScheduler: _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 @@ class OffloadingConnectorScheduler: 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 @@ class OffloadingConnectorScheduler: 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 @@ class OffloadingConnectorScheduler: 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 @@ class OffloadingConnectorScheduler: 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