__init__.py 75 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import contextlib
5
import datetime
Woosuk Kwon's avatar
Woosuk Kwon committed
6
import enum
7
import getpass
8
import hashlib
9
import importlib
10
import inspect
11
import ipaddress
12
import json
13
import multiprocessing
14
import os
15
import pickle
16
import signal
17
import socket
18
import subprocess
19
import sys
20
import tempfile
21
import textwrap
22
import threading
23
import traceback
Zhuohan Li's avatar
Zhuohan Li committed
24
import uuid
25
import warnings
26
import weakref
27
28
29
30
31
32
33
34
from argparse import (
    Action,
    ArgumentDefaultsHelpFormatter,
    ArgumentParser,
    ArgumentTypeError,
    RawDescriptionHelpFormatter,
    _ArgumentGroup,
)
35
from collections import defaultdict
36
from collections.abc import (
37
    Callable,
38
39
40
41
    Collection,
    Iterator,
    Sequence,
)
42
from concurrent.futures.process import ProcessPoolExecutor
43
from functools import cache, lru_cache, partial, wraps
44
from pathlib import Path
45
from typing import TYPE_CHECKING, Any, TextIO, TypeVar
46
from urllib.parse import urlparse
47
from uuid import uuid4
Zhuohan Li's avatar
Zhuohan Li committed
48

49
import cbor2
50
import cloudpickle
51
import numpy as np
52
import numpy.typing as npt
53
import psutil
54
import regex as re
55
import setproctitle
Zhuohan Li's avatar
Zhuohan Li committed
56
import torch
57
import yaml
58
59
import zmq
import zmq.asyncio
60
from packaging import version
61
from packaging.version import Version
62
from torch.library import Library
63

64
import vllm.envs as envs
65
from vllm.logger import enable_trace_function_call, init_logger
66
from vllm.ray.lazy_utils import is_in_ray_actor
67

68
if TYPE_CHECKING:
69
70
    from argparse import Namespace

71
    from vllm.config import ModelConfig, VllmConfig
72
    from vllm.sequence import IntermediateTensors
73
74
75
76
77
78
else:
    Namespace = object

    ModelConfig = object
    VllmConfig = object
    IntermediateTensors = object
79

80
81
logger = init_logger(__name__)

82
83
84
85
86
87
# This value is chosen to have a balance between ITL and TTFT. Note it is
# not optimized for throughput.
DEFAULT_MAX_NUM_BATCHED_TOKENS = 2048
POOLING_MODEL_MAX_NUM_BATCHED_TOKENS = 32768
MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS = 5120

88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
# Constants related to forcing the attention backend selection

# String name of register which may be set in order to
# force auto-selection of attention backend by Attention
# wrapper
STR_BACKEND_ENV_VAR: str = "VLLM_ATTENTION_BACKEND"

# Possible string values of STR_BACKEND_ENV_VAR
# register, corresponding to possible backends
STR_FLASHINFER_ATTN_VAL: str = "FLASHINFER"
STR_TORCH_SDPA_ATTN_VAL: str = "TORCH_SDPA"
STR_XFORMERS_ATTN_VAL: str = "XFORMERS"
STR_FLASH_ATTN_VAL: str = "FLASH_ATTN"
STR_INVALID_VAL: str = "INVALID"

103

104
# ANSI color codes
105
106
CYAN = "\033[1;36m"
RESET = "\033[0;0m"
107

108
STR_DTYPE_TO_TORCH_DTYPE = {
109
    "float32": torch.float32,
110
111
112
    "half": torch.half,
    "bfloat16": torch.bfloat16,
    "float": torch.float,
113
    "fp8": torch.uint8,
114
115
    "fp8_e4m3": torch.uint8,
    "fp8_e5m2": torch.uint8,
116
    "int8": torch.int8,
117
    "fp8_inc": torch.float8_e4m3fn,
118
    "fp8_ds_mla": torch.uint8,
119
}
Zhuohan Li's avatar
Zhuohan Li committed
120

121
122
123
124
125
126
127
128
129
TORCH_DTYPE_TO_NUMPY_DTYPE = {
    torch.float16: np.float16,
    torch.float32: np.float32,
    torch.float64: np.float64,
    torch.uint8: np.uint8,
    torch.int32: np.int32,
    torch.int64: np.int64,
}

130
131
132
133
134
135
136
137
138
139

@contextlib.contextmanager
def set_default_torch_num_threads(num_threads: int):
    """Sets the default number of threads for PyTorch to the given value."""
    old_num_threads = torch.get_num_threads()
    torch.set_num_threads(num_threads)
    yield
    torch.set_num_threads(old_num_threads)


140
T = TypeVar("T")
141
U = TypeVar("U")
142

143

Woosuk Kwon's avatar
Woosuk Kwon committed
144
145
146
147
148
class Device(enum.Enum):
    GPU = enum.auto()
    CPU = enum.auto()


149
150
151
152
153
class LayerBlockType(enum.Enum):
    attention = "attention"
    mamba = "mamba"


Woosuk Kwon's avatar
Woosuk Kwon committed
154
155
156
157
class Counter:
    def __init__(self, start: int = 0) -> None:
        self.counter = start

Woosuk Kwon's avatar
Woosuk Kwon committed
158
    def __next__(self) -> int:
159
        i = self.counter
Woosuk Kwon's avatar
Woosuk Kwon committed
160
        self.counter += 1
161
        return i
Woosuk Kwon's avatar
Woosuk Kwon committed
162
163
164

    def reset(self) -> None:
        self.counter = 0
Zhuohan Li's avatar
Zhuohan Li committed
165

166

Zhuohan Li's avatar
Zhuohan Li committed
167
168
def random_uuid() -> str:
    return str(uuid.uuid4().hex)
169

170

171
def close_sockets(sockets: Sequence[zmq.Socket | zmq.asyncio.Socket]):
172
173
174
175
176
    for sock in sockets:
        if sock is not None:
            sock.close(linger=0)


177
def get_ip() -> str:
178
    host_ip = envs.VLLM_HOST_IP
179
180
181
182
    if "HOST_IP" in os.environ and "VLLM_HOST_IP" not in os.environ:
        logger.warning(
            "The environment variable HOST_IP is deprecated and ignored, as"
            " it is often used by Docker and other software to"
183
            " interact with the container's network stack. Please "
184
            "use VLLM_HOST_IP instead to set the IP address for vLLM processes"
185
186
            " to communicate with each other."
        )
187
188
189
190
191
    if host_ip:
        return host_ip

    # IP is not set, try to get it from the network interface

192
    # try ipv4
193
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
194
    try:
195
        s.connect(("8.8.8.8", 80))  # Doesn't need to be reachable
196
        return s.getsockname()[0]
197
198
199
200
201
    except Exception:
        pass

    # try ipv6
    try:
202
        s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
203
204
205
        # Google's public DNS server, see
        # https://developers.google.com/speed/public-dns/docs/using#addresses
        s.connect(("2001:4860:4860::8888", 80))  # Doesn't need to be reachable
206
        return s.getsockname()[0]
207
208
209
210
211
    except Exception:
        pass

    warnings.warn(
        "Failed to get the IP address, using 0.0.0.0 by default."
212
213
        "The value can be set by the environment variable"
        " VLLM_HOST_IP or HOST_IP.",
214
215
        stacklevel=2,
    )
216
    return "0.0.0.0"
217
218


219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
def test_loopback_bind(address, family):
    try:
        s = socket.socket(family, socket.SOCK_DGRAM)
        s.bind((address, 0))  # Port 0 = auto assign
        s.close()
        return True
    except OSError:
        return False


def get_loopback_ip() -> str:
    loopback_ip = envs.VLLM_LOOPBACK_IP
    if loopback_ip:
        return loopback_ip

    # VLLM_LOOPBACK_IP is not set, try to get it based on network interface

    if test_loopback_bind("127.0.0.1", socket.AF_INET):
        return "127.0.0.1"
    elif test_loopback_bind("::1", socket.AF_INET6):
        return "::1"
    else:
        raise RuntimeError(
            "Neither 127.0.0.1 nor ::1 are bound to a local interface. "
243
244
            "Set the VLLM_LOOPBACK_IP environment variable explicitly."
        )
245
246


247
248
249
250
251
252
253
254
def is_valid_ipv6_address(address: str) -> bool:
    try:
        ipaddress.IPv6Address(address)
        return True
    except ValueError:
        return False


255
def split_host_port(host_port: str) -> tuple[str, int]:
256
    # ipv6
257
258
    if host_port.startswith("["):
        host, port = host_port.rsplit("]", 1)
259
        host = host[1:]
260
        port = port.split(":")[1]
261
262
        return host, int(port)
    else:
263
        host, port = host_port.split(":")
264
265
266
267
268
269
270
271
272
273
        return host, int(port)


def join_host_port(host: str, port: int) -> str:
    if is_valid_ipv6_address(host):
        return f"[{host}]:{port}"
    else:
        return f"{host}:{port}"


274
def get_distributed_init_method(ip: str, port: int) -> str:
275
276
277
278
    return get_tcp_uri(ip, port)


def get_tcp_uri(ip: str, port: int) -> str:
279
280
281
282
    if is_valid_ipv6_address(ip):
        return f"tcp://[{ip}]:{port}"
    else:
        return f"tcp://{ip}:{port}"
283
284


285
286
287
288
289
def get_open_zmq_ipc_path() -> str:
    base_rpc_path = envs.VLLM_RPC_BASE_PATH
    return f"ipc://{base_rpc_path}/{uuid4()}"


290
291
292
293
def get_open_zmq_inproc_path() -> str:
    return f"inproc://{uuid4()}"


294
def get_open_port() -> int:
295
296
297
298
299
300
301
302
303
    """
    Get an open port for the vLLM process to listen on.
    An edge case to handle, is when we run data parallel,
    we need to avoid ports that are potentially used by
    the data parallel master process.
    Right now we reserve 10 ports for the data parallel master
    process. Currently it uses 2 ports.
    """
    if "VLLM_DP_MASTER_PORT" in os.environ:
304
305
        dp_master_port = envs.VLLM_DP_MASTER_PORT
        reserved_port_range = range(dp_master_port, dp_master_port + 10)
306
        while True:
