[cuda] manually import the correct pynvml module (#12679)

fixes problems like https://github.com/vllm-project/vllm/pull/12635 and https://github.com/vllm-project/vllm/pull/12636 and https://github.com/vllm-project/vllm/pull/12565 --------- Signed-off-by: youkaichao <youkaichao@gmail.com>
2025-02-03 15:58:21 +08:00
parent b9986454fe
commit ad4a9dc817
3 changed files with 56 additions and 9 deletions
--- a/vllm/platforms/init.py
+++ b/vllm/platforms/init.py
@@ -33,7 +33,8 @@ def cuda_platform_plugin() -> Optional[str]:
    is_cuda = False

    try:
-        import pynvml
+        from vllm.utils import import_pynvml
+        pynvml = import_pynvml()
        pynvml.nvmlInit()
        try:
            if pynvml.nvmlDeviceGetCount() > 0: