utils.py 29.8 KB
Newer Older
1
import argparse
2
import asyncio
3
import contextlib
4
import datetime
Woosuk Kwon's avatar
Woosuk Kwon committed
5
import enum
6
import gc
7
import os
8
import socket
9
import subprocess
10
import sys
11
12
import tempfile
import threading
Zhuohan Li's avatar
Zhuohan Li committed
13
import uuid
14
import warnings
15
from collections import defaultdict
16
from functools import lru_cache, partial, wraps
17
from platform import uname
18
from typing import (Any, AsyncIterator, Awaitable, Callable, Dict, Generic,
19
                    Hashable, List, Optional, OrderedDict, Set, Tuple, TypeVar,
20
                    Union)
Zhuohan Li's avatar
Zhuohan Li committed
21

22
import numpy as np
23
import psutil
Zhuohan Li's avatar
Zhuohan Li committed
24
import torch
25
26
import torch.types
from typing_extensions import ParamSpec
27

28
import vllm.envs as envs
29
from vllm import _custom_ops as ops
30
from vllm.logger import enable_trace_function_call, init_logger
31
32
33
34
35
36
37

logger = init_logger(__name__)

STR_DTYPE_TO_TORCH_DTYPE = {
    "half": torch.half,
    "bfloat16": torch.bfloat16,
    "float": torch.float,
38
    "fp8": torch.uint8,
39
40
    "fp8_e4m3": torch.uint8,
    "fp8_e5m2": torch.uint8,
41
}
Zhuohan Li's avatar
Zhuohan Li committed
42

43
44
45
46
P = ParamSpec('P')
K = TypeVar("K")
T = TypeVar("T")

Woosuk Kwon's avatar
Woosuk Kwon committed
47

48
49
50
51
52
53
54
class _Sentinel:
    ...


ALL_PINNED_SENTINEL = _Sentinel()


Woosuk Kwon's avatar
Woosuk Kwon committed
55
56
57
58
59
60
61
62
63
64
class Device(enum.Enum):
    GPU = enum.auto()
    CPU = enum.auto()


class Counter:

    def __init__(self, start: int = 0) -> None:
        self.counter = start

Woosuk Kwon's avatar
Woosuk Kwon committed
65
    def __next__(self) -> int:
66
        i = self.counter
Woosuk Kwon's avatar
Woosuk Kwon committed
67
        self.counter += 1
68
        return i
Woosuk Kwon's avatar
Woosuk Kwon committed
69
70
71

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

73

74
class LRUCache(Generic[T]):
75
76

    def __init__(self, capacity: int):
77
        self.cache: OrderedDict[Hashable, T] = OrderedDict()
78
        self.pinned_items: Set[Hashable] = set()
79
80
81
82
83
84
85
86
        self.capacity = capacity

    def __contains__(self, key: Hashable) -> bool:
        return key in self.cache

    def __len__(self) -> int:
        return len(self.cache)

87
    def __getitem__(self, key: Hashable) -> Optional[T]:
88
89
        return self.get(key)

90
    def __setitem__(self, key: Hashable, value: T) -> None:
91
92
93
94
95
96
97
98
        self.put(key, value)

    def __delitem__(self, key: Hashable) -> None:
        self.pop(key)

    def touch(self, key: Hashable) -> None:
        self.cache.move_to_end(key)

99
100
101
    def get(self,
            key: Hashable,
            default_value: Optional[T] = None) -> Optional[T]:
102
        if key in self.cache:
103
            value: Optional[T] = self.cache[key]
104
105
106
107
108
            self.cache.move_to_end(key)
        else:
            value = default_value
        return value

109
    def put(self, key: Hashable, value: T) -> None:
110
111
112
113
        self.cache[key] = value
        self.cache.move_to_end(key)
        self._remove_old_if_needed()

