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

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

46
import vllm.envs as envs
47
from vllm.logger import enable_trace_function_call, init_logger
48
from vllm.ray.lazy_utils import is_in_ray_actor
49

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


def __getattr__(name: str) -> Any:  # noqa: D401 - short deprecation docstring
58
59
60
    """Module-level getattr to handle deprecated utilities."""
    if name in _DEPRECATED_MAPPINGS:
        submodule_name = _DEPRECATED_MAPPINGS[name]
61
62
        warnings.warn(
            f"vllm.utils.{name} is deprecated and will be removed in a future version. "
63
            f"Use vllm.utils.{submodule_name}.{name} instead.",
64
65
66
            DeprecationWarning,
            stacklevel=2,
        )
67
68
        module = __import__(f"vllm.utils.{submodule_name}", fromlist=[submodule_name])
        return getattr(module, name)
69
70
71
72
73
    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
74
    return sorted(list(globals().keys()) + list(_DEPRECATED_MAPPINGS.keys()))
75
76


77
if TYPE_CHECKING:
78
79
    from argparse import Namespace

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

    ModelConfig = object
    VllmConfig = object
86

87
88
logger = init_logger(__name__)

89
90
91
92
93
94
# 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

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

110

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

115

116
T = TypeVar("T")
117
U = TypeVar("U")
118

119

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


125
126
127
128
129
class LayerBlockType(enum.Enum):
    attention = "attention"
    mamba = "mamba"


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

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

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

142

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

146

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


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


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


171
172
173
174
175
176
177
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)


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


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


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

190
    return current_platform.is_pin_memory_available()
191
192


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


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

209
    init_hf_modules()
210
211


212
def enable_trace_function_call_for_thread(vllm_config: VllmConfig) -> None:
213
214
215
216
    """Set up function tracing for the current thread,
    if enabled via the VLLM_TRACE_FUNCTION environment variable
    """

217
    if envs.VLLM_TRACE_FUNCTION:
218
        tmp_dir = tempfile.gettempdir()
219
220
        # add username to tmp_dir to avoid permission issues
        tmp_dir = os.path.join(tmp_dir, getpass.getuser())
221
222
223
224
225
226
227
228
        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
        )
229
230
        os.makedirs(os.path.dirname(log_path), exist_ok=True)
        enable_trace_function_call(log_path)
231
232


233
234
235
236
237
238
239
def cuda_is_initialized() -> bool:
    """Check if CUDA is initialized."""
    if not torch.cuda._is_compiled():
        return False
    return torch.cuda.is_initialized()


240
241
242
243
244
245
246
def xpu_is_initialized() -> bool:
    """Check if XPU is initialized."""
    if not torch.xpu._is_compiled():
        return False
    return torch.xpu.is_initialized()


247
248
249
def cuda_get_device_properties(
    device, names: Sequence[str], init_cuda=False
) -> tuple[Any, ...]:
250
251
252
253
254
255
256
257
258
    """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:
259
        return executor.submit(cuda_get_device_properties, device, names, True).result()
260
261


262
263
264
def weak_bind(
    bound_method: Callable[..., Any],
) -> Callable[..., None]:
265
266
267
268
269
270
271
272
273
274
275
276
277
    """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


278
class StoreBoolean(Action):
279
280
281
282
283
284
    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:
285
286
287
            raise ValueError(
                f"Invalid boolean value: {values}. Expected 'true' or 'false'."
            )
288
289


290
class SortedHelpFormatter(ArgumentDefaultsHelpFormatter, RawDescriptionHelpFormatter):
291
292
    """SortedHelpFormatter that sorts arguments by their option strings."""

293
294
295
296
297
298
299
    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
300
301
        single_newline = re.compile(r"(?<!\n)\n(?!\n)\s*")
        multiple_newlines = re.compile(r"\n{2,}\s*")
302
        text = single_newline.sub(" ", text)
303
304
305
        lines = re.split(multiple_newlines, text)
        return sum([textwrap.wrap(line, width) for line in lines], [])

306
307
    def add_arguments(self, actions):
        actions = sorted(actions, key=lambda x: x.option_strings)
308
        super().add_arguments(actions)
309
310


311
class FlexibleArgumentParser(ArgumentParser):
312
313
    """ArgumentParser that allows both underscore and dash in names."""

314
    _deprecated: set[Action] = set()
