From f0fbc98c488b40a07d3b269aeba2bda3c4c3b03b Mon Sep 17 00:00:00 2001 From: Chauncey Date: Tue, 15 Sep 2026 22:50:26 +0800 Subject: [PATCH] [Frontend] Share max_num_queued_reqs across API server processes (#54746) Signed-off-by: chaunceyjiang --- .../test_api_server_process_manager.py | 39 ++++++++++++++ tests/v1/engine/test_admission_control.py | 39 ++++++++++++++ vllm/v1/engine/admission_control.py | 52 +++++++++++++++++++ vllm/v1/engine/async_llm.py | 24 +++++++-- vllm/v1/engine/output_processor.py | 14 ++++- vllm/v1/utils.py | 11 ++++ 6 files changed, 173 insertions(+), 6 deletions(-) create mode 100644 vllm/v1/engine/admission_control.py diff --git a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py index 40fefc2374..668f115448 100644 --- a/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py +++ b/tests/entrypoints/launchers/api_server/test_api_server_process_manager.py @@ -5,6 +5,7 @@ import multiprocessing import socket import threading import time +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -14,6 +15,7 @@ from tests.entrypoints.launchers.api_server._api_server_spawn_workers import ( exit_before_report_worker, ) from vllm.utils.network_utils import make_zmq_socket, split_zmq_path +from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.utils import ( APIServerProcessManager, get_engine_client_zmq_addr, @@ -32,6 +34,18 @@ def mock_run_api_server_worker(listen_address, sock, args, client_config=None): print("Mock worker completed successfully") +def update_admission_stats_worker(listen_address, sock, args, client_config): + stats = SharedAdmissionStats( + client_config, + client_config["client_count"], + client_config["client_index"], + ) + value = client_config["client_index"] + 1 + stats.set_num_requests(value) + args.barrier.wait() + args.results.put(stats.get_num_requests()) + + # Module-level stub for the gather_actual_addresses test. Must be # importable by `multiprocessing.spawn` (no closures, no nesting). def defer_addresses_stub_worker(listen_address, sock, args, client_config): @@ -131,6 +145,31 @@ def test_api_server_process_manager_init(api_server_args, with_stats_update): assert not proc.is_alive() +@pytest.mark.skip_global_cleanup +def test_api_server_processes_share_admission_stats(api_server_args): + args = api_server_args.copy() + spawn_context = multiprocessing.get_context("spawn") + worker_args = SimpleNamespace( + barrier=spawn_context.Barrier(args["num_servers"]), + results=spawn_context.Queue(), + max_num_queued_reqs=10, + ) + args["args"] = worker_args + args["target_server_fn"] = update_admission_stats_worker + manager = APIServerProcessManager(**args) + + try: + for process in manager.processes: + process.join(timeout=30) + assert process.exitcode == 0 + + assert [worker_args.results.get(timeout=5) for _ in range(3)] == [6] * 3 + finally: + manager.shutdown() + worker_args.results.close() + worker_args.results.join_thread() + + @patch("vllm.v1.utils.run_api_server_worker_proc", mock_run_api_server_worker) def test_wait_for_completion_or_failure(api_server_args): """Test that wait_for_completion_or_failure works with failures.""" diff --git a/tests/v1/engine/test_admission_control.py b/tests/v1/engine/test_admission_control.py index a03718ab5a..219850f5d4 100644 --- a/tests/v1/engine/test_admission_control.py +++ b/tests/v1/engine/test_admission_control.py @@ -14,6 +14,7 @@ These tests cover: import argparse import asyncio +import multiprocessing from http import HTTPStatus from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -34,6 +35,7 @@ from vllm.exceptions import ( 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.admission_control import SharedAdmissionStats from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.output_processor import OutputProcessor @@ -66,6 +68,7 @@ def _make_async_llm( 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 + llm.admission_stats = None return llm @@ -84,6 +87,42 @@ def _make_scheduler_config(**kwargs) -> SchedulerConfig: ) +def _make_shared_stats(client_count: int = 2) -> list[SharedAdmissionStats]: + counters = multiprocessing.RawArray( + "q", SharedAdmissionStats.num_counters(client_count) + ) + return [ + SharedAdmissionStats({"mp_admission_counters": counters}, client_count, index) + for index in range(client_count) + ] + + +@pytest.mark.skip_global_cleanup +def test_shared_admission_stats_aggregate_api_servers(): + server_0, server_1 = _make_shared_stats() + + server_0.set_num_requests(2) + assert server_0.get_num_requests() == 2 + server_1.set_num_requests(3) + assert server_0.get_num_requests() == server_1.get_num_requests() == 5 + + server_0.set_num_requests(1) + assert server_0.get_num_requests() == 4 + + +@pytest.mark.skip_global_cleanup +def test_shared_request_snapshot_enforces_global_limit(): + server_0, server_1 = _make_shared_stats() + server_0.set_num_requests(2) + llm = _make_async_llm(max_num_queued_reqs=2) + llm.admission_stats = server_1 + + with pytest.raises(QueueOverflowError): + llm.check_admission() + + assert server_0.get_num_requests() == 2 + + # --------------------------------------------------------------------------- # Exception classes # --------------------------------------------------------------------------- diff --git a/vllm/v1/engine/admission_control.py b/vllm/v1/engine/admission_control.py new file mode 100644 index 0000000000..b72d50f346 --- /dev/null +++ b/vllm/v1/engine/admission_control.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any, Protocol, cast + + +class AdmissionCounterBuffer(Protocol): + def __getitem__(self, index: int) -> int: ... + + def __setitem__(self, index: int, value: int) -> None: ... + + +class SharedAdmissionStats: + """Lock-free admission stats shared by API server processes. + + Each process writes exclusively to its own cache-line-sized slot. Readers + aggregate the slots, avoiding both inter-process locks and cache-line + contention on the request path. + """ + + # Space per-process counters 64 bytes apart to avoid false sharing on CPUs + # with 64-byte cache lines: one process's writes would otherwise invalidate + # the cache line containing another process's counter. + CACHE_LINE_BYTES = 64 + COUNTER_BYTES = 8 + COUNTERS_PER_SLOT = CACHE_LINE_BYTES // COUNTER_BYTES + + def __init__( + self, + client_addresses: dict[str, Any], + client_count: int = 1, + client_index: int = 0, + ): + self._counters = cast( + AdmissionCounterBuffer, + client_addresses.pop("mp_admission_counters"), + ) + self._client_count = client_count + self._slot_offset = client_index * self.COUNTERS_PER_SLOT + + @classmethod + def num_counters(cls, client_count: int) -> int: + return client_count * cls.COUNTERS_PER_SLOT + + def set_num_requests(self, count: int) -> None: + self._counters[self._slot_offset] = count + + def get_num_requests(self) -> int: + return sum( + self._counters[client_index * self.COUNTERS_PER_SLOT] + for client_index in range(self._client_count) + ) diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 91618c4093..bf3f5f420a 100644 --- a/vllm/v1/engine/async_llm.py +++ b/vllm/v1/engine/async_llm.py @@ -44,6 +44,7 @@ from vllm.usage.usage_lib import UsageContext from vllm.utils.async_utils import cancel_task_threadsafe from vllm.utils.collection_utils import as_list from vllm.v1.engine import EngineCoreRequest, PauseMode +from vllm.v1.engine.admission_control import SharedAdmissionStats from vllm.v1.engine.core_client import EngineCoreClient from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm.v1.engine.input_processor import InputProcessor @@ -144,12 +145,20 @@ class AsyncLLM(EngineClient): # Convert EngineInput --> EngineCoreRequest. self.input_processor = InputProcessor(self.vllm_config, renderer) + self.admission_stats = ( + SharedAdmissionStats(client_addresses, client_count, client_index) + if client_addresses is not None + and "mp_admission_counters" in client_addresses + else None + ) + # Converts EngineCoreOutputs --> RequestOutput. self.output_processor = OutputProcessor( renderer.tokenizer, log_stats=self.log_stats, stream_interval=self.vllm_config.scheduler_config.stream_interval, tracing_enabled=tracing_endpoint is not None, + admission_stats=self.admission_stats, ) # EngineCore (starts the engine in background process). @@ -291,8 +300,9 @@ class AsyncLLM(EngineClient): 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_reqs``: approximate global 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. @@ -312,13 +322,17 @@ class AsyncLLM(EngineClient): """ 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: + current_requests = ( + self.admission_stats.get_num_requests() + if self.admission_stats is not None + else self.get_num_unfinished_requests() + ) + if current_requests + n > max_num_reqs: logger.info( "Request queue full - rejecting request %s " "(current=%d, n=%d, max=%d).", request_id, - current, + current_requests, n, max_num_reqs, ) diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 419e3aa8e9..257287243d 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -5,7 +5,7 @@ import asyncio from collections import defaultdict, deque from collections.abc import Iterable from dataclasses import dataclass -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast import numpy as np import torch @@ -41,6 +41,9 @@ from vllm.v1.metrics.stats import ( ) from vllm.v1.outputs import SamplingMaskLists +if TYPE_CHECKING: + from vllm.v1.engine.admission_control import SharedAdmissionStats + # shared empty CPU tensor used as a placeholder pooling output EMPTY_CPU_TENSOR = torch.empty(0, device="cpu") @@ -455,6 +458,7 @@ class OutputProcessor: log_stats: bool, stream_interval: int = 1, tracing_enabled: bool = False, + admission_stats: "SharedAdmissionStats | None" = None, ): self.log_stats = log_stats self.tokenizer = tokenizer @@ -464,6 +468,7 @@ class OutputProcessor: self.external_req_ids: defaultdict[str, list[str]] = defaultdict(list) self.lora_states = LoRARequestStates(log_stats) self.tracing_enabled = tracing_enabled + self.admission_stats = admission_stats def get_num_unfinished_requests(self): return len(self.request_states) @@ -553,6 +558,7 @@ class OutputProcessor: child_reqs = self.abort_requests(child_reqs, internal=True) request_ids_to_abort.extend(child_reqs) self.parent_requests.pop(request_id, None) + self._update_admission_stats() return request_ids_to_abort def add_request( @@ -585,6 +591,7 @@ class OutputProcessor: # Track the external_req_id -> [internal_req_id, ...] mapping self.external_req_ids[req_state.external_req_id].append(request_id) + self._update_admission_stats() def _update_streaming_request_state( self, req_state: RequestState, request: EngineCoreRequest, prompt: str | None @@ -763,6 +770,11 @@ class OutputProcessor: parent_req = req_state.parent_req if parent_req and not parent_req.child_requests: self.parent_requests.pop(parent_req.request_id, None) + self._update_admission_stats() + + def _update_admission_stats(self) -> None: + if self.admission_stats is not None: + self.admission_stats.set_num_requests(self.get_num_unfinished_requests()) def update_scheduler_stats(self, scheduler_stats: SchedulerStats | None): self.lora_states.update_scheduler_stats(scheduler_stats) diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index b0411dc9e8..da17110d80 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -209,6 +209,15 @@ class APIServerProcessManager: self.processes: list[BaseProcess] = [] self._address_pipes: list[connection.Connection] = [] + admission_counters = None + if num_servers > 1 and getattr(args, "max_num_queued_reqs", None) is not None: + from vllm.v1.engine.admission_control import SharedAdmissionStats + + admission_counters = spawn_context.RawArray( + "q", + SharedAdmissionStats.num_counters(num_servers), + ) + for i, in_addr, out_addr in zip( range(num_servers), input_addresses, output_addresses ): @@ -218,6 +227,8 @@ class APIServerProcessManager: "client_count": num_servers, "client_index": i, } + if admission_counters is not None: + client_config["mp_admission_counters"] = admission_counters if stats_update_address is not None: client_config["stats_update_address"] = stats_update_address if tensor_queue is not None: