__init__.py 45 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 importlib
9
import inspect
10
import json
11
import multiprocessing
12
import os
13
import signal
14
import subprocess
15
import sys
16
import tempfile
17
import textwrap
18
import threading
19
import traceback
Zhuohan Li's avatar
Zhuohan Li committed
20
import uuid
21
import warnings
22
import weakref
23
24
25
26
27
28
29
30
from argparse import (
    Action,
    ArgumentDefaultsHelpFormatter,
    ArgumentParser,
    ArgumentTypeError,
    RawDescriptionHelpFormatter,
    _ArgumentGroup,
)
31
from collections import defaultdict
32
from collections.abc import (
33
    Callable,
34
35
    Sequence,
)
36
from concurrent.futures.process import ProcessPoolExecutor
37
from functools import cache, partial, wraps
38
from pathlib import Path
39
from typing import TYPE_CHECKING, Any, TextIO, TypeVar
Zhuohan Li's avatar
Zhuohan Li committed
40

41
import cloudpickle
42
import psutil
43
import regex as re
44
import setproctitle
Zhuohan Li's avatar
Zhuohan Li committed
45
import torch
46
import yaml
47

48
import vllm.envs as envs
49
from vllm.logger import enable_trace_function_call, init_logger
50
from vllm.ray.lazy_utils import is_in_ray_actor
51

52
53
54
55
56
_DEPRECATED_MAPPINGS = {
    "cprofile": "profiling",
    "cprofile_context": "profiling",
    "get_open_port": "network_utils",
}
57
58
59


def __getattr__(name: str) -> Any:  # noqa: D401 - short deprecation docstring
60
61
62
    """Module-level getattr to handle deprecated utilities."""
    if name in _DEPRECATED_MAPPINGS:
        submodule_name = _DEPRECATED_MAPPINGS[name]
63
64
        warnings.warn(
            f"vllm.utils.{name} is deprecated and will be removed in a future version. "
65
            f"Use vllm.utils.{submodule_name}.{name} instead.",
66
67
68
            DeprecationWarning,
            stacklevel=2,
        )
69
70
        module = __import__(f"vllm.utils.{submodule_name}", fromlist=[submodule_name])
        return getattr(module, name)
71
72
73
74
75
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")


def __dir__() -> list[str]:
    # expose deprecated names in dir() for better UX/tab-completion
76
    return sorted(list(globals().keys()) + list(_DEPRECATED_MAPPINGS.keys()))
77
78


79
if TYPE_CHECKING:
80
81
    from argparse import Namespace

82
    from vllm.config import ModelConfig, VllmConfig
83
84
85
86
87
else:
    Namespace = object

    ModelConfig = object
    VllmConfig = object
88

89
90
logger = init_logger(__name__)

91
92
93
94
95
96
# 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

97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
# 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"

112

113
# ANSI color codes
114
115
CYAN = "\033[1;36m"
RESET = "\033[0;0m"
116

117

118
T = TypeVar("T")
119
U = TypeVar("U")
120

121

Woosuk Kwon's avatar
Woosuk Kwon committed
122
123
124
125
126
class Device(enum.Enum):
    GPU = enum.auto()
    CPU = enum.auto()


127
128
129
130
131
class LayerBlockType(enum.Enum):
    attention = "attention"
    mamba = "mamba"


Woosuk Kwon's avatar
Woosuk Kwon committed
132
133
134
135
class Counter:
    def __init__(self, start: int = 0) -> None:
        self.counter = start

Woosuk Kwon's avatar
Woosuk Kwon committed
136
    def __next__(self) -> int:
137
        i = self.counter
Woosuk Kwon's avatar
Woosuk Kwon committed
138
        self.counter += 1
139
        return i
Woosuk Kwon's avatar
Woosuk Kwon committed
140
141
142

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

144

Zhuohan Li's avatar
Zhuohan Li committed
145
146
def random_uuid() -> str:
    return str(uuid.uuid4().hex)
