Skip to content

Commit

Permalink
[doc] fix doc build error caused by msgspec (vllm-project#7659)
Browse files Browse the repository at this point in the history
  • Loading branch information
youkaichao authored and omrishiv committed Aug 26, 2024
1 parent 9689a5d commit 71b4af3
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 8 deletions.
1 change: 1 addition & 0 deletions docs/requirements-docs.txt
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ sphinx-book-theme==1.0.1
sphinx-copybutton==0.5.2
myst-parser==2.0.0
sphinx-argparse==0.4.0
msgspec

# packages to install to build the documentation
pydantic
Expand Down
47 changes: 39 additions & 8 deletions vllm/platforms/__init__.py
Original file line number Diff line number Diff line change
@@ -1,23 +1,54 @@
import torch

from .interface import Platform, PlatformEnum, UnspecifiedPlatform

current_platform: Platform

# NOTE: we don't use `torch.version.cuda` / `torch.version.hip` because
# they only indicate the build configuration, not the runtime environment.
# For example, people can install a cuda build of pytorch but run on tpu.

is_tpu = False
try:
import torch_xla.core.xla_model as xm
xm.xla_device(devkind="TPU")
is_tpu = True
except Exception:
pass

is_cuda = False

try:
import pynvml
pynvml.nvmlInit()
try:
if pynvml.nvmlDeviceGetCount() > 0:
is_cuda = True
finally:
pynvml.nvmlShutdown()
except Exception:
pass

is_rocm = False

try:
import libtpu
except ImportError:
libtpu = None
import amdsmi
amdsmi.amdsmi_init()
try:
if len(amdsmi.amdsmi_get_processor_handles()) > 0:
is_rocm = True
finally:
amdsmi.amdsmi_shut_down()
except Exception:
pass

if libtpu is not None:
if is_tpu:
# people might install pytorch built with cuda but run on tpu
# so we need to check tpu first
from .tpu import TpuPlatform
current_platform = TpuPlatform()
elif torch.version.cuda is not None:
elif is_cuda:
from .cuda import CudaPlatform
current_platform = CudaPlatform()
elif torch.version.hip is not None:
elif is_rocm:
from .rocm import RocmPlatform
current_platform = RocmPlatform()
else:
Expand Down

0 comments on commit 71b4af3

Please sign in to comment.