@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6f0dd93801
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 7, 07:45 UTC

[Core] Remove busy loop from idle buffer readers (#28053)

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com> Signed-off-by: Travis Johnson <tsjohnso@us.ibm.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Travis Johnson <tsjohnso@us.ibm.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Joe Runde committed Mar 4, 2026 at 00:44 UTC 6f0dd93801163a6418695c6dc0b43c516261f55a
5 files changed +584 -140
tests/basic_correctness/test_basic_correctness.py
-2
@@ -124,8 +124,6 @@ def test_models(
124 [
125 ("facebook/opt-125m", "ray", "", "L4", {}),
126 ("facebook/opt-125m", "mp", "", "L4", {}),
127 - ("facebook/opt-125m", "ray", "", "L4", {"VLLM_SLEEP_WHEN_IDLE": "1"}),
128 - ("facebook/opt-125m", "mp", "", "L4", {"VLLM_SLEEP_WHEN_IDLE": "1"}),
127 ("meta-llama/Llama-3.2-1B-Instruct", "ray", "", "L4", {}),
128 ("meta-llama/Llama-3.2-1B-Instruct", "mp", "", "L4", {}),
129 ("facebook/opt-125m", "ray", "", "A100", {}),
tests/distributed/test_shm_broadcast.py
+285 -8
@@ -1,11 +1,14 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -import multiprocessing
4 import random
5 +import threading
6 import time
7 +from unittest import mock
8
9 +import multiprocess as mp
10 import numpy as np
11 +import pytest
12 import torch.distributed as dist
13
14 from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
@@ -22,7 +25,14 @@ def get_arrays(n: int, seed: int = 0) -> list[np.ndarray]:
25 return [np.random.randint(1, 100, i) for i in sizes]
26
27
25 -def distributed_run(fn, world_size):
28 +def distributed_run(fn, world_size, timeout=60):
29 + """Run a function in multiple processes with proper error handling.
30 +
31 + Args:
32 + fn: Function to run in each process
33 + world_size: Number of processes to spawn
34 + timeout: Maximum time in seconds to wait for processes (default: 60)
35 + """
36 number_of_processes = world_size
37 processes = []
38 for i in range(number_of_processes):
@@ -33,19 +43,45 @@ def distributed_run(fn, world_size):
43 env["LOCAL_WORLD_SIZE"] = str(number_of_processes)
44 env["MASTER_ADDR"] = "localhost"
45 env["MASTER_PORT"] = "12345"
36 - p = multiprocessing.Process(target=fn, args=(env,))
46 + p = mp.Process(target=fn, args=(env,))
47 processes.append(p)
48 p.start()
49
40 - for p in processes:
41 - p.join()
50 + # Monitor processes and fail fast if any process fails
51 + start_time = time.time()
52 + failed_processes = []
53 +
54 + # Wait for all processes, checking for failures
55 + while time.time() - start_time < timeout:
56 + all_done = True
57 + for i, p in enumerate(processes):
58 + if p.is_alive():
59 + all_done = False
60 + elif p.exitcode != 0:
61 + # Process failed
62 + failed_processes.append((i, p.exitcode))
63 + break
64 +
65 + if failed_processes or all_done:
66 + break
67 + time.sleep(0.1) # Check every 100ms
68
43 - for p in processes:
44 - assert p.exitcode == 0
69 + # Check for timeout if no failures detected yet
70 + for i, p in enumerate(processes):
71 + if p.is_alive():
72 + p.kill()
73 + p.join()
74 +
75 + # Report failures
76 + if failed_processes:
77 + error_msg = "Distributed test failed:\n"
78 + for rank, status in failed_processes:
79 + error_msg += f" Rank {rank}: Exit code {status}\n"
80 + raise AssertionError(error_msg)
81
82
83 def worker_fn_wrapper(fn):
48 - # `multiprocessing.Process` cannot accept environment variables directly
84 + # `mp.Process` cannot accept environment variables directly
85 # so we need to pass the environment variables as arguments
86 # and update the environment variables in the function
87 def wrapped_fn(env):
@@ -115,3 +151,244 @@ def worker_fn():
151
152 def test_shm_broadcast():
153 distributed_run(worker_fn, 4)
154 +
155 +
156 +@worker_fn_wrapper
157 +def worker_fn_test_shutdown_busy():
158 + rank = dist.get_rank()
159 + writer_rank = 2
160 + message_queue = MessageQueue.create_from_process_group(
161 + dist.group.WORLD, 40 * 1024, 2, writer_rank
162 + )
163 +
164 + if not message_queue._is_writer:
165 + # Put into busy mode
166 + message_queue._spin_condition.busy_loop_s = 9999
167 +
168 + shutdown_event = threading.Event()
169 +
170 + def shutdown_thread(mq, shutdown_event):
171 + shutdown_event.wait()
172 + mq.shutdown()
173 +
174 + threading.Thread(
175 + target=shutdown_thread, args=(message_queue, shutdown_event)
176 + ).start()
177 +
178 + with pytest.raises(TimeoutError):
179 + message_queue.dequeue(timeout=0.01)
180 +
181 + shutdown_event.set()
182 +
183 + with pytest.raises(RuntimeError, match="cancelled"):
184 + message_queue.dequeue(timeout=1)
185 +
186 + assert message_queue.shutting_down
187 +
188 + print(f"torch distributed passed the test! Rank {rank}")
189 + dist.barrier()
190 +
191 +
192 +def test_message_queue_shutdown_busy(caplog_vllm):
193 + distributed_run(worker_fn_test_shutdown_busy, 4)
194 + print(caplog_vllm.text)
195 +
196 +
197 +@worker_fn_wrapper
198 +def worker_fn_test_shutdown_idle():
199 + rank = dist.get_rank()
200 + writer_rank = 2
201 + message_queue = MessageQueue.create_from_process_group(
202 + dist.group.WORLD, 40 * 1024, 2, writer_rank
203 + )
204 +
205 + if not message_queue._is_writer:
206 + # Put into idle mode
207 + message_queue._spin_condition.last_read = 0
208 +
209 + shutdown_event = threading.Event()
210 +
211 + def shutdown_thread(mq, shutdown_event):
212 + shutdown_event.wait()
213 + mq.shutdown()
214 +
215 + threading.Thread(
216 + target=shutdown_thread, args=(message_queue, shutdown_event)
217 + ).start()
218 +
219 + with pytest.raises(TimeoutError):
220 + message_queue.dequeue(timeout=0.01)
221 +
222 + shutdown_event.set()
223 +
224 + with pytest.raises(RuntimeError, match="cancelled"):
225 + message_queue.dequeue(timeout=1)
226 +
227 + assert message_queue.shutting_down
228 +
229 + print(f"torch distributed passed the test! Rank {rank}")
230 + dist.barrier()
231 +
232 +
233 +def test_message_queue_shutdown_idle():
234 + distributed_run(worker_fn_test_shutdown_idle, 4)
235 +
236 +
237 +@worker_fn_wrapper
238 +def worker_fn_test_idle_to_busy():
239 + rank = dist.get_rank()
240 + writer_rank = 2
241 + message_queue = MessageQueue.create_from_process_group(
242 + dist.group.WORLD, 40 * 1024, 2, writer_rank
243 + )
244 +
245 + message1 = "hello world"
246 + message2 = np.random.randint(1, 100, 100)
247 + with mock.patch.object(
248 + message_queue._spin_condition, "wait", wraps=message_queue._spin_condition.wait
249 + ) as wrapped_wait:
250 + if not message_queue._is_writer:
251 + # Put into idle mode
252 + message_queue._spin_condition.last_read = 0
253 +
254 + # no messages, so expect a TimeoutError
255 + with pytest.raises(TimeoutError):
256 + message_queue.dequeue(timeout=0.01)
257 + # wait should only be called once while idle
258 + assert wrapped_wait.call_count == 1
259 +
260 + # sync with the writer and wait for message1
261 + dist.barrier()
262 + recv_message = message_queue.dequeue(timeout=5)
263 + assert recv_message == message1
264 + # second call to wait, with a message read, this puts in a busy spin
265 + assert wrapped_wait.call_count == 2
266 +
267 + # sync with the writer and wait for message2
268 + dist.barrier()
269 + recv_message = message_queue.dequeue(timeout=1)
270 + assert np.array_equal(recv_message, message2)
271 + # in busy mode, we expect wait to have been called multiple times
272 + assert wrapped_wait.call_count > 3
273 + else:
274 + # writer writes two messages in sync with the reader
275 + dist.barrier()
276 + # sleep delays the send to ensure reader enters the read loop
277 + time.sleep(0.1)
278 + message_queue.enqueue(message1)
279 +
280 + dist.barrier()
281 + time.sleep(0.1)
282 + message_queue.enqueue(message2)
283 +
284 + message_queue.shutdown()
285 + assert message_queue.shutting_down
286 + print(f"torch distributed passed the test! Rank {rank}")
287 +
288 +
289 +def test_message_queue_idle_wake():
290 + distributed_run(worker_fn_test_idle_to_busy, 4)
291 +
292 +
293 +@worker_fn_wrapper
294 +def worker_fn_test_busy_to_idle():
295 + rank = dist.get_rank()
296 + writer_rank = 2
297 + message_queue = MessageQueue.create_from_process_group(
298 + dist.group.WORLD, 40 * 1024, 2, writer_rank
299 + )
300 +
301 + message1 = 12345
302 + message2 = list(range(3))
303 + with mock.patch.object(
304 + message_queue._spin_condition, "wait", wraps=message_queue._spin_condition.wait
305 + ) as wrapped_wait:
306 + if not message_queue._is_writer:
307 + # Put into busy mode
308 + message_queue._spin_condition.busy_loop_s = 9999
309 +
310 + # sync with the writer and wait for message1
311 + dist.barrier()
312 + recv_message = message_queue.dequeue(timeout=1)
313 + assert recv_message == message1
314 + # in busy mode, we expect wait to have been called many times
315 + assert wrapped_wait.call_count > 1
316 +
317 + # simulate busy loop ending
318 + message_queue._spin_condition.busy_loop_s = 0
319 + # ensure we enter idle mode, then record call count
320 + with pytest.raises(TimeoutError):
321 + message_queue.dequeue(timeout=0.01)
322 + call_count = wrapped_wait.call_count
323 +
324 + # sync with the writer and wait for message2
325 + dist.barrier()
326 + recv_message = message_queue.dequeue(timeout=1)
327 + assert recv_message == message2
328 +
329 + # call to wait after idle should only happen once
330 + assert wrapped_wait.call_count == call_count + 1
331 + else:
332 + # writer writes two messages in sync with the reader
333 + dist.barrier()
334 + # sleep delays the send to ensure reader enters the read loop
335 + time.sleep(0.1)
336 + message_queue.enqueue(message1)
337 +
338 + dist.barrier()
339 + time.sleep(0.1)
340 + message_queue.enqueue(message2)
341 +
342 + message_queue.shutdown()
343 + assert message_queue.shutting_down
344 + print(f"torch distributed passed the test! Rank {rank}")
345 +
346 +
347 +def test_message_queue_busy_to_idle():
348 + distributed_run(worker_fn_test_busy_to_idle, 4)
349 +
350 +
351 +def test_warning_logs(caplog_vllm):
352 + """
353 + Test that warning logs are emitted at VLLM_RINGBUFFER_WARNING_INTERVAL intervals
354 + when indefinite=False, and are not emitted when indefinite=True.
355 + """
356 +
357 + # Patch the warning log interval to every 1 ms during reads
358 + with mock.patch(
359 + "vllm.distributed.device_communicators.shm_broadcast.VLLM_RINGBUFFER_WARNING_INTERVAL",
360 + new=0.001, # 1 ms
361 + ):
362 + writer = MessageQueue(
363 + n_reader=1,
364 + n_local_reader=1,
365 + max_chunk_bytes=1024 * 1024, # 1MB chunks
366 + max_chunks=10,
367 + )
368 + reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0)
369 + writer.wait_until_ready()
370 + reader.wait_until_ready()
371 +
372 + # We should have at least one warning log here
373 + # "0 seconds" expected due to rounding of 1ms test interval
374 + with pytest.raises(TimeoutError):
375 + reader.dequeue(timeout=0.01, indefinite=False)
376 + assert any(
377 + "No available shared memory broadcast block found in 0 seconds"
378 + in record.message
379 + for record in caplog_vllm.records
380 + )
381 + caplog_vllm.clear()
382 +
383 + # We should have no warnings this time
384 + with pytest.raises(TimeoutError):
385 + reader.dequeue(timeout=0.01, indefinite=True)
386 + assert all(
387 + "No available shared memory broadcast block found in 0 seconds"
388 + not in record.message
389 + for record in caplog_vllm.records
390 + )
391 +
392 + # Clean up when done
393 + writer.shutdown()
394 + reader.shutdown()
vllm/distributed/device_communicators/shm_broadcast.py
+194 -64
@@ -2,13 +2,13 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import functools
4 import pickle
5 +import sys
6 import threading
7 import time
8 from contextlib import contextmanager
9 from dataclasses import dataclass, field
10 from multiprocessing import shared_memory
11 from pickle import PickleBuffer
11 -from threading import Event
12 from typing import TYPE_CHECKING, Any, cast
13 from unittest.mock import patch
14
@@ -18,6 +18,7 @@ import zmq
18 from torch.distributed import ProcessGroup
19 from zmq import ( # type: ignore
20 IPV6, # type: ignore
21 + PUB,
22 SUB,
23 SUBSCRIBE,
24 XPUB,
@@ -32,6 +33,7 @@ from vllm.platforms import current_platform
33 from vllm.utils.network_utils import (
34 get_ip,
35 get_open_port,
36 + get_open_zmq_inproc_path,
37 get_open_zmq_ipc_path,
38 is_valid_ipv6_address,
39 )
@@ -78,50 +80,125 @@ def to_bytes_big(value: int, size: int) -> bytes:
80 logger = init_logger(__name__)
81
82
81 -def long_wait_time_msg(threshold: int) -> str:
82 - return (
83 - "No available shared memory broadcast block found "
84 - f"in {threshold} seconds. This typically happens "
85 - "when some processes are hanging or doing some "
86 - "time-consuming work (e.g. compilation, "
87 - "weight/kv cache quantization)."
88 - )
89 -
90 -
91 -class SpinTimer:
92 - def record_activity(self):
93 - pass
94 -
95 - def spin(self):
96 - sched_yield()
83 +LONG_WAIT_TIME_LOG_MSG = (
84 + "No available shared memory broadcast block found "
85 + "in %d seconds. This typically happens "
86 + "when some processes are hanging or doing some "
87 + "time-consuming work (e.g. compilation, "
88 + "weight/kv cache quantization)."
89 +)
90
91
99 -class SpinSleepTimer(SpinTimer):
92 +class SpinCondition:
93 """
101 - In setups which have long inactivity periods it is desirable to reduce
102 - system power consumption when vllm does nothing. This would lead to more
103 - CPU thermal headroom when a request eventually comes, especially when
104 - multiple GPUs are connected as each GPU would otherwise pin one thread at
105 - 100% CPU usage.
106 -
107 - The simplest solution is to reduce polling frequency when there is no
108 - activity for a certain period of time.
94 + This class implements an interface similar to a threading.Condition. It
95 + allows a writer to notify readers to wake up and read from the shared memory
96 + buffer. This notification is done over a zmq socket.
97 +
98 + For optimal performance under load we don't want the readers to need to poll
99 + the zmq socket for every read. So the `wait` method here will return
100 + immediately when reads are frequent, and will only enter "idle mode" and
101 + await a notification on the zmq socket after a period of inactivity. This
102 + allows the readers to spin quickly, hence "SpinCondition".
103 +
104 + To support clean shutdown, a separate thread in the reader's process must be
105 + able to wake the reader so that it can exit. A separate cancel() method is
106 + implemented with an in-process socket to allow this interruption.
107 """
108
111 - def __init__(self, busy_loop_s: float = 3.0, wait_sleep_s: float = 0.1):
112 - self.last_activity = time.monotonic()
113 - self.busy_loop_s = busy_loop_s
114 - self.wait_sleep_s = wait_sleep_s
115 -
116 - def record_activity(self):
117 - self.last_activity = time.monotonic()
118 -
119 - def spin(self):
120 - curr_time = time.monotonic()
121 - if curr_time >= self.last_activity + self.busy_loop_s:
122 - time.sleep(self.wait_sleep_s)
109 + def __init__(
110 + self,
111 + is_reader: bool,
112 + context: zmq.Context,
113 + notify_address: str,
114 + busy_loop_s: float = 1,
115 + ):
116 + self.is_reader = is_reader
117 +
118 + if is_reader:
119 + # Time of last shm buffer read
120 + self.last_read = time.monotonic()
121 +
122 + # Time to keep busy-looping on the shm buffer before going idle
123 + self.busy_loop_s = busy_loop_s
124 +
125 + # Readers subscribe to write notifications
126 + self.local_notify_socket: zmq.Socket = context.socket(SUB)
127 + # Set zmq.CONFLATE to only keep the last message that the socket
128 + # receives. This prevents us from piling up notification messages
129 + # under high load when we aren't polling the socket.
130 + self.local_notify_socket.setsockopt(zmq.CONFLATE, 1)
131 + # Subscribe to all messages on the socket
132 + self.local_notify_socket.setsockopt_string(SUBSCRIBE, "")
133 + self.local_notify_socket.connect(notify_address)
134 +
135 + # Readers require a process-local socket to poll for cancellation
136 + cancel_path = get_open_zmq_inproc_path()
137 + self.write_cancel_socket: zmq.Socket = context.socket(zmq.PAIR)
138 + self.write_cancel_socket.bind(cancel_path)
139 + self.read_cancel_socket: zmq.Socket = context.socket(zmq.PAIR)
140 + self.read_cancel_socket.connect(cancel_path)
141 +
142 + # Poller allows waiting on either `.notify()` or `.cancel()`
143 + self.poller = zmq.Poller()
144 + self.poller.register(self.read_cancel_socket, zmq.POLLIN)
145 + self.poller.register(self.local_notify_socket, zmq.POLLIN)
146 else:
147 + # Writer side publishes write notifications
148 + self.local_notify_socket: zmq.Socket = context.socket(PUB) # type: ignore
149 + # Set high water mark to 1 - we don't need to send a massive amount of
150 + # pings during busy operation. PUB sockets will silently drop subsequent
151 + # messages after the high water mark is reached.
152 + self.local_notify_socket.setsockopt(zmq.SNDHWM, 1)
153 + self.local_notify_socket.bind(notify_address)
154 +
155 + self.last_read = 0
156 + self.busy_loop_s = 0
157 + self.read_cancel_socket = None
158 + self.write_cancel_socket = None
159 + self.poller = None
160 +
161 + def record_read(self):
162 + self.last_read = time.monotonic()
163 +
164 + def cancel(self):
165 + # Sends cancellation ping that will cause the reader to wake up.
166 + # This is done from a monitor thread in the same process as the reader.
167 + if self.is_reader:
168 + logger.debug("Canceling waiting reads on SHM Buffer")
169 + self.write_cancel_socket.send(b"\x00")
170 +
171 + def wait(self, timeout_ms: int | None = None) -> None:
172 + """Wait for data on the shared memory buffer.
173 +
174 + Yields the scheduler then returns immediately if it has been less than
175 + self.busy_loop_s since the last read.
176 +
177 + Otherwise, enters idle mode and awaits a socket ping for at most
178 + `timeout_ms` milliseconds, or indefinitely if timeout_ms is None.
179 + """
180 + assert self.is_reader, "Only readers can wait"
181 +
182 + current_time = time.monotonic()
183 + if current_time <= self.last_read + self.busy_loop_s:
184 sched_yield()
185 + else:
186 + events = dict(self.poller.poll(timeout=timeout_ms))
187 +
188 + if self.read_cancel_socket in events:
189 + logger.debug("Poller received cancel event")
190 + elif self.local_notify_socket in events:
191 + logger.debug("Poller received notify event")
192 + # Since zmq.CONFLATE is set, there will only be one notification
193 + # to read from the socket
194 + self.local_notify_socket.recv(flags=zmq.NOBLOCK, copy=False)
195 + else:
196 + logger.debug("Poller timed out")
197 +
198 + def notify(self):
199 + """Notifies all readers to wake up"""
200 + assert not self.is_reader, "Only writers can notify"
201 + self.local_notify_socket.send(b"\x00")
202
203
204 class ShmRingBuffer:
@@ -265,6 +342,7 @@ class Handle:
342
343 buffer_handle: tuple[int, int, int, str] | None = None
344 local_subscribe_addr: str | None = None
345 + local_notify_addr: str | None = None
346 remote_subscribe_addr: str | None = None
347 remote_addr_ipv6: bool = False
348
@@ -288,7 +366,7 @@ class MessageQueue:
366 self.n_local_reader = n_local_reader
367 n_remote_reader = n_reader - n_local_reader
368 self.n_remote_reader = n_remote_reader
291 -
369 + self.shutting_down = False
370 context = Context()
371
372 if n_local_reader > 0:
@@ -310,11 +388,19 @@ class MessageQueue:
388 self.local_socket.bind(local_subscribe_addr)
389
390 self.current_idx = 0
391 +
392 + # Create the notification side of the SpinCondition
393 + local_notify_addr = get_open_zmq_ipc_path()
394 + self._spin_condition = SpinCondition(
395 + is_reader=False, context=context, notify_address=local_notify_addr
396 + )
397 else:
398 self.buffer = None # type: ignore
399 local_subscribe_addr = None
400 self.local_socket = None
401 self.current_idx = -1
402 + local_notify_addr = None
403 + self._spin_condition = None # type: ignore
404
405 remote_addr_ipv6 = False
406 if n_remote_reader > 0:
@@ -341,12 +427,12 @@ class MessageQueue:
427 self.local_reader_rank = -1
428 # rank does not matter for remote readers
429 self._is_remote_reader = False
344 - self._read_spin_timer = SpinTimer()
430
431 self.handle = Handle(
432 local_reader_ranks=local_reader_ranks,
433 buffer_handle=self.buffer.handle() if self.buffer is not None else None,
434 local_subscribe_addr=local_subscribe_addr,
435 + local_notify_addr=local_notify_addr,
436 remote_subscribe_addr=remote_subscribe_addr,
437 remote_addr_ipv6=remote_addr_ipv6,
438 )
@@ -379,9 +465,9 @@ class MessageQueue:
465 self.local_socket.connect(socket_addr)
466
467 self.remote_socket = None
382 -
383 - self._read_spin_timer = (
384 - SpinSleepTimer() if envs.VLLM_SLEEP_WHEN_IDLE else SpinTimer()
468 + assert isinstance(handle.local_notify_addr, str)
469 + self._spin_condition = SpinCondition(
470 + is_reader=True, context=context, notify_address=handle.local_notify_addr
471 )
472 else:
473 self.buffer = None # type: ignore
@@ -399,7 +485,9 @@ class MessageQueue:
485 socket_addr = handle.remote_subscribe_addr
486 logger.debug("Connecting to %s", socket_addr)
487 self.remote_socket.connect(socket_addr)
488 + self._spin_condition = None # type: ignore
489
490 + self.shutting_down = False
491 return self
492
493 def wait_until_ready(self):
@@ -435,6 +523,13 @@ class MessageQueue:
523 recv = self.remote_socket.recv()
524 assert recv == b"READY"
525
526 + def shutdown(self):
527 + """If this is an idle reader, wakes it up so it can clean up and shut
528 + down"""
529 + self.shutting_down = True
530 + if self._spin_condition is not None:
531 + self._spin_condition.cancel()
532 +
533 @contextmanager
534 def acquire_write(self, timeout: float | None = None):
535 assert self._is_writer, "Only writers can acquire write"
@@ -465,7 +560,7 @@ class MessageQueue:
560 # if we wait for a long time, log a message
561 if elapsed > VLLM_RINGBUFFER_WARNING_INTERVAL * n_warning:
562 logger.info(
468 - long_wait_time_msg(VLLM_RINGBUFFER_WARNING_INTERVAL)
563 + LONG_WAIT_TIME_LOG_MSG, VLLM_RINGBUFFER_WARNING_INTERVAL
564 )
565 n_warning += 1
566
@@ -503,16 +598,60 @@ class MessageQueue:
598 self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks
599 break
600
601 + class ReadTimeoutWithWarnings:
602 + def __init__(self, timeout: float | None, should_warn: bool) -> None:
603 + self.started = time.monotonic()
604 + self.deadline = sys.maxsize if timeout is None else self.started + timeout
605 +
606 + # if should_warn, we need to wake up periodically to log
607 + self.warning_wait_time_ms: int | None = (
608 + VLLM_RINGBUFFER_WARNING_INTERVAL * 1000 if should_warn else None
609 + )
610 +
611 + self._should_warn = should_warn
612 + self.n_warning = 1
613 + self.timeout = timeout
614 +
615 + def timeout_ms(self) -> int | None:
616 + """Returns a timeout that is:
617 + - min(time to deadline, time to next warning) if we're logging warnings
618 + - time to deadline, if we're not logging warnings
619 + - None if the timeout is None and we're not logging warnings
620 + - raise TimeoutError if we are past the deadline
621 + """
622 + warning_wait_time = self.warning_wait_time_ms
623 + if self.timeout is None:
624 + return warning_wait_time
625 +
626 + time_left_ms = int((self.deadline - time.monotonic()) * 1000)
627 + if time_left_ms <= 0:
628 + raise TimeoutError
629 +
630 + if warning_wait_time and warning_wait_time < time_left_ms:
631 + return warning_wait_time
632 +
633 + return time_left_ms
634 +
635 + def should_warn(self) -> bool:
636 + """Returns true if it's time to log a warning for a timeout that is not
637 + indefinite"""
638 + if self._should_warn:
639 + elapsed = time.monotonic() - self.started
640 + if elapsed >= VLLM_RINGBUFFER_WARNING_INTERVAL * self.n_warning:
641 + self.n_warning += 1
642 + return True
643 + return False
644 +
645 @contextmanager
646 def acquire_read(
647 self,
648 timeout: float | None = None,
510 - cancel: Event | None = None,
649 indefinite: bool = False,
650 ):
651 assert self._is_local_reader, "Only readers can acquire read"
514 - start_time = time.monotonic()
515 - n_warning = 1
652 + read_timeout = self.ReadTimeoutWithWarnings(
653 + timeout=timeout, should_warn=not indefinite
654 + )
655 with self.buffer.get_metadata(self.current_idx) as metadata_buffer:
656 while True:
657 # Memory fence ensures we see the latest writes from the writer.
@@ -529,26 +668,16 @@ class MessageQueue:
668 # for readers, `self.current_idx` is the next block to read
669 # if this block is not ready,
670 # we need to wait until it is written
671 + self._spin_condition.wait(timeout_ms=read_timeout.timeout_ms())
672
533 - # Release the processor to other threads
534 - self._read_spin_timer.spin()
535 -
536 - if cancel is not None and cancel.is_set():
673 + if self.shutting_down:
674 raise RuntimeError("cancelled")
675
539 - # if we time out, raise an exception
540 - elapsed = time.monotonic() - start_time
541 - if timeout is not None and elapsed > timeout:
542 - raise TimeoutError
543 -
676 # if we wait for a long time, log a message
545 - if not indefinite and (
546 - elapsed > VLLM_RINGBUFFER_WARNING_INTERVAL * n_warning
547 - ):
677 + if read_timeout.should_warn():
678 logger.info(
549 - long_wait_time_msg(VLLM_RINGBUFFER_WARNING_INTERVAL)
679 + LONG_WAIT_TIME_LOG_MSG, VLLM_RINGBUFFER_WARNING_INTERVAL
680 )
551 - n_warning += 1
681
682 continue
683 # found a block that is not read by this reader
@@ -565,7 +694,7 @@ class MessageQueue:
694 memory_fence()
695 self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks
696
568 - self._read_spin_timer.record_activity()
697 + self._spin_condition.record_read()
698 break
699
700 def enqueue(self, obj, timeout: float | None = None):
@@ -608,18 +737,19 @@ class MessageQueue:
737 buf[offset:buf_offset] = to_bytes_big(buf_len, 4)
738 buf[buf_offset : (offset := buf_offset + buf_len)] = buffer
739
740 + self._spin_condition.notify()
741 +
742 if self.n_remote_reader > 0:
743 self.remote_socket.send_multipart(all_buffers, copy=False)
744
745 def dequeue(
746 self,
747 timeout: float | None = None,
617 - cancel: Event | None = None,
748 indefinite: bool = False,
749 ):
750 """Read from message queue with optional timeout (in seconds)"""
751 if self._is_local_reader:
622 - with self.acquire_read(timeout, cancel, indefinite) as buf:
752 + with self.acquire_read(timeout, indefinite) as buf:
753 overflow = buf[0] == 1
754 if not overflow:
755 offset = 3
vllm/envs.py
-5
@@ -179,7 +179,6 @@ if TYPE_CHECKING:
179 VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998
180 VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE: int = 163840
181 VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1
182 - VLLM_SLEEP_WHEN_IDLE: bool = False
182 VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16
183 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300
184 VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None
@@ -1338,9 +1337,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
1337 "VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS": lambda: int(
1338 os.getenv("VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS", "1")
1339 ),
1341 - # Reduce CPU usage when vLLM is idle. Enabling this will incur small
1342 - # latency penalty when a request eventually comes.
1343 - "VLLM_SLEEP_WHEN_IDLE": lambda: bool(int(os.getenv("VLLM_SLEEP_WHEN_IDLE", "0"))),
1340 # Control the max chunk bytes (in MB) for the rpc message queue.
1341 # Object larger than this threshold will be broadcast to worker
1342 # processes via zmq.
@@ -1751,7 +1747,6 @@ def compile_factors() -> dict[str, object]:
1747 "VLLM_HTTP_TIMEOUT_KEEP_ALIVE",
1748 "VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS",
1749 "VLLM_KEEP_ALIVE_ON_ENGINE_DEATH",
1754 - "VLLM_SLEEP_WHEN_IDLE",
1750 "VLLM_IMAGE_FETCH_TIMEOUT",
1751 "VLLM_VIDEO_FETCH_TIMEOUT",
1752 "VLLM_AUDIO_FETCH_TIMEOUT",
vllm/v1/executor/multiproc_executor.py
+105 -61
@@ -104,7 +104,6 @@ class MultiprocExecutor(Executor):
104 # and ensure workers will be terminated.
105 self._finalizer = weakref.finalize(self, self.shutdown)
106 self.is_failed = False
107 - self.shutdown_event = threading.Event()
107 self.failure_callback: FailureCallback | None = None
108
109 tp_size, pp_size, pcp_size = self._get_parallel_sizes()
@@ -158,20 +157,31 @@ class MultiprocExecutor(Executor):
157 global_start_rank = (
158 self.local_world_size * self.parallel_config.node_rank_within_dp
159 )
160 + # Keep track of socket file descriptors that are inherited by the
161 + # worker when using fork, so that we can close them in subsequent
162 + # workers
163 + inherited_fds: list[int] = []
164 for local_rank in range(self.local_world_size):
165 global_rank = global_start_rank + local_rank
166 is_driver_worker = self._is_driver_worker(global_rank)
164 - unready_workers.append(
165 - WorkerProc.make_worker_process(
166 - vllm_config=self.vllm_config,
167 - local_rank=local_rank,
168 - rank=global_rank,
169 - distributed_init_method=distributed_init_method,
170 - input_shm_handle=scheduler_output_handle,
171 - shared_worker_lock=shared_worker_lock,
172 - is_driver_worker=is_driver_worker,
173 - )
167 + unready_worker_handle = WorkerProc.make_worker_process(
168 + vllm_config=self.vllm_config,
169 + local_rank=local_rank,
170 + rank=global_rank,
171 + distributed_init_method=distributed_init_method,
172 + input_shm_handle=scheduler_output_handle,
173 + shared_worker_lock=shared_worker_lock,
174 + is_driver_worker=is_driver_worker,
175 + inherited_fds=inherited_fds,
176 )
177 + unready_workers.append(unready_worker_handle)
178 + if context.get_start_method() == "fork":
179 + inherited_fds.extend(
180 + [
181 + unready_worker_handle.death_writer.fileno(),
182 + unready_worker_handle.ready_pipe.fileno(),
183 + ]
184 + )
185
186 # Workers must be created before wait_for_ready to avoid
187 # deadlock, since worker.init_device() does a device sync.
@@ -220,6 +230,7 @@ class MultiprocExecutor(Executor):
230 for uw in unready_workers:
231 if uw.death_writer is not None:
232 uw.death_writer.close()
233 + uw.death_writer = None
234 self._ensure_worker_termination([uw.proc for uw in unready_workers])
235
236 self.output_rank = self._get_output_rank()
@@ -255,6 +266,7 @@ class MultiprocExecutor(Executor):
266 died = multiprocessing.connection.wait(sentinels)
267 _self = self_ref()
268 if not _self or getattr(_self, "shutting_down", False):
269 + logger.debug("MultiprocWorkerMonitor: shutdown already initiated")
270 return
271 _self.is_failed = True
272 proc_name = next(h.proc.name for h in workers if h.proc.sentinel == died[0])
@@ -354,8 +366,6 @@ class MultiprocExecutor(Executor):
366 if output_rank is not None:
367 response_mqs = (response_mqs[output_rank],)
368
357 - shutdown_event = self.shutdown_event
358 -
369 def get_response():
370 responses = []
371 for mq in response_mqs:
@@ -363,9 +373,7 @@ class MultiprocExecutor(Executor):
373 None if deadline is None else (deadline - time.monotonic())
374 )
375 try:
366 - status, result = mq.dequeue(
367 - timeout=dequeue_timeout, cancel=shutdown_event
368 - )
376 + status, result = mq.dequeue(timeout=dequeue_timeout)
377 except TimeoutError as e:
378 raise TimeoutError(f"RPC call to {method} timed out.") from e
379 if status != WorkerProc.ResponseStatus.SUCCESS:
@@ -408,20 +416,26 @@ class MultiprocExecutor(Executor):
416
417 active_procs = lambda: [proc for proc in worker_procs if proc.is_alive()]
418 # Give processes time to clean themselves up properly first
419 + logger.debug("Worker Termination: allow workers to gracefully shutdown")
420 if wait_for_termination(active_procs(), 4):
421 return
422
423 # Send SIGTERM if still running
424 + logger.debug("Worker Termination: workers still running sending SIGTERM")
425 for p in active_procs():
426 p.terminate()
427 if not wait_for_termination(active_procs(), 4):
428 # Send SIGKILL if still running
429 + logger.debug(
430 + "Worker Termination: resorting to SIGKILL to take down workers"
431 + )
432 for p in active_procs():
433 p.kill()
434
435 def shutdown(self):
436 """Properly shut down the executor and its workers"""
437 if not getattr(self, "shutting_down", False):
438 + logger.debug("Triggering shutdown of workers")
439 self.shutting_down = True
440
441 # Make sure all the worker processes are terminated first.
@@ -431,12 +445,20 @@ class MultiprocExecutor(Executor):
445 if w.death_writer is not None:
446 w.death_writer.close()
447 w.death_writer = None
434 - w.worker_response_mq = None
448 self._ensure_worker_termination([w.proc for w in workers])
449
437 - self.shutdown_event.set()
438 -
439 - self.rpc_broadcast_mq = None
450 + for w in workers:
451 + # Shutdown response queues
452 + if w.worker_response_mq is not None:
453 + w.worker_response_mq.shutdown()
454 + w.worker_response_mq = None
455 +
456 + if self.rpc_broadcast_mq is not None:
457 + self.rpc_broadcast_mq.shutdown()
458 + self.rpc_broadcast_mq = None
459 + for mq in self.response_mqs:
460 + mq.shutdown()
461 + self.response_mqs = []
462
463 def check_health(self) -> None:
464 self.collective_rpc("check_health", timeout=10)
@@ -609,24 +631,26 @@ class WorkerProc:
631 input_shm_handle, # Receive SchedulerOutput
632 shared_worker_lock: LockType,
633 is_driver_worker: bool,
634 + inherited_fds: list[int],
635 ) -> UnreadyWorkerProcHandle:
636 context = get_mp_context()
614 - # (reader, writer)
615 - reader, writer = context.Pipe(duplex=False)
616 -
617 - # Create death pipe to detect parent process exit
637 + # Ready pipe to communicate readiness from child to parent
638 + ready_reader, ready_writer = context.Pipe(duplex=False)
639 + # Death pipe to let child detect parent process exit
640 death_reader, death_writer = context.Pipe(duplex=False)
619 -
641 process_kwargs = {
642 "vllm_config": vllm_config,
643 "local_rank": local_rank,
644 "rank": rank,
645 "distributed_init_method": distributed_init_method,
646 "input_shm_handle": input_shm_handle,
626 - "ready_pipe": (reader, writer),
647 + "ready_pipe": ready_writer,
648 "death_pipe": death_reader,
649 "shared_worker_lock": shared_worker_lock,
650 "is_driver_worker": is_driver_worker,
651 + # Have the worker close parent end of this worker's pipes too
652 + "inherited_fds": inherited_fds
653 + + [ready_reader.fileno(), death_writer.fileno()],
654 }
655 # Run EngineCore busy loop in background process.
656 proc = context.Process(
@@ -637,10 +661,12 @@ class WorkerProc:
661 )
662
663 proc.start()
640 - writer.close()
664 + # Close child ends of pipes here in the parent
665 + ready_writer.close()
666 + death_reader.close()
667 # Keep death_writer open in parent - when parent exits,
668 # death_reader in child will get EOFError
643 - return UnreadyWorkerProcHandle(proc, rank, reader, death_writer)
669 + return UnreadyWorkerProcHandle(proc, rank, ready_reader, death_writer)
670
671 @staticmethod
672 def wait_for_response_handle_ready(
@@ -703,12 +729,41 @@ class WorkerProc:
729 return cast(list[WorkerProcHandle], ready_proc_handles)
730
731 def shutdown(self):
732 + if self.rpc_broadcast_mq is not None:
733 + self.rpc_broadcast_mq.shutdown()
734 + if self.worker_response_mq is not None:
735 + self.worker_response_mq.shutdown()
736 self.worker.shutdown()
737 self.rpc_broadcast_mq = None
738 self.worker_response_mq = None
739 destroy_model_parallel()
740 destroy_distributed_environment()
741
742 + def monitor_death_pipe(self, death_pipe, shutdown_requested: threading.Event):
743 + if death_pipe is None:
744 + return
745 +
746 + def death_pipe_monitor(queues_to_shutdown: list[MessageQueue]):
747 + try:
748 + # This will block until parent process exits (pipe closes)
749 + death_pipe.recv()
750 + except EOFError:
751 + logger.info_once("Parent process exited, terminating worker queues")
752 + shutdown_requested.set()
753 + for mq in queues_to_shutdown:
754 + if mq is not None:
755 + mq.shutdown()
756 + except Exception as e:
757 + logger.warning("Death monitoring error: %s", e)
758 +
759 + # Pass queue references directly to avoid gc issues if passing self
760 + Thread(
761 + target=death_pipe_monitor,
762 + args=([self.rpc_broadcast_mq, self.worker_response_mq],),
763 + daemon=True,
764 + name="DeathPipeMonitor",
765 + ).start()
766 +
767 @staticmethod
768 def worker_main(*args, **kwargs):
769 """Worker initialization and execution loops.
@@ -717,12 +772,12 @@ class WorkerProc:
772 # Signal handler used for graceful termination.
773 # SystemExit exception is only raised once to allow this and worker
774 # processes to terminate without error
720 - shutdown_requested = False
775 + shutdown_requested = threading.Event()
776
777 def signal_handler(signum, frame):
778 nonlocal shutdown_requested
724 - if not shutdown_requested:
725 - shutdown_requested = True
779 + if not shutdown_requested.is_set():
780 + shutdown_requested.set()
781 logger.debug(
782 "WorkerProc handling signal %d, raising SystemExit", signum
783 )
@@ -733,33 +788,20 @@ class WorkerProc:
788 signal.signal(signal.SIGINT, signal_handler)
789
790 worker = None
736 - # tuple[Connection, Connection]
737 - reader, ready_writer = kwargs.pop("ready_pipe")
738 - death_pipe: Connection | None = kwargs.pop("death_pipe", None)
739 - shutdown_event = threading.Event()
740 - # Start death monitoring thread if death_pipe is provided
741 - if death_pipe is not None:
742 -
743 - def monitor_parent_death():
744 - try:
745 - # This will block until parent process exits (pipe closes)
746 - death_pipe.recv()
747 - except EOFError:
748 - # Parent process has exited, terminate this worker
749 - logger.info_once("Parent process exited, terminating worker")
750 - # Send signal to self to trigger clean shutdown
751 - shutdown_event.set()
752 - except Exception as e:
753 - logger.warning("Death monitoring error: %s", e)
754 -
755 - death_monitor = Thread(
756 - target=monitor_parent_death, daemon=True, name="WorkerDeathMonitor"
757 - )
758 - death_monitor.start()
791 + ready_writer = kwargs.pop("ready_pipe")
792 + death_pipe = kwargs.pop("death_pipe", None)
793 +
794 + # Close inherited pipes from parent (incl. other worker pipes)
795 + # Explicitly passing in existing pipes and closing them makes the pipe
796 + # behave when using fork. Otherwise, a hidden reference to the pipes
797 + # exist in the child process and prevents EOF closure.
798 + for fd in kwargs.pop("inherited_fds", []):
799 + try:
800 + os.close(fd)
801 + except Exception as e:
802 + logger.warning("Exception closing inherited connection: %s", e)
803
804 try:
761 - reader.close()
762 -
805 # Initialize tracer
806 rank = kwargs.get("rank", 0)
807 maybe_init_worker_tracer(
@@ -771,6 +813,8 @@ class WorkerProc:
813 worker = WorkerProc(*args, **kwargs)
814 assert worker.worker_response_mq is not None
815
816 + worker.monitor_death_pipe(death_pipe, shutdown_requested)
817 +
818 # Send READY once we know everything is loaded
819 ready_writer.send(
820 {
@@ -788,7 +832,7 @@ class WorkerProc:
832 ready_writer.close()
833 ready_writer = None
834
791 - worker.worker_busy_loop(cancel=shutdown_event)
835 + worker.worker_busy_loop()
836
837 except Exception:
838 # NOTE: if an Exception arises in busy_loop, we send
@@ -798,7 +842,7 @@ class WorkerProc:
842
843 if ready_writer is not None:
844 logger.exception("WorkerProc failed to start.")
801 - elif shutdown_event.is_set():
845 + elif shutdown_requested.is_set():
846 logger.info("WorkerProc shutting down.")
847 else:
848 logger.exception("WorkerProc failed.")
@@ -806,7 +850,7 @@ class WorkerProc:
850 # The parent sends a SIGTERM to all worker processes if
851 # any worker dies. Set this value so we don't re-throw
852 # SystemExit() to avoid zmq exceptions in __del__.
809 - shutdown_requested = True
853 + shutdown_requested.set()
854
855 except SystemExit as e:
856 # SystemExit is raised on SIGTERM or SIGKILL, which usually indicates that
@@ -859,12 +903,12 @@ class WorkerProc:
903 output = self.async_output_queue.get()
904 self.enqueue_output(output)
905
862 - def worker_busy_loop(self, cancel: threading.Event | None = None):
906 + def worker_busy_loop(self):
907 """Main busy loop for Multiprocessing Workers"""
908 assert self.rpc_broadcast_mq is not None
909 while True:
910 method, args, kwargs, output_rank = self.rpc_broadcast_mq.dequeue(
867 - cancel=cancel, indefinite=True
911 + indefinite=True
912 )
913 try:
914 if isinstance(method, str):