__init__.py 10.3 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 _Backend  # noqa: F401
13
from .interface import CpuArchEnum, Platform, PlatformEnum
14

15
logger = logging.getLogger(__name__)
16

17

18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
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


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

    # 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
42
43
44
    try:
        # While it's technically possible to install libtpu on a
        # non-TPU machine, this is a very uncommon scenario. Therefore,
45
        # we assume that libtpu is installed only if the machine
46
        # has TPUs.
47

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


56
57
def cuda_platform_plugin() -> Optional[str]:
    is_cuda = False
58
    logger.debug("Checking if CUDA platform is available.")
59
    try:
60
61
        from vllm.utils import import_pynvml
        pynvml = import_pynvml()
62
63
        pynvml.nvmlInit()
        try:
64
65
66
67
68
69
            # 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
70
                       and not vllm_version_matches_substr("cpu"))
71
72
73
74
75
76
77
78
            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.")
79
80
        finally:
            pynvml.nvmlShutdown()
81
    except Exception as e:
82
83
        logger.debug("Exception happens when checking CUDA platform: %s",
                     str(e))
84
85
86
87
        if "nvml" not in e.__class__.__name__.lower():
            # If the error is not related to NVML, re-raise it.
            raise e

88
89
90
91
92
93
94
95
        # 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():
96
            logger.debug("Confirmed CUDA platform is available on Jetson.")
97
            is_cuda = True
98
99
        else:
            logger.debug("CUDA platform is not available because: %s", str(e))
100

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


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

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


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

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

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


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

166
167
    except Exception as e:
        logger.debug("CPU platform is not available because: %s", str(e))
168
169
170
171
172

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


def neuron_platform_plugin() -> Optional[str]:
173
174
    tnx_installed = False
    nxd_installed = False
175
    logger.debug("Checking if Neuron platform is available.")
176
177
    try:
        import transformers_neuronx  # noqa: F401
178
        tnx_installed = True
179
180
        logger.debug("Confirmed Neuron platform is available because"
                     " transformers_neuronx is found.")
181
    except ImportError:
182
        pass
183

184
185
186
187
188
189
190
191
192
    try:
        import neuronx_distributed_inference  # noqa: F401
        nxd_installed = True
        logger.debug("Confirmed Neuron platform is available because"
                     " neuronx_distributed_inference is found.")
    except ImportError:
        pass

    is_neuron = tnx_installed or nxd_installed
193
    return "vllm.platforms.neuron.NeuronPlatform" if is_neuron else None
194
195


196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
builtin_platform_plugins = {
    'tpu': tpu_platform_plugin,
    'cuda': cuda_platform_plugin,
    'rocm': rocm_platform_plugin,
    'xpu': xpu_platform_plugin,
    'cpu': cpu_platform_plugin,
    'neuron': neuron_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:
219
            pass
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243

    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:
244
        platform_cls_qualname = "vllm.platforms.interface.UnspecifiedPlatform"
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
        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
278
    elif name in globals():
279
        return globals()[name]
280
281
282
    else:
        raise AttributeError(
            f"No attribute named '{name}' exists in {__name__}.")
283
284
285
286
287
288


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