315
316
317
318
319
320
321
    _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'
322
323
        "   --json-arg.key4+ value3 --json-arg.key4+='value4,value5'\n\n"
    )
324
    _search_keyword: str | None = None
325

326
    def __init__(self, *args, **kwargs):
327
328
329
        # Set the default "formatter_class" to SortedHelpFormatter
        if "formatter_class" not in kwargs:
            kwargs["formatter_class"] = SortedHelpFormatter
330
331
        # Pop kwarg "add_json_tip" to control whether to add the JSON tip
        self.add_json_tip = kwargs.pop("add_json_tip", True)
332
333
        super().__init__(*args, **kwargs)

334
    if sys.version_info < (3, 13):
335
        # Enable the deprecated kwarg for Python 3.12 and below
336

337
        def parse_known_args(self, args=None, namespace=None):
338
339
340
341
            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(
342
343
                    "argument '--disable-log-requests' is deprecated and "
                    "replaced with '--enable-log-requests'. This will be "
344
345
                    "removed in v0.12.0."
                )
346
347
            namespace, args = super().parse_known_args(args, namespace)
            for action in FlexibleArgumentParser._deprecated:
348
349
350
351
                if (
                    hasattr(namespace, dest := action.dest)
                    and getattr(namespace, dest) != action.default
                ):
352
                    logger.warning_once("argument '%s' is deprecated", dest)
353
354
            return namespace, args

355
356
        def add_argument(self, *args, **kwargs):
            deprecated = kwargs.pop("deprecated", False)
357
            action = super().add_argument(*args, **kwargs)
358
359
            if deprecated:
                FlexibleArgumentParser._deprecated.add(action)
360
361
            return action

362
363
364
365
366
367
368
369
370
371
372
373
        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
374

375
376
377
378
379
380
381
382
383
384
385
386
    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'
387
            if search_keyword == "all":
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
                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
406
407
408
                    if any(
                        search_keyword in opt.lower() for opt in action.option_strings
                    ):
409
410
                        matched_actions.append(action)
            if matched_actions:
411
                formatter.start_section(f"Arguments matching '{search_keyword}'")
412
413
414
415
416
417
418
419
420
                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 "
421
422
                "'--help=all' to see all available parameters."
            )
423
424
425
            return formatter.format_help()

        # usage
426
        formatter.add_usage(self.usage, self._actions, self._mutually_exclusive_groups)
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447

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

449
450
451
452
453
    def parse_args(  # type: ignore[override]
        self,
        args: list[str] | None = None,
        namespace: Namespace | None = None,
    ):
454
455
456
        if args is None:
            args = sys.argv[1:]

457
458
        # Check for --model in command line arguments first
        if args and args[0] == "serve":
459
460
            try:
                model_idx = next(
461
462
463
464
                    i
                    for i, arg in enumerate(args)
                    if arg == "--model" or arg.startswith("--model=")
                )
465
                logger.warning(
466
467
                    "With `vllm serve`, you should provide the model as a "
                    "positional argument or in a config file instead of via "
468
                    "the `--model` option. "
469
470
                    "The `--model` option will be removed in v0.13."
                )
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492

                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
493

494
        if "--config" in args:
495
            args = self._pull_args_from_config(args)
496

497
498
499
500
501
502
503
        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"(?<=--)[^\.]*")

504
        # Convert underscores to dashes and vice versa in argument names
505
        processed_args = list[str]()
506
        for i, arg in enumerate(args):
507
            if arg.startswith("--help="):
508
                FlexibleArgumentParser._search_keyword = arg.split("=", 1)[-1].lower()
509
                processed_args.append("--help")
510
511
512
            elif arg.startswith("--"):
                if "=" in arg:
                    key, value = arg.split("=", 1)
513
                    key = pattern.sub(repl, key, count=1)
514
                    processed_args.append(f"{key}={value}")
515
                else:
516
517
                    key = pattern.sub(repl, arg, count=1)
                    processed_args.append(key)
518
            elif arg.startswith("-O") and arg != "-O" and arg[2] != ".":
519
520
                # allow -O flag to be used without space, e.g. -O3 or -Odecode
                # -O.<...> handled later
521
522
523
                # also handle -O=<mode> here
                mode = arg[3:] if arg[2] == "=" else arg[2:]
                processed_args.append(f"-O.mode={mode}")
