| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import copyreg |
| 4 | import functools |
| 5 | import io |
| 6 | import os |
| 7 | import pickle |
| 8 | import shutil |
| 9 | import sys |
| 10 | import threading |
| 11 | import time |
| 12 | from contextlib import contextmanager |
| 13 | from dataclasses import dataclass, field |
| 14 | from multiprocessing import shared_memory |
| 15 | from pickle import PickleBuffer |
| 16 | from typing import TYPE_CHECKING, Any, cast |
| 17 | from unittest.mock import patch |
| 18 | |
| 19 | import torch |
| 20 | import torch.distributed as dist |
| 21 | import zmq |
| 22 | from torch.distributed import ProcessGroup |
| 23 | from zmq import ( # type: ignore |
| 24 | IPV6, # type: ignore |
| 25 | PUB, |
| 26 | SUB, |
| 27 | SUBSCRIBE, |
| 28 | XPUB, |
| 29 | XPUB_VERBOSE, |
| 30 | Context, |
| 31 | ) |
| 32 | |
| 33 | import vllm.envs as envs |
| 34 | from vllm.distributed.utils import StatelessProcessGroup, sched_yield |
| 35 | from vllm.logger import init_logger |
| 36 | from vllm.platforms import current_platform |
| 37 | from vllm.utils.cpu_resource_utils import check_cgroup_memory_available |
| 38 | from vllm.utils.network_utils import ( |
| 39 | get_ip, |
| 40 | get_open_zmq_inproc_path, |
| 41 | get_open_zmq_ipc_path, |
| 42 | is_valid_ipv6_address, |
| 43 | ) |
| 44 | |
| 45 | logger = init_logger(__name__) |
| 46 | |
| 47 | |
| 48 | SPINLOOP_EXT_ENABLED = False |
| 49 | if envs.VLLM_USE_SPINLOOP_EXT: |
| 50 | try: |
| 51 | from vllm.spinloop import spinloop |
| 52 | |
| 53 | SPINLOOP_EXT_ENABLED = True |
| 54 | except ImportError: |
| 55 | logger.warning( |
| 56 | "spinloop extension could not be loaded, disabling VLLM_USE_SPINLOOP_EXT!" |
| 57 | ) |
| 58 | SPINLOOP_TIMEOUT_SECONDS = 0.1 |
| 59 | |
| 60 | if TYPE_CHECKING: |
| 61 | from _typeshed import SizedBuffer |
| 62 | |
| 63 | VLLM_RINGBUFFER_WARNING_INTERVAL = envs.VLLM_RINGBUFFER_WARNING_INTERVAL |
| 64 | # Cap on how long an idle reader parks before re-reading the authoritative SHM |
| 65 | # written-flag. Bounds lost-notify recovery latency to ~5s while the periodic |
| 66 | # wakeup stays negligible (one flag check per reader every 5s). |
| 67 | SHM_READER_RECHECK_INTERVAL_MS = 5000 |
| 68 | |
| 69 | |
| 70 | from_bytes_big = functools.partial(int.from_bytes, byteorder="big") |
| 71 | |
| 72 | |
| 73 | # Memory fence for cross-process shared memory visibility. |
| 74 | # Required for correct producer-consumer synchronization when using |
| 75 | # shared memory without locks. |
| 76 | _memory_fence_lock = threading.Lock() |
| 77 | |
| 78 | |
| 79 | def memory_fence(): |
| 80 | """Full memory barrier for shared memory synchronization. |
| 81 | |
| 82 | Ensures all prior memory writes are visible to other processes before |
| 83 | any subsequent reads. This is critical for lock-free producer-consumer |
| 84 | patterns using shared memory. |
| 85 | |
| 86 | Implementation acquires and immediately releases a lock. Python's |
| 87 | threading.Lock provides sequentially consistent memory barrier semantics |
| 88 | across all major platforms (POSIX, Windows). This is a lightweight |
| 89 | operation (~20ns) that guarantees: |
| 90 | - All stores before the barrier are visible to other threads/processes |
| 91 | - All loads after the barrier see the latest values |
| 92 | """ |
| 93 | # Lock acquire/release provides full memory barrier semantics. |
| 94 | # Using context manager ensures lock release even on exceptions. |
| 95 | with _memory_fence_lock: |
| 96 | pass |
| 97 | |
| 98 | |
| 99 | def to_bytes_big(value: int, size: int) -> bytes: |
| 100 | return value.to_bytes(size, byteorder="big") |
| 101 | |
| 102 | |
| 103 | LONG_WAIT_TIME_LOG_MSG = ( |
| 104 | "No available shared memory broadcast block found " |
| 105 | "in %d seconds. This typically happens " |
| 106 | "when some processes are hanging or doing some " |
| 107 | "time-consuming work (e.g. compilation, " |
| 108 | "weight/kv cache quantization)." |
| 109 | ) |
| 110 | |
| 111 | |
| 112 | class SpinCondition: |
| 113 | """This class implements an interface similar to a threading.Condition. It |
| 114 | allows a writer to notify readers to wake up and read from the shared memory |
| 115 | buffer. This notification is done over a zmq socket. |
| 116 | |
| 117 | For optimal performance under load we don't want the readers to need to poll |
| 118 | the zmq socket for every read. So the `wait` method here will return |
| 119 | immediately when reads are frequent, and will only enter "idle mode" and |
| 120 | await a notification on the zmq socket after a period of inactivity. This |
| 121 | allows the readers to spin quickly, hence "SpinCondition". |
| 122 | |
| 123 | To support clean shutdown, a separate thread in the reader's process must be |
| 124 | able to wake the reader so that it can exit. A separate cancel() method is |
| 125 | implemented with an in-process socket to allow this interruption. |
| 126 | """ |
| 127 | |
| 128 | def __init__( |
| 129 | self, |
| 130 | is_reader: bool, |
| 131 | context: zmq.Context, |
| 132 | notify_address: str, |
| 133 | busy_loop_s: float = 1, |
| 134 | ): |
| 135 | self.is_reader = is_reader |
| 136 | |
| 137 | if is_reader: |
| 138 | # Time of last shm buffer read |
| 139 | self.last_read = time.monotonic() |
| 140 | |
| 141 | # Time to keep busy-looping on the shm buffer before going idle |
| 142 | self.busy_loop_s = busy_loop_s |
| 143 | |
| 144 | # Readers subscribe to write notifications |
| 145 | self.local_notify_socket: zmq.Socket = context.socket(SUB) |
| 146 | # Set zmq.CONFLATE to only keep the last message that the socket |
| 147 | # receives. This prevents us from piling up notification messages |
| 148 | # under high load when we aren't polling the socket. |
| 149 | self.local_notify_socket.setsockopt(zmq.CONFLATE, 1) |
| 150 | # Subscribe to all messages on the socket |
| 151 | self.local_notify_socket.setsockopt_string(SUBSCRIBE, "") |
| 152 | self.local_notify_socket.connect(notify_address) |
| 153 | |
| 154 | # Readers require a process-local socket to poll for cancellation |
| 155 | cancel_path = get_open_zmq_inproc_path() |
| 156 | self.write_cancel_socket: zmq.Socket = context.socket(zmq.PAIR) |
| 157 | self.write_cancel_socket.bind(cancel_path) |
| 158 | self.read_cancel_socket: zmq.Socket = context.socket(zmq.PAIR) |
| 159 | self.read_cancel_socket.connect(cancel_path) |
| 160 | |
| 161 | # Poller allows waiting on either `.notify()` or `.cancel()` |
| 162 | self.poller = zmq.Poller() |
| 163 | self.poller.register(self.read_cancel_socket, zmq.POLLIN) |
| 164 | self.poller.register(self.local_notify_socket, zmq.POLLIN) |
| 165 | else: |
| 166 | # Writer side publishes write notifications |
| 167 | self.local_notify_socket: zmq.Socket = context.socket(PUB) # type: ignore |
| 168 | # Set high water mark to 1 - we don't need to send a massive amount of |
| 169 | # pings during busy operation. PUB sockets will silently drop subsequent |
| 170 | # messages after the high water mark is reached. |
| 171 | self.local_notify_socket.setsockopt(zmq.SNDHWM, 1) |
| 172 | self.local_notify_socket.bind(notify_address) |
| 173 | |
| 174 | self.last_read = 0 |
| 175 | self.busy_loop_s = 0 |
| 176 | self.read_cancel_socket = None |
| 177 | self.write_cancel_socket = None |
| 178 | self.poller = None |
| 179 | |
| 180 | def record_read(self): |
| 181 | self.last_read = time.monotonic() |
| 182 | |
| 183 | def cancel(self): |
| 184 | # Sends cancellation ping that will cause the reader to wake up. |
| 185 | # This is done from a monitor thread in the same process as the reader. |
| 186 | if self.is_reader: |
| 187 | logger.debug("Canceling waiting reads on SHM Buffer") |
| 188 | self.write_cancel_socket.send(b"\x00") |
| 189 | |
| 190 | def wait(self, timeout_ms: int | None = None) -> None: |
| 191 | """Wait for data on the shared memory buffer. |
| 192 | |
| 193 | Yields the scheduler then returns immediately if it has been less than |
| 194 | self.busy_loop_s since the last read. |
| 195 | |
| 196 | Otherwise, enters idle mode and awaits a socket ping for at most |
| 197 | `timeout_ms` milliseconds, or indefinitely if timeout_ms is None. |
| 198 | """ |
| 199 | assert self.is_reader, "Only readers can wait" |
| 200 | |
| 201 | current_time = time.monotonic() |
| 202 | if current_time <= self.last_read + self.busy_loop_s: |
| 203 | sched_yield() |
| 204 | else: |
| 205 | events = dict(self.poller.poll(timeout=timeout_ms)) |
| 206 | |
| 207 | if self.read_cancel_socket in events: |
| 208 | logger.debug("Poller received cancel event") |
| 209 | elif self.local_notify_socket in events: |
| 210 | logger.debug("Poller received notify event") |
| 211 | # Since zmq.CONFLATE is set, there will only be one notification |
| 212 | # to read from the socket |
| 213 | self.local_notify_socket.recv(flags=zmq.NOBLOCK, copy=False) |
| 214 | else: |
| 215 | logger.debug("Poller timed out") |
| 216 | |
| 217 | def notify(self): |
| 218 | """Notifies all readers to wake up""" |
| 219 | assert not self.is_reader, "Only writers can notify" |
| 220 | self.local_notify_socket.send(b"\x00") |
| 221 | |
| 222 | |
| 223 | SHM_PATH = "/dev/shm" |
| 224 | |
| 225 | |
| 226 | def check_shm_free_space( |
| 227 | required_bytes: int, |
| 228 | shm_path: str = SHM_PATH, |
| 229 | *, |
| 230 | allocation_name: str = "shared-memory allocation", |
| 231 | ) -> None: |
| 232 | """Raise if SHM cannot fit a shared segment and log cgroup headroom. |
| 233 | |
| 234 | Args: |
| 235 | required_bytes: Size of the shared-memory segment to be created. |
| 236 | shm_path: Mount point backing POSIX shared memory; its filesystem |
| 237 | check is skipped if absent. |
| 238 | allocation_name: Human-readable name used in errors and logs. |
| 239 | |
| 240 | Raises: |
| 241 | RuntimeError: If the SHM filesystem has insufficient space. |
| 242 | |
| 243 | """ |
| 244 | if os.path.isdir(shm_path): |
| 245 | free_bytes = shutil.disk_usage(shm_path).free |
| 246 | if required_bytes > free_bytes: |
| 247 | mib = 1 << 20 |
| 248 | raise RuntimeError( |
| 249 | f"Insufficient space in {shm_path} for {allocation_name}: " |
| 250 | f"{required_bytes / mib:.0f} MiB required, " |
| 251 | f"{free_bytes / mib:.0f} MiB free. Increase {shm_path} " |
| 252 | "(e.g. --shm-size or --ipc=host)." |
| 253 | ) |
| 254 | |
| 255 | check_cgroup_memory_available( |
| 256 | required_bytes, |
| 257 | allocation_name, |
| 258 | ) |
| 259 | |
| 260 | |
| 261 | class ShmRingBuffer: |
| 262 | def __init__( |
| 263 | self, |
| 264 | n_reader: int, |
| 265 | max_chunk_bytes: int, |
| 266 | max_chunks: int, |
| 267 | name: str | None = None, |
| 268 | ): |
| 269 | """ |
| 270 | A shared memory ring buffer implementation for broadcast communication. |
| 271 | Essentially, it is a queue where only one will `enqueue` and multiple |
| 272 | will `dequeue`. The max size of each item, together with the max number |
| 273 | of items that can be stored in the buffer are known in advance. |
| 274 | In this case, we don't need to synchronize the access to |
| 275 | the buffer. |
| 276 | |
| 277 | Buffer memory layout: |
| 278 | data metadata |
| 279 | | | |
| 280 | | (current_idx) | (current_idx) |
| 281 | v v |
| 282 | +-------------------------------+----------------------------------------+ |
| 283 | | chunk0 | chunk1 | ... | chunk | metadata0 | metadata1 | ... | metadata | |
| 284 | +-------------------------------+----------------------------------------+ |
| 285 | | max_chunks x max_chunk_bytes | max_chunks x (1 + n_reader) bytes | |
| 286 | |
| 287 | metadata memory layout: each byte is a flag, the first byte is the written |
| 288 | flag, and the rest are reader flags. The flags are set to 0 by default. |
| 289 | +--------------+--------------+--------------+-----+--------------+ |
| 290 | | written_flag | reader0_flag | reader1_flag | ... | readerN_flag | |
| 291 | +--------------+--------------+--------------+-----+--------------+ |
| 292 | |
| 293 | The state of metadata is as follows: |
| 294 | |
| 295 | (case 1) 0???...???: the block is not written yet, cannot read, can write |
| 296 | (case 2) 1000...000: the block is just written, can read, cannot write |
| 297 | (case 3) 1???...???: the block is written and read by some readers, can read if not read, cannot write |
| 298 | (case 4) 1111...111: the block is written and read by all readers, cannot read, can write |
| 299 | |
| 300 | State transition for readers: |
| 301 | |
| 302 | When a reader finds a block that it can read (case 2 or 3), it can yield the block for caller to read. |
| 303 | Only after the caller finishes reading the block, the reader can mark the block as read. |
| 304 | Readers only mark the block as read (from 0 to 1), the writer marks the block as ready to read (from 1 to 0). |
| 305 | |
| 306 | State transition for writer: |
| 307 | |
| 308 | When the writer writes to a block (case 1 or 4), it first resets the written flag to 0, converting either case |
| 309 | to case 1. Then it can yield the block for caller to write. After the caller finishes writing the block, the writer |
| 310 | can reset the reader flags to 0, and mark the block as written (from 0 to 1). |
| 311 | NOTE: the order is important here, first reset the reader flags (so that we are still in case 1), then mark the block as written. The state transition is atomic. If we do it in the reverse order, it will go through case 3 and then back to case 2, and readers might read the intermediate case 3, which is not correct. |
| 312 | |
| 313 | During creation, `name` is None and the buffer is created. We can pass the |
| 314 | created object to other processes by pickling it. The other processes will |
| 315 | get the name of the shared memory and open it, so that they can access the |
| 316 | same shared memory buffer. |
| 317 | """ # noqa |
| 318 | self.n_reader = n_reader |
| 319 | self.metadata_size = 1 + n_reader |
| 320 | self.max_chunk_bytes = max_chunk_bytes |
| 321 | self.max_chunks = max_chunks |
| 322 | self.total_bytes_of_buffer = ( |
| 323 | self.max_chunk_bytes + self.metadata_size |
| 324 | ) * self.max_chunks |
| 325 | self.data_offset = 0 |
| 326 | self.metadata_offset = self.max_chunk_bytes * self.max_chunks |
| 327 | |
| 328 | if name is None: |
| 329 | # we are creating a buffer |
| 330 | self.is_creator = True |
| 331 | check_shm_free_space(self.total_bytes_of_buffer) |
| 332 | self.shared_memory = shared_memory.SharedMemory( |
| 333 | create=True, size=self.total_bytes_of_buffer |
| 334 | ) |
| 335 | assert self.shared_memory.buf is not None, "Buffer was not created" |
| 336 | # initialize the metadata section to 0 |
| 337 | with self.shared_memory.buf[self.metadata_offset :] as metadata_buffer: |
| 338 | torch.frombuffer(metadata_buffer, dtype=torch.uint8).fill_(0) |
| 339 | else: |
| 340 | # we are opening an existing buffer |
| 341 | self.is_creator = False |
| 342 | # fix to https://stackoverflow.com/q/62748654/9191338 |
| 343 | # Python incorrectly tracks shared memory even if it is not |
| 344 | # created by the process. The following patch is a workaround. |
| 345 | with patch( |
| 346 | "multiprocessing.resource_tracker.register", |
| 347 | lambda *args, **kwargs: None, |
| 348 | ): |
| 349 | try: |
| 350 | self.shared_memory = shared_memory.SharedMemory(name=name) |
| 351 | # See https://docs.python.org/3/library/multiprocessing.shared_memory.html # noqa |
| 352 | # Some platforms allocate memory based on page size, |
| 353 | # so the shared memory block size may be larger or equal |
| 354 | # to the requested size. The size parameter is ignored |
| 355 | # when attaching to an existing block. |
| 356 | assert self.shared_memory.size >= self.total_bytes_of_buffer |
| 357 | except FileNotFoundError: |
| 358 | # we might deserialize the object in a different node |
| 359 | # in this case, this object is not used, |
| 360 | # and we should suppress the error |
| 361 | pass |
| 362 | |
| 363 | def handle(self): |
| 364 | return ( |
| 365 | self.n_reader, |
| 366 | self.max_chunk_bytes, |
| 367 | self.max_chunks, |
| 368 | self.shared_memory.name, |
| 369 | ) |
| 370 | |
| 371 | def __reduce__(self): |
| 372 | return ( |
| 373 | self.__class__, |
| 374 | self.handle(), |
| 375 | ) |
| 376 | |
| 377 | def __del__(self): |
| 378 | if hasattr(self, "shared_memory"): |
| 379 | self.shared_memory.close() |
| 380 | if self.is_creator: |
| 381 | self.shared_memory.unlink() |
| 382 | |
| 383 | @contextmanager |
| 384 | def get_data(self, current_idx: int): |
| 385 | start = self.data_offset + current_idx * self.max_chunk_bytes |
| 386 | end = start + self.max_chunk_bytes |
| 387 | assert self.shared_memory.buf is not None, "Buffer has been closed" |
| 388 | with self.shared_memory.buf[start:end] as buf: |
| 389 | yield buf |
| 390 | |
| 391 | @contextmanager |
| 392 | def get_metadata(self, current_idx: int): |
| 393 | start = self.metadata_offset + current_idx * self.metadata_size |
| 394 | end = start + self.metadata_size |
| 395 | assert self.shared_memory.buf is not None, "Buffer has been closed" |
| 396 | with self.shared_memory.buf[start:end] as buf: |
| 397 | yield buf |
| 398 | |
| 399 | |
| 400 | def _rebuild_tensor(buf: Any, shape: tuple[int, ...], dtype_str: str) -> torch.Tensor: |
| 401 | """Rebuild a tensor from an out-of-band pickle buffer. |
| 402 | |
| 403 | Counterpart of `_reduce_tensor`. Note that pickle passes the original |
| 404 | buffer-providing object from `loads(buffers=...)` straight to this |
| 405 | function (no `PickleBuffer` wrapper on the receiving side), so `buf` is |
| 406 | a `zmq.Frame`, a `memoryview` of a shared-memory ring chunk, or `bytes` |
| 407 | if the buffer was serialized in-band. |
| 408 | """ |
| 409 | dtype = getattr(torch, dtype_str) |
| 410 | assert isinstance(dtype, torch.dtype) |
| 411 | if isinstance(buf, zmq.Frame): |
| 412 | # ZMQ frames own their message memory independently of any context, |
| 413 | # so the tensor can safely alias it with zero copies. The tensor's |
| 414 | # storage keeps the frame (and thus its bytes) alive via a strong |
| 415 | # reference for as long as the tensor is. |
| 416 | try: |
| 417 | return torch.frombuffer(buf, dtype=torch.uint8).view(dtype).view(shape) |
| 418 | except ValueError: |
| 419 | # Empty or read-only frame buffer; fall through to the copy path. |
| 420 | pass |
| 421 | # Shared-memory ring buffer chunks are reused by the writer once all |
| 422 | # readers have marked them read, so we must copy out of them. bytearray |
| 423 | # (vs bytes) keeps the resulting tensor writable, matching normal tensor |
| 424 | # semantics. |
| 425 | raw = bytearray(buf) |
| 426 | if not raw: |
| 427 | assert 0 in shape |
| 428 | return torch.empty(shape, dtype=dtype) |
| 429 | return torch.frombuffer(raw, dtype=torch.uint8).view(dtype).view(shape) |
| 430 | |
| 431 | |
| 432 | def _reduce_tensor(tensor: torch.Tensor): |
| 433 | """Reduce a CPU tensor to a `PickleBuffer` for out-of-band pickling. |
| 434 | |
| 435 | `torch.Tensor.__reduce_ex__` copies the tensor bytes into the pickle |
| 436 | byte stream via `torch.serialization` and never emits a `PickleBuffer`, |
| 437 | which defeats the out-of-band buffer handling in `MessageQueue.enqueue`. |
| 438 | This reducer instead exposes the tensor's memory directly, so large |
| 439 | tensors (e.g. `prompt_embeds` in `SchedulerOutput`) traverse the queue |
| 440 | without being copied into and back out of the pickled message. |
| 441 | """ |
| 442 | if ( |
| 443 | tensor.device.type == "cpu" |
| 444 | and tensor.layout == torch.strided |
| 445 | and not tensor.requires_grad |
| 446 | ): |
| 447 | try: |
| 448 | # The uint8 view exposes the raw bytes via the buffer protocol, |
| 449 | # including for dtypes numpy doesn't recognize (bfloat16, fp8, ...). |
| 450 | # reshape(-1) first so that 0-dim tensors can be viewed as well. |
| 451 | raw = tensor.contiguous().reshape(-1).view(torch.uint8).numpy() |
| 452 | except RuntimeError: |
| 453 | # Exotic tensors (e.g. with the conjugate bit set) that don't |
| 454 | # support aliasing views; let torch handle them. |
| 455 | pass |
| 456 | else: |
| 457 | dtype_str = str(tensor.dtype).removeprefix("torch.") |
| 458 | return _rebuild_tensor, (PickleBuffer(raw), tuple(tensor.shape), dtype_str) |
| 459 | |
| 460 | # Fall back to torch's default (copying) reduction. |
| 461 | return tensor.__reduce_ex__(pickle.HIGHEST_PROTOCOL) |
| 462 | |
| 463 | |
| 464 | @dataclass |
| 465 | class Handle: |
| 466 | local_reader_ranks: list[int] = field(default_factory=list) |
| 467 | |
| 468 | buffer_handle: tuple[int, int, int, str] | None = None |
| 469 | local_subscribe_addr: str | None = None |
| 470 | local_notify_addr: str | None = None |
| 471 | remote_subscribe_addr: str | None = None |
| 472 | remote_addr_ipv6: bool = False |
| 473 | |
| 474 | |
| 475 | class MessageQueue: |
| 476 | def __init__( |
| 477 | self, |
| 478 | n_reader, # number of all readers |
| 479 | n_local_reader, # number of local readers through shared memory |
| 480 | local_reader_ranks: list[int] | None = None, |
| 481 | # Default of 24MiB chosen to be large enough to accommodate grammar |
| 482 | # bitmask tensors for large batches (1024 requests). |
| 483 | max_chunk_bytes: int = 1024 * 1024 * 24, |
| 484 | max_chunks: int = 10, |
| 485 | connect_ip: str | None = None, |
| 486 | ): |
| 487 | if local_reader_ranks is None: |
| 488 | local_reader_ranks = list(range(n_local_reader)) |
| 489 | else: |
| 490 | assert len(local_reader_ranks) == n_local_reader |
| 491 | self.n_local_reader = n_local_reader |
| 492 | n_remote_reader = n_reader - n_local_reader |
| 493 | self.n_remote_reader = n_remote_reader |
| 494 | self.shutting_down = False |
| 495 | context = Context() |
| 496 | |
| 497 | if n_local_reader > 0: |
| 498 | # for local readers, we will: |
| 499 | # 1. create a shared memory ring buffer to communicate small data |
| 500 | # 2. create a publish-subscribe socket to communicate large data |
| 501 | self.buffer = ShmRingBuffer(n_local_reader, max_chunk_bytes, max_chunks) |
| 502 | |
| 503 | # XPUB is very similar to PUB, |
| 504 | # except that it can receive subscription messages |
| 505 | # to confirm the number of subscribers |
| 506 | self.local_socket = context.socket(XPUB) |
| 507 | # set the verbose option so that we can receive every subscription |
| 508 | # message. otherwise, we will only receive the first subscription |
| 509 | # see http://api.zeromq.org/3-3:zmq-setsockopt for more details |
| 510 | self.local_socket.setsockopt(XPUB_VERBOSE, True) |
| 511 | local_subscribe_addr = get_open_zmq_ipc_path() |
| 512 | logger.debug("Binding to %s", local_subscribe_addr) |
| 513 | self.local_socket.bind(local_subscribe_addr) |
| 514 | |
| 515 | self.current_idx = 0 |
| 516 | |
| 517 | # Create the notification side of the SpinCondition |
| 518 | local_notify_addr = get_open_zmq_ipc_path() |
| 519 | self._spin_condition = SpinCondition( |
| 520 | is_reader=False, context=context, notify_address=local_notify_addr |
| 521 | ) |
| 522 | else: |
| 523 | self.buffer = None # type: ignore |
| 524 | local_subscribe_addr = None |
| 525 | self.local_socket = None |
| 526 | self.current_idx = -1 |
| 527 | local_notify_addr = None |
| 528 | self._spin_condition = None # type: ignore |
| 529 | |
| 530 | remote_addr_ipv6 = False |
| 531 | if n_remote_reader > 0: |
| 532 | # for remote readers, we will: |
| 533 | # create a publish-subscribe socket to communicate large data |
| 534 | if not connect_ip: |
| 535 | connect_ip = get_ip() |
| 536 | self.remote_socket = context.socket(XPUB) |
| 537 | self.remote_socket.setsockopt(XPUB_VERBOSE, True) |
| 538 | if is_valid_ipv6_address(connect_ip): |
| 539 | self.remote_socket.setsockopt(IPV6, 1) |
| 540 | remote_addr_ipv6 = True |
| 541 | connect_ip = f"[{connect_ip}]" |
| 542 | self.remote_socket.bind(f"tcp://{connect_ip}:0") |
| 543 | last_endpoint = self.remote_socket.getsockopt(zmq.LAST_ENDPOINT) |
| 544 | remote_subscribe_port = last_endpoint.decode().rsplit(":", 1)[1] |
| 545 | remote_subscribe_addr = f"tcp://{connect_ip}:{remote_subscribe_port}" |
| 546 | else: |
| 547 | remote_subscribe_addr = None |
| 548 | self.remote_socket = None |
| 549 | |
| 550 | self._is_writer = True |
| 551 | self._is_local_reader = False |
| 552 | self.local_reader_rank = -1 |
| 553 | # rank does not matter for remote readers |
| 554 | self._is_remote_reader = False |
| 555 | |
| 556 | self.handle = Handle( |
| 557 | local_reader_ranks=local_reader_ranks, |
| 558 | buffer_handle=self.buffer.handle() if self.buffer is not None else None, |
| 559 | local_subscribe_addr=local_subscribe_addr, |
| 560 | local_notify_addr=local_notify_addr, |
| 561 | remote_subscribe_addr=remote_subscribe_addr, |
| 562 | remote_addr_ipv6=remote_addr_ipv6, |
| 563 | ) |
| 564 | |
| 565 | logger.debug("vLLM message queue communication handle: %s", self.handle) |
| 566 | |
| 567 | def export_handle(self) -> Handle: |
| 568 | return self.handle |
| 569 | |
| 570 | @staticmethod |
| 571 | def create_from_handle(handle: Handle, rank) -> "MessageQueue": |
| 572 | self = MessageQueue.__new__(MessageQueue) |
| 573 | self.handle = handle |
| 574 | self._is_writer = False |
| 575 | |
| 576 | context = Context() |
| 577 | |
| 578 | if rank in handle.local_reader_ranks: |
| 579 | assert handle.buffer_handle is not None |
| 580 | self.buffer = ShmRingBuffer(*handle.buffer_handle) |
| 581 | self.current_idx = 0 |
| 582 | self.local_reader_rank = handle.local_reader_ranks.index(rank) |
| 583 | self._is_local_reader = True |
| 584 | self._is_remote_reader = False |
| 585 | |
| 586 | self.local_socket = context.socket(SUB) |
| 587 | self.local_socket.setsockopt_string(SUBSCRIBE, "") |
| 588 | socket_addr = handle.local_subscribe_addr |
| 589 | logger.debug("Connecting to %s", socket_addr) |
| 590 | self.local_socket.connect(socket_addr) |
| 591 | |
| 592 | self.remote_socket = None |
| 593 | assert isinstance(handle.local_notify_addr, str) |
| 594 | self._spin_condition = SpinCondition( |
| 595 | is_reader=True, context=context, notify_address=handle.local_notify_addr |
| 596 | ) |
| 597 | else: |
| 598 | self.buffer = None # type: ignore |
| 599 | self.current_idx = -1 |
| 600 | self.local_reader_rank = -1 |
| 601 | self._is_local_reader = False |
| 602 | self._is_remote_reader = True |
| 603 | |
| 604 | self.local_socket = None |
| 605 | |
| 606 | self.remote_socket = context.socket(SUB) |
| 607 | self.remote_socket.setsockopt_string(SUBSCRIBE, "") |
| 608 | if handle.remote_addr_ipv6: |
| 609 | self.remote_socket.setsockopt(IPV6, 1) |
| 610 | socket_addr = handle.remote_subscribe_addr |
| 611 | logger.debug("Connecting to %s", socket_addr) |
| 612 | self.remote_socket.connect(socket_addr) |
| 613 | self._spin_condition = None # type: ignore |
| 614 | |
| 615 | self.shutting_down = False |
| 616 | return self |
| 617 | |
| 618 | def wait_until_ready(self): |
| 619 | """This is a collective operation. All processes (including the |
| 620 | readers and the writer) should call this function. |
| 621 | """ |
| 622 | if self._is_writer: |
| 623 | # wait for all readers to connect |
| 624 | |
| 625 | # local readers |
| 626 | for i in range(self.n_local_reader): |
| 627 | # wait for subscription messages from all local readers |
| 628 | self.local_socket.recv() |
| 629 | if self.n_local_reader > 0: |
| 630 | # send a message to all local readers |
| 631 | # to make sure the publish channel is working |
| 632 | self.local_socket.send(b"READY") |
| 633 | |
| 634 | # remote readers |
| 635 | for i in range(self.n_remote_reader): |
| 636 | # wait for subscription messages from all remote readers |
| 637 | self.remote_socket.recv() |
| 638 | if self.n_remote_reader > 0: |
| 639 | # send a message to all remote readers |
| 640 | # to make sure the publish channel is working |
| 641 | self.remote_socket.send(b"READY") |
| 642 | elif self._is_local_reader: |
| 643 | # wait for the writer to send a message |
| 644 | recv = self.local_socket.recv() |
| 645 | assert recv == b"READY" |
| 646 | elif self._is_remote_reader: |
| 647 | # wait for the writer to send a message |
| 648 | recv = self.remote_socket.recv() |
| 649 | assert recv == b"READY" |
| 650 | |
| 651 | def shutdown(self): |
| 652 | """If this is an idle reader, wakes it up so it can clean up and shut |
| 653 | down""" |
| 654 | self.shutting_down = True |
| 655 | if self._spin_condition is not None: |
| 656 | self._spin_condition.cancel() |
| 657 | |
| 658 | @contextmanager |
| 659 | def acquire_write(self, timeout: float | None = None): |
| 660 | assert self._is_writer, "Only writers can acquire write" |
| 661 | start_time = time.monotonic() |
| 662 | n_warning = 1 |
| 663 | while True: |
| 664 | with self.buffer.get_metadata(self.current_idx) as metadata_buffer: |
| 665 | |
| 666 | def check(): |
| 667 | memory_fence() |
| 668 | read_count = sum(metadata_buffer[1:]) |
| 669 | written_flag = metadata_buffer[0] |
| 670 | return not (written_flag and read_count != self.buffer.n_reader) |
| 671 | |
| 672 | if SPINLOOP_EXT_ENABLED and not check(): |
| 673 | spinloop(metadata_buffer, check, timeout=SPINLOOP_TIMEOUT_SECONDS) |
| 674 | |
| 675 | if not check(): |
| 676 | # this block is written and not read by all readers |
| 677 | # for writers, `self.current_idx` is the next block to write |
| 678 | # if this block is not ready to write, |
| 679 | # we need to wait until it is read by all readers |
| 680 | |
| 681 | # Release the processor to other threads |
| 682 | sched_yield() |
| 683 | |
| 684 | # if we time out, raise an exception |
| 685 | elapsed = time.monotonic() - start_time |
| 686 | if timeout is not None and elapsed > timeout: |
| 687 | raise TimeoutError |
| 688 | |
| 689 | # if we wait for a long time, log a message |
| 690 | if elapsed > VLLM_RINGBUFFER_WARNING_INTERVAL * n_warning: |
| 691 | logger.info( |
| 692 | LONG_WAIT_TIME_LOG_MSG, VLLM_RINGBUFFER_WARNING_INTERVAL |
| 693 | ) |
| 694 | n_warning += 1 |
| 695 | |
| 696 | continue |
| 697 | # found a block that is either |
| 698 | # (1) not written |
| 699 | # (2) read by all readers |
| 700 | |
| 701 | # mark the block as not written |
| 702 | metadata_buffer[0] = 0 |
| 703 | # let caller write to the buffer |
| 704 | with self.buffer.get_data(self.current_idx) as buf: |
| 705 | yield buf |
| 706 | |
| 707 | # caller has written to the buffer |
| 708 | # NOTE: order is important here |
| 709 | # first set the read flags to 0 |
| 710 | # then set the written flag to 1 |
| 711 | # otherwise, the readers may think they already read the block |
| 712 | for i in range(1, self.buffer.n_reader + 1): |
| 713 | # set read flag to 0, meaning it is not read yet |
| 714 | metadata_buffer[i] = 0 |
| 715 | # Memory fence here ensures the order of the buffer and flag |
| 716 | # writes. This guarantees that when `metadata_buffer[0] = 1` is |
| 717 | # visible to readers, `buf` can be completely ready. Without |
| 718 | # this, some CPU architectures with weak ordering may incur |
| 719 | # memory inconsistency. |
| 720 | memory_fence() |
| 721 | # mark the block as written |
| 722 | metadata_buffer[0] = 1 |
| 723 | # Memory fence ensures the write is visible to readers on other cores |
| 724 | # before we proceed. Without this, readers may spin indefinitely |
| 725 | # waiting for a write that's stuck in our CPU's store buffer. |
| 726 | memory_fence() |
| 727 | self.current_idx = (self.current_idx + 1) % self.buffer.max_chunks |
| 728 | break |
| 729 | |
| 730 | class ReadTimeoutWithWarnings: |
| 731 | def __init__(self, timeout: float | None, should_warn: bool) -> None: |
| 732 | self.started = time.monotonic() |
| 733 | self.deadline = sys.maxsize if timeout is None else self.started + timeout |
| 734 | |
| 735 | # if should_warn, we need to wake up periodically to log |
| 736 | self.warning_wait_time_ms: int | None = ( |
| 737 | VLLM_RINGBUFFER_WARNING_INTERVAL * 1000 if should_warn else None |
| 738 | ) |
| 739 | |
| 740 | self._should_warn = should_warn |
| 741 | self.n_warning = 1 |
| 742 | self.timeout = timeout |
| 743 | |
| 744 | def timeout_ms(self) -> int: |
| 745 | """Returns a timeout, capped at the recheck interval, that is: |
| 746 | - min(time to deadline, time to next warning) if we're logging warnings |
| 747 | - time to deadline, if we're not logging warnings |
| 748 | - recheck interval if the timeout is None and we're not logging warnings |
| 749 | - raise TimeoutError if we are past the deadline |
| 750 | """ |
| 751 | wait_ms = SHM_READER_RECHECK_INTERVAL_MS |
| 752 | if self.warning_wait_time_ms is not None: |
| 753 | wait_ms = min(wait_ms, self.warning_wait_time_ms) |
| 754 | if self.timeout is None: |
| 755 | return wait_ms |
| 756 | time_left_ms = int((self.deadline - time.monotonic()) * 1000) |
| 757 | if time_left_ms <= 0: |
| 758 | raise TimeoutError |
| 759 | return min(wait_ms, time_left_ms) |
| 760 | |
| 761 | def should_warn(self) -> bool: |
| 762 | """Returns true if it's time to log a warning for a timeout that is not |
| 763 | indefinite""" |
| 764 | if self._should_warn: |
| 765 | elapsed = time.monotonic() - self.started |
| 766 | if elapsed >= VLLM_RINGBUFFER_WARNING_INTERVAL * self.n_warning: |
| 767 | self.n_warning += 1 |
| 768 | return True |
| 769 | return False |
| 770 | |
| 771 | @contextmanager |
| 772 | def acquire_read( |
| 773 | self, |
| 774 | timeout: float | None = None, |
| 775 | indefinite: bool = False, |
| 776 | ): |
| 777 | assert self._is_local_reader, "Only readers can acquire read" |
| 778 | read_timeout = self.ReadTimeoutWithWarnings( |
| 779 | timeout=timeout, should_warn=not indefinite |
| 780 | ) |
| 781 | with self.buffer.get_metadata(self.current_idx) as metadata_buffer: |
| 782 | while True: |
| 783 | |
| 784 | def check(): |
| 785 | memory_fence() |
| 786 | read_flag = metadata_buffer[self.local_reader_rank + 1] |
| 787 | written_flag = metadata_buffer[0] |
| 788 | return not (not written_flag or read_flag) |
| 789 | |
| 790 | if SPINLOOP_EXT_ENABLED and not check(): |
| 791 | spinloop( |
| 792 | metadata_buffer[0 : self.local_reader_rank + 1], |
| 793 | check, |
| 794 | timeout=SPINLOOP_TIMEOUT_SECONDS, |
| 795 | ) |
| 796 | |
| 797 | if not check(): |
| 798 | # this block is either |
| 799 | # (1) not written |
| 800 | # (2) already read by this reader |
| 801 | |
| 802 | # for readers, `self.current_idx` is the next block to read |
| 803 | # if this block is not ready, |
| 804 | # we need to wait until it is written |
| 805 | self._spin_condition.wait(timeout_ms=read_timeout.timeout_ms()) |
| 806 | |
| 807 | if self.shutting_down: |
| 808 | raise RuntimeError("cancelled") |
| 809 | |
| 810 | # if we wait for a long time, log a message |
| 811 | if read_timeout.should_warn(): |
| 812 | logger.info( |
| 813 | LONG_WAIT_TIME_LOG_MSG, VLLM_RINGBUFFER_WARNING_INTERVAL |
| 814 | ) |
| 815 | |
| 816 | continue |
| 817 | # found a block that is not read by this reader |
| 818 | # let caller read from the buffer |
| 819 | with self.buffer.get_data(self.current_idx) as buf: |
| 820 | try: |
| 821 | yield buf |
| 822 | finally: |
| 823 | # caller has read from the buffer; set the read flag. |
| 824 | metadata_buffer[self.local_reader_rank + 1] = 1 |
| 825 | # Memory fence ensures the read flag is visible to the writer. |
| 826 | # Without this, writer may not see our read completion and |
| 827 | # could wait indefinitely for all readers to finish. |
| 828 | memory_fence() |
| 829 | next_idx = self.current_idx + 1 |
| 830 | self.current_idx = next_idx % self.buffer.max_chunks |
| 831 | self._spin_condition.record_read() |
| 832 | break |
| 833 | |
| 834 | def enqueue(self, obj, timeout: float | None = None): |
| 835 | """Write to message queue with optional timeout (in seconds)""" |
| 836 | assert self._is_writer, "Only writers can enqueue" |
| 837 | all_buffers: list[SizedBuffer] = [b""] |
| 838 | total_bytes = 6 # 2 bytes for oob buffer count, 4 for main buffer size |
| 839 | |
| 840 | def oob_callback(buf: PickleBuffer) -> bool: |
| 841 | raw_buf = buf.raw() |
| 842 | if len(raw_buf) < 1024 * 1024: |
| 843 | # In-line buffers smaller than 1MiB. |
| 844 | return True |
| 845 | all_buffers.append(raw_buf) |
| 846 | nonlocal total_bytes |
| 847 | total_bytes += len(raw_buf) + 4 |
| 848 | return False |
| 849 | |
| 850 | # CPU tensors are routed through `_reduce_tensor` so that their |
| 851 | # bytes are emitted as out-of-band buffers instead of being |
| 852 | # copied into the pickle stream by torch's default reducer. |
| 853 | # Start from `copyreg.dispatch_table` to preserve globally |
| 854 | # registered reducers (e.g. `re.Pattern`); the per-pickler |
| 855 | # dispatch table would otherwise shadow them. |
| 856 | dispatch_table = dict(copyreg.dispatch_table) |
| 857 | dispatch_table[torch.Tensor] = _reduce_tensor |
| 858 | with io.BytesIO() as bio: |
| 859 | pickler = pickle.Pickler( |
| 860 | bio, |
| 861 | protocol=pickle.HIGHEST_PROTOCOL, |
| 862 | buffer_callback=oob_callback, |
| 863 | ) |
| 864 | pickler.dispatch_table = dispatch_table |
| 865 | pickler.dump(obj) |
| 866 | all_buffers[0] = bio.getvalue() |
| 867 | if self.n_local_reader > 0: |
| 868 | if total_bytes + len(all_buffers[0]) >= self.buffer.max_chunk_bytes: |
| 869 | with self.acquire_write(timeout) as buf: |
| 870 | buf[0] = 1 # overflow |
| 871 | self.local_socket.send_multipart(all_buffers, copy=False) |
| 872 | else: |
| 873 | # Byte 0: 0 |
| 874 | # Bytes 1-2: Count of buffers |
| 875 | # Then each buffer follows, preceded by 4 bytes containing its length: |
| 876 | # [4 byte int L][L bytes of buffer content] ... |
| 877 | with self.acquire_write(timeout) as buf: |
| 878 | buf[0] = 0 # not overflow |
| 879 | offset = 3 |
| 880 | buf[1:offset] = to_bytes_big(len(all_buffers), 2) # oob buf count |
| 881 | for buffer in all_buffers: |
| 882 | buf_len = len(buffer) |
| 883 | # prepend each buffer with 4 bytes containing its size. |
| 884 | buf_offset = offset + 4 |
| 885 | buf[offset:buf_offset] = to_bytes_big(buf_len, 4) |
| 886 | buf[buf_offset : (offset := buf_offset + buf_len)] = buffer |
| 887 | |
| 888 | self._spin_condition.notify() |
| 889 | |
| 890 | if self.n_remote_reader > 0: |
| 891 | self.remote_socket.send_multipart(all_buffers, copy=False) |
| 892 | |
| 893 | def dequeue( |
| 894 | self, |
| 895 | timeout: float | None = None, |
| 896 | indefinite: bool = False, |
| 897 | ): |
| 898 | """Read from message queue with optional timeout (in seconds)""" |
| 899 | if self._is_local_reader: |
| 900 | with self.acquire_read(timeout, indefinite) as buf: |
| 901 | overflow = buf[0] == 1 |
| 902 | if not overflow: |
| 903 | offset = 3 |
| 904 | buf_count = from_bytes_big(buf[1:offset]) |
| 905 | all_buffers = [] |
| 906 | for i in range(buf_count): |
| 907 | buf_offset = offset + 4 |
| 908 | buf_len = from_bytes_big(buf[offset:buf_offset]) |
| 909 | offset = buf_offset + buf_len |
| 910 | all_buffers.append(buf[buf_offset:offset]) |
| 911 | obj = pickle.loads(all_buffers[0], buffers=all_buffers[1:]) |
| 912 | if overflow: |
| 913 | obj = MessageQueue.recv(self.local_socket, timeout) |
| 914 | elif self._is_remote_reader: |
| 915 | obj = MessageQueue.recv(self.remote_socket, timeout) |
| 916 | else: |
| 917 | raise RuntimeError("Only readers can dequeue") |
| 918 | return obj |
| 919 | |
| 920 | @staticmethod |
| 921 | def recv(socket: zmq.Socket, timeout: float | None) -> Any: |
| 922 | # Ensure non-negative timeout passed to zmq poll. |
| 923 | timeout_ms = None if timeout is None else max(0, int(timeout * 1000)) |
| 924 | if not socket.poll(timeout=timeout_ms): |
| 925 | raise TimeoutError |
| 926 | recv, *recv_oob = socket.recv_multipart(copy=False) |
| 927 | return pickle.loads(recv, buffers=recv_oob) |
| 928 | |
| 929 | def broadcast_object(self, obj=None): |
| 930 | if self._is_writer: |
| 931 | self.enqueue(obj) |
| 932 | return obj |
| 933 | return self.dequeue() |
| 934 | |
| 935 | @staticmethod |
| 936 | def create_from_process_group_single_reader( |
| 937 | pg: ProcessGroup, |
| 938 | max_chunk_bytes, |
| 939 | max_chunks, |
| 940 | reader_rank: int = 0, |
| 941 | blocking: bool = False, |
| 942 | ) -> tuple["MessageQueue", list[Handle]]: |
| 943 | """Creates a MessageQueue for a process group with a single reader. |
| 944 | |
| 945 | This method is designed for scenarios where only one process (the reader) |
| 946 | will consume messages, and all other processes are writers. It sets up |
| 947 | the shared memory buffer and communication handles accordingly, and |
| 948 | gathers the handles from all processes to the reader. |
| 949 | |
| 950 | Args: |
| 951 | pg (ProcessGroup): The torch distributed process group. |
| 952 | max_chunk_bytes (int): Maximum size in bytes for each chunk in the buffer. |
| 953 | max_chunks (int): Maximum number of chunks in the buffer. |
| 954 | reader_rank (int, optional): The global rank that will act as the reader. |
| 955 | Defaults to 0. |
| 956 | blocking (bool, optional): If True, blocks until all processes are ready. |
| 957 | Defaults to False. |
| 958 | |
| 959 | Returns: |
| 960 | tuple[MessageQueue, list[Handle]]: |
| 961 | The MessageQueue instance for the calling process, |
| 962 | and a list of handles (only non-empty for the reader process). |
| 963 | |
| 964 | """ |
| 965 | from vllm.platforms.interface import get_assigned_physical_gpu_ids |
| 966 | |
| 967 | assigned_physical_gpu_ids = get_assigned_physical_gpu_ids() |
| 968 | if assigned_physical_gpu_ids is not None: |
| 969 | local_size = len(assigned_physical_gpu_ids) |
| 970 | else: |
| 971 | local_size = current_platform.device_count() |
| 972 | rank = dist.get_rank() |
| 973 | same_node = rank // local_size == reader_rank // local_size |
| 974 | buffer_io = MessageQueue( |
| 975 | n_reader=1, |
| 976 | n_local_reader=1 if same_node else 0, |
| 977 | max_chunk_bytes=max_chunk_bytes, |
| 978 | max_chunks=max_chunks, |
| 979 | ) |
| 980 | handle = buffer_io.export_handle() |
| 981 | handles = [None] * dist.get_world_size(pg) if rank == reader_rank else None |
| 982 | dist.gather_object(handle, handles, dst=reader_rank, group=pg) |
| 983 | if blocking: |
| 984 | buffer_io.wait_until_ready() |
| 985 | return buffer_io, cast(list[Handle], handles or []) |
| 986 | |
| 987 | @staticmethod |
| 988 | def create_from_process_group( |
| 989 | pg: ProcessGroup | StatelessProcessGroup, |
| 990 | max_chunk_bytes, |
| 991 | max_chunks, |
| 992 | writer_rank: int = 0, |
| 993 | external_writer_handle=None, |
| 994 | blocking: bool = True, |
| 995 | ) -> "MessageQueue": |
| 996 | """Creates a MessageQueue for a distributed process group with one writer and |
| 997 | multiple readers. |
| 998 | |
| 999 | This method is designed for scenarios where one process (the writer) sends |
| 1000 | messages, and all other processes (the readers) receive messages. It sets up |
| 1001 | the shared memory buffer and socket communication handles accordingly, and |
| 1002 | broadcasts the handle from the writer to all readers. |
| 1003 | |
| 1004 | Args: |
| 1005 | pg (ProcessGroup | StatelessProcessGroup): The torch distributed process |
| 1006 | group. |
| 1007 | max_chunk_bytes (int): Maximum size in bytes for each chunk in the buffer. |
| 1008 | max_chunks (int): Maximum number of chunks in the buffer. |
| 1009 | writer_rank (int, optional): The global rank that will act as the writer. |
| 1010 | Defaults to 0. |
| 1011 | external_writer_handle (Handle, optional): Used when there is a handle |
| 1012 | from an external Message Queue. If provided, use this handle to init |
| 1013 | PG writer message queue instead of creating a new one. Defaults to None. |
| 1014 | blocking (bool, optional): If True, blocks until all processes are ready. |
| 1015 | Defaults to True. |
| 1016 | |
| 1017 | Returns: |
| 1018 | MessageQueue: The MessageQueue instance for the calling process. |
| 1019 | |
| 1020 | """ |
| 1021 | if isinstance(pg, ProcessGroup): |
| 1022 | group_rank = dist.get_rank(pg) |
| 1023 | group_world_size = dist.get_world_size(pg) |
| 1024 | global_ranks = dist.get_process_group_ranks(pg) |
| 1025 | else: |
| 1026 | group_rank = pg.rank |
| 1027 | group_world_size = pg.world_size |
| 1028 | global_ranks = list(range(pg.world_size)) |
| 1029 | from vllm.distributed.parallel_state import in_the_same_node_as |
| 1030 | |
| 1031 | status = in_the_same_node_as(pg, source_rank=writer_rank) |
| 1032 | if group_rank == writer_rank: |
| 1033 | if external_writer_handle is not None: |
| 1034 | buffer_io = MessageQueue.create_from_handle( |
| 1035 | external_writer_handle, group_rank |
| 1036 | ) |
| 1037 | else: |
| 1038 | same_node_ranks = [i for i, s in enumerate(status) if s] |
| 1039 | n_reader = group_world_size - 1 |
| 1040 | n_local_reader = len(same_node_ranks) - 1 |
| 1041 | local_reader_ranks = [i for i in same_node_ranks if i != writer_rank] |
| 1042 | buffer_io = MessageQueue( |
| 1043 | n_reader=n_reader, |
| 1044 | n_local_reader=n_local_reader, |
| 1045 | local_reader_ranks=local_reader_ranks, |
| 1046 | max_chunk_bytes=max_chunk_bytes, |
| 1047 | max_chunks=max_chunks, |
| 1048 | ) |
| 1049 | handle = buffer_io.export_handle() |
| 1050 | if isinstance(pg, ProcessGroup): |
| 1051 | dist.broadcast_object_list( |
| 1052 | [handle], src=global_ranks[writer_rank], group=pg |
| 1053 | ) |
| 1054 | else: |
| 1055 | pg.broadcast_obj(handle, writer_rank) |
| 1056 | else: |
| 1057 | if isinstance(pg, ProcessGroup): |
| 1058 | recv = [None] |
| 1059 | dist.broadcast_object_list( |
| 1060 | recv, src=global_ranks[writer_rank], group=pg |
| 1061 | ) |
| 1062 | handle = recv[0] # type: ignore |
| 1063 | else: |
| 1064 | handle = pg.broadcast_obj(None, writer_rank) |
| 1065 | buffer_io = MessageQueue.create_from_handle(handle, group_rank) |
| 1066 | if blocking: |
| 1067 | buffer_io.wait_until_ready() |
| 1068 | return buffer_io |