| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import argparse |
| 4 | import contextlib |
| 5 | import json |
| 6 | import multiprocessing |
| 7 | import threading |
| 8 | import time |
| 9 | import weakref |
| 10 | from collections.abc import Callable, Sequence |
| 11 | from contextlib import AbstractContextManager |
| 12 | from dataclasses import dataclass |
| 13 | from multiprocessing import connection |
| 14 | from multiprocessing.process import BaseProcess |
| 15 | from multiprocessing.queues import Queue |
| 16 | from typing import ( |
| 17 | TYPE_CHECKING, |
| 18 | Any, |
| 19 | Generic, |
| 20 | TypeVar, |
| 21 | Union, |
| 22 | overload, |
| 23 | ) |
| 24 | |
| 25 | import torch |
| 26 | import uvloop |
| 27 | from torch.autograd.profiler import record_function |
| 28 | |
| 29 | import vllm.envs as envs |
| 30 | from vllm.logger import init_logger |
| 31 | from vllm.usage.usage_lib import UsageContext, is_usage_stats_enabled, usage_message |
| 32 | from vllm.utils.network_utils import get_open_zmq_ipc_path, get_tcp_uri |
| 33 | from vllm.utils.system_utils import decorate_logs, kill_process_tree, set_process_title |
| 34 | from vllm.utils.torch_utils import PIN_MEMORY |
| 35 | from vllm.v1.core.sched.output import SchedulerOutput |
| 36 | |
| 37 | if TYPE_CHECKING: |
| 38 | import numpy as np |
| 39 | |
| 40 | from vllm.v1.engine.coordinator import DPCoordinator |
| 41 | from vllm.v1.engine.utils import CoreEngineActorManager, CoreEngineProcManager |
| 42 | |
| 43 | logger = init_logger(__name__) |
| 44 | |
| 45 | T = TypeVar("T") |
| 46 | |
| 47 | |
| 48 | class ConstantList(Generic[T], Sequence): |
| 49 | def __init__(self, x: list[T]) -> None: |
| 50 | self._x = x |
| 51 | |
| 52 | def append(self, item): |
| 53 | raise TypeError("Cannot append to a constant list") |
| 54 | |
| 55 | def extend(self, item): |
| 56 | raise TypeError("Cannot extend a constant list") |
| 57 | |
| 58 | def insert(self, item): |
| 59 | raise TypeError("Cannot insert into a constant list") |
| 60 | |
| 61 | def pop(self, item): |
| 62 | raise TypeError("Cannot pop from a constant list") |
| 63 | |
| 64 | def remove(self, item): |
| 65 | raise TypeError("Cannot remove from a constant list") |
| 66 | |
| 67 | def clear(self): |
| 68 | raise TypeError("Cannot clear a constant list") |
| 69 | |
| 70 | def index(self, item: T, start: int = 0, stop: int | None = None) -> int: |
| 71 | return self._x.index(item, start, stop if stop is not None else len(self._x)) |
| 72 | |
| 73 | @overload |
| 74 | def __getitem__(self, item: int) -> T: ... |
| 75 | |
| 76 | @overload |
| 77 | def __getitem__(self, s: slice, /) -> list[T]: ... |
| 78 | |
| 79 | def __getitem__(self, item: int | slice) -> T | list[T]: |
| 80 | return self._x[item] |
| 81 | |
| 82 | @overload |
| 83 | def __setitem__(self, item: int, value: T): ... |
| 84 | |
| 85 | @overload |
| 86 | def __setitem__(self, s: slice, value: T, /): ... |
| 87 | |
| 88 | def __setitem__(self, item: int | slice, value: T | list[T]): |
| 89 | raise TypeError("Cannot set item in a constant list") |
| 90 | |
| 91 | def __delitem__(self, item): |
| 92 | raise TypeError("Cannot delete item from a constant list") |
| 93 | |
| 94 | def __iter__(self): |
| 95 | return iter(self._x) |
| 96 | |
| 97 | def __contains__(self, item): |
| 98 | return item in self._x |
| 99 | |
| 100 | def __len__(self): |
| 101 | return len(self._x) |
| 102 | |
| 103 | def __repr__(self): |
| 104 | return f"ConstantList({self._x})" |
| 105 | |
| 106 | def copy(self) -> list[T]: |
| 107 | return self._x.copy() |
| 108 | |
| 109 | |
| 110 | class CpuGpuBuffer: |
| 111 | """Buffer to easily copy tensors between CPU and GPU.""" |
| 112 | |
| 113 | def __init__( |
| 114 | self, |
| 115 | *size: int | torch.SymInt, |
| 116 | dtype: torch.dtype, |
| 117 | device: torch.device, |
| 118 | pin_memory: bool = PIN_MEMORY, |
| 119 | with_numpy: bool = True, |
| 120 | ) -> None: |
| 121 | # these buffers are mutable runtime state, so allocate them as normal |
| 122 | with torch.inference_mode(False): |
| 123 | self.cpu = torch.zeros( |
| 124 | *size, dtype=dtype, device="cpu", pin_memory=pin_memory |
| 125 | ) |
| 126 | self.gpu = torch.zeros_like(self.cpu, device=device) |
| 127 | self.np: np.ndarray |
| 128 | # To keep type hints simple (avoiding generics and subclasses), we |
| 129 | # only conditionally create the numpy array attribute. This can cause |
| 130 | # AttributeError if `self.np` is accessed when `with_numpy=False`. |
| 131 | if with_numpy: |
| 132 | if dtype == torch.bfloat16: |
| 133 | raise ValueError( |
| 134 | "Bfloat16 torch tensors cannot be directly cast to a " |
| 135 | "numpy array, so call CpuGpuBuffer with with_numpy=False" |
| 136 | ) |
| 137 | self.np = self.cpu.numpy() |
| 138 | |
| 139 | def copy_to_gpu(self, n: int | None = None) -> torch.Tensor: |
| 140 | cpu, gpu = self.cpu, self.gpu |
| 141 | if n is not None: |
| 142 | cpu, gpu = cpu[:n], gpu[:n] |
| 143 | return gpu.copy_(cpu.pin_memory() if PIN_MEMORY else cpu, non_blocking=True) |
| 144 | |
| 145 | def copy_to_cpu(self, n: int | None = None) -> torch.Tensor: |
| 146 | """NOTE: Because this method is non-blocking, explicit synchronization |
| 147 | is needed to ensure the data is copied to CPU.""" |
| 148 | if n is None: |
| 149 | return self.cpu.copy_(self.gpu, non_blocking=True) |
| 150 | return self.cpu[:n].copy_(self.gpu[:n], non_blocking=True) |
| 151 | |
| 152 | |
| 153 | def get_engine_client_zmq_addr( |
| 154 | local_only: bool, |
| 155 | host: str, |
| 156 | port: int = 0, |
| 157 | ) -> str: |
| 158 | """Return an IPC path (``local_only=True``) or ``tcp://host:port``. |
| 159 | |
| 160 | ``port=0`` lets the kernel assign the port at ``bind()`` time; the |
| 161 | caller must recover it via ``getsockopt(zmq.LAST_ENDPOINT)``.""" |
| 162 | if local_only: |
| 163 | return get_open_zmq_ipc_path() |
| 164 | return get_tcp_uri(host, port) |
| 165 | |
| 166 | |
| 167 | class APIServerProcessManager: |
| 168 | """Manages a group of API server processes. |
| 169 | |
| 170 | Handles creation, monitoring, and termination of API server worker |
| 171 | processes. Also monitors extra processes to check if they are healthy. |
| 172 | """ |
| 173 | |
| 174 | def __init__( |
| 175 | self, |
| 176 | listen_address: str, |
| 177 | sock: Any, |
| 178 | args: argparse.Namespace, |
| 179 | num_servers: int, |
| 180 | input_addresses: list[str], |
| 181 | output_addresses: list[str], |
| 182 | target_server_fn: Callable | None = None, |
| 183 | stats_update_address: str | None = None, |
| 184 | tensor_queue: Queue | None = None, |
| 185 | ): |
| 186 | """Initialize and start API server worker processes. |
| 187 | |
| 188 | ``input_addresses``/``output_addresses`` may contain |
| 189 | ``tcp://host:0`` placeholders; each child must report the actual |
| 190 | bound endpoint over its ``actual_address_pipe`` in ``client_config`` |
| 191 | and the parent collects them via |
| 192 | :py:meth:`gather_actual_addresses`. |
| 193 | |
| 194 | Args: |
| 195 | target_server_fn: Override function to call for each API server process |
| 196 | listen_address: Address to listen for client connections |
| 197 | sock: Socket for client connections |
| 198 | args: Command line arguments |
| 199 | num_servers: Number of API server processes to start |
| 200 | input_addresses: Input addresses for each API server |
| 201 | output_addresses: Output addresses for each API server |
| 202 | stats_update_address: Optional stats update address |
| 203 | tensor_queue: Optional tensor IPC queue for sharing MM tensors |
| 204 | |
| 205 | """ |
| 206 | self.listen_address = listen_address |
| 207 | self.sock = sock |
| 208 | self.args = args |
| 209 | |
| 210 | spawn_context = multiprocessing.get_context("spawn") |
| 211 | self.processes: list[BaseProcess] = [] |
| 212 | self._address_pipes: list[connection.Connection] = [] |
| 213 | |
| 214 | admission_counters = None |
| 215 | if num_servers > 1 and getattr(args, "max_num_queued_reqs", None) is not None: |
| 216 | from vllm.v1.engine.admission_control import SharedAdmissionStats |
| 217 | |
| 218 | admission_counters = spawn_context.RawArray( |
| 219 | "q", |
| 220 | SharedAdmissionStats.num_counters(num_servers), |
| 221 | ) |
| 222 | |
| 223 | for i, in_addr, out_addr in zip( |
| 224 | range(num_servers), input_addresses, output_addresses |
| 225 | ): |
| 226 | client_config: dict[str, Any] = { |
| 227 | "input_address": in_addr, |
| 228 | "output_address": out_addr, |
| 229 | "client_count": num_servers, |
| 230 | "client_index": i, |
| 231 | } |
| 232 | if admission_counters is not None: |
| 233 | client_config["mp_admission_counters"] = admission_counters |
| 234 | if stats_update_address is not None: |
| 235 | client_config["stats_update_address"] = stats_update_address |
| 236 | if tensor_queue is not None: |
| 237 | client_config["tensor_queue"] = tensor_queue |
| 238 | |
| 239 | parent_recv, child_send = spawn_context.Pipe(duplex=False) |
| 240 | self._address_pipes.append(parent_recv) |
| 241 | client_config["actual_address_pipe"] = child_send |
| 242 | |
| 243 | proc = spawn_context.Process( |
| 244 | target=target_server_fn or run_api_server_worker_proc, |
| 245 | name=f"ApiServer_{i}", |
| 246 | args=(listen_address, sock, args, client_config), |
| 247 | ) |
| 248 | self.processes.append(proc) |
| 249 | proc.start() |
| 250 | |
| 251 | # Drop parent's write end so reader sees EOF on child death. |
| 252 | child_send.close() |
| 253 | |
| 254 | logger.info("Started %d API server processes", len(self.processes)) |
| 255 | |
| 256 | # Shutdown only the API server processes on garbage collection |
| 257 | # The extra processes are managed by their owners |
| 258 | self._finalizer = weakref.finalize(self, shutdown, self.processes) |
| 259 | |
| 260 | def gather_actual_addresses( |
| 261 | self, |
| 262 | timeout: float = envs.VLLM_ENGINE_READY_TIMEOUT_S, |
| 263 | ) -> tuple[list[str], list[str]]: |
| 264 | """Return (inputs, outputs) reported by each child, indexed by |
| 265 | ``client_index``. Raises ``RuntimeError`` on timeout or premature |
| 266 | child exit.""" |
| 267 | n = len(self._address_pipes) |
| 268 | inputs: list[str | None] = [None] * n |
| 269 | outputs: list[str | None] = [None] * n |
| 270 | pending: dict[connection.Connection, int] = { |
| 271 | pipe: i for i, pipe in enumerate(self._address_pipes) |
| 272 | } |
| 273 | sentinel_to_idx: dict[Any, int] = { |
| 274 | proc.sentinel: i for i, proc in enumerate(self.processes) |
| 275 | } |
| 276 | |
| 277 | deadline = time.monotonic() + timeout |
| 278 | try: |
| 279 | while pending: |
| 280 | remaining = deadline - time.monotonic() |
| 281 | if remaining <= 0: |
| 282 | missing = [self.processes[i].name for i in pending.values()] |
| 283 | raise RuntimeError( |
| 284 | f"Timed out after {timeout:.1f}s waiting for " |
| 285 | f"API server(s) to report bound ZMQ addresses: " |
| 286 | f"{missing}" |
| 287 | ) |
| 288 | waitables: list[Any] = list(pending.keys()) + list( |
| 289 | sentinel_to_idx.keys() |
| 290 | ) |
| 291 | ready = connection.wait(waitables, timeout=remaining) |
| 292 | # Drain pipes before checking sentinels: a child that sent |
| 293 | # its message and then exited can surface both events in |
| 294 | # the same poll, and we must record the success first. |
| 295 | for item in ready: |
| 296 | if isinstance(item, connection.Connection) and item in pending: |
| 297 | idx = pending.pop(item) |
| 298 | try: |
| 299 | msg: dict[str, str] = item.recv() |
| 300 | except EOFError as e: |
| 301 | raise RuntimeError( |
| 302 | f"API server {self.processes[idx].name} " |
| 303 | f"closed its address pipe without " |
| 304 | f"reporting its bound ZMQ addresses" |
| 305 | ) from e |
| 306 | inputs[idx] = msg["input_address"] |
| 307 | outputs[idx] = msg["output_address"] |
| 308 | item.close() |
| 309 | for item in ready: |
| 310 | if item in sentinel_to_idx: |
| 311 | idx = sentinel_to_idx.pop(item) |
| 312 | pipe = self._address_pipes[idx] |
| 313 | if pipe in pending: |
| 314 | proc = self.processes[idx] |
| 315 | raise RuntimeError( |
| 316 | f"API server process {proc.name} exited " |
| 317 | f"(code={proc.exitcode}) before reporting " |
| 318 | f"its bound ZMQ addresses" |
| 319 | ) |
| 320 | finally: |
| 321 | for pipe in pending: |
| 322 | with contextlib.suppress(Exception): |
| 323 | pipe.close() |
| 324 | |
| 325 | return inputs, outputs # type: ignore[return-value] |
| 326 | |
| 327 | def shutdown(self, timeout: float | None = None) -> None: |
| 328 | """Shutdown API server processes with configurable timeout""" |
| 329 | for pipe in self._address_pipes: |
| 330 | with contextlib.suppress(Exception): |
| 331 | pipe.close() |
| 332 | self._address_pipes = [] |
| 333 | |
| 334 | if self._finalizer.detach() is not None: |
| 335 | shutdown(self.processes, timeout=timeout) |
| 336 | |
| 337 | |
| 338 | class RustFrontendProcessManager: |
| 339 | """Manages a single Rust frontend subprocess. |
| 340 | |
| 341 | Launches the Rust vllm-rs binary in 'frontend' mode, passing the |
| 342 | listening socket fd and ZMQ transport addresses. Provides the same |
| 343 | interface as APIServerProcessManager for process monitoring. |
| 344 | """ |
| 345 | |
| 346 | def __init__( |
| 347 | self, |
| 348 | binary_path: str, |
| 349 | sock: Any, |
| 350 | args: argparse.Namespace, |
| 351 | input_address: str, |
| 352 | output_address: str, |
| 353 | engine_start_index: int, |
| 354 | engine_count: int, |
| 355 | data_parallel_size: int, |
| 356 | stats_update_address: str | None = None, |
| 357 | grpc_sock: Any | None = None, |
| 358 | ): |
| 359 | import os |
| 360 | import subprocess |
| 361 | |
| 362 | fd = sock.fileno() |
| 363 | os.set_inheritable(fd, True) |
| 364 | |
| 365 | cmd = [ |
| 366 | binary_path, |
| 367 | "frontend", |
| 368 | "--listen-fd", |
| 369 | str(fd), |
| 370 | "--input-address", |
| 371 | input_address, |
| 372 | "--output-address", |
| 373 | output_address, |
| 374 | "--engine-start-index", |
| 375 | str(engine_start_index), |
| 376 | "--engine-count", |
| 377 | str(engine_count), |
| 378 | "--data-parallel-size", |
| 379 | str(data_parallel_size), |
| 380 | ] |
| 381 | pass_fds = [fd] |
| 382 | if grpc_sock is not None: |
| 383 | grpc_fd = grpc_sock.fileno() |
| 384 | os.set_inheritable(grpc_fd, True) |
| 385 | cmd.extend(["--grpc-listen-fd", str(grpc_fd)]) |
| 386 | pass_fds.append(grpc_fd) |
| 387 | if stats_update_address is not None: |
| 388 | cmd.extend(["--coordinator-address", stats_update_address]) |
| 389 | from vllm.entrypoints.serve.utils.api_utils import jsonify_non_default_args |
| 390 | |
| 391 | args_dict = jsonify_non_default_args( |
| 392 | args, |
| 393 | exclude={ |
| 394 | "api_server_count", |
| 395 | # Python passes the bootstrapped engine range explicitly. |
| 396 | "data_parallel_rank", |
| 397 | "data_parallel_external_lb", |
| 398 | "data_parallel_hybrid_lb", |
| 399 | }, |
| 400 | ) |
| 401 | |
| 402 | # `model_tag` is the positional `vllm serve` model argument. When the |
| 403 | # model is supplied only through `--config`, argparse populates `model` |
| 404 | # while leaving `model_tag` unset. The Rust frontend requires |
| 405 | # `model_tag` in its JSON bootstrap payload, so use the resolved model |
| 406 | # as a fallback. |
| 407 | model_tag = getattr(args, "model_tag", None) or getattr(args, "model", None) |
| 408 | if model_tag is not None: |
| 409 | args_dict["model_tag"] = model_tag |
| 410 | |
| 411 | # The Rust `frontend` subcommand parses --args-json via serde_json, |
| 412 | # which bypasses clap and therefore ignores any `#[arg(env = ...)]` |
| 413 | # declarations on SharedRuntimeArgs fields. Forward the env-driven |
| 414 | # values explicitly so VLLM_ENGINE_READY_TIMEOUT_S and |
| 415 | # VLLM_HTTP_TIMEOUT_KEEP_ALIVE behave the same on both Python and Rust |
| 416 | # frontends. |
| 417 | args_dict["engine_ready_timeout_secs"] = envs.VLLM_ENGINE_READY_TIMEOUT_S |
| 418 | args_dict["http_timeout_keep_alive"] = envs.VLLM_HTTP_TIMEOUT_KEEP_ALIVE |
| 419 | args_json = json.dumps(args_dict, sort_keys=True) |
| 420 | cmd.extend(["--args-json", args_json]) |
| 421 | |
| 422 | # The subprocess needs the real values, but the log must not carry |
| 423 | # credentials such as api_key or hf_token. |
| 424 | from vllm.entrypoints.serve.utils.api_utils import redact_sensitive_args |
| 425 | |
| 426 | redacted_json = json.dumps(redact_sensitive_args(args_dict), sort_keys=True) |
| 427 | logger.info("Launching Rust frontend: %s", " ".join(cmd[:-1] + [redacted_json])) |
| 428 | self._proc = subprocess.Popen(cmd, pass_fds=pass_fds) |
| 429 | |
| 430 | # Create a process wrapper with a sentinel fd for monitoring |
| 431 | self.processes: list[_SubprocessWrapper] = [ |
| 432 | _SubprocessWrapper(self._proc, "RustFrontend") |
| 433 | ] |
| 434 | |
| 435 | self._finalizer = weakref.finalize(self, _shutdown_subprocesses, self.processes) |
| 436 | |
| 437 | def shutdown(self, timeout: float | None = None) -> None: |
| 438 | if self._finalizer.detach() is not None: |
| 439 | _shutdown_subprocesses(self.processes, timeout=timeout) |
| 440 | |
| 441 | |
| 442 | class _SubprocessWrapper: |
| 443 | """Wraps subprocess.Popen to provide the BaseProcess-like interface |
| 444 | needed by wait_for_completion_or_failure.""" |
| 445 | |
| 446 | def __init__(self, proc, name: str): |
| 447 | self._proc = proc |
| 448 | self.name = name |
| 449 | self.pid = proc.pid |
| 450 | self._sentinel_conn: connection.Connection | None = None |
| 451 | self._sentinel_send: connection.Connection | None = None |
| 452 | |
| 453 | # Use a Pipe-based sentinel so subprocess monitoring works uniformly |
| 454 | # across platforms with multiprocessing.connection.wait(). |
| 455 | recv, send = connection.Pipe(duplex=False) |
| 456 | self._sentinel_conn = recv |
| 457 | self._sentinel_send = send |
| 458 | |
| 459 | def monitor_subprocess() -> None: |
| 460 | try: |
| 461 | proc.wait() |
| 462 | finally: |
| 463 | with contextlib.suppress(Exception): |
| 464 | send.close() |
| 465 | |
| 466 | threading.Thread( |
| 467 | target=monitor_subprocess, daemon=True, name=f"{name}Monitor" |
| 468 | ).start() |
| 469 | |
| 470 | @property |
| 471 | def sentinel(self): |
| 472 | return self._sentinel_conn |
| 473 | |
| 474 | @property |
| 475 | def exitcode(self) -> int | None: |
| 476 | return self._proc.returncode if self._proc.poll() is not None else None |
| 477 | |
| 478 | def is_alive(self) -> bool: |
| 479 | return self._proc.poll() is None |
| 480 | |
| 481 | def terminate(self): |
| 482 | self._proc.terminate() |
| 483 | |
| 484 | def join(self, timeout=None): |
| 485 | with contextlib.suppress(Exception): |
| 486 | self._proc.wait(timeout=timeout) |
| 487 | |
| 488 | def __del__(self): |
| 489 | with contextlib.suppress(Exception): |
| 490 | if self._sentinel_conn is not None: |
| 491 | self._sentinel_conn.close() |
| 492 | if self._sentinel_send is not None: |
| 493 | self._sentinel_send.close() |
| 494 | |
| 495 | |
| 496 | def _shutdown_subprocesses( |
| 497 | procs: list[_SubprocessWrapper], timeout: float | None = None |
| 498 | ) -> None: |
| 499 | """Shutdown subprocess wrappers (mirrors the shutdown() function).""" |
| 500 | if timeout is None: |
| 501 | timeout = 0.0 |
| 502 | timeout = max(timeout, 5.0) |
| 503 | |
| 504 | logger.debug( |
| 505 | "[shutdown] Subprocess manager: start process_count=%d timeout=%ss", |
| 506 | len(procs), |
| 507 | timeout, |
| 508 | ) |
| 509 | |
| 510 | for proc in procs: |
| 511 | if proc.is_alive(): |
| 512 | proc.terminate() |
| 513 | |
| 514 | deadline = time.monotonic() + timeout |
| 515 | for proc in procs: |
| 516 | remaining = deadline - time.monotonic() |
| 517 | if remaining <= 0: |
| 518 | break |
| 519 | if proc.is_alive(): |
| 520 | proc.join(remaining) |
| 521 | |
| 522 | remaining_pids = [ |
| 523 | proc.pid for proc in procs if proc.is_alive() and proc.pid is not None |
| 524 | ] |
| 525 | if remaining_pids: |
| 526 | logger.warning( |
| 527 | "[shutdown] Subprocess manager: force killing remaining processes count=%d", |
| 528 | len(remaining_pids), |
| 529 | ) |
| 530 | for pid in remaining_pids: |
| 531 | kill_process_tree(pid) |
| 532 | |
| 533 | logger.debug_once("[shutdown] Subprocess manager: complete") |
| 534 | |
| 535 | |
| 536 | def run_api_server_worker_proc( |
| 537 | listen_address, sock, args, client_config=None, **uvicorn_kwargs |
| 538 | ) -> None: |
| 539 | """Entrypoint for individual API server worker processes.""" |
| 540 | if logging_config := getattr(args, "logging_config", None): |
| 541 | from vllm.logger import configure_logging |
| 542 | |
| 543 | configure_logging(logging_config) |
| 544 | |
| 545 | from vllm.entrypoints.launchers.api_server.entry import run_server_worker |
| 546 | |
| 547 | client_config = client_config or {} |
| 548 | server_index = client_config.get("client_index", 0) |
| 549 | |
| 550 | # Set process title and process-specific log metadata. |
| 551 | set_process_title("APIServer", str(server_index)) |
| 552 | decorate_logs() |
| 553 | |
| 554 | uvloop.run( |
| 555 | run_server_worker(listen_address, sock, args, client_config, **uvicorn_kwargs) |
| 556 | ) |
| 557 | |
| 558 | |
| 559 | def wait_for_completion_or_failure( |
| 560 | api_server_manager: "APIServerProcessManager | RustFrontendProcessManager", |
| 561 | engine_manager: Union["CoreEngineProcManager", "CoreEngineActorManager"] |
| 562 | | None = None, |
| 563 | coordinator: "DPCoordinator | None" = None, |
| 564 | ) -> None: |
| 565 | """Wait for all processes to complete or detect if any fail. |
| 566 | |
| 567 | Raises an exception if any process exits with a non-zero status. |
| 568 | |
| 569 | Args: |
| 570 | api_server_manager: The manager for API servers. |
| 571 | engine_manager: The manager for engine processes. |
| 572 | If CoreEngineProcManager, it manages local engines; |
| 573 | if CoreEngineActorManager, it manages all engines. |
| 574 | coordinator: The coordinator for data parallel. |
| 575 | |
| 576 | """ |
| 577 | try: |
| 578 | logger.info("Waiting for API servers to complete ...") |
| 579 | # Create a mapping of sentinels to their corresponding processes |
| 580 | # for efficient lookup |
| 581 | sentinel_to_proc: dict[Any, BaseProcess | _SubprocessWrapper | None] = { |
| 582 | proc.sentinel: proc for proc in api_server_manager.processes |
| 583 | } |
| 584 | |
| 585 | if coordinator: |
| 586 | sentinel_to_proc[coordinator.proc.sentinel] = coordinator.proc |
| 587 | |
| 588 | if engine_manager: |
| 589 | core_shutdown_recv, core_shutdown_send = connection.Pipe(duplex=False) |
| 590 | |
| 591 | def monitor_engines(): |
| 592 | try: |
| 593 | engine_manager.monitor_engine_liveness() |
| 594 | finally: |
| 595 | core_shutdown_send.close() |
| 596 | core_shutdown_recv.close() |
| 597 | |
| 598 | # start monitor for engine liveness |
| 599 | threading.Thread(target=monitor_engines, daemon=True).start() |
| 600 | sentinel_to_proc[core_shutdown_recv] = None # type: ignore[assignment] |
| 601 | |
| 602 | # Check if any process terminates |
| 603 | while sentinel_to_proc: |
| 604 | # Wait for any process to terminate (or engine shutdown signal) |
| 605 | ready_sentinels: list[Any] = connection.wait(sentinel_to_proc) |
| 606 | |
| 607 | # Process any terminated processes |
| 608 | for sentinel in ready_sentinels: |
| 609 | proc = sentinel_to_proc.pop(sentinel) |
| 610 | |
| 611 | # Check if process exited with error |
| 612 | if proc is not None and proc.exitcode != 0: |
| 613 | raise RuntimeError( |
| 614 | f"Process {proc.name} (PID: {proc.pid}) " |
| 615 | f"died with exit code {proc.exitcode}" |
| 616 | ) |
| 617 | if engine_manager and engine_manager.failed_proc_name is not None: |
| 618 | raise RuntimeError( |
| 619 | f"Engine core process {engine_manager.failed_proc_name} " |
| 620 | "died unexpectedly." |
| 621 | ) |
| 622 | |
| 623 | except KeyboardInterrupt: |
| 624 | logger.info("Received KeyboardInterrupt, shutting down API servers...") |
| 625 | except Exception as e: |
| 626 | logger.exception("Exception occurred while running API servers: %s", str(e)) |
| 627 | raise |
| 628 | |
| 629 | |
| 630 | # Note(rob): shutdown function cannot be a bound method, |
| 631 | # else the gc cannot collect the object. |
| 632 | def shutdown(procs: list[BaseProcess], timeout: float | None = None) -> None: |
| 633 | """Shutdown processes with timeout. |
| 634 | |
| 635 | Args: |
| 636 | procs: List of processes to shutdown |
| 637 | timeout: Maximum time in seconds to wait for graceful shutdown |
| 638 | |
| 639 | """ |
| 640 | if timeout is None: |
| 641 | # Keep a small grace period for best-effort cleanup paths that do not |
| 642 | # have a user-configured shutdown timeout. |
| 643 | timeout = 5.0 |
| 644 | |
| 645 | logger.debug( |
| 646 | "[shutdown] Process manager: start process_count=%d timeout=%ss names=%s", |
| 647 | len(procs), |
| 648 | timeout, |
| 649 | (",").join([proc.name for proc in procs]), |
| 650 | ) |
| 651 | |
| 652 | # Shutdown the process. |
| 653 | for proc in procs: |
| 654 | if proc.is_alive(): |
| 655 | logger.info( |
| 656 | "[shutdown] Process manager: send sigterm to process %s", proc.name |
| 657 | ) |
| 658 | proc.terminate() |
| 659 | |
| 660 | # Allow time for remaining procs to terminate. |
| 661 | deadline = time.monotonic() + timeout |
| 662 | for proc in procs: |
| 663 | remaining = deadline - time.monotonic() |
| 664 | if remaining <= 0: |
| 665 | break |
| 666 | if proc.is_alive(): |
| 667 | proc.join(remaining) |
| 668 | |
| 669 | remaining_procs = [ |
| 670 | (proc.pid, proc.name) |
| 671 | for proc in procs |
| 672 | if proc.is_alive() and proc.pid is not None |
| 673 | ] |
| 674 | if remaining_procs: |
| 675 | logger.warning( |
| 676 | "[shutdown] Process manager: force killing remaining processes count=%d", |
| 677 | len(remaining_procs), |
| 678 | ) |
| 679 | for pid, proc_name in remaining_procs: |
| 680 | logger.warning( |
| 681 | "[shutdown] Process manager: force killing remaining process %s pid %d", |
| 682 | proc_name, |
| 683 | pid, |
| 684 | ) |
| 685 | kill_process_tree(pid) |
| 686 | |
| 687 | logger.debug_once("[shutdown] Process manager: complete") |
| 688 | |
| 689 | |
| 690 | def copy_slice( |
| 691 | from_tensor: torch.Tensor, to_tensor: torch.Tensor, length: int |
| 692 | ) -> torch.Tensor: |
| 693 | """Copy the first length elements of a tensor into another tensor in a |
| 694 | non-blocking manner. |
| 695 | |
| 696 | Used to copy pinned CPU tensor data to pre-allocated GPU tensors. |
| 697 | |
| 698 | Returns the sliced target tensor. |
| 699 | """ |
| 700 | return to_tensor[:length].copy_(from_tensor[:length], non_blocking=True) |
| 701 | |
| 702 | |
| 703 | def report_usage_stats( |
| 704 | vllm_config, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT |
| 705 | ) -> None: |
| 706 | """Report usage statistics if enabled.""" |
| 707 | if not is_usage_stats_enabled(): |
| 708 | return |
| 709 | |
| 710 | from vllm.model_executor.model_loader import get_architecture_class_name |
| 711 | |
| 712 | model_config = vllm_config.model_config |
| 713 | scheduler_config = vllm_config.scheduler_config |
| 714 | parallel_config = vllm_config.parallel_config |
| 715 | attention_config = vllm_config.attention_config |
| 716 | compilation_config = vllm_config.compilation_config |
| 717 | speculative_config = vllm_config.speculative_config |
| 718 | |
| 719 | # Prepare KV connector string if applicable |
| 720 | kv_connector = None |
| 721 | if vllm_config.kv_transfer_config is not None: |
| 722 | kv_connector = vllm_config.kv_transfer_config.kv_connector |
| 723 | |
| 724 | # Attention backend is None when set to "auto" (resolved at runtime per platform). |
| 725 | attention_backend = ( |
| 726 | attention_config.backend.name if attention_config.backend is not None else None |
| 727 | ) |
| 728 | |
| 729 | # CompilationMode is an IntEnum; report the name for readability in dashboards. |
| 730 | compilation_mode = ( |
| 731 | compilation_config.mode.name if compilation_config.mode is not None else None |
| 732 | ) |
| 733 | |
| 734 | # Speculative decoding fields default to None when spec decode is disabled. |
| 735 | spec_decode_method = ( |
| 736 | speculative_config.method if speculative_config is not None else None |
| 737 | ) |
| 738 | num_speculative_tokens = ( |
| 739 | speculative_config.num_speculative_tokens |
| 740 | if speculative_config is not None |
| 741 | else None |
| 742 | ) |
| 743 | |
| 744 | if model_config.using_transformers_backend(): |
| 745 | backend_cls = model_config._model_info.architecture |
| 746 | # Show what was wrapped e.g. TransformersForCausalLM(Starcoder2ForCausalLM) |
| 747 | architecture = f"{backend_cls}({model_config.architectures[0]})" |
| 748 | else: |
| 749 | architecture = get_architecture_class_name(model_config) |
| 750 | |
| 751 | usage_message.report_usage( |
| 752 | architecture, |
| 753 | usage_context, |
| 754 | extra_kvs={ |
| 755 | # Common configuration |
| 756 | "dtype": str(model_config.dtype), |
| 757 | "block_size": vllm_config.cache_config.block_size, |
| 758 | "gpu_memory_utilization": vllm_config.cache_config.gpu_memory_utilization, |
| 759 | "kv_cache_memory_bytes": vllm_config.cache_config.kv_cache_memory_bytes, |
| 760 | # Quantization |
| 761 | "quantization": model_config.quantization, |
| 762 | "kv_cache_dtype": str(vllm_config.cache_config.cache_dtype), |
| 763 | # Feature flags |
| 764 | "enable_lora": bool(vllm_config.lora_config), |
| 765 | "enable_prefix_caching": vllm_config.cache_config.enable_prefix_caching, |
| 766 | "enforce_eager": model_config.enforce_eager, |
| 767 | "disable_custom_all_reduce": parallel_config.disable_custom_all_reduce, |
| 768 | # Distributed parallelism settings |
| 769 | "tensor_parallel_size": parallel_config.tensor_parallel_size, |
| 770 | "data_parallel_size": parallel_config.data_parallel_size, |
| 771 | "pipeline_parallel_size": parallel_config.pipeline_parallel_size, |
| 772 | "enable_expert_parallel": parallel_config.enable_expert_parallel, |
| 773 | # All2All backend for MoE expert parallel |
| 774 | "all2all_backend": parallel_config.all2all_backend, |
| 775 | # KV connector used |
| 776 | "kv_connector": kv_connector, |
| 777 | # Batching limits — tuning knobs operators commonly override |
| 778 | "max_model_len": model_config.max_model_len, |
| 779 | "max_num_seqs": scheduler_config.max_num_seqs, |
| 780 | "max_num_batched_tokens": scheduler_config.max_num_batched_tokens, |
| 781 | # Attention backend (user-requested; None = auto-selected at runtime) |
| 782 | "attention_backend": attention_backend, |
| 783 | # torch.compile mode (e.g. NONE, STOCK_TORCH_COMPILE, VLLM_COMPILE) |
| 784 | "compilation_mode": compilation_mode, |
| 785 | # Speculative decoding configuration |
| 786 | "spec_decode_method": spec_decode_method, |
| 787 | "num_speculative_tokens": num_speculative_tokens, |
| 788 | # Wide expert parallel: load balancer + redundant/total expert counts |
| 789 | "enable_eplb": parallel_config.enable_eplb, |
| 790 | "num_redundant_experts": parallel_config.eplb_config.num_redundant_experts, |
| 791 | "num_experts": model_config.get_num_experts(), |
| 792 | }, |
| 793 | ) |
| 794 | |
| 795 | |
| 796 | _PROFILER_FUNC = None |
| 797 | |
| 798 | |
| 799 | def record_function_or_nullcontext(name: str) -> AbstractContextManager: |
| 800 | global _PROFILER_FUNC |
| 801 | |
| 802 | # fast path assume it is set |
| 803 | if _PROFILER_FUNC is not None: |
| 804 | return _PROFILER_FUNC(name) |
| 805 | |
| 806 | func = contextlib.nullcontext |
| 807 | if envs.VLLM_CUSTOM_SCOPES_FOR_PROFILING: |
| 808 | func = record_function |
| 809 | elif envs.VLLM_NVTX_SCOPES_FOR_PROFILING: |
| 810 | import nvtx |
| 811 | |
| 812 | func = nvtx.annotate |
| 813 | |
| 814 | _PROFILER_FUNC = func |
| 815 | return func(name) |
| 816 | |
| 817 | |
| 818 | def tensor_data(tensor: torch.Tensor) -> memoryview: |
| 819 | """Get the raw data of a tensor as a uint8 memoryview, useful for |
| 820 | serializing and hashing. |
| 821 | |
| 822 | Args: |
| 823 | tensor: The input tensor. |
| 824 | |
| 825 | Returns: |
| 826 | A memoryview of the tensor data as uint8. |
| 827 | |
| 828 | """ |
| 829 | return tensor.flatten().cpu().contiguous().view(torch.uint8).numpy().data |
| 830 | |
| 831 | |
| 832 | @dataclass |
| 833 | class IterationDetails: |
| 834 | num_ctx_requests: int |
| 835 | num_ctx_tokens: int |
| 836 | num_generation_requests: int |
| 837 | num_generation_tokens: int |
| 838 | num_encoder_inputs: int = 0 |
| 839 | num_encoder_output_tokens: int = 0 |
| 840 | |
| 841 | def __repr__(self) -> str: |
| 842 | return f"IterationDetails(num_ctx_requests={self.num_ctx_requests},\ |
| 843 | num_ctx_tokens={self.num_ctx_tokens}, \ |
| 844 | num_generation_requests={self.num_generation_requests}, \ |
| 845 | num_generation_tokens={self.num_generation_tokens}, \ |
| 846 | num_encoder_inputs={self.num_encoder_inputs}, \ |
| 847 | num_encoder_output_tokens={self.num_encoder_output_tokens})" |
| 848 | |
| 849 | |
| 850 | def compute_iteration_details(scheduler_output: SchedulerOutput) -> IterationDetails: |
| 851 | """Compute the number of context/generation requests and tokens |
| 852 | for the current iteration's scheduler output. A requests is regarded |
| 853 | as a context request if its output tokens are still 0, an extended chunk |
| 854 | of chunked prefill falls into this category. |
| 855 | |
| 856 | Args: |
| 857 | scheduler_output: The scheduler output for the current iteration. |
| 858 | |
| 859 | Returns: |
| 860 | An IterationDetails object containing the number of |
| 861 | context/generation requests and tokens. |
| 862 | |
| 863 | """ |
| 864 | num_context_requests = 0 |
| 865 | num_context_tokens = 0 |
| 866 | num_generation_requests = 0 |
| 867 | num_generation_tokens = 0 |
| 868 | new_req_ids = {new_req.req_id for new_req in scheduler_output.scheduled_new_reqs} |
| 869 | for req_id, num_tokens in scheduler_output.num_scheduled_tokens.items(): |
| 870 | if scheduler_output.scheduled_cached_reqs.is_context_phase(req_id) or ( |
| 871 | req_id in new_req_ids |
| 872 | ): |
| 873 | num_context_requests += 1 |
| 874 | num_context_tokens += num_tokens |
| 875 | else: |
| 876 | num_generation_requests += 1 |
| 877 | num_generation_tokens += num_tokens |
| 878 | scheduled_encoder_input_stats = scheduler_output.scheduled_encoder_input_stats |
| 879 | num_encoder_inputs = 0 |
| 880 | num_encoder_output_tokens = 0 |
| 881 | if scheduled_encoder_input_stats is not None: |
| 882 | num_encoder_inputs = scheduled_encoder_input_stats.num_inputs |
| 883 | num_encoder_output_tokens = scheduled_encoder_input_stats.output_tokens |
| 884 | |
| 885 | return IterationDetails( |
| 886 | num_context_requests, |
| 887 | num_context_tokens, |
| 888 | num_generation_requests, |
| 889 | num_generation_tokens, |
| 890 | num_encoder_inputs, |
| 891 | num_encoder_output_tokens, |
| 892 | ) |