| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import multiprocessing |
| 5 | from collections.abc import Sequence |
| 6 | from concurrent.futures.process import ProcessPoolExecutor |
| 7 | from functools import cache |
| 8 | from typing import Any |
| 9 | |
| 10 | import regex as re |
| 11 | import torch |
| 12 | |
| 13 | |
| 14 | def cuda_is_initialized() -> bool: |
| 15 | """Check if CUDA is initialized.""" |
| 16 | if not torch.cuda._is_compiled(): |
| 17 | return False |
| 18 | return torch.cuda.is_initialized() |
| 19 | |
| 20 | |
| 21 | def xpu_is_initialized() -> bool: |
| 22 | """Check if XPU is initialized.""" |
| 23 | if not torch.xpu._is_compiled(): |
| 24 | return False |
| 25 | return torch.xpu.is_initialized() |
| 26 | |
| 27 | |
| 28 | def cuda_get_device_properties( |
| 29 | device, names: Sequence[str], init_cuda=False |
| 30 | ) -> tuple[Any, ...]: |
| 31 | """Get specified CUDA device property values without initializing CUDA in |
| 32 | the current process.""" |
| 33 | if init_cuda or cuda_is_initialized(): |
| 34 | props = torch.cuda.get_device_properties(device) |
| 35 | return tuple(getattr(props, name) for name in names) |
| 36 | |
| 37 | # Run in subprocess to avoid initializing CUDA as a side effect. |
| 38 | mp_ctx = multiprocessing.get_context("fork") |
| 39 | with ProcessPoolExecutor(max_workers=1, mp_context=mp_ctx) as executor: |
| 40 | return executor.submit(cuda_get_device_properties, device, names, True).result() |
| 41 | |
| 42 | |
| 43 | @cache |
| 44 | def is_pin_memory_available() -> bool: |
| 45 | from vllm.platforms import current_platform |
| 46 | |
| 47 | return current_platform.is_pin_memory_available() |
| 48 | |
| 49 | |
| 50 | @cache |
| 51 | def is_uva_available() -> bool: |
| 52 | """Check if Unified Virtual Addressing (UVA) is available.""" |
| 53 | # UVA requires pinned memory. |
| 54 | from vllm.platforms import current_platform |
| 55 | |
| 56 | # TODO: Add more requirements for UVA if needed. |
| 57 | return is_pin_memory_available() or current_platform.is_cpu() |
| 58 | |
| 59 | |
| 60 | @cache |
| 61 | def num_compute_units(device_id: int = 0) -> int: |
| 62 | """Get the number of compute units of the current device.""" |
| 63 | from vllm.platforms import current_platform |
| 64 | |
| 65 | return current_platform.num_compute_units(device_id) |
| 66 | |
| 67 | |
| 68 | @cache |
| 69 | def get_device_name_as_file_name(device_id: int = 0) -> str: |
| 70 | from vllm.platforms import current_platform |
| 71 | |
| 72 | name = current_platform.get_device_name(device_id) |
| 73 | name = re.sub(r"[\s/]+", "_", name) |
| 74 | return name |