[Frontend] Share max_num_queued_reqs across API server processes (#54746)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Chauncey committed
Sep 15, 2026 at 22:50 UTC
f0fbc98c488b40a07d3b269aeba2bda3c4c3b03b
6 files changed
+173
-6
tests/entrypoints/launchers/api_server/test_api_server_process_manager.py
+39
@@ -5,6 +5,7 @@ import multiprocessing
5
import socket
6
import threading
7
import time
8
+from types import SimpleNamespace
9
from unittest.mock import patch
10
11
import pytest
@@ -14,6 +15,7 @@ from tests.entrypoints.launchers.api_server._api_server_spawn_workers import (
15
exit_before_report_worker,
16
)
17
from vllm.utils.network_utils import make_zmq_socket, split_zmq_path
18
+from vllm.v1.engine.admission_control import SharedAdmissionStats
19
from vllm.v1.utils import (
20
APIServerProcessManager,
21
get_engine_client_zmq_addr,
@@ -32,6 +34,18 @@ def mock_run_api_server_worker(listen_address, sock, args, client_config=None):
34
print("Mock worker completed successfully")
35
36
37
+def update_admission_stats_worker(listen_address, sock, args, client_config):
38
+ stats = SharedAdmissionStats(
39
+ client_config,
40
+ client_config["client_count"],
41
+ client_config["client_index"],
42
+ )
43
+ value = client_config["client_index"] + 1
44
+ stats.set_num_requests(value)
45
+ args.barrier.wait()
46
+ args.results.put(stats.get_num_requests())
47
+
48
+
49
# Module-level stub for the gather_actual_addresses test. Must be
50
# importable by `multiprocessing.spawn` (no closures, no nesting).
51
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):
145
assert not proc.is_alive()
146
147
148
+@pytest.mark.skip_global_cleanup
149
+def test_api_server_processes_share_admission_stats(api_server_args):
150
+ args = api_server_args.copy()
151
+ spawn_context = multiprocessing.get_context("spawn")
152
+ worker_args = SimpleNamespace(
153
+ barrier=spawn_context.Barrier(args["num_servers"]),
154
+ results=spawn_context.Queue(),
155
+ max_num_queued_reqs=10,
156
+ )
157
+ args["args"] = worker_args
158
+ args["target_server_fn"] = update_admission_stats_worker
159
+ manager = APIServerProcessManager(**args)
160
+
161
+ try:
162
+ for process in manager.processes:
163
+ process.join(timeout=30)
164
+ assert process.exitcode == 0
165
+
166
+ assert [worker_args.results.get(timeout=5) for _ in range(3)] == [6] * 3
167
+ finally:
168
+ manager.shutdown()
169
+ worker_args.results.close()
170
+ worker_args.results.join_thread()
171
+
172
+
173
@patch("vllm.v1.utils.run_api_server_worker_proc", mock_run_api_server_worker)
174
def test_wait_for_completion_or_failure(api_server_args):
175
"""Test that wait_for_completion_or_failure works with failures."""
tests/v1/engine/test_admission_control.py
+39
@@ -14,6 +14,7 @@ These tests cover:
14
15
import argparse
16
import asyncio
17
+import multiprocessing
18
from http import HTTPStatus
19
from types import SimpleNamespace
20
from unittest.mock import AsyncMock, MagicMock
@@ -34,6 +35,7 @@ from vllm.exceptions import (
35
from vllm.pooling_params import PoolingParams
36
from vllm.sampling_params import SamplingParams
37
from vllm.utils.argparse_utils import human_readable_int
38
+from vllm.v1.engine.admission_control import SharedAdmissionStats
39
from vllm.v1.engine.async_llm import AsyncLLM
40
from vllm.v1.engine.output_processor import OutputProcessor
41
@@ -66,6 +68,7 @@ def _make_async_llm(
68
llm.output_processor = MagicMock()
69
llm.output_processor.get_num_unfinished_requests.return_value = num_unfinished
70
llm.output_processor.get_num_queued_tokens.return_value = num_queued_tokens
71
+ llm.admission_stats = None
72
return llm
73
74
@@ -84,6 +87,42 @@ def _make_scheduler_config(**kwargs) -> SchedulerConfig:
87
)
88
89
90
+def _make_shared_stats(client_count: int = 2) -> list[SharedAdmissionStats]:
91
+ counters = multiprocessing.RawArray(
92
+ "q", SharedAdmissionStats.num_counters(client_count)
93
+ )
94
+ return [
95
+ SharedAdmissionStats({"mp_admission_counters": counters}, client_count, index)
96
+ for index in range(client_count)
97
+ ]
98
+
99
+
100
+@pytest.mark.skip_global_cleanup
101
+def test_shared_admission_stats_aggregate_api_servers():
102
+ server_0, server_1 = _make_shared_stats()
103
+
104
+ server_0.set_num_requests(2)
105
+ assert server_0.get_num_requests() == 2
106
+ server_1.set_num_requests(3)
107
+ assert server_0.get_num_requests() == server_1.get_num_requests() == 5
108
+
109
+ server_0.set_num_requests(1)
110
+ assert server_0.get_num_requests() == 4
111
+
112
+
113
+@pytest.mark.skip_global_cleanup
114
+def test_shared_request_snapshot_enforces_global_limit():
115
+ server_0, server_1 = _make_shared_stats()
116
+ server_0.set_num_requests(2)
117
+ llm = _make_async_llm(max_num_queued_reqs=2)
118
+ llm.admission_stats = server_1
119
+
120
+ with pytest.raises(QueueOverflowError):
121
+ llm.check_admission()
122
+
123
+ assert server_0.get_num_requests() == 2
124
+
125
+
126
# ---------------------------------------------------------------------------
127
# Exception classes
128
# ---------------------------------------------------------------------------
vllm/v1/engine/admission_control.py
new
+52
@@ -0,0 +1,52 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from typing import Any, Protocol, cast
5
+
6
+
7
+class AdmissionCounterBuffer(Protocol):
8
+ def __getitem__(self, index: int) -> int: ...
9
+
10
+ def __setitem__(self, index: int, value: int) -> None: ...
11
+
12
+
13
+class SharedAdmissionStats:
14
+ """Lock-free admission stats shared by API server processes.
15
+
16
+ Each process writes exclusively to its own cache-line-sized slot. Readers
17
+ aggregate the slots, avoiding both inter-process locks and cache-line
18
+ contention on the request path.
19
+ """
20
+
21
+ # Space per-process counters 64 bytes apart to avoid false sharing on CPUs
22
+ # with 64-byte cache lines: one process's writes would otherwise invalidate
23
+ # the cache line containing another process's counter.
24
+ CACHE_LINE_BYTES = 64
25
+ COUNTER_BYTES = 8
26
+ COUNTERS_PER_SLOT = CACHE_LINE_BYTES // COUNTER_BYTES
27
+
28
+ def __init__(
29
+ self,
30
+ client_addresses: dict[str, Any],
31
+ client_count: int = 1,
32
+ client_index: int = 0,
33
+ ):
34
+ self._counters = cast(
35
+ AdmissionCounterBuffer,
36
+ client_addresses.pop("mp_admission_counters"),
37
+ )
38
+ self._client_count = client_count
39
+ self._slot_offset = client_index * self.COUNTERS_PER_SLOT
40
+
41
+ @classmethod
42
+ def num_counters(cls, client_count: int) -> int:
43
+ return client_count * cls.COUNTERS_PER_SLOT
44
+
45
+ def set_num_requests(self, count: int) -> None:
46
+ self._counters[self._slot_offset] = count
47
+
48
+ def get_num_requests(self) -> int:
49
+ return sum(
50
+ self._counters[client_index * self.COUNTERS_PER_SLOT]
51
+ for client_index in range(self._client_count)
52
+ )
vllm/v1/engine/async_llm.py
+19
-5
@@ -44,6 +44,7 @@ from vllm.usage.usage_lib import UsageContext
44
from vllm.utils.async_utils import cancel_task_threadsafe
45
from vllm.utils.collection_utils import as_list
46
from vllm.v1.engine import EngineCoreRequest, PauseMode
47
+from vllm.v1.engine.admission_control import SharedAdmissionStats
48
from vllm.v1.engine.core_client import EngineCoreClient
49
from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
50
from vllm.v1.engine.input_processor import InputProcessor
@@ -144,12 +145,20 @@ class AsyncLLM(EngineClient):
145
# Convert EngineInput --> EngineCoreRequest.
146
self.input_processor = InputProcessor(self.vllm_config, renderer)
147
148
+ self.admission_stats = (
149
+ SharedAdmissionStats(client_addresses, client_count, client_index)
150
+ if client_addresses is not None
151
+ and "mp_admission_counters" in client_addresses
152
+ else None
153
+ )
154
+
155
# Converts EngineCoreOutputs --> RequestOutput.
156
self.output_processor = OutputProcessor(
157
renderer.tokenizer,
158
log_stats=self.log_stats,
159
stream_interval=self.vllm_config.scheduler_config.stream_interval,
160
tracing_enabled=tracing_endpoint is not None,
161
+ admission_stats=self.admission_stats,
162
)
163
164
# EngineCore (starts the engine in background process).
@@ -291,8 +300,9 @@ class AsyncLLM(EngineClient):
300
Both limits return HTTP 503 (Service Unavailable) so that load
301
balancers and client SDKs retry on a different instance.
302
294
- - ``max_num_queued_reqs``: hard cap on the number of unfinished requests
295
- (waiting + running). A request with ``n > 1`` counts as ``n`` slots.
303
+ - ``max_num_queued_reqs``: approximate global cap on the number of
304
+ unfinished requests (waiting + running). A request with ``n > 1``
305
+ counts as ``n`` slots.
306
- ``max_num_queued_tokens``: TTFT QoS — cap on the total prompt
307
tokens of requests still in prefill.
308
@@ -312,13 +322,17 @@ class AsyncLLM(EngineClient):
322
"""
323
max_num_reqs = self.scheduler_config.max_num_queued_reqs
324
if max_num_reqs is not None:
315
- current = self.get_num_unfinished_requests()
316
- if current + n > max_num_reqs:
325
+ current_requests = (
326
+ self.admission_stats.get_num_requests()
327
+ if self.admission_stats is not None
328
+ else self.get_num_unfinished_requests()
329
+ )
330
+ if current_requests + n > max_num_reqs:
331
logger.info(
332
"Request queue full - rejecting request %s "
333
"(current=%d, n=%d, max=%d).",
334
request_id,
321
- current,
335
+ current_requests,
336
n,
337
max_num_reqs,
338
)
vllm/v1/engine/output_processor.py
+13
-1
@@ -5,7 +5,7 @@ import asyncio
5
from collections import defaultdict, deque
6
from collections.abc import Iterable
7
from dataclasses import dataclass
8
-from typing import Any, cast
8
+from typing import TYPE_CHECKING, Any, cast
9
10
import numpy as np
11
import torch
@@ -41,6 +41,9 @@ from vllm.v1.metrics.stats import (
41
)
42
from vllm.v1.outputs import SamplingMaskLists
43
44
+if TYPE_CHECKING:
45
+ from vllm.v1.engine.admission_control import SharedAdmissionStats
46
+
47
# shared empty CPU tensor used as a placeholder pooling output
48
EMPTY_CPU_TENSOR = torch.empty(0, device="cpu")
49
@@ -455,6 +458,7 @@ class OutputProcessor:
458
log_stats: bool,
459
stream_interval: int = 1,
460
tracing_enabled: bool = False,
461
+ admission_stats: "SharedAdmissionStats | None" = None,
462
):
463
self.log_stats = log_stats
464
self.tokenizer = tokenizer
@@ -464,6 +468,7 @@ class OutputProcessor:
468
self.external_req_ids: defaultdict[str, list[str]] = defaultdict(list)
469
self.lora_states = LoRARequestStates(log_stats)
470
self.tracing_enabled = tracing_enabled
471
+ self.admission_stats = admission_stats
472
473
def get_num_unfinished_requests(self):
474
return len(self.request_states)
@@ -553,6 +558,7 @@ class OutputProcessor:
558
child_reqs = self.abort_requests(child_reqs, internal=True)
559
request_ids_to_abort.extend(child_reqs)
560
self.parent_requests.pop(request_id, None)
561
+ self._update_admission_stats()
562
return request_ids_to_abort
563
564
def add_request(
@@ -585,6 +591,7 @@ class OutputProcessor:
591
592
# Track the external_req_id -> [internal_req_id, ...] mapping
593
self.external_req_ids[req_state.external_req_id].append(request_id)
594
+ self._update_admission_stats()
595
596
def _update_streaming_request_state(
597
self, req_state: RequestState, request: EngineCoreRequest, prompt: str | None
@@ -763,6 +770,11 @@ class OutputProcessor:
770
parent_req = req_state.parent_req
771
if parent_req and not parent_req.child_requests:
772
self.parent_requests.pop(parent_req.request_id, None)
773
+ self._update_admission_stats()
774
+
775
+ def _update_admission_stats(self) -> None:
776
+ if self.admission_stats is not None:
777
+ self.admission_stats.set_num_requests(self.get_num_unfinished_requests())
778
779
def update_scheduler_stats(self, scheduler_stats: SchedulerStats | None):
780
self.lora_states.update_scheduler_stats(scheduler_stats)
vllm/v1/utils.py
+11
@@ -209,6 +209,15 @@ class APIServerProcessManager:
209
self.processes: list[BaseProcess] = []
210
self._address_pipes: list[connection.Connection] = []
211
212
+ admission_counters = None
213
+ if num_servers > 1 and getattr(args, "max_num_queued_reqs", None) is not None:
214
+ from vllm.v1.engine.admission_control import SharedAdmissionStats
215
+
216
+ admission_counters = spawn_context.RawArray(
217
+ "q",
218
+ SharedAdmissionStats.num_counters(num_servers),
219
+ )
220
+
221
for i, in_addr, out_addr in zip(
222
range(num_servers), input_addresses, output_addresses
223
):
@@ -218,6 +227,8 @@ class APIServerProcessManager:
227
"client_count": num_servers,
228
"client_index": i,
229
}
230
+ if admission_counters is not None:
231
+ client_config["mp_admission_counters"] = admission_counters
232
if stats_update_address is not None:
233
client_config["stats_update_address"] = stats_update_address
234
if tensor_queue is not None: