[Bugfix][Platform][CPU] Fix cuda platform detection on CPU backend edge case (#13358)

Signed-off-by: Isotr0py <2037008807@qq.com>
This commit is contained in:
Isotr0py 2025-02-17 02:55:27 +08:00 committed by GitHub
parent e18227b04a
commit d67cc21b78
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -33,12 +33,19 @@ def cuda_platform_plugin() -> Optional[str]:
is_cuda = False is_cuda = False
try: try:
from importlib.metadata import version
from vllm.utils import import_pynvml from vllm.utils import import_pynvml
pynvml = import_pynvml() pynvml = import_pynvml()
pynvml.nvmlInit() pynvml.nvmlInit()
try: try:
if pynvml.nvmlDeviceGetCount() > 0: # NOTE: Edge case: vllm cpu build on a GPU machine.
is_cuda = True # 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
and "cpu" not in version("vllm"))
finally: finally:
pynvml.nvmlShutdown() pynvml.nvmlShutdown()
except Exception as e: except Exception as e: