iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][Core] Keep resumable sessions that stop while preempted (#58778)

Browse Source 
Signed-off-by: David Cheung <d7cheung@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Yifan Zong <yzong@redhat.com>
This commit is contained in:
David Cheung
2026-10-17 09:08:14 +00:00
committed by GitHub
co-authored by Claude Opus 5.5 Yifan Zong
parent 2067297b2f
commit f98ae6a9fd
4 changed files with 102 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/streaming_input/test_async_llm_streaming.py
+23
View File
@@ -9,6 +9,7 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from vllm.engine.protocol import StreamingInput
from vllm.exceptions import VLLMValidationError
from vllm.outputs import STREAM_FINISHED, RequestOutput
from vllm.sampling_params import RequestOutputKind, SamplingParams
from vllm.v1.engine.async_llm import AsyncLLM
@@ -107,6 +108,25 @@ def make_output(request_id: str, finished: bool) -> RequestOutput:
)
@pytest.mark.parametrize(
("is_mm_encoder_only", "is_diffusion"),
[(True, False), (False, True)],
ids=["encoder-only", "diffusion"],
)
def test_resumable_requests_rejected_for_non_decoder_models(
is_mm_encoder_only: bool, is_diffusion: bool
):
llm = MagicMock()
llm.vllm_config.is_mm_encoder_only = is_mm_encoder_only
llm.model_config.is_diffusion = is_diffusion
with pytest.raises(
VLLMValidationError,
match="Resumable requests are not supported",
):
AsyncLLM._validate_resumable_request(llm)
@pytest.mark.asyncio
async def test_generate_with_async_generator():
"""Test generate with an async input generator.
@@ -180,6 +200,8 @@ async def test_empty_input_stream_finishes_without_engine_request():
llm.get_supported_tasks = AsyncMock(return_value=("generate",))
llm._run_output_handler = MagicMock()
llm._add_request = AsyncMock()
llm.vllm_config = MagicMock(is_mm_encoder_only=False)
llm.model_config = MagicMock(is_diffusion=False)
final_req = MagicMock()
final_req.request_id = "request-int"
@@ -204,6 +226,7 @@ async def test_empty_input_stream_finishes_without_engine_request():
def _make_streaming_llm(num_chunks: int, has_request: bool = True):
llm = AsyncLLM.__new__(AsyncLLM)
llm._validate_streaming_input_sampling_params = MagicMock()
llm._validate_resumable_request = MagicMock()
llm.get_supported_tasks = AsyncMock(return_value=("generate",))
llm._run_output_handler = MagicMock()
llm._add_request = AsyncMock()
tests/v1/streaming_input/test_scheduler_streaming.py
+57
View File
@@ -9,6 +9,7 @@ import pytest
import torch
import vllm.v1.core.kv_cache_utils as kv_cache_utils
from tests.v1.core.utils import create_scheduler as create_async_scheduler
from tests.v1.core.utils import create_scheduler as create_spec_scheduler
from vllm.config import DeviceConfig, VllmConfig
from vllm.multimodal.inputs import (
@@ -179,6 +180,7 @@ class TestStreamingScheduler(unittest.TestCase):
request_id="session",
prompt_token_ids=[1, 2, 3],
)
session.append_output_token_ids(10)
session.num_computed_tokens = len(session.prompt_token_ids)
session.max_tokens = 2
@@ -195,6 +197,7 @@ class TestStreamingScheduler(unittest.TestCase):
assert session.sampling_params.max_tokens == 10
assert session.max_tokens == 10
session.append_output_token_ids(11)
session.num_computed_tokens = len(session.prompt_token_ids)
new_request2 = DummyRequest(
@@ -216,6 +219,7 @@ class TestStreamingScheduler(unittest.TestCase):
request_id="session",
prompt_token_ids=[1, 2, 3],
)
session.append_output_token_ids(10)
session.num_computed_tokens = len(session.prompt_token_ids)
new_request = DummyRequest(
@@ -246,6 +250,7 @@ class TestStreamingScheduler(unittest.TestCase):
prompt_token_ids=[1, 2, 3],
mm_features=[mm_feature],
)
session.append_output_token_ids(10)
session.num_computed_tokens = len(session.prompt_token_ids)
mm_feature = MultiModalFeatureSpec(
@@ -316,6 +321,7 @@ class TestStreamingScheduler(unittest.TestCase):
)
scheduler.add_request(session)
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
session.append_output_token_ids(10)
session.num_computed_tokens = len(session.prompt_token_ids)
next_request = DummyRequest(
@@ -716,3 +722,54 @@ def test_next_chunk_fully_computed_after_stop_on_accepted_draft(async_scheduling
assert list(session.prompt_token_ids) == [1, 2, 3, 10, 7, 8]
assert output.num_scheduled_tokens["session"] == 2
@pytest.mark.parametrize("chunk_queued_before_stop", [False, True])
def test_session_resumes_after_stop_while_preempted(chunk_queued_before_stop):
"""A session preempted with a step in flight can stop when that step's
output is delivered. It must stay schedulable and keep its tokens."""
scheduler = create_async_scheduler(async_scheduling=True)
session = DummyRequest("session", prompt_token_ids=list(range(10)))
scheduler.add_request(session)
scheduler.update_from_output(scheduler.schedule(), _model_output("session", [10]))
in_flight = scheduler.schedule()
next_chunk = DummyRequest("session", prompt_token_ids=[20, 21])
if chunk_queued_before_stop:
scheduler.add_request(next_chunk)
scheduler.running.remove(session)
scheduler._preempt_request(session, timestamp=0.0)
scheduler.update_from_output(in_flight, _model_output("session", [STOP_TOKEN]))
# The parked variant holds a streaming-input slot until the next chunk.
assert scheduler.num_waiting_for_streaming_input == (
0 if chunk_queued_before_stop else 1
)
if not chunk_queued_before_stop:
scheduler.add_request(next_chunk)
output = scheduler.schedule()
assert scheduler.num_waiting_for_streaming_input == 0
expected = list(range(10)) + [10, 20, 21]
assert list(session.all_token_ids) == expected
assert list(session.prompt_token_ids) == expected
assert output.num_scheduled_tokens["session"] == len(expected)
def test_non_resumable_request_stays_dequeued_after_stop_while_preempted():
"""Only resumable requests continue after a stop; a request that finishes
while preempted must not be put back in a queue."""
scheduler = create_async_scheduler(async_scheduling=True)
request = DummyRequest("req", resumable=False, prompt_token_ids=list(range(10)))
scheduler.add_request(request)
scheduler.update_from_output(scheduler.schedule(), _model_output("req", [10]))
in_flight = scheduler.schedule()
scheduler.running.remove(request)
scheduler._preempt_request(request, timestamp=0.0)
scheduler.update_from_output(in_flight, _model_output("req", [STOP_TOKEN]))
assert request.status == RequestStatus.FINISHED_STOPPED
assert request not in scheduler.waiting
assert request not in scheduler.kv_holding_waiting
assert request not in scheduler.deferred_waiting
assert "req" not in scheduler.schedule().num_scheduled_tokens
vllm/v1/core/sched/scheduler.py
+14 -9
View File
@@ -1732,14 +1732,14 @@ class Scheduler(SchedulerInterface):
Discards the last sampled output token from the prior input chunk.
"""
# Current streaming input behaviour: Keep only computed output tokens
# (discard final sampled output token).
num_computed_tokens = session.num_computed_tokens
# Keep all output tokens but the final sampled one.
keep_end = session.num_tokens - 1
assert keep_end >= session.num_computed_tokens
kept_output_tokens = session._all_token_ids[
session.num_prompt_tokens : num_computed_tokens
session.num_prompt_tokens : keep_end
]
del session._all_token_ids[num_computed_tokens:]
del session.block_hashes[num_computed_tokens // self.hash_block_size :]
del session._all_token_ids[keep_end:]
del session.block_hashes[keep_end // self.hash_block_size :]
session._output_token_ids.clear()
assert session.prompt_token_ids is not None
# Extend prompt with kept output tokens.
@@ -2155,6 +2155,8 @@ class Scheduler(SchedulerInterface):
# to avoid expensive operations inside the loop.
stopped_running_reqs: set[Request] = set()
stopped_preempted_reqs: set[Request] = set()
# Stopped resumable requests to re-enqueue, in output order.
continuing_reqs: list[Request] = []
for req_id, num_tokens_scheduled in num_scheduled_tokens.items():
assert num_tokens_scheduled > 0
request = self.requests.get(req_id)
@@ -2320,6 +2322,8 @@ class Scheduler(SchedulerInterface):
kv_transfer_params["remote_prefill_cached_tokens"] = (
prefill_stats.num_cached_tokens
)
else:
continuing_reqs.append(request)
if status_before_stop == RequestStatus.RUNNING:
stopped_running_reqs.add(request)
@@ -2371,14 +2375,16 @@ class Scheduler(SchedulerInterface):
assert not prompt_logprobs_tensors
assert prompt_token_id_logprobs is None
# Remove the stopped requests from the running and waiting queues.
# Remove stopped requests from their queues, then enqueue the ones
# that continue, so each ends up in exactly one queue.
if stopped_running_reqs:
self.running = remove_all(self.running, stopped_running_reqs)
if stopped_preempted_reqs:
# This is a rare case and unlikely to impact performance.
self.waiting.remove_requests(stopped_preempted_reqs)
self.kv_holding_waiting.remove_requests(stopped_preempted_reqs)
self.deferred_waiting.difference_update(stopped_preempted_reqs)
for request in continuing_reqs:
self._enqueue_waiting_request(request)
error_req_ids = set(self.grammar_compile_error_reqs)
self.grammar_compile_error_reqs.clear()
@@ -2568,7 +2574,6 @@ class Scheduler(SchedulerInterface):
request.status = RequestStatus.WAITING_FOR_STREAMING_REQ
self.num_waiting_for_streaming_input += 1
self._enqueue_waiting_request(request)
return False
def _update_request_with_output(
vllm/v1/engine/async_llm.py
+8
View File
@@ -599,6 +599,7 @@ class AsyncLLM(EngineClient):
session_id: str | None = None,
kv_hints: KvHintsEnvelope | None = None,
) -> RequestOutputCollector:
self._validate_resumable_request()
self._validate_streaming_input_sampling_params(sampling_params)
inputs = dict(
@@ -693,6 +694,13 @@ class AsyncLLM(EngineClient):
queue._input_stream_task = asyncio.create_task(handle_inputs())
return queue
def _validate_resumable_request(self) -> None:
if self.vllm_config.is_mm_encoder_only or self.model_config.is_diffusion:
raise VLLMValidationError(
"Resumable requests are not supported for encoder-only or "
"diffusion models."
)
@staticmethod
def _validate_streaming_input_sampling_params(
params: SamplingParams | PoolingParams,