307
308
309
            candidate_port = _get_open_port()
            if candidate_port not in reserved_port_range:
                return candidate_port
310
311
    return _get_open_port()

youkaichao's avatar
youkaichao committed
312

313
314
def get_open_ports_list(count: int = 5) -> list[int]:
    """Get a list of open ports."""
315
    ports = set[int]()
316
317
318
319
320
    while len(ports) < count:
        ports.add(get_open_port())
    return list(ports)


321
def _get_open_port() -> int:
322
    port = envs.VLLM_PORT
323
    if port is not None:
324
325
326
327
328
329
330
        while True:
            try:
                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                    s.bind(("", port))
                    return port
            except OSError:
                port += 1  # Increment port number if already in use
331
                logger.info("Port %d is already in use, trying port %d", port - 1, port)
332
333
334
335
336
337
338
339
340
341
    # try ipv4
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("", 0))
            return s.getsockname()[1]
    except OSError:
        # try ipv6
        with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
            s.bind(("", 0))
            return s.getsockname()[1]
342
343


344
def find_process_using_port(port: int) -> psutil.Process | None:
345
346
347
348
349
350
351
    # TODO: We can not check for running processes with network
    # port on macOS. Therefore, we can not have a full graceful shutdown
    # of vLLM. For now, let's not look for processes in this case.
    # Ref: https://www.florianreinhard.de/accessdenied-in-psutil/
    if sys.platform.startswith("darwin"):
        return None

352
    our_pid = os.getpid()
353
    for conn in psutil.net_connections():
354
        if conn.laddr.port == port and (conn.pid is not None and conn.pid != our_pid):
355
356
357
358
359
360
361
            try:
                return psutil.Process(conn.pid)
            except psutil.NoSuchProcess:
                return None
    return None


362
def update_environment_variables(envs: dict[str, str]):
363
    for k, v in envs.items():
364
        if k in os.environ and os.environ[k] != v:
365
            logger.warning(
366
367
368
369
370
                "Overwriting environment variable %s from '%s' to '%s'",
                k,
                os.environ[k],
                v,
            )
371
        os.environ[k] = v
372
373