114
115
116
117
118
119
120
121
122
123
124
125
    def pin(self, key: Hashable) -> None:
        """
        Pins a key in the cache preventing it from being
        evicted in the LRU order.
        """
        if key not in self.cache:
            raise ValueError(f"Cannot pin key: {key} not in cache.")
        self.pinned_items.add(key)

    def _unpin(self, key: Hashable) -> None:
        self.pinned_items.remove(key)

126
    def _on_remove(self, key: Hashable, value: Optional[T]):
127
128
        pass

129
    def remove_oldest(self, remove_pinned=False):
130
131
        if not self.cache:
            return
132
133
134
135
136
137
138
139
140
141
142
143

        if not remove_pinned:
            # pop the oldest item in the cache that is not pinned
            lru_key = next(
                (key for key in self.cache if key not in self.pinned_items),
                ALL_PINNED_SENTINEL)
            if lru_key is ALL_PINNED_SENTINEL:
                raise RuntimeError("All items are pinned, "
                                   "cannot remove oldest from the cache.")
        else:
            lru_key = next(iter(self.cache))
        self.pop(lru_key)
144
145
146
147
148

    def _remove_old_if_needed(self) -> None:
        while len(self.cache) > self.capacity:
            self.remove_oldest()

149
150
151
    def pop(self,
            key: Hashable,
            default_value: Optional[T] = None) -> Optional[T]:
152
        run_on_remove = key in self.cache
153
        value: Optional[T] = self.cache.pop(key, default_value)
154
155
156
        # remove from pinned items
        if key in self.pinned_items:
            self._unpin(key)
157
158
159
160
161
162
        if run_on_remove:
            self._on_remove(key, value)
        return value

    def clear(self):
        while len(self.cache) > 0:
163
            self.remove_oldest(remove_pinned=True)
164
165
166
        self.cache.clear()


167
168
169
170
def is_hip() -> bool:
    return torch.version.hip is not None


171
172
@lru_cache(maxsize=None)
def is_cpu() -> bool:
173
174
175
176
177
    from importlib.metadata import PackageNotFoundError, version
    try:
        return "cpu" in version("vllm")
    except PackageNotFoundError:
        return False
178
179


180
181
182
183
184
185
186
187
188
@lru_cache(maxsize=None)
def is_openvino() -> bool:
    from importlib.metadata import PackageNotFoundError, version
    try:
        return "openvino" in version("vllm")
    except PackageNotFoundError:
        return False


189
@lru_cache(maxsize=None)
190
191
192
193
194
195
196
197
def is_neuron() -> bool:
    try:
        import transformers_neuronx
    except ImportError:
        transformers_neuronx = None
    return transformers_neuronx is not None


198
199
200
201
202
203
204
205
206
@lru_cache(maxsize=None)
def is_tpu() -> bool:
    try:
        import libtpu
    except ImportError:
        libtpu = None
    return libtpu is not None


207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
@lru_cache(maxsize=None)
def is_xpu() -> bool:
    from importlib.metadata import version
    is_xpu_flag = "xpu" in version("vllm")
    # vllm is not build with xpu
    if not is_xpu_flag:
        return False
    try:
        import intel_extension_for_pytorch as ipex  # noqa: F401
        _import_ipex = True
    except ImportError as e:
        logger.warning("Import Error for IPEX: %s", e.msg)
        _import_ipex = False
    # ipex dependency is not ready
    if not _import_ipex:
        logger.warning("not found ipex lib")
        return False
    return hasattr(torch, "xpu") and torch.xpu.is_available()


227
@lru_cache(maxsize=None)
228
229
def get_max_shared_memory_bytes(gpu: int = 0) -> int:
    """Returns the maximum shared memory per thread block in bytes."""
230
    max_shared_mem = (
231
        ops.get_max_shared_memory_per_block_device_attribute(gpu))
232
233
    # value 0 will cause MAX_SEQ_LEN become negative and test_attention.py
    # will fail
234
    assert max_shared_mem > 0, "max_shared_mem can not be zero"