524
525
526
527
528
            elif (
                arg == "-O"
                and i + 1 < len(args)
                and args[i + 1] in {"0", "1", "2", "3"}
            ):
529
530
                # Convert -O <n> to -O.mode <n>
                processed_args.append("-O.mode")
531
532
533
            else:
                processed_args.append(arg)

534
        def create_nested_dict(keys: list[str], value: str) -> dict[str, Any]:
535
536
537
538
539
540
541
542
543
544
            """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

545
546
547
        def recursive_dict_update(
            original: dict[str, Any],
            update: dict[str, Any],
548
549
550
551
552
        ) -> set[str]:
            """Recursively updates a dictionary with another dictionary.
            Returns a set of duplicate keys that were overwritten.
            """
            duplicates = set[str]()
553
554
            for k, v in update.items():
                if isinstance(v, dict) and isinstance(original.get(k), dict):
555
556
557
558
                    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
559
                else:
560
561
                    if k in original:
                        duplicates.add(k)
562
                    original[k] = v
563
            return duplicates
564

565
566
        delete = set[int]()
        dict_args = defaultdict[str, dict[str, Any]](dict)
567
        duplicates = set[str]()
568
        for i, processed_arg in enumerate(processed_args):
569
570
571
572
            if i in delete:  # skip if value from previous arg
                continue

            if processed_arg.startswith("-") and "." in processed_arg:
573
                if "=" in processed_arg:
574
                    processed_arg, value_str = processed_arg.split("=", 1)
575
                    if "." not in processed_arg:
576
                        # False positive, '.' was only in the value
577
578
                        continue
                else:
579
                    value_str = processed_args[i + 1]
580
                    delete.add(i + 1)
581

582
583
584
585
                if processed_arg.endswith("+"):
                    processed_arg = processed_arg[:-1]
                    value_str = json.dumps(list(value_str.split(",")))

586
                key, *keys = processed_arg.split(".")
587
588
589
590
591
                try:
                    value = json.loads(value_str)
                except json.decoder.JSONDecodeError:
                    value = value_str

592
593
                # Merge all values with the same key into a single dict
                arg_dict = create_nested_dict(keys, value)
594
595
                arg_duplicates = recursive_dict_update(dict_args[key], arg_dict)
                duplicates |= {f"{key}.{d}" for d in arg_duplicates}
596
597
                delete.add(i)
        # Filter out the dict args we set to None
598
        processed_args = [a for i, a in enumerate(processed_args) if i not in delete]
599
600
601
        if duplicates:
            logger.warning("Found duplicate keys %s", ", ".join(duplicates))

602
603
604
605
606
        # 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))

607
        return super().parse_args(processed_args, namespace)
608

609
610
611
612
    def check_port(self, value):
        try:
            value = int(value)
        except ValueError:
613
            msg = "Port must be an integer"
614
            raise ArgumentTypeError(msg) from None
615
616

        if not (1024 <= value <= 65535):
617
            raise ArgumentTypeError("Port must be between 1024 and 65535")
618
619
620

        return value

621
    def _pull_args_from_config(self, args: list[str]) -> list[str]:
622
623
        """Method to pull arguments specified in the config file
        into the command-line args variable.
624
625

        The arguments in config file will be inserted between
626
        the argument list.
627

628
629
630
631
632
633
634
635
636
637
        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",
638
639
            "facebook/opt-12B",
            '--config', 'config.yaml',
