diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 7a8a1a3e4e..615dd82a1c 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -2573,6 +2573,96 @@ def _step_until_kv_transfer_finished(scheduler: Scheduler, req_ids: list[str]): return initial_ecos +def _runner_output( + req_ids: list[str], finished_recving: set[str] | None = None +) -> ModelRunnerOutput: + return ModelRunnerOutput( + req_ids=req_ids, + req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)}, + sampled_token_ids=[[1] for _ in req_ids], + kv_connector_output=KVConnectorOutput(finished_recving=finished_recving), + ) + + +@pytest.mark.parametrize( + "load_done_before_pause", [False, True], ids=["load-in-flight", "load-done"] +) +def test_pause_new_drains_async_kv_loads(load_done_before_pause: bool): + """A PAUSED_NEW (pause mode="wait") drain must finish requests that hold KV + blocks for an async load, keep blockless queued requests queued, and leave + no blocks behind, so the post-drain cache reset succeeds.""" + block_size = 16 + scheduler = create_scheduler( + enable_prefix_caching=True, + use_kv_connector=mock_kv(matched_tokens=2 * block_size, is_async=True), + block_size=block_size, + ) + loading, queued = create_requests( + num_requests=2, + num_tokens=3 * block_size, + max_tokens=2, + block_size=block_size, + req_ids=["loading", "queued"], + ) + scheduler.add_request(loading) + output = scheduler.schedule() + assert loading.status == RequestStatus.WAITING_FOR_REMOTE_KVS + pending_recv = {"loading"} + if load_done_before_pause: + scheduler.update_from_output(output, _runner_output([], pending_recv)) + pending_recv = None + else: + scheduler.update_from_output(output, _runner_output([])) + + scheduler.set_pause_state(PauseState.PAUSED_NEW) + scheduler.add_request(queued) + + num_steps = 0 + while scheduler.has_requests(): + num_steps += 1 + assert num_steps <= 4, "pause drain did not finish" + output = scheduler.schedule() + assert set(output.num_scheduled_tokens) <= {"loading"} + scheduler.update_from_output( + output, _runner_output(list(output.num_scheduled_tokens), pending_recv) + ) + pending_recv = None + + # What pause runs once drained; it raises if the load still holds blocks. + assert scheduler.reset_prefix_cache(reset_running_requests=True) + assert loading.status == RequestStatus.FINISHED_LENGTH_CAPPED + assert loading.request_id not in scheduler.requests + assert list(scheduler.waiting) == [queued] + assert queued.status == RequestStatus.WAITING + + scheduler.set_pause_state(PauseState.UNPAUSED) + scheduler.schedule() + assert queued.status == RequestStatus.WAITING_FOR_REMOTE_KVS + + +def test_pause_new_ignores_streaming_session_waiting_for_input(): + """A session idle in WAITING_FOR_STREAMING_REQ cannot progress, so it must + not keep a PAUSED_NEW drain from completing.""" + stop_token = 7 + scheduler = create_scheduler() + (session,) = create_requests( + num_requests=1, num_tokens=4, stop_token_ids=[stop_token], req_ids=["s"] + ) + session.resumable = True + scheduler.add_request(session) + output = scheduler.schedule() + scheduler.update_from_output( + output, + ModelRunnerOutput( + req_ids=["s"], req_id_to_index={"s": 0}, sampled_token_ids=[[stop_token]] + ), + ) + assert session in scheduler.kv_holding_waiting + + scheduler.set_pause_state(PauseState.PAUSED_NEW) + assert scheduler.get_num_unfinished_requests() == 0 + + @pytest.mark.parametrize( ("load_modes", "expected_has_sync_loads"), [ diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index fb17ccb1d0..f4b4ca2a79 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -747,12 +747,13 @@ async def test_dp_pause_wait_mode_drains_in_flight(): prompt=DP_PAUSE_PROMPT, params=SamplingParams(max_tokens=64), ) + # Wait for the first output, so drain has something to observe: a + # request still queued holds no blocks and is not waited on. + out = await asyncio.wait_for(collector.get(), timeout=30) await engine.pause_generation(mode="wait") assert await engine.is_paused() - while True: + while not out.finished: out = await asyncio.wait_for(collector.get(), timeout=30) - if out.finished: - break await engine.resume_generation() async for out in engine.generate( diff --git a/vllm/v1/core/sched/interface.py b/vllm/v1/core/sched/interface.py index c868d249c8..dd88d99c1f 100644 --- a/vllm/v1/core/sched/interface.py +++ b/vllm/v1/core/sched/interface.py @@ -26,7 +26,7 @@ class PauseState(enum.IntEnum): - UNPAUSED: Normal operation - PAUSE_NEW: No new requests are scheduled, requests already in - running state are scheduled. + running state or holding KV blocks are scheduled. - PAUSE_ALL: No requests are scheduled """ diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 91ea8e93e4..9739dbb3d5 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -887,7 +887,7 @@ class Scheduler(SchedulerInterface): assert len(scheduled_loras) <= self.lora_config.max_loras # Next, schedule the WAITING requests. - if not preempted_reqs and self._pause_state == PauseState.UNPAUSED: + if not preempted_reqs and self._pause_state != PauseState.PAUSED_ALL: step_skipped_waiting: deque[Request] = deque() step_skipped_kv_holding: deque[Request] = deque() @@ -900,8 +900,11 @@ class Scheduler(SchedulerInterface): step_skipped_waiting.appendleft(request_to_skip) while token_budget > 0: - # Requests holding KV blocks are always drained first. - request_queue = self.kv_holding_waiting or self.waiting + # Requests holding KV blocks are always drained first, and are + # the only ones admitted under PAUSED_NEW (e.g. async KV loads). + request_queue = self.kv_holding_waiting + if not request_queue and self._pause_state == PauseState.UNPAUSED: + request_queue = self.waiting if not request_queue or input_budget <= draft_slots: break # Paused streaming sessions (WAITING_FOR_STREAMING_REQ) are not @@ -2802,7 +2805,8 @@ class Scheduler(SchedulerInterface): return 0 num_running, num_waiting = self.get_request_counts() if self._pause_state == PauseState.PAUSED_NEW: - return num_running + # Requests holding KV blocks still drain; queued ones stay queued. + num_waiting = len(self.kv_holding_waiting) num_waiting -= self.num_waiting_for_streaming_input return num_waiting + num_running