235
236
237
    return int(max_shared_mem)


238
def get_cpu_memory() -> int:
239
    """Returns the total CPU memory of the node in bytes."""
240
    return psutil.virtual_memory().total
Zhuohan Li's avatar
Zhuohan Li committed
241
242
243
244


def random_uuid() -> str:
    return str(uuid.uuid4().hex)
245

246

247
@lru_cache(maxsize=None)
248
def get_vllm_instance_id() -> str:
249
250
251
252
253
254
    """
    If the environment variable VLLM_INSTANCE_ID is set, return it.
    Otherwise, return a random UUID.
    Instance id represents an instance of the VLLM. All processes in the same
    instance should have the same instance id.
    """
255
    return envs.VLLM_INSTANCE_ID or f"vllm-instance-{random_uuid()}"
256
257


258
@lru_cache(maxsize=None)
259
260
261
def in_wsl() -> bool:
    # Reference: https://github.com/microsoft/WSL/issues/4071
    return "microsoft" in " ".join(uname()).lower()
262
263


264
def make_async(func: Callable[P, T]) -> Callable[P, Awaitable[T]]:
265
266
267
268
269
270
271
    """Take a blocking function, and run it on in an executor thread.

    This function prevents the blocking function from blocking the
    asyncio event loop.
    The code in this function needs to be thread safe.
    """

272
    def _async_wrapper(*args: P.args, **kwargs: P.kwargs) -> asyncio.Future:
273
274
275
276
277
278
279
        loop = asyncio.get_event_loop()
        p_func = partial(func, *args, **kwargs)
        return loop.run_in_executor(executor=None, func=p_func)

    return _async_wrapper


280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
def merge_async_iterators(
        *iterators: AsyncIterator[T]) -> AsyncIterator[Tuple[int, T]]:
    """Merge multiple asynchronous iterators into a single iterator.

    This method handle the case where some iterators finish before others.
    When it yields, it yields a tuple (i, item) where i is the index of the
    iterator that yields the item.
    """
    queue: asyncio.Queue[Union[Tuple[int, T], Exception]] = asyncio.Queue()

    finished = [False] * len(iterators)

    async def producer(i: int, iterator: AsyncIterator[T]):
        try:
            async for item in iterator:
                await queue.put((i, item))
        except Exception as e:
            await queue.put(e)
        finished[i] = True

    _tasks = [
        asyncio.create_task(producer(i, iterator))
        for i, iterator in enumerate(iterators)
    ]

    async def consumer():
306
307
308
309
310
311
312
313
        try:
            while not all(finished) or not queue.empty():
                item = await queue.get()
                if isinstance(item, Exception):
                    raise item
                yield item
        except (Exception, asyncio.CancelledError) as e:
            for task in _tasks:
314
315
316
317
318
                if sys.version_info >= (3, 9):
                    # msg parameter only supported in Python 3.9+
                    task.cancel(e)
                else:
                    task.cancel()
319
            raise e
320
321
322
323
324
        await asyncio.gather(*_tasks)

    return consumer()


325
def get_ip() -> str:
326
    host_ip = envs.VLLM_HOST_IP
327
328
329
330
331
    if host_ip:
        return host_ip

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

332
    # try ipv4
333
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
334
    try:
335
        s.connect(("8.8.8.8", 80))  # Doesn't need to be reachable
336
        return s.getsockname()[0]
337
338
339
340
341
    except Exception:
        pass

    # try ipv6
    try:
342
        s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM)
343
344
345
        # Google's public DNS server, see
        # https://developers.google.com/speed/public-dns/docs/using#addresses
        s.connect(("2001:4860:4860::8888", 80))  # Doesn't need to be reachable
346
        return s.getsockname()[0]
347
348
349
350
351
    except Exception:
        pass

    warnings.warn(
        "Failed to get the IP address, using 0.0.0.0 by default."
352
353
        "The value can be set by the environment variable"
        " VLLM_HOST_IP or HOST_IP.",
354
355
        stacklevel=2)
    return "0.0.0.0"
356
357


358
def get_distributed_init_method(ip: str, port: int) -> str:
359
360
361
    # Brackets are not permitted in ipv4 addresses,
    # see https://github.com/python/cpython/issues/103848
    return f"tcp://[{ip}]:{port}" if ":" in ip else f"tcp://{ip}:{port}"
362
363


364
def get_open_port() -> int:
365
366
    port = envs.VLLM_PORT
    if port is not None:
367
368
369
370
371
372
373
374
375
        while True:
            try:
                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                    s.bind(("", port))
                    return port
            except OSError:
                port += 1  # Increment port number if already in use
                logger.info("Port %d is already in use, trying port %d",
                            port - 1, port)
376
377
378
379
380
381
382
383
384
385
    # try ipv4
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("", 0))
            return s.getsockname()[1]
    except OSError:
        # try ipv6
        with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s:
            s.bind(("", 0))
            return s.getsockname()[1]
386
387


388
389
def update_environment_variables(envs: Dict[str, str]):
    for k, v in envs.items():
390
        if k in os.environ and os.environ[k] != v:
391
392
393
            logger.warning(
                "Overwriting environment variable %s "
                "from '%s' to '%s'", k, os.environ[k], v)
394
        os.environ[k] = v
395
396


397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
def init_kmp_env():
    if not is_cpu():
        return

    ld_prealod_str = os.getenv("LD_PRELOAD", "")
    if "libiomp5.so" not in ld_prealod_str:
        return

    # The time(milliseconds) that a thread should wait after completing the
    # execution of a parallel region, before sleeping.
    os.environ['KMP_BLOCKTIME'] = "1"
    # dump settings on start up
    os.environ['KMP_SETTINGS'] = "1"
    # Prevents the CPU to run into low performance state
    os.environ['KMP_TPAUSE'] = "0"
    # Provides fine granularity parallelism
    os.environ['KMP_FORKJOIN_BARRIER_PATTERN'] = "dist,dist"
    os.environ['KMP_PLAIN_BARRIER_PATTERN'] = "dist,dist"
    os.environ['KMP_REDUCTION_BARRIER_PATTERN'] = "dist,dist"


418
def chunk_list(lst: List[T], chunk_size: int) -> List[List[T]]:
419
420
421
422
423
424
425
426
427
    """Yield successive chunk_size chunks from lst."""
    return [lst[i:i + chunk_size] for i in range(0, len(lst), chunk_size)]


def cdiv(a: int, b: int) -> int:
    """Ceiling division."""
    return -(a // -b)


428
def _generate_random_fp8(
429
    tensor: torch.Tensor,
430
431
432
433
434
435
    low: float,
    high: float,
) -> None:
    # NOTE(zhaoyang): Due to NaN and Inf representation for fp8 data type,
    # it may occur Inf or NaN if we directly use torch.randint
    # to generate random data for fp8 data.
436
    # For example, s.11111.00 in fp8e5m2 format represents Inf.
437
438
439
440
    #     | E4M3        | E5M2
    #-----|-------------|-------------------
    # Inf | N/A         | s.11111.00
    # NaN | s.1111.111  | s.11111.{01,10,11}
441
    from vllm import _custom_ops as ops
442
443
    tensor_tmp = torch.empty_like(tensor, dtype=torch.float16)
    tensor_tmp.uniform_(low, high)
444
    ops.convert_fp8(tensor, tensor_tmp)
445
446
447
    del tensor_tmp


448
449
450
def get_kv_cache_torch_dtype(
        cache_dtype: Optional[Union[str, torch.dtype]],
        model_dtype: Optional[Union[str, torch.dtype]] = None) -> torch.dtype:
451
452
453
454
455
456
457
458
459
460
    if isinstance(cache_dtype, str):
        if cache_dtype == "auto":
            if isinstance(model_dtype, str):
                torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[model_dtype]
            elif isinstance(model_dtype, torch.dtype):
                torch_dtype = model_dtype
            else:
                raise ValueError(f"Invalid model dtype: {model_dtype}")
        elif cache_dtype in ["half", "bfloat16", "float"]:
            torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype]
