Unverified Commit d67cc21b authored by Isotr0py's avatar Isotr0py Committed by GitHub
Browse files

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


Signed-off-by: default avatarIsotr0py <2037008807@qq.com>
parent e18227b0
...@@ -33,12 +33,19 @@ def cuda_platform_plugin() -> Optional[str]: ...@@ -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:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment