__init__.py 9.44 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
7
import logging
import traceback
from itertools import chain
from typing import TYPE_CHECKING, Optional

8
from vllm import envs
9
from vllm.plugins import load_plugins_by_group
10
from vllm.utils import resolve_obj_by_qualname, supports_xccl
11
12

from .interface import CpuArchEnum, Platform, PlatformEnum
13

14
logger = logging.getLogger(__name__)
15

16

17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def vllm_version_matches_substr(substr: str) -> bool:
    """
    Check to see if the vLLM version matches a substring.
    """
    from importlib.metadata import PackageNotFoundError, version
    try:
        vllm_version = version("vllm")
    except PackageNotFoundError as e:
        logger.warning(
            "The vLLM package was not found, so its version could not be "
            "inspected. This may cause platform detection to fail.")
        raise e
    return substr in vllm_version


32
def tpu_platform_plugin() -> Optional[str]:
33
    logger.debug("Checking if TPU platform is available.")
34
35
36
37
38
39
40

    # Check for Pathways TPU proxy
    if envs.VLLM_TPU_USING_PATHWAYS:
        logger.debug("Confirmed TPU platform is available via Pathways proxy.")
        return "tpu_commons.platforms.tpu_jax.TpuPlatform"

    # Check for libtpu installation
41
42
43
    try:
        # While it's technically possible to install libtpu on a
        # non-TPU machine, this is a very uncommon scenario. Therefore,
44
        # we assume that libtpu is installed only if the machine
45
        # has TPUs.
46

47
        import libtpu  # noqa: F401
48
        logger.debug("Confirmed TPU platform is available.")
49
        return "vllm.platforms.tpu.TpuPlatform"
50
51
    except Exception as e:
        logger.debug("TPU platform is not available because: %s", str(e))
52
        return None
53
54


55
56
def cuda_platform_plugin() -> Optional[str]:
    is_cuda = False
57
    logger.debug("Checking if CUDA platform is available.")
58
    try:
59
60
        from vllm.utils import import_pynvml
        pynvml = import_pynvml()
61
62
        pynvml.nvmlInit()
        try:
63
64
65
66
67
68
            # NOTE: Edge case: vllm cpu build on a GPU machine.
            # Third-party pynvml can be imported in cpu build,
            # we need to check if vllm is built with cpu too.
            # Otherwise, vllm will always activate cuda plugin
            # on a GPU machine, even if in a cpu build.
            is_cuda = (pynvml.nvmlDeviceGetCount() > 0
69
                       and not vllm_version_matches_substr("cpu"))
70
71
72
73
74
75
76
77
            if pynvml.nvmlDeviceGetCount() <= 0:
                logger.debug(
                    "CUDA platform is not available because no GPU is found.")
            if vllm_version_matches_substr("cpu"):
                logger.debug("CUDA platform is not available because"
                             " vLLM is built with CPU.")
            if is_cuda:
                logger.debug("Confirmed CUDA platform is available.")
78
79
        finally:
            pynvml.nvmlShutdown()
80
    except Exception as e:
81
82
        logger.debug("Exception happens when checking CUDA platform: %s",
                     str(e))
83
84
85
86
        if "nvml" not in e.__class__.__name__.lower():
            # If the error is not related to NVML, re-raise it.
            raise e

87
88
89
90
91
92
93
94
        # CUDA is supported on Jetson, but NVML may not be.
        import os

        def cuda_is_jetson() -> bool:
            return os.path.isfile("/etc/nv_tegra_release") \
                or os.path.exists("/sys/class/tegra-firmware")

        if cuda_is_jetson():
95
            logger.debug("Confirmed CUDA platform is available on Jetson.")
96
            is_cuda = True
97
98
        else:
            logger.debug("CUDA platform is not available because: %s", str(e))
99

100
101
102
103
104
    return "vllm.platforms.cuda.CudaPlatform" if is_cuda else None


def rocm_platform_plugin() -> Optional[str]:
    is_rocm = False
105
    logger.debug("Checking if ROCm platform is available.")
106
107
108
109
110
111
    try:
        import amdsmi
        amdsmi.amdsmi_init()
        try:
            if len(amdsmi.amdsmi_get_processor_handles()) > 0:
                is_rocm = True
112
                logger.debug("Confirmed ROCm platform is available.")
113
114
115
            else:
                logger.debug("ROCm platform is not available because"
                             " no GPU is found.")
116
117
        finally:
            amdsmi.amdsmi_shut_down()
118
119
    except Exception as e:
        logger.debug("ROCm platform is not available because: %s", str(e))
120
121
122
123
124
125

    return "vllm.platforms.rocm.RocmPlatform" if is_rocm else None


def xpu_platform_plugin() -> Optional[str]:
    is_xpu = False
126
    logger.debug("Checking if XPU platform is available.")
127
128
129
130
    try:
        # installed IPEX if the machine has XPUs.
        import intel_extension_for_pytorch  # noqa: F401
        import torch
131
132
133
134
135
136
        if supports_xccl():
            dist_backend = "xccl"
        else:
            dist_backend = "ccl"
            import oneccl_bindings_for_pytorch  # noqa: F401

137
138
        if hasattr(torch, 'xpu') and torch.xpu.is_available():
            is_xpu = True
139
140
141
142
            from vllm.platforms.xpu import XPUPlatform
            XPUPlatform.dist_backend = dist_backend
            logger.debug("Confirmed %s backend is available.",
                         XPUPlatform.dist_backend)
143
144
145
            logger.debug("Confirmed XPU platform is available.")
    except Exception as e:
        logger.debug("XPU platform is not available because: %s", str(e))
146
147
148
149
150
151

    return "vllm.platforms.xpu.XPUPlatform" if is_xpu else None


def cpu_platform_plugin() -> Optional[str]:
    is_cpu = False
152
    logger.debug("Checking if CPU platform is available.")
153
    try:
154
        is_cpu = vllm_version_matches_substr("cpu")
155
156
157
        if is_cpu:
            logger.debug("Confirmed CPU platform is available because"
                         " vLLM is built with CPU.")
158
        if not is_cpu:
159
160
            import sys
            is_cpu = sys.platform.startswith("darwin")
161
162
            if is_cpu:
                logger.debug("Confirmed CPU platform is available"
163
                             " because the machine is MacOS.")
164

165
166
    except Exception as e:
        logger.debug("CPU platform is not available because: %s", str(e))
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192

    return "vllm.platforms.cpu.CpuPlatform" if is_cpu else None


builtin_platform_plugins = {
    'tpu': tpu_platform_plugin,
    'cuda': cuda_platform_plugin,
    'rocm': rocm_platform_plugin,
    'xpu': xpu_platform_plugin,
    'cpu': cpu_platform_plugin,
}


def resolve_current_platform_cls_qualname() -> str:
    platform_plugins = load_plugins_by_group('vllm.platform_plugins')

    activated_plugins = []

    for name, func in chain(builtin_platform_plugins.items(),
                            platform_plugins.items()):
        try:
            assert callable(func)
            platform_cls_qualname = func()
            if platform_cls_qualname is not None:
                activated_plugins.append(name)
        except Exception:
193
            pass
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217

    activated_builtin_plugins = list(
        set(activated_plugins) & set(builtin_platform_plugins.keys()))
    activated_oot_plugins = list(
        set(activated_plugins) & set(platform_plugins.keys()))

    if len(activated_oot_plugins) >= 2:
        raise RuntimeError(
            "Only one platform plugin can be activated, but got: "
            f"{activated_oot_plugins}")
    elif len(activated_oot_plugins) == 1:
        platform_cls_qualname = platform_plugins[activated_oot_plugins[0]]()
        logger.info("Platform plugin %s is activated",
                    activated_oot_plugins[0])
    elif len(activated_builtin_plugins) >= 2:
        raise RuntimeError(
            "Only one platform plugin can be activated, but got: "
            f"{activated_builtin_plugins}")
    elif len(activated_builtin_plugins) == 1:
        platform_cls_qualname = builtin_platform_plugins[
            activated_builtin_plugins[0]]()
        logger.info("Automatically detected platform %s.",
                    activated_builtin_plugins[0])
    else:
218
        platform_cls_qualname = "vllm.platforms.interface.UnspecifiedPlatform"
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
        logger.info(
            "No platform detected, vLLM is running on UnspecifiedPlatform")
    return platform_cls_qualname


_current_platform = None
_init_trace: str = ''

if TYPE_CHECKING:
    current_platform: Platform


def __getattr__(name: str):
    if name == 'current_platform':
        # lazy init current_platform.
        # 1. out-of-tree platform plugins need `from vllm.platforms import
        #    Platform` so that they can inherit `Platform` class. Therefore,
        #    we cannot resolve `current_platform` during the import of
        #    `vllm.platforms`.
        # 2. when users use out-of-tree platform plugins, they might run
        #    `import vllm`, some vllm internal code might access
        #    `current_platform` during the import, and we need to make sure
        #    `current_platform` is only resolved after the plugins are loaded
        #    (we have tests for this, if any developer violate this, they will
        #    see the test failures).
        global _current_platform
        if _current_platform is None:
            platform_cls_qualname = resolve_current_platform_cls_qualname()
            _current_platform = resolve_obj_by_qualname(
                platform_cls_qualname)()
            global _init_trace
            _init_trace = "".join(traceback.format_stack())
        return _current_platform
252
    elif name in globals():
253
        return globals()[name]
254
255
256
    else:
        raise AttributeError(
            f"No attribute named '{name}' exists in {__name__}.")
257
258
259
260
261
262


__all__ = [
    'Platform', 'PlatformEnum', 'current_platform', 'CpuArchEnum',
    "_init_trace"
]