461
        elif cache_dtype == "fp8":
462
463
464
465
466
467
468
            torch_dtype = torch.uint8
        else:
            raise ValueError(f"Invalid kv cache dtype: {cache_dtype}")
    elif isinstance(cache_dtype, torch.dtype):
        torch_dtype = cache_dtype
    else:
        raise ValueError(f"Invalid kv cache dtype: {cache_dtype}")
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
    return torch_dtype


def create_kv_caches_with_random_flash(
    num_blocks: int,
    block_size: int,
    num_layers: int,
    num_heads: int,
    head_size: int,
    cache_dtype: Optional[Union[str, torch.dtype]],
    model_dtype: Optional[Union[str, torch.dtype]] = None,
    seed: int = 0,
    device: Optional[str] = "cuda",
) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
    assert cache_dtype != "fp8"
    torch.random.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)

    torch_dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
    key_value_cache_shape = (num_blocks, 2, block_size, num_heads, head_size)
    scale = head_size**-0.5
491
492
493
494

    key_caches: List[torch.Tensor] = []
    value_caches: List[torch.Tensor] = []

495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
    for _ in range(num_layers):
        key_value_cache = torch.empty(size=key_value_cache_shape,
                                      dtype=torch_dtype,
                                      device=device)
        key_value_cache.uniform_(-scale, scale)
        key_caches.append(key_value_cache[:, 0])
        value_caches.append(key_value_cache[:, 1])
    return key_caches, value_caches


def create_kv_caches_with_random(
    num_blocks: int,
    block_size: int,
    num_layers: int,
    num_heads: int,
    head_size: int,
    cache_dtype: Optional[Union[str, torch.dtype]],
    model_dtype: Optional[Union[str, torch.dtype]] = None,
    seed: int = 0,
    device: Optional[str] = "cuda",
) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
    torch.random.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)

    torch_dtype = get_kv_cache_torch_dtype(cache_dtype, model_dtype)
