[Frontend] Share max_num_queued_reqs across API server processes (#54746)

Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>

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."""

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
 # ---------------------------------------------------------------------------

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)
+        )

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,
                 )

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)

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: