diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 5ad8269b37..b7b02efd12 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -42,7 +42,7 @@ from vllm.entrypoints.openai.models.serving import ( ) from vllm.entrypoints.openai.parser.harmony_utils import get_encoding from vllm.entrypoints.serve.engine.protocol import ErrorResponse -from vllm.exceptions import VLLMValidationError +from vllm.exceptions import QueueOverflowError, VLLMValidationError from vllm.inputs import TokensPrompt from vllm.multimodal.inputs import PlaceholderRange from vllm.outputs import CompletionOutput, RequestOutput @@ -848,6 +848,9 @@ class MockEngine: renderer: MagicMock = field(default_factory=MagicMock) errored: bool = False + def check_admission(self, n: int = 1, request_id: str | None = None) -> None: + pass + async def _async_serving_chat_init(): engine = MockEngine() @@ -944,6 +947,35 @@ async def test_serving_chat_returns_correct_model_name(): assert await serving_chat.create_chat_completion(req) == MODEL_NAME +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [True, False]) +async def test_admission_rejection_escapes_before_response_starts(stream): + """Overload rejections must propagate out of create_chat_completion. + + For streaming this is the only chance to return a real HTTP status: once + StreamingResponse is constructed the 200 has already been sent and the + rejection would degrade into an in-band SSE error chunk. + """ + mock_engine = MagicMock(spec=AsyncLLM) + mock_engine.errored = False + mock_engine.model_config = MockModelConfig() + mock_engine.input_processor = MagicMock() + mock_engine.renderer = _build_renderer(mock_engine.model_config) + mock_engine.check_admission.side_effect = QueueOverflowError() + + serving_chat = _build_serving_chat(mock_engine) + req = ChatCompletionRequest( + model=MODEL_NAME, + messages=[{"role": "user", "content": "what is 1+1?"}], + stream=stream, + ) + + with pytest.raises(QueueOverflowError): + await serving_chat.create_chat_completion(req) + + mock_engine.generate.assert_not_called() + + @pytest.mark.asyncio async def test_serving_chat_should_set_correct_max_tokens(): mock_engine = MagicMock(spec=AsyncLLM) diff --git a/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py b/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py index 45772708d9..0ff8adb6b5 100644 --- a/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py +++ b/tests/entrypoints/speech_to_text/test_speech_to_text_cancellation.py @@ -38,6 +38,7 @@ async def test_non_streaming_cancel_aborts_engine_requests( ): engine_client = SimpleNamespace( errored=False, + check_admission=Mock(), generate=Mock(side_effect=lambda *_args, **_kwargs: _never_finishes()), abort=AsyncMock(), is_tracing_enabled=AsyncMock(return_value=False), @@ -104,6 +105,7 @@ async def test_non_streaming_cancel_advances_all_chunk_generators(): started_request_ids: list[str] = [] engine_client = SimpleNamespace( errored=False, + check_admission=Mock(), generate=Mock( side_effect=lambda *_args, **_kwargs: _records_start_then_never_finishes( started_request_ids, _args[2] diff --git a/tests/v1/engine/test_admission_control.py b/tests/v1/engine/test_admission_control.py new file mode 100644 index 0000000000..a03718ab5a --- /dev/null +++ b/tests/v1/engine/test_admission_control.py @@ -0,0 +1,384 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for admission control (max_num_queued_reqs / max_num_queued_tokens). + +These tests cover: +- OutputProcessor.get_num_queued_tokens() token counting +- AsyncLLM.check_admission() admission control logic +- Exception classes (GracefulHTTPError, QueueOverflowError, MaxQueuedTokensError) +- create_error_response() mapping GracefulHTTPError to HTTP 503 +- SchedulerConfig field defaults and validation +- human_readable_int CLI notation for max_num_queued_tokens +""" + +import argparse +import asyncio +from http import HTTPStatus +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from pydantic import ValidationError + +from vllm.config.scheduler import SchedulerConfig +from vllm.entrypoints.serve.exception_handling.error_response import ( + create_error_response, +) +from vllm.exceptions import ( + GracefulHTTPError, + MaxQueuedTokensError, + QueueOverflowError, + VLLMError, +) +from vllm.pooling_params import PoolingParams +from vllm.sampling_params import SamplingParams +from vllm.utils.argparse_utils import human_readable_int +from vllm.v1.engine.async_llm import AsyncLLM +from vllm.v1.engine.output_processor import OutputProcessor + +pytestmark = pytest.mark.cpu_test + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_req_state(prompt_len: int, is_prefilling: bool = True): + return SimpleNamespace( + prompt_len=prompt_len, + is_prefilling=is_prefilling, + ) + + +def _make_async_llm( + max_num_queued_reqs: int | None = None, + max_num_queued_tokens: int | None = None, + num_unfinished: int = 0, + num_queued_tokens: int = 0, +) -> AsyncLLM: + """Create a bare AsyncLLM with just the attributes needed for scheduling.""" + llm = AsyncLLM.__new__(AsyncLLM) + llm.scheduler_config = SimpleNamespace( + max_num_queued_reqs=max_num_queued_reqs, + max_num_queued_tokens=max_num_queued_tokens, + ) + llm.output_processor = MagicMock() + llm.output_processor.get_num_unfinished_requests.return_value = num_unfinished + llm.output_processor.get_num_queued_tokens.return_value = num_queued_tokens + return llm + + +def _make_output_processor(**request_states) -> OutputProcessor: + op = OutputProcessor.__new__(OutputProcessor) + op.request_states = request_states + return op + + +def _make_scheduler_config(**kwargs) -> SchedulerConfig: + return SchedulerConfig( + runner_type="generate", + max_model_len=4096, + is_encoder_decoder=False, + **kwargs, + ) + + +# --------------------------------------------------------------------------- +# Exception classes +# --------------------------------------------------------------------------- + + +def test_graceful_http_error_carries_status_and_message(): + err = GracefulHTTPError("custom message", HTTPStatus.SERVICE_UNAVAILABLE) + assert err.message == "custom message" + assert err.http_status == HTTPStatus.SERVICE_UNAVAILABLE + assert str(err) == "custom message" + + +def test_graceful_http_error_is_vllm_error(): + err = GracefulHTTPError("msg", HTTPStatus.TOO_MANY_REQUESTS) + assert isinstance(err, VLLMError) + + +def test_queue_overflow_error(): + err = QueueOverflowError() + assert err.http_status == HTTPStatus.SERVICE_UNAVAILABLE + assert isinstance(err, GracefulHTTPError) + assert "busy" in err.message.lower() or "try again" in err.message.lower() + + +def test_max_queued_tokens_error(): + err = MaxQueuedTokensError() + assert err.http_status == HTTPStatus.SERVICE_UNAVAILABLE + assert isinstance(err, GracefulHTTPError) + assert "backlog" in err.message.lower() or "try again" in err.message.lower() + + +@pytest.mark.parametrize("exc_cls", [QueueOverflowError, MaxQueuedTokensError]) +def test_admission_exceptions_are_vllm_errors(exc_cls): + """Admission rejections must reach the VLLMError HTTP handler.""" + assert issubclass(exc_cls, VLLMError) + + +# --------------------------------------------------------------------------- +# OutputProcessor.get_num_queued_tokens +# --------------------------------------------------------------------------- + + +def test_queued_tokens_empty(): + assert _make_output_processor().get_num_queued_tokens() == 0 + + +def test_queued_tokens_sums_prefilling_requests(): + op = _make_output_processor(r1=_make_req_state(100), r2=_make_req_state(200)) + assert op.get_num_queued_tokens() == 300 + + +def test_queued_tokens_excludes_non_prefilling(): + op = _make_output_processor( + r1=_make_req_state(100, is_prefilling=True), + r2=_make_req_state(200, is_prefilling=False), + r3=_make_req_state(50, is_prefilling=True), + ) + assert op.get_num_queued_tokens() == 150 + + +def test_queued_tokens_all_non_prefilling(): + op = _make_output_processor( + r1=_make_req_state(100, is_prefilling=False), + r2=_make_req_state(200, is_prefilling=False), + ) + assert op.get_num_queued_tokens() == 0 + + +# --------------------------------------------------------------------------- +# AsyncLLM.check_admission +# --------------------------------------------------------------------------- + + +def test_admission_no_limits_allows_everything(): + llm = _make_async_llm(num_unfinished=999, num_queued_tokens=999) + llm.check_admission() + + +# -- max_num_queued_reqs ---------------------------------------------------- + + +def test_admission_reqs_allows_when_under_limit(): + llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=5) + llm.check_admission() + + +def test_admission_reqs_rejects_at_limit(): + llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=10) + with pytest.raises(QueueOverflowError): + llm.check_admission() + + +def test_admission_reqs_rejects_with_n(): + llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=8) + with pytest.raises(QueueOverflowError): + llm.check_admission(3) + + +def test_admission_reqs_allows_n_at_boundary(): + llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=7) + llm.check_admission(3) + + +def test_admission_reqs_rejects_when_zero_limit(): + llm = _make_async_llm(max_num_queued_reqs=0, num_unfinished=0) + with pytest.raises(QueueOverflowError): + llm.check_admission() + + +# -- max_num_queued_tokens -------------------------------------------------- + + +def test_admission_tokens_allows_when_under_limit(): + llm = _make_async_llm(max_num_queued_tokens=1000, num_queued_tokens=500) + llm.check_admission() + + +def test_admission_tokens_rejects_at_limit(): + llm = _make_async_llm(max_num_queued_tokens=1000, num_queued_tokens=1000) + with pytest.raises(MaxQueuedTokensError): + llm.check_admission() + + +def test_admission_tokens_rejects_over_limit(): + llm = _make_async_llm(max_num_queued_tokens=1000, num_queued_tokens=1500) + with pytest.raises(MaxQueuedTokensError): + llm.check_admission() + + +def test_admission_tokens_rejects_when_zero_limit(): + llm = _make_async_llm(max_num_queued_tokens=0, num_queued_tokens=0) + with pytest.raises(MaxQueuedTokensError): + llm.check_admission() + + +# -- interaction between both limits ---------------------------------------- + + +def test_admission_both_limits_checked_independently(): + llm = _make_async_llm( + max_num_queued_reqs=100, + max_num_queued_tokens=1000, + num_unfinished=5, + num_queued_tokens=1000, + ) + with pytest.raises(MaxQueuedTokensError): + llm.check_admission() + + +def test_admission_req_limit_checked_before_token_limit(): + llm = _make_async_llm( + max_num_queued_reqs=10, + max_num_queued_tokens=1000, + num_unfinished=10, + num_queued_tokens=1000, + ) + with pytest.raises(QueueOverflowError): + llm.check_admission() + + +# -- n derived from params at the add_request call site --------------------- + + +@pytest.mark.parametrize("params", [SamplingParams(), PoolingParams()]) +def test_admission_params_without_explicit_n_count_as_one_slot(params): + """PoolingParams has no ``.n``; add_request must fall back to 1.""" + llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=9) + llm.check_admission(getattr(params, "n", 1) or 1) + + llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=10) + with pytest.raises(QueueOverflowError): + llm.check_admission(getattr(params, "n", 1) or 1) + + +@pytest.mark.asyncio +async def test_concurrent_single_request_admission_respects_limit(): + """Concurrent n=1 admissions cannot consume the same slot.""" + llm = _make_async_llm(max_num_queued_reqs=1) + request_states: dict[str, SimpleNamespace] = {} + llm.output_processor.get_num_unfinished_requests.side_effect = lambda: len( + request_states + ) + llm.output_processor.has_request.side_effect = request_states.__contains__ + + def add_request(request, *_args): + request_states[request.request_id] = _make_req_state( + len(request.prompt_token_ids) + ) + + llm.output_processor.add_request.side_effect = add_request + + async def add_request_async(_request): + await asyncio.sleep(0) + + llm.engine_core = SimpleNamespace( + add_request_async=AsyncMock(side_effect=add_request_async), + shutdown=MagicMock(), + ) + llm.log_requests = False + + requests = [ + SimpleNamespace(request_id=f"request-{idx}", prompt_token_ids=[idx]) + for idx in range(2) + ] + results = await asyncio.gather( + *( + llm._add_request(request, None, None, 0, MagicMock()) + for request in requests + ), + return_exceptions=True, + ) + + assert sum(result is None for result in results) == 1 + assert sum(isinstance(result, QueueOverflowError) for result in results) == 1 + assert len(request_states) == 1 + + +# --------------------------------------------------------------------------- +# create_error_response integration +# --------------------------------------------------------------------------- + + +def test_queue_overflow_maps_to_503(): + resp = create_error_response(QueueOverflowError()) + assert resp.error.code == HTTPStatus.SERVICE_UNAVAILABLE.value + assert resp.error.type == HTTPStatus.SERVICE_UNAVAILABLE.phrase + assert resp.error.param is None + msg = resp.error.message.lower() + assert "busy" in msg or "try again" in msg + + +def test_max_queued_tokens_maps_to_503(): + resp = create_error_response(MaxQueuedTokensError()) + assert resp.error.code == HTTPStatus.SERVICE_UNAVAILABLE.value + assert resp.error.type == HTTPStatus.SERVICE_UNAVAILABLE.phrase + msg = resp.error.message.lower() + assert "backlog" in msg or "try again" in msg + + +def test_custom_graceful_error_maps_to_its_status(): + err = GracefulHTTPError("custom", HTTPStatus.SERVICE_UNAVAILABLE) + resp = create_error_response(err) + assert resp.error.code == HTTPStatus.SERVICE_UNAVAILABLE.value + assert resp.error.type == HTTPStatus.SERVICE_UNAVAILABLE.phrase + + +# --------------------------------------------------------------------------- +# SchedulerConfig field defaults +# --------------------------------------------------------------------------- + + +def test_scheduler_config_defaults_are_none(): + config = _make_scheduler_config() + assert config.max_num_queued_reqs is None + assert config.max_num_queued_tokens is None + + +def test_scheduler_config_accepts_explicit_values(): + config = _make_scheduler_config( + max_num_queued_reqs=100, + max_num_queued_tokens=32000, + ) + assert config.max_num_queued_reqs == 100 + assert config.max_num_queued_tokens == 32000 + + +@pytest.mark.parametrize("field", ["max_num_queued_reqs", "max_num_queued_tokens"]) +def test_scheduler_config_rejects_negative(field): + with pytest.raises(ValidationError): + _make_scheduler_config(**{field: -1}) + + +# --------------------------------------------------------------------------- +# human_readable_int for CLI notation +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "input_str, expected", + [ + ("32k", 32_000), + ("1k", 1_000), + ("1K", 1_024), + ("1m", 1_000_000), + ("1M", 1_048_576), + ("100", 100), + ("2.5k", 2_500), + ("0", 0), + ], +) +def test_human_readable_int_parses_notation(input_str: str, expected: int): + assert human_readable_int(input_str) == expected + + +@pytest.mark.parametrize("invalid", ["abc", "1x", "", "k", "1.5K"]) +def test_human_readable_int_rejects_invalid(invalid: str): + with pytest.raises((argparse.ArgumentTypeError, ValueError)): + human_readable_int(invalid) diff --git a/vllm/config/scheduler.py b/vllm/config/scheduler.py index bcff098c34..7ac0f8f837 100644 --- a/vllm/config/scheduler.py +++ b/vllm/config/scheduler.py @@ -71,6 +71,37 @@ class SchedulerConfig: """For chunked prefill, a request is considered long if the prompt is longer than this number of tokens. 0 disables the cap (default).""" + max_num_queued_reqs: int | None = Field(default=None, ge=0) + """Maximum number of requests that can be in-flight (waiting or running) + at the same time, or None for no limit. When the limit is reached, new + requests are rejected with HTTP 503 so the client can retry on another + instance. This bounds vLLM's otherwise unbounded request queue and is + primarily a coarse capacity valve.""" + + max_num_queued_tokens: int | None = Field(default=None, ge=0) + """Maximum total prompt tokens of requests currently in the prefill + phase, or None for no limit. When the limit is reached, new requests + are rejected with HTTP 503. + + This is a TTFT QoS mechanism: by setting it to + ``target_TTFT * prefill_throughput`` you reject requests when the + prefill backlog would exceed the latency target. In a disaggregated + prefill-decode setup this maps directly to the prefill pool's + capacity. + + Note: the count is conservative. A partially prefilled request + still contributes its full ``prompt_len`` until it transitions out + of the prefill phase, because the scheduler's per-iteration + ``num_computed_tokens`` progress is not propagated to the API + server process during prefill (``EngineCoreOutput`` is only + emitted once the request starts producing tokens). Similarly, + prefix-cache hits (``num_cached_tokens``) are only known to the + OutputProcessor after prefill completes. This overestimates the + real backlog, causing earlier rejection than strictly necessary + — the safe direction for QoS. The impact is limited to long + prompts under chunked prefill; short prompts that prefill in a + single iteration are unaffected.""" + enable_chunked_prefill: bool = True """If True, prefill requests can be chunked based on the remaining `max_num_batched_tokens`. diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index d117fedf37..08ec2962f7 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -367,6 +367,7 @@ def _compute_kwargs(cls: ConfigType) -> dict[str, dict[str, Any]]: "max_num_scheduled_tokens", "kv_cache_memory_bytes", "safetensors_prefetch_block_size", + "max_num_queued_tokens", } if name == "max_model_len": kwargs[name]["type"] = human_readable_int_or_auto @@ -544,6 +545,8 @@ class EngineArgs: max_num_scheduled_tokens: int | None = None long_prefill_token_threshold: int = SchedulerConfig.long_prefill_token_threshold max_num_seqs: int | None = None + max_num_queued_reqs: int | None = None + max_num_queued_tokens: int | None = None max_logprobs: int = ModelConfig.max_logprobs logprobs_mode: LogprobsMode = ModelConfig.logprobs_mode use_fp64_gumbel: bool = ModelConfig.use_fp64_gumbel @@ -1554,6 +1557,13 @@ class EngineArgs: "default": None, }, ) + scheduler_group.add_argument( + "--max-num-queued-reqs", **scheduler_kwargs["max_num_queued_reqs"] + ) + scheduler_group.add_argument( + "--max-num-queued-tokens", + **scheduler_kwargs["max_num_queued_tokens"], + ) scheduler_group.add_argument( "--long-prefill-token-threshold", **scheduler_kwargs["long_prefill_token_threshold"], @@ -2348,6 +2358,8 @@ class EngineArgs: max_num_batched_tokens=self.max_num_batched_tokens, max_num_scheduled_tokens=self.max_num_scheduled_tokens, max_num_seqs=self.max_num_seqs, + max_num_queued_reqs=self.max_num_queued_reqs, + max_num_queued_tokens=self.max_num_queued_tokens, max_model_len=model_config.max_model_len, enable_chunked_prefill=self.enable_chunked_prefill, disable_chunked_mm_input=self.disable_chunked_mm_input, diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 0f61b73839..63f6970e05 100644 --- a/vllm/engine/protocol.py +++ b/vllm/engine/protocol.py @@ -62,6 +62,23 @@ class EngineClient(ABC): @abstractmethod def dead_error(self) -> BaseException: ... + def check_admission( # noqa: B027 + self, n: int = 1, request_id: str | None = None + ) -> None: + """Reject the request up front if it would exceed queue limits. + + Called before a response is started so that overload rejections can + carry an HTTP status, which is not possible once a streaming response + has begun. Engines without admission control accept everything. + + Args: + n: Number of sequences the request will occupy. + request_id: Request id, used for logging only. + + Raises: + GracefulHTTPError: If the request cannot be admitted. + """ + @abstractmethod def generate( self, diff --git a/vllm/entrypoints/generate/base/serving.py b/vllm/entrypoints/generate/base/serving.py index 0532828011..dc41d3f067 100644 --- a/vllm/entrypoints/generate/base/serving.py +++ b/vllm/entrypoints/generate/base/serving.py @@ -171,6 +171,19 @@ class GenerateBaseServing(BaseServing, BeamSearchOnlineMixin): # Never fail server startup over the fingerprint. self.system_fingerprint = None + def _preflight(self, n: int = 1) -> None: + """Engine checks that must run before a response is started. + + This is required for the streaming case, where we return a success + status before we actually start generating text :). + + Args: + n: Number of sequences the request will occupy. + """ + if self.engine_client.errored: + raise self.engine_client.dead_error + self.engine_client.check_admission(n) + def create_streaming_error_response( self, message: str | Exception, diff --git a/vllm/entrypoints/generate/generative_scoring/serving.py b/vllm/entrypoints/generate/generative_scoring/serving.py index 55b695184e..59db919225 100644 --- a/vllm/entrypoints/generate/generative_scoring/serving.py +++ b/vllm/entrypoints/generate/generative_scoring/serving.py @@ -195,6 +195,7 @@ class ServingGenerativeScoring(BaseServing): # Check if engine is alive if self.engine_client.errored: raise self.engine_client.dead_error + self.engine_client.check_admission(len(request.items)) # Get tokenizer tokenizer = self.renderer.tokenizer diff --git a/vllm/entrypoints/openai/chat_completion/batch_serving.py b/vllm/entrypoints/openai/chat_completion/batch_serving.py index ab26d2a549..06b3dbfd4a 100644 --- a/vllm/entrypoints/openai/chat_completion/batch_serving.py +++ b/vllm/entrypoints/openai/chat_completion/batch_serving.py @@ -56,8 +56,7 @@ class OpenAIServingChatBatch(OpenAIServingChat): logger.error("Error with model %s", error_check_ret) return error_check_ret - if self.engine_client.errored: - raise self.engine_client.dead_error + self._preflight() renderer = self.online_renderer diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index d8755ed0a1..06e7315c24 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -235,11 +235,7 @@ class OpenAIServingChat(GenerateBaseServing): logger.error("Error with model %s", error_check_ret) return error_check_ret - # If the engine is dead, raise the engine's DEAD_ERROR. - # This is required for the streaming case, where we return a - # success status before we actually start generating text :). - if self.engine_client.errored: - raise self.engine_client.dead_error + self._preflight(request.n or 1) return await self.online_renderer.render_chat(request) diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index d5fb3557fe..2fa01478ce 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -102,11 +102,7 @@ class OpenAIServingCompletion(GenerateBaseServing): if error_check_ret is not None: return error_check_ret - # If the engine is dead, raise the engine's DEAD_ERROR. - # This is required for the streaming case, where we return a - # success status before we actually start generating text :). - if self.engine_client.errored: - raise self.engine_client.dead_error + self._preflight(request.n or 1) return await self.online_renderer.render_completion(request) diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index 5ac5a69cdb..7053aba292 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -350,11 +350,7 @@ class OpenAIServingResponses(GenerateBaseServing): if maybe_validation_error is not None: return maybe_validation_error - # If the engine is dead, raise the engine's DEAD_ERROR. - # This is required for the streaming case, where we return a - # success status before we actually start generating text :). - if self.engine_client.errored: - raise self.engine_client.dead_error + self._preflight() if request.store and not self.enable_store: # Disable the store option. diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index d63f904b03..7fd8865ade 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -110,11 +110,7 @@ class ServingTokens(GenerateBaseServing): logger.error("Error with model %s", error_check_ret) return error_check_ret - # If the engine is dead, raise the engine's DEAD_ERROR. - # This is required for the streaming case, where we return a - # success status before we actually start generating text :). - if self.engine_client.errored: - raise self.engine_client.dead_error + self._preflight() lora_request = None lora_request = self._maybe_get_adapters(request, supports_default_mm_loras=True) diff --git a/vllm/entrypoints/serve/exception_handling/error_response.py b/vllm/entrypoints/serve/exception_handling/error_response.py index 6412a30840..505c5ac35c 100644 --- a/vllm/entrypoints/serve/exception_handling/error_response.py +++ b/vllm/entrypoints/serve/exception_handling/error_response.py @@ -26,6 +26,7 @@ def create_error_response( ) from vllm.exceptions import ( + GracefulHTTPError, VLLMClientError, VLLMNotFoundError, VLLMServerError, @@ -33,7 +34,11 @@ def create_error_response( VLLMValidationError, ) - if isinstance(exc, VLLMValidationError): + if isinstance(exc, GracefulHTTPError): + err_type = HTTPStatus(exc.http_status).phrase + status_code = exc.http_status + param = None + elif isinstance(exc, VLLMValidationError): err_type = "BadRequestError" status_code = HTTPStatus.BAD_REQUEST param = exc.parameter diff --git a/vllm/entrypoints/speech_to_text/base/serving.py b/vllm/entrypoints/speech_to_text/base/serving.py index e046230a56..dc09db6aa3 100644 --- a/vllm/entrypoints/speech_to_text/base/serving.py +++ b/vllm/entrypoints/speech_to_text/base/serving.py @@ -440,11 +440,7 @@ class SpeechToTextBaseServing(GenerateBaseServing): if not request.model: request.model = self.models.model_name() - # If the engine is dead, raise the engine's DEAD_ERROR. - # This is required for the streaming case, where we return a - # success status before we actually start generating text :). - if self.engine_client.errored: - raise self.engine_client.dead_error + self._preflight() if request.response_format not in [ "text", diff --git a/vllm/exceptions.py b/vllm/exceptions.py index 0d4f1e7104..582e83ec4a 100644 --- a/vllm/exceptions.py +++ b/vllm/exceptions.py @@ -125,3 +125,48 @@ class GenerationError(VLLMServerError): def __init__(self, message: str = "Internal server error"): super().__init__(message) self.status_code = HTTPStatus.INTERNAL_SERVER_ERROR + + +class GracefulHTTPError(VLLMError): + """Exception that should be translated into an HTTP error response. + + These are expected to occur during normal operation (e.g. admission + control rejections) and should be surfaced to the client with the + explicit HTTP status code they carry, rather than being mapped to a + generic 4xx/5xx by the client/server split. + """ + + def __init__(self, message: str, http_status: HTTPStatus): + super().__init__(message) + self.message = message + self.http_status = http_status + + +class QueueOverflowError(GracefulHTTPError): + """Raised when admitting a request would exceed the request queue limit. + + Returns HTTP 503 (Service Unavailable) so that load balancers and + client SDKs retry the request on a different instance. + """ + + def __init__(self): + super().__init__( + "The engine is currently busy and cannot accept new requests. " + "Please try again later or on a different instance.", + HTTPStatus.SERVICE_UNAVAILABLE, + ) + + +class MaxQueuedTokensError(GracefulHTTPError): + """Raised when the pending prefill tokens exceed the configured limit. + + Returns HTTP 503 (Service Unavailable) so that load balancers and + client SDKs retry the request on a different instance. + """ + + def __init__(self): + super().__init__( + "The engine has reached its prefill token backlog limit. " + "Please try again later or on a different instance.", + HTTPStatus.SERVICE_UNAVAILABLE, + ) diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index dd5cb2c262..0a7263ee2f 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -21,7 +21,13 @@ from vllm.distributed.weight_transfer.base import ( from vllm.engine.arg_utils import AsyncEngineArgs from vllm.engine.protocol import EngineClient, StreamingInput from vllm.entrypoints.serve.elastic_ep.middleware import set_scaling_elastic_ep -from vllm.exceptions import VLLMClientError, VLLMValidationError +from vllm.exceptions import ( + GracefulHTTPError, + MaxQueuedTokensError, + QueueOverflowError, + VLLMClientError, + VLLMValidationError, +) from vllm.inputs import EngineInput, PromptType from vllm.logger import init_logger from vllm.lora.request import LoRARequest @@ -112,6 +118,7 @@ class AsyncLLM(EngineClient): self.vllm_config = vllm_config self._elastic_ep_lock = asyncio.Lock() self.model_config = vllm_config.model_config + self.scheduler_config = vllm_config.scheduler_config self.observability_config = vllm_config.observability_config tracing_endpoint = self.observability_config.otlp_traces_endpoint @@ -273,6 +280,64 @@ class AsyncLLM(EngineClient): if handler is not None: cancel_task_threadsafe(handler) + def get_num_unfinished_requests(self) -> int: + return self.output_processor.get_num_unfinished_requests() + + def get_num_queued_tokens(self) -> int: + return self.output_processor.get_num_queued_tokens() + + def check_admission(self, n: int = 1, request_id: str | None = None) -> None: + """Reject the request if it would exceed queue limits. + + Both limits return HTTP 503 (Service Unavailable) so that load + balancers and client SDKs retry on a different instance. + + - ``max_num_queued_reqs``: hard cap on the number of unfinished requests + (waiting + running). A request with ``n > 1`` counts as ``n`` slots. + - ``max_num_queued_tokens``: TTFT QoS — cap on the total prompt + tokens of requests still in prefill. + + Note: ``get_num_queued_tokens`` uses ``prompt_len`` for all prefilling requests. + Chunked prefill progress and prefix-cache hits are not subtracted because the + scheduler's ``num_computed_tokens`` and ``num_cached_tokens`` are only + propagated to the API server after prefill completes. The overestimation is + conservative — earlier rejection, preserving TTFT targets. + + Args: + n: Number of sequences the request will occupy. + request_id: Request id, used for logging only. + + Raises: + QueueOverflowError: If ``max_num_queued_reqs`` would be exceeded. + MaxQueuedTokensError: If ``max_num_queued_tokens`` would be exceeded. + """ + max_num_reqs = self.scheduler_config.max_num_queued_reqs + if max_num_reqs is not None: + current = self.get_num_unfinished_requests() + if current + n > max_num_reqs: + logger.info( + "Request queue full - rejecting request %s " + "(current=%d, n=%d, max=%d).", + request_id, + current, + n, + max_num_reqs, + ) + raise QueueOverflowError() + + max_queued_tokens = self.scheduler_config.max_num_queued_tokens + if max_queued_tokens is not None: + current_tokens = self.get_num_queued_tokens() + if current_tokens >= max_queued_tokens: + logger.info( + "Max queued tokens reached - rejecting request %s " + "(current_tokens=%d, max=%d).", + request_id, + current_tokens, + max_queued_tokens, + ) + raise MaxQueuedTokensError() + async def get_supported_tasks(self) -> tuple[SupportedTask, ...]: if not hasattr(self, "_supported_tasks"): # Cache the result @@ -317,6 +382,10 @@ class AsyncLLM(EngineClient): "prompt logprobs" ) + if isinstance(params, SamplingParams) and params.n > 1: + # TODO (NickLucche) Batch check admission check for all n requests + self.check_admission(params.n, request_id) + if isinstance(prompt, AsyncGenerator): if reasoning_ended is not None or reasoning_parser_kwargs is not None: raise NotImplementedError @@ -430,7 +499,12 @@ class AsyncLLM(EngineClient): index: int, queue: RequestOutputCollector, ): - # Add the request to OutputProcessor (this process). + if parent_req is None and not self.output_processor.has_request( + request.request_id + ): + self.check_admission(request_id=request.request_id) + + # Register locally before the first await so concurrent tasks see this request. self.output_processor.add_request(request, prompt, parent_req, index, queue) # Add the EngineCoreRequest to EngineCore (separate process). @@ -630,8 +704,8 @@ class AsyncLLM(EngineClient): logger.info("Request %s failed (engine dead).", request_id) raise - # Request validation error. - except VLLMClientError as e: + # Request validation error or admission control rejection. + except (VLLMClientError, GracefulHTTPError) as e: if self.log_requests: logger.info("Request %s failed (bad request): %s.", request_id, e) raise @@ -907,8 +981,8 @@ class AsyncLLM(EngineClient): logger.info("Request %s failed (engine dead).", request_id) raise - # Request validation error. - except VLLMClientError: + # Request validation error or admission control rejection. + except (VLLMClientError, GracefulHTTPError): if self.log_requests: logger.info("Request %s failed (bad request).", request_id) raise diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 6238fbc817..7bbb1285be 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -467,6 +467,20 @@ class OutputProcessor: def get_num_unfinished_requests(self): return len(self.request_states) + def has_request(self, request_id: str) -> bool: + return request_id in self.request_states + + def get_num_queued_tokens(self) -> int: + """Total prompt tokens of requests currently in the prefill phase. + + Uses ``prompt_len`` rather than remaining prefill work because the + scheduler's ``num_computed_tokens`` is not propagated to the API + server until prefill completes. See ``SchedulerConfig`` docs. + """ + return sum( + req.prompt_len for req in self.request_states.values() if req.is_prefilling + ) + def has_unfinished_requests(self) -> bool: return len(self.request_states) > 0