640
641
642
643
            '-tp', '2'
        ]
        $: args = [
            "serve,chat,complete",
644
645
646
            "facebook/opt-12B",
            '--port', '12323',
            '--tensor-parallel-size', '4',
647
648
649
650
651
            '-tp', '2'
            ]
        ```

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

657
        index = args.index("--config")
658
        if index == len(args) - 1:
659
660
661
662
            raise ValueError(
                "No config file specified! \
                             Please check your command-line arguments."
            )
663
664
665

        file_path = args[index + 1]

666
        config_args = self.load_config_file(file_path)
667

668
        # 0th index might be the sub command {serve,chat,complete,...}
669
        # optionally followed by model_tag (only for serve)
670
671
672
673
        # followed by config args
        # followed by rest of cli args.
        # maintaining this order will enforce the precedence
        # of cli > config > defaults
674
        if args[0].startswith("-"):
675
            # No sub command (e.g., api_server entry point)
676
            args = config_args + args[0:index] + args[index + 2 :]
677
        elif args[0] == "serve":
678
679
            model_in_cli = len(args) > 1 and not args[1].startswith("-")
            model_in_config = any(arg == "--model" for arg in config_args)
680
681

            if not model_in_cli and not model_in_config:
682
                raise ValueError(
683
                    "No model specified! Please specify model either "
684
685
                    "as a positional argument or in a config file."
                )
686
687
688

            if model_in_cli:
                # Model specified as positional arg, keep CLI version
689
690
691
692
693
694
695
                args = (
                    [args[0]]
                    + [args[1]]
                    + config_args
                    + args[2:index]
                    + args[index + 2 :]
                )
696
697
            else:
                # No model in CLI, use config if available
698
                args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
699
        else:
700
            args = [args[0]] + config_args + args[1:index] + args[index + 2 :]
701
702
703

        return args

704
    def load_config_file(self, file_path: str) -> list[str]:
705
        """Loads a yaml file and returns the key value pairs as a
706
707
708
709
710
711
712
713
714
715
716
        flattened list with argparse like pattern
        ```yaml
            port: 12323
            tensor-parallel-size: 4
        ```
        returns:
            processed_args: list[str] = [
                '--port': '12323',
                '--tensor-parallel-size': '4'
            ]
        """
717
718
        extension: str = file_path.split(".")[-1]
        if extension not in ("yaml", "yml"):
719
720
            raise ValueError(
                "Config file must be of a yaml/yml type.\
721
722
723
                              %s supplied",
                extension,
            )
724
725

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

728
        config: dict[str, int | str] = {}
729
        try:
730
            with open(file_path) as config_file:
731
732
733
734
                config = yaml.safe_load(config_file)
        except Exception as ex:
            logger.error(
                "Unable to read the config file at %s. \
735
736
737
                Make sure path is correct",
                file_path,
            )
738
739
            raise ex

740
        store_boolean_arguments = [
741
            action.dest for action in self._actions if isinstance(action, StoreBoolean)
742
743
        ]

744
        for key, value in config.items():
745
746
            if isinstance(value, bool) and key not in store_boolean_arguments:
                if value:
747
                    processed_args.append("--" + key)
748
749
            elif isinstance(value, list):
                if value:
750
                    processed_args.append("--" + key)
751
752
                    for item in value:
                        processed_args.append(str(item))
753
            else:
754
                processed_args.append("--" + key)
755
                processed_args.append(str(value))
756
757
758

        return processed_args

759

760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
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
783
784


785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
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)
808
809


810
# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/srt/utils.py#L630 # noqa: E501
811
def set_ulimit(target_soft_limit=65535):
812
    if sys.platform.startswith("win"):
813
814
815
816
        logger.info("Windows detected, skipping ulimit adjustment.")
        return

    import resource
817

818
819
820
821
822
    resource_type = resource.RLIMIT_NOFILE
    current_soft, current_hard = resource.getrlimit(resource_type)

    if current_soft < target_soft_limit:
        try:
823
            resource.setrlimit(resource_type, (target_soft_limit, current_hard))
824
825
        except ValueError as e:
            logger.warning(
826
827
                "Found ulimit of %s and failed to automatically increase "
                "with error %s. This can cause fd limit errors like "
828
                "`OSError: [Errno 24] Too many open files`. Consider "
829
830
831
832
                "increasing with ulimit -n",
                current_soft,
                e,
            )
833
834
835
836
837
838
839
840
841


# 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


842
843
844
845
846
847
848
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

849
850
    reasons = []
    if is_in_ray_actor():
851
852
853
854
        # 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
855

856
        os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address
857
858
859
860
861
862
        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")
863

864
    if reasons:
865
866
867
        logger.warning(
            "We must use the `spawn` multiprocessing start method. "
            "Overriding VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. "
868
            "See https://docs.vllm.ai/en/latest/usage/"
869
            "troubleshooting.html#python-multiprocessing "
870
871
872
            "for more information. Reasons: %s",
            "; ".join(reasons),
        )
873
874
875
876
        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"


def get_mp_context():
877
878
879
880
881
882
883
    """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()
884
885
    mp_method = envs.VLLM_WORKER_MULTIPROC_METHOD
    return multiprocessing.get_context(mp_method)
886
887
888


def bind_kv_cache(
889
890
    ctx: dict[str, Any],
    kv_cache: list[list[torch.Tensor]],  # [virtual_engine][layer_index]
891
    shared_kv_cache_layers: dict[str, str] | None = None,
892
893
894
895
896
897
898
899
900
901
902
) -> 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
903
904
905
906
    # 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 = {}
907
908
    from vllm.attention import AttentionType
    from vllm.model_executor.models.utils import extract_layer_index
909

910
    layer_need_kv_cache = [
911
912
913
914
915
916
917
918
        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
919
920
    ]
    layer_index_sorted = sorted(
921
922
        set(extract_layer_index(layer_name) for layer_name in layer_need_kv_cache)
    )
923
    for layer_name in layer_need_kv_cache:
924
        kv_cache_idx = layer_index_sorted.index(extract_layer_index(layer_name))
925
926
927
928
        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]
929
930
    if shared_kv_cache_layers is not None:
        for layer_name, target_layer_name in shared_kv_cache_layers.items():
931
932
933
            assert extract_layer_index(target_layer_name) < extract_layer_index(
                layer_name
            ), "v0 doesn't support interleaving kv sharing"
934
            ctx[layer_name].kv_cache = ctx[target_layer_name].kv_cache
935
936


937
938
def run_method(
    obj: Any,
939
    method: str | bytes | Callable,
940
941
942
    args: tuple[Any],
    kwargs: dict[str, Any],
) -> Any:
943
944
945
946
947
948
949
950
951
952
953
954
955
    """
    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:
956
957
958
            raise NotImplementedError(
                f"Method {method!r} is not implemented."
            ) from None
959
960
961
    else:
        func = partial(method, obj)  # type: ignore
    return func(*args, **kwargs)
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982


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.
983
984
985
986
987
988
989
    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.
990
    """
991
    import vllm.third_party.pynvml as pynvml
992

993
    return pynvml
994
995


996
def warn_for_unimplemented_methods(cls: type[T]) -> type[T]:
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
    """
    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
1012
            if attr_name.startswith("_"):
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
                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:
1026
1027
            method_names = ",".join(unimplemented_methods)
            msg = f"Methods {method_names} not implemented in {self}"
1028
            logger.debug(msg)
1029
1030
1031
1032
1033
1034

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

1035
    type.__setattr__(cls, "__init__", wrapped_init)
1036
    return cls
1037
1038


1039
1040
# Only relevant for models using ALiBi (e.g, MPT)
def check_use_alibi(model_config: ModelConfig) -> bool:
1041
    cfg = model_config.hf_text_config
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
    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)
                )
            )
        )
    )
1062
1063


1064
1065
1066
def set_process_title(
    name: str, suffix: str = "", prefix: str = envs.VLLM_PROCESS_NAME_PREFIX
) -> None:
1067
1068
1069
    """
    Set the current process title to a specific name with an
    optional suffix.
1070
1071

    Args:
1072
        name: The title to assign to the current process.
1073
        suffix: An optional suffix to append to the base name.
1074
        prefix: A prefix to prepend to the front separated by `::`.
1075
1076
1077
    """
    if suffix:
        name = f"{name}_{suffix}"
1078
    setproctitle.setproctitle(f"{prefix}::{name}")
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092


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
1093
        while (next_idx := s.find("\n", idx)) != -1:
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
            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]


1108
def decorate_logs(process_name: str | None = None) -> None:
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
    """
    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)
1128
1129
1130


def length_from_prompt_token_ids_or_embeds(
1131
1132
    prompt_token_ids: list[int] | None,
    prompt_embeds: torch.Tensor | None,
1133
) -> int:
1134
    """Calculate the request length (in number of tokens) give either
1135
1136
    prompt_token_ids or prompt_embeds.
    """
1137
1138
    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)
1139
1140
1141

    if prompt_token_len is None:
        if prompt_embeds_len is None:
1142
            raise ValueError("Neither prompt_token_ids nor prompt_embeds were defined.")
1143
1144
        return prompt_embeds_len
    else:
1145
        if prompt_embeds_len is not None and prompt_embeds_len != prompt_token_len:
1146
1147
1148
            raise ValueError(
                "Prompt token ids and prompt embeds had different lengths"
                f" prompt_token_ids={prompt_token_len}"
1149
1150
                f" prompt_embeds={prompt_embeds_len}"
            )
1151
        return prompt_token_len
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164


@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
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184


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