521
522
523
524

    scale = head_size**-0.5
    x = 16 // torch.tensor([], dtype=torch_dtype).element_size()
    key_cache_shape = (num_blocks, num_heads, head_size // x, block_size, x)
525
    key_caches: List[torch.Tensor] = []
526
527
528
529
    for _ in range(num_layers):
        key_cache = torch.empty(size=key_cache_shape,
                                dtype=torch_dtype,
                                device=device)
530
        if cache_dtype in ["auto", "half", "bfloat16", "float"]:
531
            key_cache.uniform_(-scale, scale)
532
533
        elif cache_dtype == 'fp8':
            _generate_random_fp8(key_cache, -scale, scale)
534
535
536
        else:
            raise ValueError(
                f"Does not support key cache of type {cache_dtype}")
537
538
539
        key_caches.append(key_cache)

    value_cache_shape = (num_blocks, num_heads, head_size, block_size)
540
    value_caches: List[torch.Tensor] = []
541
542
543
544
    for _ in range(num_layers):
        value_cache = torch.empty(size=value_cache_shape,
                                  dtype=torch_dtype,
                                  device=device)
545
        if cache_dtype in ["auto", "half", "bfloat16", "float"]:
546
            value_cache.uniform_(-scale, scale)
547
548
        elif cache_dtype == 'fp8':
            _generate_random_fp8(value_cache, -scale, scale)
549
550
551
        else:
            raise ValueError(
                f"Does not support value cache of type {cache_dtype}")
552
553
        value_caches.append(value_cache)
    return key_caches, value_caches
554
555


556
557
558
559
560
561
562
563
564
565
566
567
568
569
@lru_cache
def print_warning_once(msg: str) -> None:
    logger.warning(msg)


@lru_cache(maxsize=None)
def is_pin_memory_available() -> bool:

    if in_wsl():
        # Pinning memory in WSL is not supported.
        # https://docs.nvidia.com/cuda/wsl-user-guide/index.html#known-limitations-for-linux-cuda-applications
        print_warning_once("Using 'pin_memory=False' as WSL is detected. "
                           "This may slow down the performance.")
        return False
570
571
572
    elif is_xpu():
        print_warning_once("Pin memory is not supported on XPU.")
        return False
573
574
575
    elif is_neuron():
        print_warning_once("Pin memory is not supported on Neuron.")
        return False
576
    elif is_cpu() or is_openvino():
577
        return False
578
579
580
581
    return True


class CudaMemoryProfiler:
582

583
    def __init__(self, device: Optional[torch.types.Device] = None):
584
585
586
587
        self.device = device

    def current_memory_usage(self) -> float:
        # Return the memory usage in bytes.
588
589
590
591
592
593
        if torch.cuda.is_available():
            torch.cuda.reset_peak_memory_stats(self.device)
            mem = torch.cuda.max_memory_allocated(self.device)
        elif is_xpu():
            torch.xpu.reset_peak_memory_stats(self.device)
            mem = torch.xpu.max_memory_allocated(self.device)
594
595
596
597
598
599
600
601
602
603
604
605
606
        return mem

    def __enter__(self):
        self.initial_memory = self.current_memory_usage()
        # This allows us to call methods of the context manager if needed
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.final_memory = self.current_memory_usage()
        self.consumed_memory = self.final_memory - self.initial_memory

        # Force garbage collection
        gc.collect()
607
608


609
def str_to_int_tuple(s: str) -> Tuple[int, ...]:
610
611
612
613
614
615
616
617
618
    """Convert a string to a tuple of integers."""
    try:
        return tuple(map(int, s.split(",")))
    except ValueError as e:
        raise ValueError(
            "String must be a series of integers separated by commas "
            f"(e.g., 1, 2, 3). Given input: {s}") from e


619
620
621
622
623
624
625
626
627
628
629
630
def make_tensor_with_pad(
    x: List[List[int]],
    max_len: int,
    pad: int,
    dtype: torch.dtype,
    device: Optional[Union[str, torch.device]],
) -> torch.Tensor:
    """Make a padded tensor of a 2D inputs.

    The padding is applied to the end of each inner list until it reaches
    `max_len`.
    """
631
632
633
634
    padded_x = np.zeros([len(x), max_len], dtype=np.int32) + pad
    for ind, blocktb in enumerate(x):
        assert len(blocktb) <= max_len
        padded_x[ind, :len(blocktb)] = blocktb
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
    return torch.tensor(padded_x, dtype=dtype, device=device)


def async_tensor_h2d(
    data: list,
    dtype: torch.dtype,
    target_device: Union[str, torch.device],
    pin_memory: bool,
) -> torch.Tensor:
    """Asynchronously create a tensor and copy it from host to device."""
    t = torch.tensor(data, dtype=dtype, pin_memory=pin_memory, device="cpu")
    return t.to(device=target_device, non_blocking=True)


def maybe_expand_dim(tensor: torch.Tensor,
                     target_dims: int,
                     size: int = 1) -> torch.Tensor:
    """Expand the tensor to the target_dims."""
    if tensor.ndim < target_dims:
        tensor = tensor.view(-1, *([size] * (target_dims - tensor.ndim)))
    return tensor
656
657


658
659
660
661
662
def get_dtype_size(dtype: torch.dtype) -> int:
    """Get the size of the data type in bytes."""
    return torch.tensor([], dtype=dtype).element_size()


663
664
def merge_dicts(dict1: Dict[K, List[T]],
                dict2: Dict[K, List[T]]) -> Dict[K, List[T]]:
665
    """Merge 2 dicts that have key -> List of items.
666

667
668
    When a key conflicts, the values in dict1 is prioritized.
    """
669
    merged_dict: Dict[K, List[T]] = defaultdict(list)
670
671
672
673
674
675
676
677

    for key, value in dict1.items():
        merged_dict[key].extend(value)

    for key, value in dict2.items():
        merged_dict[key].extend(value)

    return dict(merged_dict)
678
679


680
def init_cached_hf_modules() -> None:
681
682
683
684
685
    """
    Lazy initialization of the Hugging Face modules.
    """
    from transformers.dynamic_module_utils import init_hf_modules
    init_hf_modules()
686
687


688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
@lru_cache(maxsize=None)
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
704
    env_ld_library_path = envs.LD_LIBRARY_PATH
705
706
707
708
709
710
711
712
713
714
715
    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]


