| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Unit tests for admission control (max_num_queued_reqs / max_num_queued_tokens). |
| 5 | |
| 6 | These tests cover: |
| 7 | - OutputProcessor.get_num_queued_tokens() token counting |
| 8 | - AsyncLLM.check_admission() admission control logic |
| 9 | - Exception classes (GracefulHTTPError, QueueOverflowError, MaxQueuedTokensError) |
| 10 | - create_error_response() mapping GracefulHTTPError to HTTP 503 |
| 11 | - SchedulerConfig field defaults and validation |
| 12 | - human_readable_int CLI notation for max_num_queued_tokens |
| 13 | """ |
| 14 | |
| 15 | import argparse |
| 16 | import asyncio |
| 17 | import multiprocessing |
| 18 | from collections.abc import Awaitable, Callable |
| 19 | from http import HTTPStatus |
| 20 | from types import SimpleNamespace |
| 21 | from unittest.mock import AsyncMock, MagicMock |
| 22 | |
| 23 | import pytest |
| 24 | from pydantic import ValidationError |
| 25 | |
| 26 | from vllm.config.scheduler import SchedulerConfig |
| 27 | from vllm.entrypoints.serve.exception_handling.error_response import ( |
| 28 | create_error_response, |
| 29 | ) |
| 30 | from vllm.exceptions import ( |
| 31 | GracefulHTTPError, |
| 32 | MaxQueuedTokensError, |
| 33 | QueueOverflowError, |
| 34 | VLLMError, |
| 35 | ) |
| 36 | from vllm.pooling_params import PoolingParams |
| 37 | from vllm.sampling_params import SamplingParams |
| 38 | from vllm.utils.argparse_utils import human_readable_int |
| 39 | from vllm.v1.engine import EngineCoreRequest |
| 40 | from vllm.v1.engine.admission_control import SharedAdmissionStats |
| 41 | from vllm.v1.engine.async_llm import AsyncLLM |
| 42 | from vllm.v1.engine.output_processor import OutputProcessor |
| 43 | |
| 44 | pytestmark = pytest.mark.cpu_test |
| 45 | |
| 46 | # --------------------------------------------------------------------------- |
| 47 | # Helpers |
| 48 | # --------------------------------------------------------------------------- |
| 49 | |
| 50 | |
| 51 | def _make_req_state(prompt_len: int, is_prefilling: bool = True): |
| 52 | return SimpleNamespace( |
| 53 | prompt_len=prompt_len, |
| 54 | is_prefilling=is_prefilling, |
| 55 | ) |
| 56 | |
| 57 | |
| 58 | def _make_async_llm( |
| 59 | max_num_queued_reqs: int | None = None, |
| 60 | max_num_queued_tokens: int | None = None, |
| 61 | num_unfinished: int = 0, |
| 62 | num_queued_tokens: int = 0, |
| 63 | ) -> AsyncLLM: |
| 64 | """Create a bare AsyncLLM with just the attributes needed for scheduling.""" |
| 65 | llm = AsyncLLM.__new__(AsyncLLM) |
| 66 | llm.scheduler_config = SimpleNamespace( |
| 67 | max_num_queued_reqs=max_num_queued_reqs, |
| 68 | max_num_queued_tokens=max_num_queued_tokens, |
| 69 | ) |
| 70 | llm.output_processor = MagicMock() |
| 71 | llm.output_processor.get_num_unfinished_requests.return_value = num_unfinished |
| 72 | llm.output_processor.get_num_queued_tokens.return_value = num_queued_tokens |
| 73 | llm.admission_stats = None |
| 74 | return llm |
| 75 | |
| 76 | |
| 77 | def _make_output_processor(**request_states) -> OutputProcessor: |
| 78 | op = OutputProcessor.__new__(OutputProcessor) |
| 79 | op.request_states = request_states |
| 80 | return op |
| 81 | |
| 82 | |
| 83 | def _make_engine_request(request_id: str, n: int) -> EngineCoreRequest: |
| 84 | return EngineCoreRequest( |
| 85 | request_id=request_id, |
| 86 | prompt_token_ids=[1], |
| 87 | mm_features=None, |
| 88 | sampling_params=SamplingParams(n=n), |
| 89 | pooling_params=None, |
| 90 | arrival_time=0, |
| 91 | lora_request=None, |
| 92 | cache_salt=None, |
| 93 | data_parallel_rank=None, |
| 94 | ) |
| 95 | |
| 96 | |
| 97 | def _make_request_test_llm( |
| 98 | max_num_queued_reqs: int, |
| 99 | add_request_async: Callable[[EngineCoreRequest], Awaitable[None]], |
| 100 | ) -> AsyncLLM: |
| 101 | llm = _make_async_llm(max_num_queued_reqs=max_num_queued_reqs) |
| 102 | llm.output_processor = OutputProcessor(None, log_stats=False) |
| 103 | llm.engine_core = SimpleNamespace( |
| 104 | resources=SimpleNamespace(engine_dead=False), |
| 105 | add_request_async=AsyncMock(side_effect=add_request_async), |
| 106 | abort_requests_async=AsyncMock(), |
| 107 | shutdown=MagicMock(), |
| 108 | ) |
| 109 | llm.vllm_config = SimpleNamespace( |
| 110 | cache_config=SimpleNamespace(kv_sharing_fast_prefill=False) |
| 111 | ) |
| 112 | llm.output_handler = None |
| 113 | llm.log_requests = False |
| 114 | llm._run_output_handler = MagicMock() |
| 115 | llm.input_processor = MagicMock() |
| 116 | llm.input_processor.assign_request_id.side_effect = lambda request: setattr( |
| 117 | request, "external_req_id", request.request_id |
| 118 | ) |
| 119 | return llm |
| 120 | |
| 121 | |
| 122 | def _make_scheduler_config(**kwargs) -> SchedulerConfig: |
| 123 | return SchedulerConfig( |
| 124 | runner_type="generate", |
| 125 | max_model_len=4096, |
| 126 | is_encoder_decoder=False, |
| 127 | **kwargs, |
| 128 | ) |
| 129 | |
| 130 | |
| 131 | def _make_shared_stats(client_count: int = 2) -> list[SharedAdmissionStats]: |
| 132 | counters = multiprocessing.RawArray( |
| 133 | "q", SharedAdmissionStats.num_counters(client_count) |
| 134 | ) |
| 135 | return [ |
| 136 | SharedAdmissionStats({"mp_admission_counters": counters}, client_count, index) |
| 137 | for index in range(client_count) |
| 138 | ] |
| 139 | |
| 140 | |
| 141 | @pytest.mark.skip_global_cleanup |
| 142 | def test_shared_admission_stats_aggregate_api_servers(): |
| 143 | server_0, server_1 = _make_shared_stats() |
| 144 | |
| 145 | server_0.set_num_requests(2) |
| 146 | assert server_0.get_num_requests() == 2 |
| 147 | server_1.set_num_requests(3) |
| 148 | assert server_0.get_num_requests() == server_1.get_num_requests() == 5 |
| 149 | |
| 150 | server_0.set_num_requests(1) |
| 151 | assert server_0.get_num_requests() == 4 |
| 152 | |
| 153 | |
| 154 | @pytest.mark.skip_global_cleanup |
| 155 | def test_shared_request_snapshot_enforces_global_limit(): |
| 156 | server_0, server_1 = _make_shared_stats() |
| 157 | server_0.set_num_requests(2) |
| 158 | llm = _make_async_llm(max_num_queued_reqs=2) |
| 159 | llm.admission_stats = server_1 |
| 160 | |
| 161 | with pytest.raises(QueueOverflowError): |
| 162 | llm.check_admission() |
| 163 | |
| 164 | assert server_0.get_num_requests() == 2 |
| 165 | |
| 166 | |
| 167 | # --------------------------------------------------------------------------- |
| 168 | # Exception classes |
| 169 | # --------------------------------------------------------------------------- |
| 170 | |
| 171 | |
| 172 | def test_graceful_http_error_carries_status_and_message(): |
| 173 | err = GracefulHTTPError("custom message", HTTPStatus.SERVICE_UNAVAILABLE) |
| 174 | assert err.message == "custom message" |
| 175 | assert err.http_status == HTTPStatus.SERVICE_UNAVAILABLE |
| 176 | assert str(err) == "custom message" |
| 177 | |
| 178 | |
| 179 | def test_graceful_http_error_is_vllm_error(): |
| 180 | err = GracefulHTTPError("msg", HTTPStatus.TOO_MANY_REQUESTS) |
| 181 | assert isinstance(err, VLLMError) |
| 182 | |
| 183 | |
| 184 | def test_queue_overflow_error(): |
| 185 | err = QueueOverflowError() |
| 186 | assert err.http_status == HTTPStatus.SERVICE_UNAVAILABLE |
| 187 | assert isinstance(err, GracefulHTTPError) |
| 188 | assert "busy" in err.message.lower() or "try again" in err.message.lower() |
| 189 | |
| 190 | |
| 191 | def test_max_queued_tokens_error(): |
| 192 | err = MaxQueuedTokensError() |
| 193 | assert err.http_status == HTTPStatus.SERVICE_UNAVAILABLE |
| 194 | assert isinstance(err, GracefulHTTPError) |
| 195 | assert "backlog" in err.message.lower() or "try again" in err.message.lower() |
| 196 | |
| 197 | |
| 198 | @pytest.mark.parametrize("exc_cls", [QueueOverflowError, MaxQueuedTokensError]) |
| 199 | def test_admission_exceptions_are_vllm_errors(exc_cls): |
| 200 | """Admission rejections must reach the VLLMError HTTP handler.""" |
| 201 | assert issubclass(exc_cls, VLLMError) |
| 202 | |
| 203 | |
| 204 | # --------------------------------------------------------------------------- |
| 205 | # OutputProcessor.get_num_queued_tokens |
| 206 | # --------------------------------------------------------------------------- |
| 207 | |
| 208 | |
| 209 | def test_queued_tokens_empty(): |
| 210 | assert _make_output_processor().get_num_queued_tokens() == 0 |
| 211 | |
| 212 | |
| 213 | def test_queued_tokens_sums_prefilling_requests(): |
| 214 | op = _make_output_processor(r1=_make_req_state(100), r2=_make_req_state(200)) |
| 215 | assert op.get_num_queued_tokens() == 300 |
| 216 | |
| 217 | |
| 218 | def test_queued_tokens_excludes_non_prefilling(): |
| 219 | op = _make_output_processor( |
| 220 | r1=_make_req_state(100, is_prefilling=True), |
| 221 | r2=_make_req_state(200, is_prefilling=False), |
| 222 | r3=_make_req_state(50, is_prefilling=True), |
| 223 | ) |
| 224 | assert op.get_num_queued_tokens() == 150 |
| 225 | |
| 226 | |
| 227 | def test_queued_tokens_all_non_prefilling(): |
| 228 | op = _make_output_processor( |
| 229 | r1=_make_req_state(100, is_prefilling=False), |
| 230 | r2=_make_req_state(200, is_prefilling=False), |
| 231 | ) |
| 232 | assert op.get_num_queued_tokens() == 0 |
| 233 | |
| 234 | |
| 235 | # --------------------------------------------------------------------------- |
| 236 | # AsyncLLM.check_admission |
| 237 | # --------------------------------------------------------------------------- |
| 238 | |
| 239 | |
| 240 | def test_admission_no_limits_allows_everything(): |
| 241 | llm = _make_async_llm(num_unfinished=999, num_queued_tokens=999) |
| 242 | llm.check_admission() |
| 243 | |
| 244 | |
| 245 | # -- max_num_queued_reqs ---------------------------------------------------- |
| 246 | |
| 247 | |
| 248 | def test_admission_reqs_allows_when_under_limit(): |
| 249 | llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=5) |
| 250 | llm.check_admission() |
| 251 | |
| 252 | |
| 253 | def test_admission_reqs_rejects_at_limit(): |
| 254 | llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=10) |
| 255 | with pytest.raises(QueueOverflowError): |
| 256 | llm.check_admission() |
| 257 | |
| 258 | |
| 259 | def test_admission_reqs_rejects_with_n(): |
| 260 | llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=8) |
| 261 | with pytest.raises(QueueOverflowError): |
| 262 | llm.check_admission(3) |
| 263 | |
| 264 | |
| 265 | def test_admission_reqs_allows_n_at_boundary(): |
| 266 | llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=7) |
| 267 | llm.check_admission(3) |
| 268 | |
| 269 | |
| 270 | def test_admission_reqs_rejects_when_zero_limit(): |
| 271 | llm = _make_async_llm(max_num_queued_reqs=0, num_unfinished=0) |
| 272 | with pytest.raises(QueueOverflowError): |
| 273 | llm.check_admission() |
| 274 | |
| 275 | |
| 276 | # -- max_num_queued_tokens -------------------------------------------------- |
| 277 | |
| 278 | |
| 279 | def test_admission_tokens_allows_when_under_limit(): |
| 280 | llm = _make_async_llm(max_num_queued_tokens=1000, num_queued_tokens=500) |
| 281 | llm.check_admission() |
| 282 | |
| 283 | |
| 284 | def test_admission_tokens_rejects_at_limit(): |
| 285 | llm = _make_async_llm(max_num_queued_tokens=1000, num_queued_tokens=1000) |
| 286 | with pytest.raises(MaxQueuedTokensError): |
| 287 | llm.check_admission() |
| 288 | |
| 289 | |
| 290 | def test_admission_tokens_rejects_over_limit(): |
| 291 | llm = _make_async_llm(max_num_queued_tokens=1000, num_queued_tokens=1500) |
| 292 | with pytest.raises(MaxQueuedTokensError): |
| 293 | llm.check_admission() |
| 294 | |
| 295 | |
| 296 | def test_admission_tokens_rejects_when_zero_limit(): |
| 297 | llm = _make_async_llm(max_num_queued_tokens=0, num_queued_tokens=0) |
| 298 | with pytest.raises(MaxQueuedTokensError): |
| 299 | llm.check_admission() |
| 300 | |
| 301 | |
| 302 | # -- interaction between both limits ---------------------------------------- |
| 303 | |
| 304 | |
| 305 | def test_admission_both_limits_checked_independently(): |
| 306 | llm = _make_async_llm( |
| 307 | max_num_queued_reqs=100, |
| 308 | max_num_queued_tokens=1000, |
| 309 | num_unfinished=5, |
| 310 | num_queued_tokens=1000, |
| 311 | ) |
| 312 | with pytest.raises(MaxQueuedTokensError): |
| 313 | llm.check_admission() |
| 314 | |
| 315 | |
| 316 | def test_admission_req_limit_checked_before_token_limit(): |
| 317 | llm = _make_async_llm( |
| 318 | max_num_queued_reqs=10, |
| 319 | max_num_queued_tokens=1000, |
| 320 | num_unfinished=10, |
| 321 | num_queued_tokens=1000, |
| 322 | ) |
| 323 | with pytest.raises(QueueOverflowError): |
| 324 | llm.check_admission() |
| 325 | |
| 326 | |
| 327 | # -- n derived from params at the add_request call site --------------------- |
| 328 | |
| 329 | |
| 330 | @pytest.mark.parametrize("params", [SamplingParams(), PoolingParams()]) |
| 331 | def test_admission_params_without_explicit_n_count_as_one_slot(params): |
| 332 | """PoolingParams has no ``.n``; add_request must fall back to 1.""" |
| 333 | llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=9) |
| 334 | llm.check_admission(getattr(params, "n", 1) or 1) |
| 335 | |
| 336 | llm = _make_async_llm(max_num_queued_reqs=10, num_unfinished=10) |
| 337 | with pytest.raises(QueueOverflowError): |
| 338 | llm.check_admission(getattr(params, "n", 1) or 1) |
| 339 | |
| 340 | |
| 341 | @pytest.mark.asyncio |
| 342 | async def test_concurrent_single_request_admission_respects_limit(): |
| 343 | """Concurrent n=1 admissions cannot consume the same slot.""" |
| 344 | llm = _make_async_llm(max_num_queued_reqs=1) |
| 345 | request_states: dict[str, SimpleNamespace] = {} |
| 346 | llm.output_processor.get_num_unfinished_requests.side_effect = lambda: len( |
| 347 | request_states |
| 348 | ) |
| 349 | llm.output_processor.has_request.side_effect = request_states.__contains__ |
| 350 | |
| 351 | def add_request(request, *_args): |
| 352 | request_states[request.request_id] = _make_req_state( |
| 353 | len(request.prompt_token_ids) |
| 354 | ) |
| 355 | |
| 356 | llm.output_processor.add_request.side_effect = add_request |
| 357 | |
| 358 | async def add_request_async(_request): |
| 359 | await asyncio.sleep(0) |
| 360 | |
| 361 | llm.engine_core = SimpleNamespace( |
| 362 | add_request_async=AsyncMock(side_effect=add_request_async), |
| 363 | shutdown=MagicMock(), |
| 364 | ) |
| 365 | llm.log_requests = False |
| 366 | |
| 367 | requests = [ |
| 368 | SimpleNamespace(request_id=f"request-{idx}", prompt_token_ids=[idx]) |
| 369 | for idx in range(2) |
| 370 | ] |
| 371 | results = await asyncio.gather( |
| 372 | *( |
| 373 | llm._add_request(request, None, None, 0, MagicMock()) |
| 374 | for request in requests |
| 375 | ), |
| 376 | return_exceptions=True, |
| 377 | ) |
| 378 | |
| 379 | assert sum(result is None for result in results) == 1 |
| 380 | assert sum(isinstance(result, QueueOverflowError) for result in results) == 1 |
| 381 | assert len(request_states) == 1 |
| 382 | |
| 383 | |
| 384 | @pytest.mark.asyncio |
| 385 | @pytest.mark.parametrize("first_n", [1, 3]) |
| 386 | async def test_parallel_admission_is_all_or_nothing(first_n: int): |
| 387 | """An n>1 request reserves either all its capacity or none of it.""" |
| 388 | |
| 389 | async def add_request_async(_request): |
| 390 | await asyncio.sleep(0) |
| 391 | |
| 392 | llm = _make_request_test_llm(3, add_request_async) |
| 393 | second_n = 3 if first_n == 1 else 1 |
| 394 | first = _make_engine_request("first", first_n) |
| 395 | second = _make_engine_request("second", second_n) |
| 396 | results = await asyncio.gather( |
| 397 | llm.add_request("first", first, first.params), |
| 398 | llm.add_request("second", second, second.params), |
| 399 | return_exceptions=True, |
| 400 | ) |
| 401 | |
| 402 | assert not isinstance(results[0], Exception) |
| 403 | assert isinstance(results[1], QueueOverflowError) |
| 404 | assert llm.output_processor.get_num_unfinished_requests() == first_n |
| 405 | assert llm.engine_core.add_request_async.await_count == first_n |
| 406 | |
| 407 | |
| 408 | @pytest.mark.asyncio |
| 409 | async def test_parallel_admission_cancellation_cleans_up_all_children(): |
| 410 | """Cancellation during core submission must release every reserved slot.""" |
| 411 | first_submission_started = asyncio.Event() |
| 412 | |
| 413 | async def add_request_async(_request): |
| 414 | first_submission_started.set() |
| 415 | await asyncio.Event().wait() |
| 416 | |
| 417 | llm = _make_request_test_llm(3, add_request_async) |
| 418 | request = _make_engine_request("parallel", 3) |
| 419 | output = llm.generate(request, request.params, request.request_id) |
| 420 | generate_task = asyncio.create_task(anext(output)) |
| 421 | |
| 422 | await first_submission_started.wait() |
| 423 | generate_task.cancel() |
| 424 | with pytest.raises(asyncio.CancelledError): |
| 425 | await generate_task |
| 426 | |
| 427 | assert llm.engine_core.add_request_async.await_count == 1 |
| 428 | assert llm.output_processor.get_num_unfinished_requests() == 0 |
| 429 | aborted_ids = llm.engine_core.abort_requests_async.await_args.args[0] |
| 430 | assert set(aborted_ids) == {"0_parallel", "1_parallel", "2_parallel"} |
| 431 | |
| 432 | |
| 433 | # --------------------------------------------------------------------------- |
| 434 | # create_error_response integration |
| 435 | # --------------------------------------------------------------------------- |
| 436 | |
| 437 | |
| 438 | def test_queue_overflow_maps_to_503(): |
| 439 | resp = create_error_response(QueueOverflowError()) |
| 440 | assert resp.error.code == HTTPStatus.SERVICE_UNAVAILABLE.value |
| 441 | assert resp.error.type == HTTPStatus.SERVICE_UNAVAILABLE.phrase |
| 442 | assert resp.error.param is None |
| 443 | msg = resp.error.message.lower() |
| 444 | assert "busy" in msg or "try again" in msg |
| 445 | |
| 446 | |
| 447 | def test_max_queued_tokens_maps_to_503(): |
| 448 | resp = create_error_response(MaxQueuedTokensError()) |
| 449 | assert resp.error.code == HTTPStatus.SERVICE_UNAVAILABLE.value |
| 450 | assert resp.error.type == HTTPStatus.SERVICE_UNAVAILABLE.phrase |
| 451 | msg = resp.error.message.lower() |
| 452 | assert "backlog" in msg or "try again" in msg |
| 453 | |
| 454 | |
| 455 | def test_custom_graceful_error_maps_to_its_status(): |
| 456 | err = GracefulHTTPError("custom", HTTPStatus.SERVICE_UNAVAILABLE) |
| 457 | resp = create_error_response(err) |
| 458 | assert resp.error.code == HTTPStatus.SERVICE_UNAVAILABLE.value |
| 459 | assert resp.error.type == HTTPStatus.SERVICE_UNAVAILABLE.phrase |
| 460 | |
| 461 | |
| 462 | # --------------------------------------------------------------------------- |
| 463 | # SchedulerConfig field defaults |
| 464 | # --------------------------------------------------------------------------- |
| 465 | |
| 466 | |
| 467 | def test_scheduler_config_defaults_are_none(): |
| 468 | config = _make_scheduler_config() |
| 469 | assert config.max_num_queued_reqs is None |
| 470 | assert config.max_num_queued_tokens is None |
| 471 | |
| 472 | |
| 473 | def test_scheduler_config_accepts_explicit_values(): |
| 474 | config = _make_scheduler_config( |
| 475 | max_num_queued_reqs=100, |
| 476 | max_num_queued_tokens=32000, |
| 477 | ) |
| 478 | assert config.max_num_queued_reqs == 100 |
| 479 | assert config.max_num_queued_tokens == 32000 |
| 480 | |
| 481 | |
| 482 | @pytest.mark.parametrize("field", ["max_num_queued_reqs", "max_num_queued_tokens"]) |
| 483 | def test_scheduler_config_rejects_negative(field): |
| 484 | with pytest.raises(ValidationError): |
| 485 | _make_scheduler_config(**{field: -1}) |
| 486 | |
| 487 | |
| 488 | # --------------------------------------------------------------------------- |
| 489 | # human_readable_int for CLI notation |
| 490 | # --------------------------------------------------------------------------- |
| 491 | |
| 492 | |
| 493 | @pytest.mark.parametrize( |
| 494 | "input_str, expected", |
| 495 | [ |
| 496 | ("32k", 32_000), |
| 497 | ("1k", 1_000), |
| 498 | ("1K", 1_024), |
| 499 | ("1m", 1_000_000), |
| 500 | ("1M", 1_048_576), |
| 501 | ("100", 100), |
| 502 | ("2.5k", 2_500), |
| 503 | ("0", 0), |
| 504 | ], |
| 505 | ) |
| 506 | def test_human_readable_int_parses_notation(input_str: str, expected: int): |
| 507 | assert human_readable_int(input_str) == expected |
| 508 | |
| 509 | |
| 510 | @pytest.mark.parametrize("invalid", ["abc", "1x", "", "k", "1.5K"]) |
| 511 | def test_human_readable_int_rejects_invalid(invalid: str): |
| 512 | with pytest.raises((argparse.ArgumentTypeError, ValueError)): |
| 513 | human_readable_int(invalid) |