147

148

149
def update_environment_variables(envs: dict[str, str]):
150
    for k, v in envs.items():
151
        if k in os.environ and os.environ[k] != v:
152
            logger.warning(
153
154
155
156
157
                "Overwriting environment variable %s from '%s' to '%s'",
                k,
                os.environ[k],
                v,
            )
158
        os.environ[k] = v
159
160


161
162
163
164
165
def cdiv(a: int, b: int) -> int:
    """Ceiling division."""
    return -(a // -b)


166
167
168
169
170
171
172
def next_power_of_2(n) -> int:
    """The next power of 2 (inclusive)"""
    if n < 1:
        return 1
    return 1 << (n - 1).bit_length()


173
174
175
176
177
178
179
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)


180
181
182
183
def round_up(x: int, y: int) -> int:
    return ((x + y - 1) // y) * y


184
185
186
187
def round_down(x: int, y: int) -> int:
    return (x // y) * y


188
@cache
189
def is_pin_memory_available() -> bool:
190
    from vllm.platforms import current_platform
191

192
    return current_platform.is_pin_memory_available()
193
194


195
196
197
198
199
200
201
202
@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()


203
204
# TODO: This function can be removed if transformer_modules classes are
# serialized by value when communicating between processes
205
def init_cached_hf_modules() -> None:
206
207
208
209
    """
    Lazy initialization of the Hugging Face modules.
    """
    from transformers.dynamic_module_utils import init_hf_modules
210

211
    init_hf_modules()
212
213


214
@cache
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
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
230
    env_ld_library_path = envs.LD_LIBRARY_PATH
231
232
233
234
235
236
237
238
239
240
241
    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]


242
def find_nccl_library() -> str:
243
244
245
246
247
248
    """
    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.
    """
249
    so_file = envs.VLLM_NCCL_SO_PATH
250
251
252
253

    # manually load the nccl library
    if so_file:
        logger.info(
254
255
            "Found nccl from environment variable VLLM_NCCL_SO_PATH=%s", so_file
        )
256
257
    else:
        if torch.version.cuda is not None:
258
            so_file = "libnccl.so.2"
259
        elif torch.version.hip is not None:
260
            so_file = "librccl.so.1"
261
262
        else:
            raise ValueError("NCCL only supports CUDA and ROCm backends.")
263
        logger.debug_once("Found nccl from library %s", so_file)
264
    return so_file
265
266


267
def find_nccl_include_paths() -> list[str] | None:
268
269
    """
    We either use the nccl.h specified by the `VLLM_NCCL_INCLUDE_PATH`
270
271
    environment variable, or we find the library file brought by
    nvidia-nccl-cuXX. load_inline by default uses
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
    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


298
def enable_trace_function_call_for_thread(vllm_config: VllmConfig) -> None:
299
300
301
302
    """Set up function tracing for the current thread,
    if enabled via the VLLM_TRACE_FUNCTION environment variable
    """

303
    if envs.VLLM_TRACE_FUNCTION:
304
        tmp_dir = tempfile.gettempdir()
305
306
        # add username to tmp_dir to avoid permission issues
        tmp_dir = os.path.join(tmp_dir, getpass.getuser())
307
308
309
310
311
312
313
314
        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
        )
315
316
        os.makedirs(os.path.dirname(log_path), exist_ok=True)
        enable_trace_function_call(log_path)
317
318


319
320
321
322
323
324
325
def cuda_is_initialized() -> bool:
    """Check if CUDA is initialized."""
    if not torch.cuda._is_compiled():
        return False
    return torch.cuda.is_initialized()


326
327
328
329
330
331
332
def xpu_is_initialized() -> bool:
    """Check if XPU is initialized."""
    if not torch.xpu._is_compiled():
        return False
    return torch.xpu.is_initialized()


333
334
335
def cuda_get_device_properties(
    device, names: Sequence[str], init_cuda=False
) -> tuple[Any, ...]:
336
337
338
339
340
341
342
343
344
    """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:
345
        return executor.submit(cuda_get_device_properties, device, names, True).result()
346
347


348
349
350
def weak_bind(
    bound_method: Callable[..., Any],
) -> Callable[..., None]:
351
352
353
354
355
356
357
358
359
360
361
362
363
    """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


364
class StoreBoolean(Action):
365
366
367
368
369
370
    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:
371
372
373
            raise ValueError(
                f"Invalid boolean value: {values}. Expected 'true' or 'false'."
            )
374
375


376
class SortedHelpFormatter(ArgumentDefaultsHelpFormatter, RawDescriptionHelpFormatter):
377
378
    """SortedHelpFormatter that sorts arguments by their option strings."""

379
380
381
382
383
384
385
    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
386
387
        single_newline = re.compile(r"(?<!\n)\n(?!\n)\s*")
        multiple_newlines = re.compile(r"\n{2,}\s*")
388
        text = single_newline.sub(" ", text)
389
390
391
        lines = re.split(multiple_newlines, text)
        return sum([textwrap.wrap(line, width) for line in lines], [])

392
393
    def add_arguments(self, actions):
        actions = sorted(actions, key=lambda x: x.option_strings)
394
        super().add_arguments(actions)
395
396


397
class FlexibleArgumentParser(ArgumentParser):
398
399
    """ArgumentParser that allows both underscore and dash in names."""

400
    _deprecated: set[Action] = set()
401
402
403
404
405
406
407
    _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'
408
409
        "   --json-arg.key4+ value3 --json-arg.key4+='value4,value5'\n\n"
    )
410
    _search_keyword: str | None = None
411

412
    def __init__(self, *args, **kwargs):
413
414
415
        # Set the default "formatter_class" to SortedHelpFormatter
        if "formatter_class" not in kwargs:
            kwargs["formatter_class"] = SortedHelpFormatter
416
417
        # Pop kwarg "add_json_tip" to control whether to add the JSON tip
        self.add_json_tip = kwargs.pop("add_json_tip", True)
418
419
        super().__init__(*args, **kwargs)

420
    if sys.version_info < (3, 13):
421
        # Enable the deprecated kwarg for Python 3.12 and below
422

423
        def parse_known_args(self, args=None, namespace=None):
424
425
426
427
            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(
428
429
                    "argument '--disable-log-requests' is deprecated and "
                    "replaced with '--enable-log-requests'. This will be "
430
431
                    "removed in v0.12.0."
                )
432
433
            namespace, args = super().parse_known_args(args, namespace)
            for action in FlexibleArgumentParser._deprecated:
434
435
436
437
                if (
                    hasattr(namespace, dest := action.dest)
                    and getattr(namespace, dest) != action.default
                ):
438
                    logger.warning_once("argument '%s' is deprecated", dest)
439
440
            return namespace, args

441
442
        def add_argument(self, *args, **kwargs):
            deprecated = kwargs.pop("deprecated", False)
443
            action = super().add_argument(*args, **kwargs)
444
445
            if deprecated:
                FlexibleArgumentParser._deprecated.add(action)
446
447
            return action

448
449
450
451
452
453
454
455
456
457
458
459
        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
460

461
462
463
464
465
466
467
468
469
470
471
472
    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'
473
            if search_keyword == "all":
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
                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
492
493
494
                    if any(
                        search_keyword in opt.lower() for opt in action.option_strings
                    ):
495
496
                        matched_actions.append(action)
            if matched_actions:
497
                formatter.start_section(f"Arguments matching '{search_keyword}'")
498
499
500
501
502
503
504
505
506
                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 "
507
508
                "'--help=all' to see all available parameters."
            )
509
510
511
            return formatter.format_help()

        # usage
512
        formatter.add_usage(self.usage, self._actions, self._mutually_exclusive_groups)
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533

        # 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()
534

535
536
537
538
539
    def parse_args(  # type: ignore[override]
        self,
        args: list[str] | None = None,
        namespace: Namespace | None = None,
    ):
540
541
542
        if args is None:
            args = sys.argv[1:]

543
544
        # Check for --model in command line arguments first
        if args and args[0] == "serve":
545
546
            try:
                model_idx = next(
547
548
549
550
                    i
                    for i, arg in enumerate(args)
                    if arg == "--model" or arg.startswith("--model=")
                )
551
                logger.warning(
552
553
                    "With `vllm serve`, you should provide the model as a "
                    "positional argument or in a config file instead of via "
554
                    "the `--model` option. "
555
556
                    "The `--model` option will be removed in v0.13."
                )
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578

                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
579

580
        if "--config" in args:
581
            args = self._pull_args_from_config(args)
582

583
584
585
586
587
588
589
        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"(?<=--)[^\.]*")

590
        # Convert underscores to dashes and vice versa in argument names
591
        processed_args = list[str]()
592
        for i, arg in enumerate(args):
593
            if arg.startswith("--help="):
594
                FlexibleArgumentParser._search_keyword = arg.split("=", 1)[-1].lower()
595
                processed_args.append("--help")
596
597
598
            elif arg.startswith("--"):
                if "=" in arg:
                    key, value = arg.split("=", 1)
599
                    key = pattern.sub(repl, key, count=1)
600
                    processed_args.append(f"{key}={value}")
601
                else:
602
603
                    key = pattern.sub(repl, arg, count=1)
                    processed_args.append(key)
604
            elif arg.startswith("-O") and arg != "-O" and arg[2] != ".":
605
606
                # allow -O flag to be used without space, e.g. -O3 or -Odecode
                # -O.<...> handled later
607
608
609
                # also handle -O=<mode> here
                mode = arg[3:] if arg[2] == "=" else arg[2:]
                processed_args.append(f"-O.mode={mode}")
610
611
612
613
614
            elif (
                arg == "-O"
                and i + 1 < len(args)
                and args[i + 1] in {"0", "1", "2", "3"}
            ):
615
616
                # Convert -O <n> to -O.mode <n>
                processed_args.append("-O.mode")
617
618
619
            else:
                processed_args.append(arg)

620
        def create_nested_dict(keys: list[str], value: str) -> dict[str, Any]:
621
622
623
624
625
626
627
628
629
630
            """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

631
632
633
        def recursive_dict_update(
            original: dict[str, Any],
            update: dict[str, Any],
634
635
636
637
638
        ) -> set[str]:
            """Recursively updates a dictionary with another dictionary.
            Returns a set of duplicate keys that were overwritten.
            """
            duplicates = set[str]()
639
640
            for k, v in update.items():
                if isinstance(v, dict) and isinstance(original.get(k), dict):
641
642
643
644
                    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
645
                else:
646
647
                    if k in original:
                        duplicates.add(k)
648
                    original[k] = v
649
            return duplicates
650

651
652
        delete = set[int]()
        dict_args = defaultdict[str, dict[str, Any]](dict)
653
        duplicates = set[str]()
654
        for i, processed_arg in enumerate(processed_args):
655
656
657
658
            if i in delete:  # skip if value from previous arg
                continue

            if processed_arg.startswith("-") and "." in processed_arg:
659
                if "=" in processed_arg:
660
                    processed_arg, value_str = processed_arg.split("=", 1)
661
                    if "." not in processed_arg:
662
                        # False positive, '.' was only in the value
663
664
                        continue
                else:
665
                    value_str = processed_args[i + 1]
666
                    delete.add(i + 1)
667

668
669
670
671
                if processed_arg.endswith("+"):
                    processed_arg = processed_arg[:-1]
                    value_str = json.dumps(list(value_str.split(",")))

672
                key, *keys = processed_arg.split(".")
673
674
675
676
677
                try:
                    value = json.loads(value_str)
                except json.decoder.JSONDecodeError:
                    value = value_str

678
679
                # Merge all values with the same key into a single dict
                arg_dict = create_nested_dict(keys, value)
680
681
                arg_duplicates = recursive_dict_update(dict_args[key], arg_dict)
                duplicates |= {f"{key}.{d}" for d in arg_duplicates}
682
683
                delete.add(i)
        # Filter out the dict args we set to None
684
        processed_args = [a for i, a in enumerate(processed_args) if i not in delete]
685
686
687
        if duplicates:
            logger.warning("Found duplicate keys %s", ", ".join(duplicates))

688
689
690
691
692
        # 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))

693
        return super().parse_args(processed_args, namespace)
694

695
696
697
698
    def check_port(self, value):
        try:
            value = int(value)
        except ValueError:
699
            msg = "Port must be an integer"
700
            raise ArgumentTypeError(msg) from None
701
702

        if not (1024 <= value <= 65535):
703
            raise ArgumentTypeError("Port must be between 1024 and 65535")
704
705
706

        return value

707
    def _pull_args_from_config(self, args: list[str]) -> list[str]:
708
709
        """Method to pull arguments specified in the config file
        into the command-line args variable.
710
711

        The arguments in config file will be inserted between
712
        the argument list.
713

714
715
716
717
718
719
720
721
722
723
        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",
724
725
            "facebook/opt-12B",
            '--config', 'config.yaml',
726
727
728
729
            '-tp', '2'
        ]
        $: args = [
            "serve,chat,complete",
730
731
732
            "facebook/opt-12B",
            '--port', '12323',
            '--tensor-parallel-size', '4',
733
734
735
736
737
            '-tp', '2'
            ]
        ```

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

743
        index = args.index("--config")
744
        if index == len(args) - 1:
745
746
747
748
            raise ValueError(
                "No config file specified! \
                             Please check your command-line arguments."
            )
749
750
751

        file_path = args[index + 1]

752
        config_args = self.load_config_file(file_path)
753

754
        # 0th index might be the sub command {serve,chat,complete,...}
755
        # optionally followed by model_tag (only for serve)
756
757
758
759
        # followed by config args
        # followed by rest of cli args.
        # maintaining this order will enforce the precedence
        # of cli > config > defaults
760
        if args[0].startswith("-"):
761
            # No sub command (e.g., api_server entry point)
762
            args = config_args + args[0:index] + args[index + 2 :]
763
        elif args[0] == "serve":
764
765
            model_in_cli = len(args) > 1 and not args[1].startswith("-")
            model_in_config = any(arg == "--model" for arg in config_args)
766
767

            if not model_in_cli and not model_in_config:
768
                raise ValueError(
769
                    "No model specified! Please specify model either "
770
771
                    "as a positional argument or in a config file."
                )
772
773
774

            if model_in_cli:
                # Model specified as positional arg, keep CLI version
775
776
777
778
779
780
781
                args = (
                    [args[0]]
                    + [args[1]]
                    + config_args
                    + args[2:index]
                    + args[index + 2 :]
                )
782
783
            else:
                # No model in CLI, use config if available
784
                args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
785
        else:
786
            args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
787
788
789

        return args

790
    def load_config_file(self, file_path: str) -> list[str]:
791
        """Loads a yaml file and returns the key value pairs as a
792
793
794
795
796
797
798
799
800
801
802
        flattened list with argparse like pattern
        ```yaml
            port: 12323
            tensor-parallel-size: 4
        ```
        returns:
            processed_args: list[str] = [
                '--port': '12323',
                '--tensor-parallel-size': '4'
            ]
        """
803
804
        extension: str = file_path.split(".")[-1]
        if extension not in ("yaml", "yml"):
805
806
            raise ValueError(
                "Config file must be of a yaml/yml type.\
807
808
809
                              %s supplied",
                extension,
            )
810
811

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

814
        config: dict[str, int | str] = {}
815
        try:
816
            with open(file_path) as config_file:
817
818
819
820
                config = yaml.safe_load(config_file)
        except Exception as ex:
            logger.error(
                "Unable to read the config file at %s. \
821
822
823
                Make sure path is correct",
                file_path,
            )
824
825
            raise ex

826
        store_boolean_arguments = [
827
            action.dest for action in self._actions if isinstance(action, StoreBoolean)
828
829
        ]

830
        for key, value in config.items():
831
832
            if isinstance(value, bool) and key not in store_boolean_arguments:
                if value:
833
                    processed_args.append("--" + key)
834
835
            elif isinstance(value, list):
                if value:
836
                    processed_args.append("--" + key)
837
838
                    for item in value:
                        processed_args.append(str(item))
839
            else:
840
                processed_args.append("--" + key)
841
                processed_args.append(str(value))
842
843
844

        return processed_args

845

846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
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
869
870


871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
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)
894
895


896
# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L630 # noqa: E501
897
def set_ulimit(target_soft_limit=65535):
898
    if sys.platform.startswith("win"):
899
900
901
902
        logger.info("Windows detected, skipping ulimit adjustment.")
        return

    import resource
903

904
905
906
907
908
    resource_type = resource.RLIMIT_NOFILE
    current_soft, current_hard = resource.getrlimit(resource_type)

    if current_soft < target_soft_limit:
        try:
909
            resource.setrlimit(resource_type, (target_soft_limit, current_hard))
910
911
        except ValueError as e:
            logger.warning(
912
913
                "Found ulimit of %s and failed to automatically increase "
                "with error %s. This can cause fd limit errors like "
914
                "`OSError: [Errno 24] Too many open files`. Consider "
915
916
917
918
                "increasing with ulimit -n",
                current_soft,
                e,
            )
919
920
921
922
923
924
925
926
927


# 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


928
929
930
931
932
933
934
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

935
936
    reasons = []
    if is_in_ray_actor():
937
938
939
940
        # 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
941

942
        os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address
943
944
945
946
947
948
        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")
949

950
    if reasons:
951
952
953
        logger.warning(
            "We must use the `spawn` multiprocessing start method. "
            "Overriding VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. "
954
            "See https://docs.vllm.ai/en/latest/usage/"
955
            "troubleshooting.html#python-multiprocessing "
956
957
958
            "for more information. Reasons: %s",
            "; ".join(reasons),
        )
959
960
961
962
        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"


def get_mp_context():
963
964
965
966
967
968
969
    """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()
970
971
    mp_method = envs.VLLM_WORKER_MULTIPROC_METHOD
    return multiprocessing.get_context(mp_method)
972
973
974


def bind_kv_cache(
975
976
    ctx: dict[str, Any],
    kv_cache: list[list[torch.Tensor]],  # [virtual_engine][layer_index]
977
    shared_kv_cache_layers: dict[str, str] | None = None,
978
979
980
981
982
983
984
985
986
987
988
) -> 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
989
990
991
992
    # 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 = {}
993
994
    from vllm.attention import AttentionType
    from vllm.model_executor.models.utils import extract_layer_index
995

996
    layer_need_kv_cache = [
997
998
999
1000
1001
1002
1003
1004
        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
1005
1006
    ]
    layer_index_sorted = sorted(
1007
1008
        set(extract_layer_index(layer_name) for layer_name in layer_need_kv_cache)
    )
1009
    for layer_name in layer_need_kv_cache:
1010
        kv_cache_idx = layer_index_sorted.index(extract_layer_index(layer_name))
1011
1012
1013
1014
        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]