716
def find_nccl_library() -> str:
717
718
719
720
721
722
    """
    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.
    """
723
    so_file = envs.VLLM_NCCL_SO_PATH
724
725
726
727

    # manually load the nccl library
    if so_file:
        logger.info(
728
729
            "Found nccl from environment variable VLLM_NCCL_SO_PATH=%s",
            so_file)
730
731
    else:
        if torch.version.cuda is not None:
732
            so_file = "libnccl.so.2"
733
        elif torch.version.hip is not None:
734
            so_file = "librccl.so.1"
735
736
        else:
            raise ValueError("NCCL only supports CUDA and ROCm backends.")
737
        logger.info("Found nccl from library %s", so_file)
738
    return so_file
739
740
741
742
743
744
745


def enable_trace_function_call_for_thread() -> None:
    """Set up function tracing for the current thread,
    if enabled via the VLLM_TRACE_FUNCTION environment variable
    """

746
    if envs.VLLM_TRACE_FUNCTION:
747
748
749
750
751
752
753
754
        tmp_dir = tempfile.gettempdir()
        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", get_vllm_instance_id(),
                                filename)
        os.makedirs(os.path.dirname(log_path), exist_ok=True)
        enable_trace_function_call(log_path)
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795


def identity(value: T) -> T:
    return value


F = TypeVar('F', bound=Callable[..., Any])


def deprecate_kwargs(
        *kws: str,
        is_deprecated: Union[bool, Callable[[], bool]] = True,
        additional_message: Optional[str] = None) -> Callable[[F], F]:
    deprecated_kws = set(kws)

    if not callable(is_deprecated):
        is_deprecated = partial(identity, is_deprecated)

    def wrapper(fn: F) -> F:

        @wraps(fn)
        def inner(*args, **kwargs):
            if is_deprecated():
                deprecated_kwargs = kwargs.keys() & deprecated_kws
                if deprecated_kwargs:
                    msg = (
                        f"The keyword arguments {deprecated_kwargs} are "
                        "deprecated and will be removed in a future update.")
                    if additional_message is not None:
                        msg += f" {additional_message}"

                    warnings.warn(
                        DeprecationWarning(msg),
                        stacklevel=3,  # The inner function takes up one level
                    )

            return fn(*args, **kwargs)

        return inner  # type: ignore

    return wrapper
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812


@lru_cache(maxsize=8)
def _cuda_device_count_stateless(
        cuda_visible_devices: Optional[str] = None) -> int:
    # Note: cuda_visible_devices is not used, but we keep it as an argument for
    # LRU Cache purposes.

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

    if not torch.cuda._is_compiled():
        return 0
813
814
815
816
817
818
819
820
    if is_hip():
        # ROCm uses amdsmi instead of nvml for stateless device count
        # This requires a sufficiently modern version of Torch 2.4.0
        raw_count = torch.cuda._device_count_amdsmi() if (hasattr(
            torch.cuda, "_device_count_amdsmi")) else -1
    else:
        raw_count = torch.cuda._device_count_nvml()
    r = torch._C._cuda_getDeviceCount() if raw_count < 0 else raw_count