374
375
376
377
378
def cdiv(a: int, b: int) -> int:
    """Ceiling division."""
    return -(a // -b)


379
380
381
382
383
384
385
def next_power_of_2(n) -> int:
    """The next power of 2 (inclusive)"""
    if n < 1:
        return 1
    return 1 << (n - 1).bit_length()


386
387
388
389
390
391
392
def prev_power_of_2(n: int) -> int:
    """The previous power of 2 (inclusive)"""
    if n <= 0:
        return 0
    return 1 << (n.bit_length() - 1)


393
394
395
396
def round_up(x: int, y: int) -> int:
    return ((x + y - 1) // y) * y


397
398
399
400
def round_down(x: int, y: int) -> int:
    return (x // y) * y


401
def _generate_random_fp8(
402
    tensor: torch.Tensor,
403
404
405
406
407
408
    low: float,
    high: float,
) -> None:
    # NOTE(zhaoyang): Due to NaN and Inf representation for fp8 data type,
    # it may occur Inf or NaN if we directly use torch.randint
    # to generate random data for fp8 data.
409
    # For example, s.11111.00 in fp8e5m2 format represents Inf.
410
    #     | E4M3        | E5M2
411
    # -----|-------------|-------------------
412
413
    # Inf | N/A         | s.11111.00
    # NaN | s.1111.111  | s.11111.{01,10,11}
414
    from vllm import _custom_ops as ops
415

416
417
    tensor_tmp = torch.empty_like(tensor, dtype=torch.float16)
    tensor_tmp.uniform_(low, high)
418
    ops.convert_fp8(tensor, tensor_tmp)
419
420
421
    del tensor_tmp


422
def get_kv_cache_torch_dtype(
423
424
    cache_dtype: str | torch.dtype | None,
    model_dtype: str | torch.dtype | None = None,
425
) -> torch.dtype:
426
427
    if isinstance(cache_dtype, str):
        if cache_dtype == "auto":
428
            if isinstance(model_dtype, str) and model_dtype in STR_DTYPE_TO_TORCH_DTYPE:
429
430
431
432
433
                torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[model_dtype]
            elif isinstance(model_dtype, torch.dtype):
                torch_dtype = model_dtype
            else:
                raise ValueError(f"Invalid model dtype: {model_dtype}")
434
        elif cache_dtype in STR_DTYPE_TO_TORCH_DTYPE:
435
436
437
438
439
440
441
            torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype]
        else:
            raise ValueError(f"Invalid kv cache dtype: {cache_dtype}")
    elif isinstance(cache_dtype, torch.dtype):
        torch_dtype = cache_dtype
    else:
        raise ValueError(f"Invalid kv cache dtype: {cache_dtype}")
442
443
444
445
446
447
448
449
450
    return torch_dtype


def create_kv_caches_with_random_flash(
    num_blocks: int,
    block_size: int,
    num_layers: int,
    num_heads: int,
    head_size: int,
451
452
    cache_dtype: str | torch.dtype | None,
    model_dtype: str | torch.dtype | None = None,
453
454
455
    seed: int | None = None,
    device: str | None = "cuda",
    cache_layout: str | None = "NHD",
456
) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
457
    from vllm.platforms import current_platform
458

459
    current_platform.seed_everything(seed)
460

461
    dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
462
463
    generic_kv_cache_shape = (num_blocks, 2, block_size, num_heads, head_size)
    assert cache_layout in ("NHD", "HND")
464
    stride_order = (0, 1, 2, 3, 4) if cache_layout == "NHD" else (0, 1, 3, 2, 4)
465

466
    kv_cache_allocation_shape = tuple(generic_kv_cache_shape[i] for i in stride_order)
467
    scale = head_size**-0.5
468

469
470
    key_caches: list[torch.Tensor] = []
    value_caches: list[torch.Tensor] = []
471

472
    for _ in range(num_layers):
473
        key_value_cache = torch.empty(
474
            size=kv_cache_allocation_shape, dtype=dtype, device=device
475
        ).permute(*stride_order)
476
477
        if cache_dtype in ["auto", "half", "bfloat16", "float"]:
            key_value_cache.uniform_(-scale, scale)
478
        elif cache_dtype == "fp8":
479
480
            _generate_random_fp8(key_value_cache, -scale, scale)
        else:
481
            raise ValueError(f"Does not support key cache of type {cache_dtype}")
482
483
484
485
486
487
488
489
490
491
492
        key_caches.append(key_value_cache[:, 0])
        value_caches.append(key_value_cache[:, 1])
    return key_caches, value_caches


def create_kv_caches_with_random(
    num_blocks: int,
    block_size: int,
    num_layers: int,
    num_heads: int,
    head_size: int,
493
494
    cache_dtype: str | torch.dtype | None,
    model_dtype: str | torch.dtype | None = None,
495
496
    seed: int | None = None,
    device: str | None = "cuda",
497
) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
Joe's avatar
Joe committed
498
499
500
501
    if cache_dtype == "fp8" and head_size % 16:
        raise ValueError(
            f"Does not support key cache of type fp8 with head_size {head_size}"
        )
502
    from vllm.platforms import current_platform
503

504
    current_platform.seed_everything(seed)
505

506
    dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
507
508

    scale = head_size**-0.5
509
    x = 16 // torch.tensor([], dtype=dtype).element_size()
510
    key_cache_shape = (num_blocks, num_heads, head_size // x, block_size, x)
511
    key_caches: list[torch.Tensor] = []
512
    for _ in range(num_layers):
513
        key_cache = torch.empty(size=key_cache_shape, dtype=dtype, device=device)
514
        if cache_dtype in ["auto", "half", "bfloat16", "float"]:
515
            key_cache.uniform_(-scale, scale)
516
        elif cache_dtype == "fp8":
517
            _generate_random_fp8(key_cache, -scale, scale)
518
        else:
519
            raise ValueError(f"Does not support key cache of type {cache_dtype}")
520
521
522
        key_caches.append(key_cache)

    value_cache_shape = (num_blocks, num_heads, head_size, block_size)
523
    value_caches: list[torch.Tensor] = []
524
    for _ in range(num_layers):
525
        value_cache = torch.empty(size=value_cache_shape, dtype=dtype, device=device)
526
        if cache_dtype in ["auto", "half", "bfloat16", "float"]:
527
            value_cache.uniform_(-scale, scale)
528
        elif cache_dtype == "fp8":
529
            _generate_random_fp8(value_cache, -scale, scale)
530
        else:
531
            raise ValueError(f"Does not support value cache of type {cache_dtype}")
532
533
        value_caches.append(value_cache)
    return key_caches, value_caches
534
535


536
@cache
537
def is_pin_memory_available() -> bool:
538
    from vllm.platforms import current_platform
539

540
    return current_platform.is_pin_memory_available()
541
542


543
544
545
546
547
548
549
550
@cache
def is_uva_available() -> bool:
    """Check if Unified Virtual Addressing (UVA) is available."""
    # UVA requires pinned memory.
    # TODO: Add more requirements for UVA if needed.
    return is_pin_memory_available()


551
def make_ndarray_with_pad(
552
    x: list[list[T]],
553
554
555
    pad: T,
    dtype: npt.DTypeLike,
    *,
556
    max_len: int | None = None,
557
558
559
) -> npt.NDArray:
    """
    Make a padded array from 2D inputs.
560
561
562
563

    The padding is applied to the end of each inner list until it reaches
    `max_len`.
    """
564
565
566
567
568
    if max_len is None:
        # Unlike for most functions, map is faster than a genexpr over `len`
        max_len = max(map(len, x), default=0)

    padded_x = np.full((len(x), max_len), pad, dtype=dtype)
569
570
    for ind, blocktb in enumerate(x):
        assert len(blocktb) <= max_len
571
        padded_x[ind, : len(blocktb)] = blocktb
572
573
574
575
576

    return padded_x


def make_tensor_with_pad(
577
    x: list[list[T]],
578
579
580
    pad: T,
    dtype: torch.dtype,
    *,
581
    max_len: int | None = None,
582
    device: str | torch.device | None = None,
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
    pin_memory: bool = False,
) -> torch.Tensor:
    """
    Make a padded tensor from 2D inputs.

    The padding is applied to the end of each inner list until it reaches
    `max_len`.
    """
    np_dtype = TORCH_DTYPE_TO_NUMPY_DTYPE[dtype]
    padded_x = make_ndarray_with_pad(x, pad, np_dtype, max_len=max_len)

    tensor = torch.from_numpy(padded_x).to(device)
    if pin_memory:
        tensor = tensor.pin_memory()

    return tensor
599
600
601
602
603


def async_tensor_h2d(
    data: list,
    dtype: torch.dtype,
604
    target_device: str | torch.device,
605
606
607
608
609
610
611
    pin_memory: bool,
) -> torch.Tensor:
    """Asynchronously create a tensor and copy it from host to device."""
    t = torch.tensor(data, dtype=dtype, pin_memory=pin_memory, device="cpu")
    return t.to(device=target_device, non_blocking=True)


612
613
614
615
616
def get_dtype_size(dtype: torch.dtype) -> int:
    """Get the size of the data type in bytes."""
    return torch.tensor([], dtype=dtype).element_size()


617
618
619
# bool = 0, int = 1, float = 2, complex = 3
def _get_precision_level(dtype: torch.dtype) -> int:
    # NOTE: Complex dtypes return `is_floating_point=False`
620
    return (dtype != torch.bool) + dtype.is_floating_point + dtype.is_complex * 2
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647


def is_lossless_cast(src_dtype: torch.dtype, tgt_dtype: torch.dtype):
    """
    Test whether it is lossless to cast a tensor from
    `src_dtype` to `tgt_dtype`.
    """
    if src_dtype == tgt_dtype:
        return True

    src_level = _get_precision_level(src_dtype)
    tgt_level = _get_precision_level(tgt_dtype)

    if src_level < tgt_level:
        return True
    if src_level > tgt_level:
        return False

    # Compare integral types
    if not src_dtype.is_floating_point and not src_dtype.is_complex:
        src_info = torch.iinfo(src_dtype)
        tgt_info = torch.iinfo(tgt_dtype)
        return src_info.min >= tgt_info.min and src_info.max <= tgt_info.max

    # Compare floating-point types
    src_info = torch.finfo(src_dtype)
    tgt_info = torch.finfo(tgt_dtype)
648
649
650
651
652
    return (
        src_info.min >= tgt_info.min
        and src_info.max <= tgt_info.max
        and src_info.resolution >= tgt_info.resolution
    )
653
654
655
656
657
658
659
660
661
662
663
664
665


def common_broadcastable_dtype(dtypes: Collection[torch.dtype]):
    """
    Get the common `dtype` where all of the other `dtypes` can be
    cast to it without losing any information.
    """
    return max(
        dtypes,
        key=lambda dtype: sum(is_lossless_cast(dt, dtype) for dt in dtypes),
    )


666
667
# TODO: This function can be removed if transformer_modules classes are
# serialized by value when communicating between processes
668
def init_cached_hf_modules() -> None:
669
670
671
672
    """
    Lazy initialization of the Hugging Face modules.
    """
    from transformers.dynamic_module_utils import init_hf_modules
673

674
    init_hf_modules()
675
676


677
@cache
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
def find_library(lib_name: str) -> str:
    """
    Find the library file in the system.
    `lib_name` is full filename, with both prefix and suffix.
    This function resolves `lib_name` to the full path of the library.
    """
    # Adapted from https://github.com/openai/triton/blob/main/third_party/nvidia/backend/driver.py#L19 # noqa
    # According to https://en.wikipedia.org/wiki/Filesystem_Hierarchy_Standard
    # `/sbin/ldconfig` should exist in all Linux systems.
    # `/sbin/ldconfig` searches the library in the system
    libs = subprocess.check_output(["/sbin/ldconfig", "-p"]).decode()
    # each line looks like the following:
    # libcuda.so.1 (libc6,x86-64) => /lib/x86_64-linux-gnu/libcuda.so.1
    locs = [line.split()[-1] for line in libs.splitlines() if lib_name in line]
    # `LD_LIBRARY_PATH` searches the library in the user-defined paths
693
    env_ld_library_path = envs.LD_LIBRARY_PATH
694
695
696
697
698
699
700
701
702
703
704
    if not locs and env_ld_library_path:
        locs = [
            os.path.join(dir, lib_name)
            for dir in env_ld_library_path.split(":")
            if os.path.exists(os.path.join(dir, lib_name))
        ]
    if not locs:
        raise ValueError(f"Cannot find {lib_name} in the system.")
    return locs[0]


705
def find_nccl_library() -> str:
706
707
708
709
710
711
    """
    We either use the library file specified by the `VLLM_NCCL_SO_PATH`
    environment variable, or we find the library file brought by PyTorch.
    After importing `torch`, `libnccl.so.2` or `librccl.so.1` can be
    found by `ctypes` automatically.
    """
712
    so_file = envs.VLLM_NCCL_SO_PATH
713
714
715
716

    # manually load the nccl library
    if so_file:
        logger.info(
717
718
            "Found nccl from environment variable VLLM_NCCL_SO_PATH=%s", so_file
        )
719
720
    else:
        if torch.version.cuda is not None:
721
            so_file = "libnccl.so.2"
722
        elif torch.version.hip is not None:
723
            so_file = "librccl.so.1"
724
725
        else:
            raise ValueError("NCCL only supports CUDA and ROCm backends.")
726
        logger.debug_once("Found nccl from library %s", so_file)
727
    return so_file
728
729


730
def find_nccl_include_paths() -> list[str] | None:
731
732
    """
    We either use the nccl.h specified by the `VLLM_NCCL_INCLUDE_PATH`
733
734
    environment variable, or we find the library file brought by
    nvidia-nccl-cuXX. load_inline by default uses
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
    torch.utils.cpp_extension.include_paths
    """
    paths: list[str] = []
    inc = envs.VLLM_NCCL_INCLUDE_PATH
    if inc and os.path.isdir(inc):
        paths.append(inc)

    try:
        spec = importlib.util.find_spec("nvidia.nccl")
        if spec and getattr(spec, "submodule_search_locations", None):
            for loc in spec.submodule_search_locations:
                inc_dir = os.path.join(loc, "include")
                if os.path.exists(os.path.join(inc_dir, "nccl.h")):
                    paths.append(inc_dir)
    except Exception:
        pass

    seen = set()
    out: list[str] = []
    for p in paths:
        if p and p not in seen:
            out.append(p)
            seen.add(p)
    return out or None


youkaichao's avatar
youkaichao committed
761
762
prev_set_stream = torch.cuda.set_stream

763
_current_stream_tls = threading.local()
youkaichao's avatar
youkaichao committed
764
765
766


def _patched_set_stream(stream: torch.cuda.Stream) -> None:
767
    _current_stream_tls.value = stream
youkaichao's avatar
youkaichao committed
768
769
770
771
772
773
    prev_set_stream(stream)


torch.cuda.set_stream = _patched_set_stream


774
775
776
777
778
class _StreamPlaceholder:
    def __init__(self):
        self.synchronize = lambda: None


youkaichao's avatar
youkaichao committed
779
780
781
782
783
784
785
786
787
788
789
def current_stream() -> torch.cuda.Stream:
    """
    replace `torch.cuda.current_stream()` with `vllm.utils.current_stream()`.
    it turns out that `torch.cuda.current_stream()` is quite expensive,
    as it will construct a new stream object at each call.
    here we patch `torch.cuda.set_stream` to keep track of the current stream
    directly, so that we can avoid calling `torch.cuda.current_stream()`.

    the underlying hypothesis is that we do not call `torch._C._cuda_setStream`
    from C/C++ code.
    """
790
    from vllm.platforms import current_platform
791
792

    if not hasattr(_current_stream_tls, "value") or _current_stream_tls.value is None:
youkaichao's avatar
youkaichao committed
793
794
        # when this function is called before any stream is set,
        # we return the default stream.
795
796
797
        # On ROCm using the default 0 stream in combination with RCCL
        # is hurting performance. Therefore creating a dedicated stream
        # per process
798
        if current_platform.is_rocm():
799
800
            # torch.cuda.set_stream here is the alias of _pathed_set_stream
            torch.cuda.set_stream(torch.cuda.Stream())
801
802
803
804
805
806
807
808
809
        elif current_platform.is_cpu():
            _current_stream_tls.value = _StreamPlaceholder()
        else:
            current_stream = current_platform.current_stream
            if current_stream is not None:
                _current_stream_tls.value = current_stream()
            else:
                raise ValueError(
                    "Fail to set current stream, current platform "
810
811
                    "may not support current_stream with torch API"
                )
812
    return _current_stream_tls.value
youkaichao's avatar
youkaichao committed
813
814


815
def enable_trace_function_call_for_thread(vllm_config: VllmConfig) -> None:
816
817
818
819
    """Set up function tracing for the current thread,
    if enabled via the VLLM_TRACE_FUNCTION environment variable
    """

820
    if envs.VLLM_TRACE_FUNCTION:
821
        tmp_dir = tempfile.gettempdir()
822
823
        # add username to tmp_dir to avoid permission issues
        tmp_dir = os.path.join(tmp_dir, getpass.getuser())
824
825
826
827
828
829
830
831
        filename = (
            f"VLLM_TRACE_FUNCTION_for_process_{os.getpid()}"
            f"_thread_{threading.get_ident()}_"
            f"at_{datetime.datetime.now()}.log"
        ).replace(" ", "_")
        log_path = os.path.join(
            tmp_dir, "vllm", f"vllm-instance-{vllm_config.instance_id}", filename
        )
832
833
        os.makedirs(os.path.dirname(log_path), exist_ok=True)
        enable_trace_function_call(log_path)
834
835


836
@lru_cache(maxsize=8)
837
def _cuda_device_count_stateless(cuda_visible_devices: str | None = None) -> int:
838
839
840
841
842
843
844
845
846
    # Note: cuda_visible_devices is not used, but we keep it as an argument for
    # LRU Cache purposes.

    # Code below is based on
    # https://github.com/pytorch/pytorch/blob/
    # c1cd946818442aca8c7f812b16d187ce1586c3bc/
    # torch/cuda/__init__.py#L831C1-L831C17
    import torch.cuda

847
    from vllm.platforms import current_platform
848

849
850
    if not torch.cuda._is_compiled():
        return 0
851
    if current_platform.is_rocm():
852
853
        # ROCm uses amdsmi instead of nvml for stateless device count
        # This requires a sufficiently modern version of Torch 2.4.0
854
855
856
857
858
        raw_count = (
            torch.cuda._device_count_amdsmi()
            if (hasattr(torch.cuda, "_device_count_amdsmi"))
            else -1
        )
859
860
861
    else:
        raw_count = torch.cuda._device_count_nvml()
    r = torch._C._cuda_getDeviceCount() if raw_count < 0 else raw_count
862
863
864
865
866
867
    return r


def cuda_device_count_stateless() -> int:
    """Get number of CUDA devices, caching based on the value of
    CUDA_VISIBLE_DEVICES at the time of call.
868

869
870
871
872
873
874
875
    This should be used instead of torch.cuda.device_count()
    unless CUDA_VISIBLE_DEVICES has already been set to the desired
    value."""

    # This can be removed and simply replaced with torch.cuda.get_device_count
    # after https://github.com/pytorch/pytorch/pull/122815 is released.
    return _cuda_device_count_stateless(envs.CUDA_VISIBLE_DEVICES)
876
877


878
879
880
881
882
883
884
def cuda_is_initialized() -> bool:
    """Check if CUDA is initialized."""
    if not torch.cuda._is_compiled():
        return False
    return torch.cuda.is_initialized()


885
886
887
888
889
890
891
def xpu_is_initialized() -> bool:
    """Check if XPU is initialized."""
    if not torch.xpu._is_compiled():
        return False
    return torch.xpu.is_initialized()


892
893
894
def cuda_get_device_properties(
    device, names: Sequence[str], init_cuda=False
) -> tuple[Any, ...]:
895
896
897
898
899
900
901
902
903
    """Get specified CUDA device property values without initializing CUDA in
    the current process."""
    if init_cuda or cuda_is_initialized():
        props = torch.cuda.get_device_properties(device)
        return tuple(getattr(props, name) for name in names)

    # Run in subprocess to avoid initializing CUDA as a side effect.
    mp_ctx = multiprocessing.get_context("fork")
    with ProcessPoolExecutor(max_workers=1, mp_context=mp_ctx) as executor:
904
        return executor.submit(cuda_get_device_properties, device, names, True).result()
905
906


907
908
909
def weak_bind(
    bound_method: Callable[..., Any],
) -> Callable[..., None]:
910
911
912
913
914
915
916
917
918
919
920
921
922
    """Make an instance method that weakly references
    its associated instance and no-ops once that
    instance is collected."""
    ref = weakref.ref(bound_method.__self__)  # type: ignore[attr-defined]
    unbound = bound_method.__func__  # type: ignore[attr-defined]

    def weak_bound(*args, **kwargs) -> None:
        if inst := ref():
            unbound(inst, *args, **kwargs)

    return weak_bound


923
class StoreBoolean(Action):
924
925
926
927
928
929
    def __call__(self, parser, namespace, values, option_string=None):
        if values.lower() == "true":
            setattr(namespace, self.dest, True)
        elif values.lower() == "false":
            setattr(namespace, self.dest, False)
        else:
930
931
932
            raise ValueError(
                f"Invalid boolean value: {values}. Expected 'true' or 'false'."
            )
933
934


935
class SortedHelpFormatter(ArgumentDefaultsHelpFormatter, RawDescriptionHelpFormatter):
936
937
    """SortedHelpFormatter that sorts arguments by their option strings."""

938
939
940
941
942
943
944
    def _split_lines(self, text, width):
        """
        1. Sentences split across lines have their single newlines removed.
        2. Paragraphs and explicit newlines are split into separate lines.
        3. Each line is wrapped to the specified width (width of terminal).
        """
        # The patterns also include whitespace after the newline
945
946
        single_newline = re.compile(r"(?<!\n)\n(?!\n)\s*")
        multiple_newlines = re.compile(r"\n{2,}\s*")
947
        text = single_newline.sub(" ", text)
948
949
950
        lines = re.split(multiple_newlines, text)
        return sum([textwrap.wrap(line, width) for line in lines], [])

951
952
    def add_arguments(self, actions):
        actions = sorted(actions, key=lambda x: x.option_strings)
953
        super().add_arguments(actions)
954
955


956
class FlexibleArgumentParser(ArgumentParser):
957
958
    """ArgumentParser that allows both underscore and dash in names."""

959
    _deprecated: set[Action] = set()
960
961
962
963
964
965
966
    _json_tip: str = (
        "When passing JSON CLI arguments, the following sets of arguments "
        "are equivalent:\n"
        '   --json-arg \'{"key1": "value1", "key2": {"key3": "value2"}}\'\n'
        "   --json-arg.key1 value1 --json-arg.key2.key3 value2\n\n"
        "Additionally, list elements can be passed individually using +:\n"
        '   --json-arg \'{"key4": ["value3", "value4", "value5"]}\'\n'
967
968
        "   --json-arg.key4+ value3 --json-arg.key4+='value4,value5'\n\n"
    )
969
    _search_keyword: str | None = None
970

971
    def __init__(self, *args, **kwargs):
972
973
974
        # Set the default "formatter_class" to SortedHelpFormatter
        if "formatter_class" not in kwargs:
            kwargs["formatter_class"] = SortedHelpFormatter
975
976
        # Pop kwarg "add_json_tip" to control whether to add the JSON tip
        self.add_json_tip = kwargs.pop("add_json_tip", True)
977
978
        super().__init__(*args, **kwargs)

979
    if sys.version_info < (3, 13):
980
        # Enable the deprecated kwarg for Python 3.12 and below
981

982
        def parse_known_args(self, args=None, namespace=None):
983
984
985
986
            if args is not None and "--disable-log-requests" in args:
                # Special case warning because the warning below won't trigger
                # if –-disable-log-requests because its value is default.
                logger.warning_once(
987
988
                    "argument '--disable-log-requests' is deprecated and "
                    "replaced with '--enable-log-requests'. This will be "
989
990
                    "removed in v0.12.0."
                )
991
992
            namespace, args = super().parse_known_args(args, namespace)
            for action in FlexibleArgumentParser._deprecated:
993
994
995
996
                if (
                    hasattr(namespace, dest := action.dest)
                    and getattr(namespace, dest) != action.default
                ):
997
                    logger.warning_once("argument '%s' is deprecated", dest)
998
999
            return namespace, args

1000
1001
        def add_argument(self, *args, **kwargs):
            deprecated = kwargs.pop("deprecated", False)
1002
            action = super().add_argument(*args, **kwargs)
1003
1004
            if deprecated:
                FlexibleArgumentParser._deprecated.add(action)
1005
1006
            return action

1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
        class _FlexibleArgumentGroup(_ArgumentGroup):
            def add_argument(self, *args, **kwargs):
                deprecated = kwargs.pop("deprecated", False)
                action = super().add_argument(*args, **kwargs)
                if deprecated:
                    FlexibleArgumentParser._deprecated.add(action)
                return action

        def add_argument_group(self, *args, **kwargs):
            group = self._FlexibleArgumentGroup(self, *args, **kwargs)
            self._action_groups.append(group)
            return group
1019

1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
    def format_help(self):
        # Only use custom help formatting for bottom level parsers
        if self._subparsers is not None:
            return super().format_help()

        formatter = self._get_formatter()

        # Handle keyword search of the args
        if (search_keyword := self._search_keyword) is not None:
            # Normalise the search keyword
            search_keyword = search_keyword.lower().replace("_", "-")
            # Return full help if searching for 'all'
1032
            if search_keyword == "all":
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
                self.epilog = self._json_tip
                return super().format_help()

            # Return group help if searching for a group title
            for group in self._action_groups:
                if group.title and group.title.lower() == search_keyword:
                    formatter.start_section(group.title)
                    formatter.add_text(group.description)
                    formatter.add_arguments(group._group_actions)
                    formatter.end_section()
                    formatter.add_text(self._json_tip)
                    return formatter.format_help()

            # Return matched args if searching for an arg name
            matched_actions = []
            for group in self._action_groups:
                for action in group._group_actions:
                    # search option name
1051
1052
1053
                    if any(
                        search_keyword in opt.lower() for opt in action.option_strings
                    ):
1054
1055
                        matched_actions.append(action)
            if matched_actions:
1056
                formatter.start_section(f"Arguments matching '{search_keyword}'")
1057
1058
1059
1060
1061
1062
1063
1064
1065
                formatter.add_arguments(matched_actions)
                formatter.end_section()
                formatter.add_text(self._json_tip)
                return formatter.format_help()

            # No match found
            formatter.add_text(
                f"No group or arguments matching '{search_keyword}'.\n"
                "Use '--help' to see available groups or "
1066
1067
                "'--help=all' to see all available parameters."
            )
1068
1069
1070
            return formatter.format_help()

        # usage
1071
        formatter.add_usage(self.usage, self._actions, self._mutually_exclusive_groups)
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092

        # description
        formatter.add_text(self.description)

        # positionals, optionals and user-defined groups
        formatter.start_section("Config Groups")
        config_groups = ""
        for group in self._action_groups:
            if not group._group_actions:
                continue
            title = group.title
            description = group.description or ""
            config_groups += f"{title: <24}{description}\n"
        formatter.add_text(config_groups)
        formatter.end_section()

        # epilog
        formatter.add_text(self.epilog)

        # determine help from format above
        return formatter.format_help()
1093

1094
1095
1096
1097
1098
    def parse_args(  # type: ignore[override]
        self,
        args: list[str] | None = None,
        namespace: Namespace | None = None,
    ):
1099
1100
1101
        if args is None:
            args = sys.argv[1:]

1102
1103
        # Check for --model in command line arguments first
        if args and args[0] == "serve":
1104
1105
            try:
                model_idx = next(
1106
1107
1108
1109
                    i
                    for i, arg in enumerate(args)
                    if arg == "--model" or arg.startswith("--model=")
                )
1110
                logger.warning(
1111
1112
                    "With `vllm serve`, you should provide the model as a "
                    "positional argument or in a config file instead of via "
1113
                    "the `--model` option. "
1114
1115
                    "The `--model` option will be removed in v0.13."
                )
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137

                if args[model_idx] == "--model":
                    model_tag = args[model_idx + 1]
                    rest_start_idx = model_idx + 2
                else:
                    model_tag = args[model_idx].removeprefix("--model=")
                    rest_start_idx = model_idx + 1

                # Move <model> to the front, e,g:
                # [Before]
                # vllm serve -tp 2 --model <model> --enforce-eager --port 8001
                # [After]
                # vllm serve <model> -tp 2 --enforce-eager --port 8001
                args = [
                    "serve",
                    model_tag,
                    *args[1:model_idx],
                    *args[rest_start_idx:],
                ]
                print("args", args)
            except StopIteration:
                pass
1138

1139
        if "--config" in args:
1140
            args = self._pull_args_from_config(args)
1141

1142
1143
1144
1145
1146
1147
1148
        def repl(match: re.Match) -> str:
            """Replaces underscores with dashes in the matched string."""
            return match.group(0).replace("_", "-")

        # Everything between the first -- and the first .
        pattern = re.compile(r"(?<=--)[^\.]*")

1149
        # Convert underscores to dashes and vice versa in argument names
1150
        processed_args = list[str]()
1151
        for i, arg in enumerate(args):
1152
            if arg.startswith("--help="):
1153
                FlexibleArgumentParser._search_keyword = arg.split("=", 1)[-1].lower()
1154
                processed_args.append("--help")
1155
1156
1157
            elif arg.startswith("--"):
                if "=" in arg:
                    key, value = arg.split("=", 1)
1158
                    key = pattern.sub(repl, key, count=1)
1159
                    processed_args.append(f"{key}={value}")
1160
                else:
1161
1162
                    key = pattern.sub(repl, arg, count=1)
                    processed_args.append(key)
1163
            elif arg.startswith("-O") and arg != "-O" and arg[2] != ".":
1164
1165
                # allow -O flag to be used without space, e.g. -O3 or -Odecode
                # -O.<...> handled later
1166
1167
1168
                # also handle -O=<mode> here
                mode = arg[3:] if arg[2] == "=" else arg[2:]
                processed_args.append(f"-O.mode={mode}")
1169
1170
1171
1172
1173
            elif (
                arg == "-O"
                and i + 1 < len(args)
                and args[i + 1] in {"0", "1", "2", "3"}
            ):
1174
1175
                # Convert -O <n> to -O.mode <n>
                processed_args.append("-O.mode")
1176
1177
1178
            else:
                processed_args.append(arg)

1179
        def create_nested_dict(keys: list[str], value: str) -> dict[str, Any]:
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
            """Creates a nested dictionary from a list of keys and a value.

            For example, `keys = ["a", "b", "c"]` and `value = 1` will create:
            `{"a": {"b": {"c": 1}}}`
            """
            nested_dict: Any = value
            for key in reversed(keys):
                nested_dict = {key: nested_dict}
            return nested_dict

1190
1191
1192
        def recursive_dict_update(
            original: dict[str, Any],
            update: dict[str, Any],
1193
1194
1195
1196
1197
        ) -> set[str]:
            """Recursively updates a dictionary with another dictionary.
            Returns a set of duplicate keys that were overwritten.
            """
            duplicates = set[str]()
1198
1199
            for k, v in update.items():
                if isinstance(v, dict) and isinstance(original.get(k), dict):
1200
1201
1202
1203
                    nested_duplicates = recursive_dict_update(original[k], v)
                    duplicates |= {f"{k}.{d}" for d in nested_duplicates}
                elif isinstance(v, list) and isinstance(original.get(k), list):
                    original[k] += v
1204
                else:
1205
1206
                    if k in original:
                        duplicates.add(k)
1207
                    original[k] = v
1208
            return duplicates
1209

1210
1211
        delete = set[int]()
        dict_args = defaultdict[str, dict[str, Any]](dict)
1212
        duplicates = set[str]()
1213
        for i, processed_arg in enumerate(processed_args):
1214
1215
1216
1217
            if i in delete:  # skip if value from previous arg
                continue

            if processed_arg.startswith("-") and "." in processed_arg:
1218
                if "=" in processed_arg:
1219
                    processed_arg, value_str = processed_arg.split("=", 1)
1220
                    if "." not in processed_arg:
1221
                        # False positive, '.' was only in the value
1222
1223
                        continue
                else:
1224
                    value_str = processed_args[i + 1]
1225
                    delete.add(i + 1)
1226

1227
1228
1229
1230
                if processed_arg.endswith("+"):
                    processed_arg = processed_arg[:-1]
                    value_str = json.dumps(list(value_str.split(",")))

1231
                key, *keys = processed_arg.split(".")
1232
1233
1234
1235
1236
                try:
                    value = json.loads(value_str)
                except json.decoder.JSONDecodeError:
                    value = value_str

1237
1238
                # Merge all values with the same key into a single dict
                arg_dict = create_nested_dict(keys, value)
1239
1240
                arg_duplicates = recursive_dict_update(dict_args[key], arg_dict)
                duplicates |= {f"{key}.{d}" for d in arg_duplicates}
1241
1242
                delete.add(i)
        # Filter out the dict args we set to None
1243
        processed_args = [a for i, a in enumerate(processed_args) if i not in delete]
1244
1245
1246
        if duplicates:
            logger.warning("Found duplicate keys %s", ", ".join(duplicates))

1247
1248
1249
1250
1251
        # Add the dict args back as if they were originally passed as JSON
        for dict_arg, dict_value in dict_args.items():
            processed_args.append(dict_arg)
            processed_args.append(json.dumps(dict_value))

1252
        return super().parse_args(processed_args, namespace)
1253

1254
1255
1256
1257
    def check_port(self, value):
        try:
            value = int(value)
        except ValueError:
1258
            msg = "Port must be an integer"
1259
            raise ArgumentTypeError(msg) from None
1260
1261

        if not (1024 <= value <= 65535):
1262
            raise ArgumentTypeError("Port must be between 1024 and 65535")
1263
1264
1265

        return value

1266
    def _pull_args_from_config(self, args: list[str]) -> list[str]:
1267
1268
        """Method to pull arguments specified in the config file
        into the command-line args variable.
1269
1270

        The arguments in config file will be inserted between
1271
        the argument list.
1272

1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
        example:
        ```yaml
            port: 12323
            tensor-parallel-size: 4
        ```
        ```python
        $: vllm {serve,chat,complete} "facebook/opt-12B" \
            --config config.yaml -tp 2
        $: args = [
            "serve,chat,complete",
1283
1284
            "facebook/opt-12B",
            '--config', 'config.yaml',
1285
1286
1287
1288
            '-tp', '2'
        ]
        $: args = [
            "serve,chat,complete",
1289
1290
1291
            "facebook/opt-12B",
            '--port', '12323',
            '--tensor-parallel-size', '4',
1292
1293
1294
1295
1296
            '-tp', '2'
            ]
        ```

        Please note how the config args are inserted after the sub command.
1297
        this way the order of priorities is maintained when these are args
1298
1299
        parsed by super().
        """
1300
        assert args.count("--config") <= 1, "More than one config file specified!"
1301

1302
        index = args.index("--config")
1303
        if index == len(args) - 1:
1304
1305
1306
1307
            raise ValueError(
                "No config file specified! \
                             Please check your command-line arguments."
            )
1308
1309
1310

        file_path = args[index + 1]

1311
        config_args = self.load_config_file(file_path)
1312

1313
        # 0th index might be the sub command {serve,chat,complete,...}
1314
        # optionally followed by model_tag (only for serve)
1315
1316
1317
1318
        # followed by config args
        # followed by rest of cli args.
        # maintaining this order will enforce the precedence
        # of cli > config > defaults
1319
        if args[0].startswith("-"):
1320
            # No sub command (e.g., api_server entry point)
1321
            args = config_args + args[0:index] + args[index + 2 :]
1322
        elif args[0] == "serve":
1323
1324
            model_in_cli = len(args) > 1 and not args[1].startswith("-")
            model_in_config = any(arg == "--model" for arg in config_args)
1325
1326

            if not model_in_cli and not model_in_config:
1327
                raise ValueError(
1328
                    "No model specified! Please specify model either "
1329
1330
                    "as a positional argument or in a config file."
                )
1331
1332
1333

            if model_in_cli:
                # Model specified as positional arg, keep CLI version
1334
1335
1336
1337
1338
1339
1340
                args = (
                    [args[0]]
                    + [args[1]]
                    + config_args
                    + args[2:index]
                    + args[index + 2 :]
                )
1341
1342
            else:
                # No model in CLI, use config if available
1343
                args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
1344
        else:
1345
            args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
1346
1347
1348

        return args

1349
    def load_config_file(self, file_path: str) -> list[str]:
1350
        """Loads a yaml file and returns the key value pairs as a
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
        flattened list with argparse like pattern
        ```yaml
            port: 12323
            tensor-parallel-size: 4
        ```
        returns:
            processed_args: list[str] = [
                '--port': '12323',
                '--tensor-parallel-size': '4'
            ]
        """
1362
1363
        extension: str = file_path.split(".")[-1]
        if extension not in ("yaml", "yml"):
1364
1365
            raise ValueError(
                "Config file must be of a yaml/yml type.\
1366
1367
1368
                              %s supplied",
                extension,
            )
1369
1370

        # only expecting a flat dictionary of atomic types
1371
        processed_args: list[str] = []
1372

1373
        config: dict[str, int | str] = {}
1374
        try:
1375
            with open(file_path) as config_file:
1376
1377
1378
1379
                config = yaml.safe_load(config_file)
        except Exception as ex:
            logger.error(
                "Unable to read the config file at %s. \
1380
1381
1382
                Make sure path is correct",
                file_path,
            )
1383
1384
            raise ex

1385
        store_boolean_arguments = [
1386
            action.dest for action in self._actions if isinstance(action, StoreBoolean)
1387
1388
        ]

1389
        for key, value in config.items():
1390
1391
            if isinstance(value, bool) and key not in store_boolean_arguments:
                if value:
1392
                    processed_args.append("--" + key)
1393
1394
            elif isinstance(value, list):
                if value:
1395
                    processed_args.append("--" + key)
1396
1397
                    for item in value:
                        processed_args.append(str(item))
1398
            else:
1399
                processed_args.append("--" + key)
1400
                processed_args.append(str(value))
1401
1402
1403

        return processed_args

1404

1405
1406
1407
1408
1409
1410
# Using dynamo with vLLM doesn't really work well with PyTorch versions < 2.4.0.
# In particular, the FakeScalarType is not supported for earlier versions of
# PyTorch which breaks dynamo for any ops registered using ScalarType.
def supports_dynamo() -> bool:
    base_torch_version = Version(Version(torch.__version__).base_version)
    return base_torch_version >= Version("2.4.0")
1411
1412


1413
# Supports xccl with PyTorch versions >= 2.8.0.dev for XPU platform
1414
def supports_xccl() -> bool:
1415
1416
1417
    return (
        is_torch_equal_or_newer("2.8.0.dev") and torch.distributed.is_xccl_available()
    )
1418
1419


1420
1421
1422
1423
1424
1425
# Some backends use pytorch version < 2.4.0 which doesn't
# support `torch.library.custom_op`.
def supports_custom_op() -> bool:
    return hasattr(torch.library, "custom_op")


1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
class AtomicCounter:
    """An atomic, thread-safe counter"""

    def __init__(self, initial=0):
        """Initialize a new atomic counter to given initial value"""
        self._value = initial
        self._lock = threading.Lock()

    def inc(self, num=1):
        """Atomically increment the counter by num and return the new value"""
        with self._lock:
            self._value += num
            return self._value

    def dec(self, num=1):
        """Atomically decrement the counter by num and return the new value"""
        with self._lock:
            self._value -= num
            return self._value

    @property
    def value(self):
        return self._value
1449
1450


1451
def weak_ref_tensor(tensor: Any) -> Any:
1452
1453
1454
1455
1456
    """
    Create a weak reference to a tensor.
    The new tensor will share the same data as the original tensor,
    but will not keep the original tensor alive.
    """
1457
1458
1459
1460
    if isinstance(tensor, torch.Tensor):
        return torch.ops._C.weak_ref_tensor(tensor)
    else:
        return tensor
1461
1462
1463


def weak_ref_tensors(
1464
1465
1466
1467
1468
    tensors: torch.Tensor
    | list[torch.Tensor]
    | tuple[torch.Tensor]
    | IntermediateTensors,
) -> torch.Tensor | list[Any] | tuple[Any] | Any:
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
    """
    Convenience function to create weak references to tensors,
    for single tensor, list of tensors or tuple of tensors.
    """
    if isinstance(tensors, torch.Tensor):
        return weak_ref_tensor(tensors)
    if isinstance(tensors, list):
        return [weak_ref_tensor(t) for t in tensors]
    if isinstance(tensors, tuple):
        return tuple(weak_ref_tensor(t) for t in tensors)
1479
1480
1481

    # For IntermediateTensors used in pipeline parallelism
    from vllm.sequence import IntermediateTensors
1482

1483
    if isinstance(tensors, IntermediateTensors):
1484
1485
1486
        ret = IntermediateTensors(
            {key: weak_ref_tensor(val) for key, val in tensors.tensors.items()}
        )
1487
        return ret
1488
    raise ValueError("Invalid type for tensors")
1489
1490


1491
1492
1493
1494
1495
1496
1497
1498
def get_cuda_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tensor:
    """
    Get a CUDA view of a CPU tensor using Unified Virtual Addressing (UVA).
    """
    assert cpu_tensor.is_pinned(), "CPU tensor must be pinned"
    return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor)


1499
1500
1501
1502
1503
# create a library to hold the custom op
vllm_lib = Library("vllm", "FRAGMENT")  # noqa


def direct_register_custom_op(
1504
1505
    op_name: str,
    op_func: Callable,
1506
1507
1508
1509
    mutates_args: list[str] | None = None,
    fake_impl: Callable | None = None,
    target_lib: Library | None = None,
    dispatch_key: str | None = None,
1510
    tags: tuple[torch.Tag, ...] = (),
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
):
    """
    `torch.library.custom_op` can have significant overhead because it
    needs to consider complicated dispatching logic. This function
    directly registers a custom op and dispatches it to the CUDA backend.
    See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5
    for more details.

    By default, the custom op is registered to the vLLM library. If you
    want to register it to a different library, you can pass the library
    object to the `target_lib` argument.

    IMPORTANT: the lifetime of the operator is tied to the lifetime of the
    library object. If you want to bind the operator to a different library,
    make sure the library object is alive when the operator is used.
    """
1527
    if not supports_custom_op():
1528
        from vllm.platforms import current_platform
1529

1530
1531
1532
1533
1534
        assert not current_platform.is_cuda_alike(), (
            "cuda platform needs torch>=2.4 to support custom op, "
            "chances are you are using an old version of pytorch "
            "or a custom build of pytorch. It is recommended to "
            "use vLLM in a fresh new environment and let it install "
1535
1536
            "the required dependencies."
        )
1537
1538
        return

1539
1540
1541
1542
1543
    if mutates_args is None:
        mutates_args = []

    if dispatch_key is None:
        from vllm.platforms import current_platform
1544

1545
1546
        dispatch_key = current_platform.dispatch_key

1547
    import torch.library
1548

1549
    if hasattr(torch.library, "infer_schema"):
1550
        schema_str = torch.library.infer_schema(op_func, mutates_args=mutates_args)
1551
1552
1553
    else:
        # for pytorch 2.4
        import torch._custom_op.impl
1554

1555
        schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args)
1556
    my_lib = target_lib or vllm_lib
1557
    my_lib.define(op_name + schema_str, tags=tags)
1558
    my_lib.impl(op_name, op_func, dispatch_key=dispatch_key)
1559
1560
    if fake_impl is not None:
        my_lib._register_fake(op_name, fake_impl)
1561
1562


1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
def kill_process_tree(pid: int):
    """
    Kills all descendant processes of the given pid by sending SIGKILL.

    Args:
        pid (int): Process ID of the parent process
    """
    try:
        parent = psutil.Process(pid)
    except psutil.NoSuchProcess:
        return

    # Get all children recursively
    children = parent.children(recursive=True)

    # Send SIGKILL to all children first
    for child in children:
        with contextlib.suppress(ProcessLookupError):
            os.kill(child.pid, signal.SIGKILL)

    # Finally kill the parent
    with contextlib.suppress(ProcessLookupError):
        os.kill(pid, signal.SIGKILL)
1586
1587


1588
# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L630 # noqa: E501
1589
def set_ulimit(target_soft_limit=65535):
1590
    if sys.platform.startswith("win"):
1591
1592
1593
1594
        logger.info("Windows detected, skipping ulimit adjustment.")
        return

    import resource
1595

1596
1597
1598
1599
1600
    resource_type = resource.RLIMIT_NOFILE
    current_soft, current_hard = resource.getrlimit(resource_type)

    if current_soft < target_soft_limit:
        try:
1601
            resource.setrlimit(resource_type, (target_soft_limit, current_hard))
1602
1603
        except ValueError as e:
            logger.warning(
1604
1605
                "Found ulimit of %s and failed to automatically increase "
                "with error %s. This can cause fd limit errors like "
1606
                "`OSError: [Errno 24] Too many open files`. Consider "
1607
1608
1609
1610
                "increasing with ulimit -n",
                current_soft,
                e,
            )
1611
1612
1613
1614
1615
1616
1617
1618
1619


# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/utils.py#L28 # noqa: E501
def get_exception_traceback():
    etype, value, tb = sys.exc_info()
    err_str = "".join(traceback.format_exception(etype, value, tb))
    return err_str


1620
def split_zmq_path(path: str) -> tuple[str, str, str]:
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
    """Split a zmq path into its parts."""
    parsed = urlparse(path)
    if not parsed.scheme:
        raise ValueError(f"Invalid zmq path: {path}")

    scheme = parsed.scheme
    host = parsed.hostname or ""
    port = str(parsed.port or "")

    if scheme == "tcp" and not all((host, port)):
        # The host and port fields are required for tcp
        raise ValueError(f"Invalid zmq path: {path}")

    if scheme != "tcp" and port:
        # port only makes sense with tcp
        raise ValueError(f"Invalid zmq path: {path}")

    return scheme, host, port


1641
def make_zmq_path(scheme: str, host: str, port: int | None = None) -> str:
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
    """Make a ZMQ path from its parts.

    Args:
        scheme: The ZMQ transport scheme (e.g. tcp, ipc, inproc).
        host: The host - can be an IPv4 address, IPv6 address, or hostname.
        port: Optional port number, only used for TCP sockets.

    Returns:
        A properly formatted ZMQ path string.
    """
1652
    if port is None:
1653
1654
1655
1656
1657
1658
        return f"{scheme}://{host}"
    if is_valid_ipv6_address(host):
        return f"{scheme}://[{host}]:{port}"
    return f"{scheme}://{host}:{port}"


1659
1660
# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L783 # noqa: E501
def make_zmq_socket(
1661
    ctx: zmq.asyncio.Context | zmq.Context,  # type: ignore[name-defined]
1662
    path: str,
1663
    socket_type: Any,
1664
1665
1666
    bind: bool | None = None,
    identity: bytes | None = None,
    linger: int | None = None,
1667
) -> zmq.Socket | zmq.asyncio.Socket:  # type: ignore[name-defined]
1668
1669
1670
    """Make a ZMQ socket with the proper bind/connect semantics."""

    mem = psutil.virtual_memory()
1671
    socket = ctx.socket(socket_type)
1672
1673
1674
1675
1676
1677
1678
1679

    # Calculate buffer size based on system memory
    total_mem = mem.total / 1024**3
    available_mem = mem.available / 1024**3
    # For systems with substantial memory (>32GB total, >16GB available):
    # - Set a large 0.5GB buffer to improve throughput
    # For systems with less memory:
    # - Use system default (-1) to avoid excessive memory consumption
1680
    buf_size = int(0.5 * 1024**3) if total_mem > 32 and available_mem > 16 else -1
1681

1682
    if bind is None:
1683
        bind = socket_type not in (zmq.PUSH, zmq.SUB, zmq.XSUB)
1684
1685
1686
1687
1688
1689
1690
1691
1692
1693
1694
1695

    if socket_type in (zmq.PULL, zmq.DEALER, zmq.ROUTER):
        socket.setsockopt(zmq.RCVHWM, 0)
        socket.setsockopt(zmq.RCVBUF, buf_size)

    if socket_type in (zmq.PUSH, zmq.DEALER, zmq.ROUTER):
        socket.setsockopt(zmq.SNDHWM, 0)
        socket.setsockopt(zmq.SNDBUF, buf_size)

    if identity is not None:
        socket.setsockopt(zmq.IDENTITY, identity)

1696
1697
1698
    if linger is not None:
        socket.setsockopt(zmq.LINGER, linger)

1699
1700
1701
    if socket_type == zmq.XPUB:
        socket.setsockopt(zmq.XPUB_VERBOSE, True)

1702
1703
1704
1705
1706
1707
    # Determine if the path is a TCP socket with an IPv6 address.
    # Enable IPv6 on the zmq socket if so.
    scheme, host, _ = split_zmq_path(path)
    if scheme == "tcp" and is_valid_ipv6_address(host):
        socket.setsockopt(zmq.IPV6, 1)

1708
    if bind:
1709
        socket.bind(path)
1710
    else:
1711
        socket.connect(path)
1712
1713
1714
1715
1716

    return socket


@contextlib.contextmanager
1717
1718
1719
def zmq_socket_ctx(
    path: str,
    socket_type: Any,
1720
    bind: bool | None = None,
1721
    linger: int = 0,
1722
    identity: bytes | None = None,
1723
) -> Iterator[zmq.Socket]:
1724
1725
    """Context manager for a ZMQ socket"""

1726
    ctx = zmq.Context()  # type: ignore[attr-defined]
1727
    try:
1728
        yield make_zmq_socket(ctx, path, socket_type, bind=bind, identity=identity)
1729
1730
1731
1732
    except KeyboardInterrupt:
        logger.debug("Got Keyboard Interrupt.")

    finally:
1733
        ctx.destroy(linger=linger)
1734
1735


1736
1737
1738
1739
1740
1741
1742
def _maybe_force_spawn():
    """Check if we need to force the use of the `spawn` multiprocessing start
    method.
    """
    if os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") == "spawn":
        return

1743
1744
    reasons = []
    if is_in_ray_actor():
1745
1746
1747
1748
        # even if we choose to spawn, we need to pass the ray address
        # to the subprocess so that it knows how to connect to the ray cluster.
        # env vars are inherited by subprocesses, even if we use spawn.
        import ray
1749

1750
        os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address
1751
1752
1753
1754
1755
1756
        reasons.append("In a Ray actor and can only be spawned")

    if cuda_is_initialized():
        reasons.append("CUDA is initialized")
    elif xpu_is_initialized():
        reasons.append("XPU is initialized")
1757

1758
    if reasons:
1759
1760
1761
        logger.warning(
            "We must use the `spawn` multiprocessing start method. "
            "Overriding VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. "
1762
            "See https://docs.vllm.ai/en/latest/usage/"
1763
            "troubleshooting.html#python-multiprocessing "
1764
1765
1766
            "for more information. Reasons: %s",
            "; ".join(reasons),
        )
1767
1768
1769
1770
        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"


def get_mp_context():
1771
1772
1773
1774
1775
1776
1777
    """Get a multiprocessing context with a particular method (spawn or fork).
    By default we follow the value of the VLLM_WORKER_MULTIPROC_METHOD to
    determine the multiprocessing method (default is fork). However, under
    certain conditions, we may enforce spawn and override the value of
    VLLM_WORKER_MULTIPROC_METHOD.
    """
    _maybe_force_spawn()
1778
1779
    mp_method = envs.VLLM_WORKER_MULTIPROC_METHOD
    return multiprocessing.get_context(mp_method)
1780
1781
1782


def bind_kv_cache(
1783
1784
    ctx: dict[str, Any],
    kv_cache: list[list[torch.Tensor]],  # [virtual_engine][layer_index]
1785
    shared_kv_cache_layers: dict[str, str] | None = None,
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
) -> None:
    # Bind the kv_cache tensor to Attention modules, similar to
    # ctx[layer_name].kv_cache[ve]=kv_cache[ve][extract_layer_index(layer_name)]
    # Special things handled here:
    # 1. Some models have non-attention layers, e.g., Jamba
    # 2. Pipeline parallelism, each rank only has a subset of layers
    # 3. Encoder attention has no kv cache
    # 4. Encoder-decoder models, encoder-decoder attention and decoder-only
    #    attention of the same layer (e.g., bart's decoder.layers.1.self_attn
    #    and decoder.layers.1.encoder_attn) is mapped to the same kv cache
    #    tensor
1797
1798
1799
1800
    # 5. Some models have attention layers that share kv cache with previous
    #    layers, this is specified through shared_kv_cache_layers
    if shared_kv_cache_layers is None:
        shared_kv_cache_layers = {}
1801
1802
    from vllm.attention import AttentionType
    from vllm.model_executor.models.utils import extract_layer_index
1803

1804
    layer_need_kv_cache = [
1805
1806
1807
1808
1809
1810
1811
1812
        layer_name
        for layer_name in ctx
        if (
            hasattr(ctx[layer_name], "attn_type")
            and ctx[layer_name].attn_type
            in (AttentionType.DECODER, AttentionType.ENCODER_DECODER)
        )
        and ctx[layer_name].kv_sharing_target_layer_name is None
1813
1814
    ]
    layer_index_sorted = sorted(
1815
1816
        set(extract_layer_index(layer_name) for layer_name in layer_need_kv_cache)
    )
1817
    for layer_name in layer_need_kv_cache:
1818
        kv_cache_idx = layer_index_sorted.index(extract_layer_index(layer_name))
1819
1820
1821
1822
        forward_ctx = ctx[layer_name]
        assert len(forward_ctx.kv_cache) == len(kv_cache)
        for ve, ve_kv_cache in enumerate(kv_cache):
            forward_ctx.kv_cache[ve] = ve_kv_cache[kv_cache_idx]
1823
1824
    if shared_kv_cache_layers is not None:
        for layer_name, target_layer_name in shared_kv_cache_layers.items():
1825
1826
1827
            assert extract_layer_index(target_layer_name) < extract_layer_index(
                layer_name
            ), "v0 doesn't support interleaving kv sharing"
1828
            ctx[layer_name].kv_cache = ctx[target_layer_name].kv_cache
1829
1830


1831
1832
def run_method(
    obj: Any,
1833
    method: str | bytes | Callable,
1834
1835
1836
    args: tuple[Any],
    kwargs: dict[str, Any],
) -> Any:
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
    """
    Run a method of an object with the given arguments and keyword arguments.
    If the method is string, it will be converted to a method using getattr.
    If the method is serialized bytes and will be deserialized using
    cloudpickle.
    If the method is a callable, it will be called directly.
    """
    if isinstance(method, bytes):
        func = partial(cloudpickle.loads(method), obj)
    elif isinstance(method, str):
        try:
            func = getattr(obj, method)
        except AttributeError:
1850
1851
1852
            raise NotImplementedError(
                f"Method {method!r} is not implemented."
            ) from None
1853
1854
1855
    else:
        func = partial(method, obj)  # type: ignore
    return func(*args, **kwargs)
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
1873
1874
1875
1876


def import_pynvml():
    """
    Historical comments:

    libnvml.so is the library behind nvidia-smi, and
    pynvml is a Python wrapper around it. We use it to get GPU
    status without initializing CUDA context in the current process.
    Historically, there are two packages that provide pynvml:
    - `nvidia-ml-py` (https://pypi.org/project/nvidia-ml-py/): The official
        wrapper. It is a dependency of vLLM, and is installed when users
        install vLLM. It provides a Python module named `pynvml`.
    - `pynvml` (https://pypi.org/project/pynvml/): An unofficial wrapper.
        Prior to version 12.0, it also provides a Python module `pynvml`,
        and therefore conflicts with the official one. What's worse,
        the module is a Python package, and has higher priority than
        the official one which is a standalone Python file.
        This causes errors when both of them are installed.
        Starting from version 12.0, it migrates to a new module
        named `pynvml_utils` to avoid the conflict.
1877
1878
1879
1880
1881
1882
1883
    It is so confusing that many packages in the community use the
    unofficial one by mistake, and we have to handle this case.
    For example, `nvcr.io/nvidia/pytorch:24.12-py3` uses the unofficial
    one, and it will cause errors, see the issue
    https://github.com/vllm-project/vllm/issues/12847 for example.
    After all the troubles, we decide to copy the official `pynvml`
    module to our codebase, and use it directly.
1884
    """
1885
    import vllm.third_party.pynvml as pynvml
1886

1887
    return pynvml
1888
1889


1890
def warn_for_unimplemented_methods(cls: type[T]) -> type[T]:
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
    """
    A replacement for `abc.ABC`.
    When we use `abc.ABC`, subclasses will fail to instantiate
    if they do not implement all abstract methods.
    Here, we only require `raise NotImplementedError` in the
    base class, and log a warning if the method is not implemented
    in the subclass.
    """

    original_init = cls.__init__

    def find_unimplemented_methods(self: object):
        unimplemented_methods = []
        for attr_name in dir(self):
            # bypass inner method
1906
            if attr_name.startswith("_"):
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
                continue

            try:
                attr = getattr(self, attr_name)
                # get the func of callable method
                if callable(attr):
                    attr_func = attr.__func__
            except AttributeError:
                continue
            src = inspect.getsource(attr_func)
            if "NotImplementedError" in src:
                unimplemented_methods.append(attr_name)
        if unimplemented_methods:
1920
1921
            method_names = ",".join(unimplemented_methods)
            msg = f"Methods {method_names} not implemented in {self}"
1922
            logger.debug(msg)
1923
1924
1925
1926
1927
1928

    @wraps(original_init)
    def wrapped_init(self, *args, **kwargs) -> None:
        original_init(self, *args, **kwargs)
        find_unimplemented_methods(self)

1929
    type.__setattr__(cls, "__init__", wrapped_init)
1930
    return cls
1931
1932


1933
@contextlib.contextmanager
1934
def cprofile_context(save_file: str | None = None):
1935
1936
1937
1938
    """Run a cprofile

    Args:
        save_file: path to save the profile result. "1" or
1939
            None will result in printing to stdout.
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
    """
    import cProfile

    prof = cProfile.Profile()
    prof.enable()

    try:
        yield
    finally:
        prof.disable()
        if save_file and save_file != "1":
            prof.dump_stats(save_file)
        else:
            prof.print_stats(sort="cumtime")


1956
def cprofile(save_file: str | None = None, enabled: bool = True):
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
    """Decorator to profile a Python method using cProfile.

    Args:
        save_file: Path to save the profile result.
            If "1", None, or "", results will be printed to stdout.
        enabled: Set to false to turn this into a no-op
    """

    def decorator(func: Callable):
        @wraps(func)
        def wrapper(*args, **kwargs):
            if not enabled:
                # If profiling is disabled, just call the function directly.
                return func(*args, **kwargs)

            with cprofile_context(save_file):
                return func(*args, **kwargs)

        return wrapper

    return decorator
1978
1979


1980
1981
# Only relevant for models using ALiBi (e.g, MPT)
def check_use_alibi(model_config: ModelConfig) -> bool:
1982
    cfg = model_config.hf_text_config
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
    return (
        getattr(cfg, "alibi", False)  # Falcon
        or (
            "BloomForCausalLM" in getattr(model_config.hf_config, "architectures", [])
        )  # Bloom
        or getattr(cfg, "position_encoding_type", "") == "alibi"  # codellm_1b_alibi
        or (
            hasattr(cfg, "attn_config")  # MPT
            and (
                (
                    isinstance(cfg.attn_config, dict)
                    and cfg.attn_config.get("alibi", False)
                )
                or (
                    not isinstance(cfg.attn_config, dict)
                    and getattr(cfg.attn_config, "alibi", False)
                )
            )
        )
    )
2003
2004


2005
def sha256(input: Any) -> bytes:
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
    """Hash any picklable Python object using SHA-256.

    The input is serialized using pickle before hashing, which allows
    arbitrary Python objects to be used. Note that this function does
    not use a hash seed—if you need one, prepend it explicitly to the input.

    Args:
        input: Any picklable Python object.

    Returns:
2016
        Bytes representing the SHA-256 hash of the serialized input.
2017
2018
    """
    input_bytes = pickle.dumps(input, protocol=pickle.HIGHEST_PROTOCOL)
2019
    return hashlib.sha256(input_bytes).digest()
2020
2021


2022
def sha256_cbor(input: Any) -> bytes:
2023
    """
2024
    Hash objects using CBOR serialization and SHA-256.
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034

    This option is useful for non-Python-dependent serialization and hashing.

    Args:
        input: Object to be serialized and hashed. Supported types include
            basic Python types and complex structures like lists, tuples, and
            dictionaries.
            Custom classes must implement CBOR serialization methods.

    Returns:
2035
        Bytes representing the SHA-256 hash of the CBOR serialized input.
2036
2037
    """
    input_bytes = cbor2.dumps(input, canonical=True)
2038
    return hashlib.sha256(input_bytes).digest()
2039
2040


2041
def get_hash_fn_by_name(hash_fn_name: str) -> Callable[[Any], bytes]:
2042
2043
2044
2045
2046
2047
2048
2049
2050
    """Get a hash function by name, or raise an error if
    the function is not found.
    Args:
        hash_fn_name: Name of the hash function.
    Returns:
        A hash function.
    """
    if hash_fn_name == "sha256":
        return sha256
2051
2052
    if hash_fn_name == "sha256_cbor":
        return sha256_cbor
2053
2054
2055
2056

    raise ValueError(f"Unsupported hash function: {hash_fn_name}")


2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
def is_torch_equal_or_newer(target: str) -> bool:
    """Check if the installed torch version is >= the target version.

    Args:
        target: a version string, like "2.6.0".

    Returns:
        Whether the condition meets.
    """
    try:
2067
        return _is_torch_equal_or_newer(str(torch.__version__), target)
2068
2069
    except Exception:
        # Fallback to PKG-INFO to load the package info, needed by the doc gen.
2070
        return Version(importlib.metadata.version("torch")) >= Version(target)
2071
2072
2073
2074
2075
2076


# Helper function used in testing.
def _is_torch_equal_or_newer(torch_version: str, target: str) -> bool:
    torch_version = version.parse(torch_version)
    return torch_version >= version.parse(target)
2077
2078


2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
def _is_torch_equal(target: str) -> bool:
    assert target.count(".") == 2
    torch_version = str(torch.__version__)
    torch_version = version.parse(torch_version)
    # torch version is like "2.6.0.dev20240101" or "2.6.0.dev20240101+cpu"
    # or "2.6.0+cu128" but never "2.6.0.1"
    return (
        torch_version >= version.parse(target)
        and version.parse(target + ".1") > torch_version
    )


def is_torch_equal(target: str) -> bool:
    """Check if the installed torch version is == the target version.

    Args:
        target: a version string, like "2.6.0".

    Returns:
        Whether the condition meets.
    """
    try:
        return _is_torch_equal(target)
    except Exception:
        return Version(importlib.metadata.version("torch")) == Version(target)


2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
@cache
def _has_module(module_name: str) -> bool:
    """Return True if *module_name* can be found in the current environment.

    The result is cached so that subsequent queries for the same module incur
    no additional overhead.
    """
    return importlib.util.find_spec(module_name) is not None


def has_pplx() -> bool:
    """Whether the optional `pplx_kernels` package is available."""

    return _has_module("pplx_kernels")


def has_deep_ep() -> bool:
    """Whether the optional `deep_ep` package is available."""

    return _has_module("deep_ep")


def has_deep_gemm() -> bool:
    """Whether the optional `deep_gemm` package is available."""

2131
    return _has_module("deep_gemm")
2132
2133


2134
2135
2136
2137
2138
2139
def has_triton_kernels() -> bool:
    """Whether the optional `triton_kernels` package is available."""

    return _has_module("triton_kernels")


2140
2141
2142
2143
2144
2145
def has_tilelang() -> bool:
    """Whether the optional `tilelang` package is available."""

    return _has_module("tilelang")


2146
2147
2148
def set_process_title(
    name: str, suffix: str = "", prefix: str = envs.VLLM_PROCESS_NAME_PREFIX
) -> None:
2149
2150
2151
    """
    Set the current process title to a specific name with an
    optional suffix.
2152
2153

    Args:
2154
        name: The title to assign to the current process.
2155
        suffix: An optional suffix to append to the base name.
2156
        prefix: A prefix to prepend to the front separated by `::`.
2157
2158
2159
    """
    if suffix:
        name = f"{name}_{suffix}"
2160
    setproctitle.setproctitle(f"{prefix}::{name}")
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174


def _add_prefix(file: TextIO, worker_name: str, pid: int) -> None:
    """Prepend each output line with process-specific prefix"""

    prefix = f"{CYAN}({worker_name} pid={pid}){RESET} "
    file_write = file.write

    def write_with_prefix(s: str):
        if not s:
            return
        if file.start_new_line:  # type: ignore[attr-defined]
            file_write(prefix)
        idx = 0
2175
        while (next_idx := s.find("\n", idx)) != -1:
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
            next_idx += 1
            file_write(s[idx:next_idx])
            if next_idx == len(s):
                file.start_new_line = True  # type: ignore[attr-defined]
                return
            file_write(prefix)
            idx = next_idx
        file_write(s[idx:])
        file.start_new_line = False  # type: ignore[attr-defined]

    file.start_new_line = True  # type: ignore[attr-defined]
    file.write = write_with_prefix  # type: ignore[method-assign]


2190
def decorate_logs(process_name: str | None = None) -> None:
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
    """
    Adds a process-specific prefix to each line of output written to stdout and
    stderr.

    This function is intended to be called before initializing the api_server,
    engine_core, or worker classes, so that all subsequent output from the
    process is prefixed with the process name and PID. This helps distinguish
    log output from different processes in multi-process environments.

    Args:
        process_name: Optional; the name of the process to use in the prefix.
            If not provided, the current process name from the multiprocessing
            context is used.
    """
    if process_name is None:
        process_name = get_mp_context().current_process().name
    pid = os.getpid()
    _add_prefix(sys.stdout, process_name, pid)
    _add_prefix(sys.stderr, process_name, pid)
2210
2211
2212


def length_from_prompt_token_ids_or_embeds(
2213
2214
    prompt_token_ids: list[int] | None,
    prompt_embeds: torch.Tensor | None,
2215
) -> int:
2216
    """Calculate the request length (in number of tokens) give either
2217
2218
    prompt_token_ids or prompt_embeds.
    """
2219
2220
    prompt_token_len = None if prompt_token_ids is None else len(prompt_token_ids)
    prompt_embeds_len = None if prompt_embeds is None else len(prompt_embeds)
2221
2222
2223

    if prompt_token_len is None:
        if prompt_embeds_len is None:
2224
            raise ValueError("Neither prompt_token_ids nor prompt_embeds were defined.")
2225
2226
        return prompt_embeds_len
    else:
2227
        if prompt_embeds_len is not None and prompt_embeds_len != prompt_token_len:
2228
2229
2230
            raise ValueError(
                "Prompt token ids and prompt embeds had different lengths"
                f" prompt_token_ids={prompt_token_len}"
2231
2232
                f" prompt_embeds={prompt_embeds_len}"
            )
2233
        return prompt_token_len
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246


@contextlib.contextmanager
def set_env_var(key, value):
    old = os.environ.get(key)
    os.environ[key] = value
    try:
        yield
    finally:
        if old is None:
            del os.environ[key]
        else:
            os.environ[key] = old
2247
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266


def unique_filepath(fn: Callable[[int], Path]) -> Path:
    """
    unique_filepath returns a unique path by trying
    to include an integer in increasing order.

    fn should be a callable that returns a path that
    includes the passed int at a fixed location.

    Note: This function has a TOCTOU race condition.
    Caller should use atomic operations (e.g., open with 'x' mode)
    when creating the file to ensure thread safety.
    """
    i = 0
    while True:
        p = fn(i)
        if not p.exists():
            return p
        i += 1