1015
1016
    if shared_kv_cache_layers is not None:
        for layer_name, target_layer_name in shared_kv_cache_layers.items():
1017
1018
1019
            assert extract_layer_index(target_layer_name) < extract_layer_index(
                layer_name
            ), "v0 doesn't support interleaving kv sharing"
1020
            ctx[layer_name].kv_cache = ctx[target_layer_name].kv_cache
1021
1022


1023
1024
def run_method(
    obj: Any,
1025
    method: str | bytes | Callable,
1026
1027
1028
    args: tuple[Any],
    kwargs: dict[str, Any],
) -> Any:
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
    """
    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:
1042
1043
1044
            raise NotImplementedError(
                f"Method {method!r} is not implemented."
            ) from None
1045
1046
1047
    else:
        func = partial(method, obj)  # type: ignore
    return func(*args, **kwargs)
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068


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.
1069
1070
1071
1072
1073
1074
1075
    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.
1076
    """
1077
    import vllm.third_party.pynvml as pynvml
1078

1079
    return pynvml
1080
1081


1082
def warn_for_unimplemented_methods(cls: type[T]) -> type[T]:
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
    """
    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
1098
            if attr_name.startswith("_"):
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
                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:
1112
1113
            method_names = ",".join(unimplemented_methods)
            msg = f"Methods {method_names} not implemented in {self}"
1114
            logger.debug(msg)
1115
1116
1117
1118
1119
1120

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

1121
    type.__setattr__(cls, "__init__", wrapped_init)
1122
    return cls
1123
1124


1125
1126
# Only relevant for models using ALiBi (e.g, MPT)
def check_use_alibi(model_config: ModelConfig) -> bool:
1127
    cfg = model_config.hf_text_config
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
    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)
                )
            )
        )
    )
1148
1149


1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
@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."""