821
822
823
824
825
826
827
828
829
830
831
832
833
834
    return r


def cuda_device_count_stateless() -> int:
    """Get number of CUDA devices, caching based on the value of
    CUDA_VISIBLE_DEVICES at the time of call.
    
    This should be used instead of torch.cuda.device_count()
    unless CUDA_VISIBLE_DEVICES has already been set to the desired
    value."""

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


837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
def error_on_invalid_device_count_status():
    cache_entries = 0
    with contextlib.suppress(Exception):
        # future pytorch will fix the issue, device_count will not be cached
        # at that time, `.cache_info().currsize` will error out
        cache_entries = torch.cuda.device_count.cache_info().currsize
    if cache_entries != 0:
        # the function is already called, and the result is cached
        remembered = torch.cuda.device_count()
        current = cuda_device_count_stateless()
        if remembered > current:
            raise RuntimeError(
                "The number of CUDA devices has changed since the first "
                "call to torch.cuda.device_count(). This is not allowed "
                "and may result in undefined behavior. Please check out "
                "https://github.com/vllm-project/vllm/issues/6056 to "
                "find the first call to torch.cuda.device_count() "
                "and defer it until the engine is up. Or you can set "
                "CUDA_VISIBLE_DEVICES to the GPUs you want to use.")


858
859
860
861
862
863
864
865
866
867
868
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
894
895
896
897
898
899
900
901
902
903
904
905
906
907
# NVML utils
# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
# all the related functions work on real physical device ids.
# the major benefit of using NVML is that it will not initialize CUDA

try:
    import pynvml
except ImportError:
    # For non-NV devices
    pynvml = None


def with_nvml_context(fn):

    @wraps(fn)
    def wrapper(*args, **kwargs):
        if pynvml is not None:
            pynvml.nvmlInit()
        try:
            return fn(*args, **kwargs)
        finally:
            if pynvml is not None:
                pynvml.nvmlShutdown()

    return wrapper


@with_nvml_context
def is_full_nvlink(device_ids: List[int]) -> bool:
    """
    query if the set of gpus are fully connected by nvlink (1 hop)
    """
    handles = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in device_ids]
    for i, handle in enumerate(handles):
        for j, peer_handle in enumerate(handles):
            if i < j:
                try:
                    p2p_status = pynvml.nvmlDeviceGetP2PStatus(
                        handle, peer_handle, pynvml.NVML_P2P_CAPS_INDEX_NVLINK)
                    if p2p_status != pynvml.NVML_P2P_STATUS_OK:
                        return False
                except pynvml.NVMLError as error:
                    logger.error(
                        "NVLink detection failed. This is normal if your"
                        " machine has no NVLink equipped.",
                        exc_info=error)
                    return False
    return True


908
909
910
911
912
913
914
915
916
917
#From: https://stackoverflow.com/a/4104188/2749989
def run_once(f):

    def wrapper(*args, **kwargs) -> Any:
        if not wrapper.has_run:  # type: ignore[attr-defined]
            wrapper.has_run = True  # type: ignore[attr-defined]
            return f(*args, **kwargs)

    wrapper.has_run = False  # type: ignore[attr-defined]
    return wrapper
918
919
920
921
922
923
924
925
926
927
928
929
930


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

    def parse_args(self, args=None, namespace=None):
        if args is None:
            args = sys.argv[1:]

        # Convert underscores to dashes and vice versa in argument names
        processed_args = []
        for arg in args:
            if arg.startswith('--'):
931
932
933
934
935
936
937
                if '=' in arg:
                    key, value = arg.split('=', 1)
                    key = '--' + key[len('--'):].replace('_', '-')
                    processed_args.append(f'{key}={value}')
                else:
                    processed_args.append('--' +
                                          arg[len('--'):].replace('_', '-'))
938
939
940
941
            else:
                processed_args.append(arg)

        return super().parse_args(processed_args, namespace)