1175
    return _has_module("deep_gemm")
1176
1177


1178
1179
1180
1181
1182
1183
def has_triton_kernels() -> bool:
    """Whether the optional `triton_kernels` package is available."""

    return _has_module("triton_kernels")


1184
1185
1186
1187
1188
1189
def has_tilelang() -> bool:
    """Whether the optional `tilelang` package is available."""

    return _has_module("tilelang")


1190
1191
1192
def set_process_title(
    name: str, suffix: str = "", prefix: str = envs.VLLM_PROCESS_NAME_PREFIX
) -> None:
1193
1194
1195
    """
    Set the current process title to a specific name with an
    optional suffix.
1196
1197

    Args:
1198
        name: The title to assign to the current process.
1199
        suffix: An optional suffix to append to the base name.
1200
        prefix: A prefix to prepend to the front separated by `::`.
1201
1202
1203
    """
    if suffix:
        name = f"{name}_{suffix}"
1204
    setproctitle.setproctitle(f"{prefix}::{name}")
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218


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
1219
        while (next_idx := s.find("\n", idx)) != -1:
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
            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]


1234
def decorate_logs(process_name: str | None = None) -> None:
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
    """
    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)
1254
1255
1256


def length_from_prompt_token_ids_or_embeds(
1257
1258
    prompt_token_ids: list[int] | None,
    prompt_embeds: torch.Tensor | None,
1259
) -> int:
1260
    """Calculate the request length (in number of tokens) give either
1261
1262
    prompt_token_ids or prompt_embeds.
    """
1263
1264
    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)
1265
1266
1267

    if prompt_token_len is None:
        if prompt_embeds_len is None:
1268
            raise ValueError("Neither prompt_token_ids nor prompt_embeds were defined.")
1269
1270
        return prompt_embeds_len
    else:
1271
        if prompt_embeds_len is not None and prompt_embeds_len != prompt_token_len:
1272
1273
1274
            raise ValueError(
                "Prompt token ids and prompt embeds had different lengths"
                f" prompt_token_ids={prompt_token_len}"
1275
1276
                f" prompt_embeds={prompt_embeds_len}"
            )
1277
        return prompt_token_len
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290


@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
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310


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