[Misc] Improve error message for incorrect pynvml (vllm-project#12809)

Signed-off-by: youkaichao <youkaichao@gmail.com>
EmbeddedLLM · Feb 6, 2025 · 0408efc · 0408efc
1 parent 449d1bc
commit 0408efc
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/vllm/platforms/__init__.py b/vllm/platforms/__init__.py
@@ -41,7 +41,11 @@ def cuda_platform_plugin() -> Optional[str]:
                 is_cuda = True
         finally:
             pynvml.nvmlShutdown()
-    except Exception:
+    except Exception as e:
+        if "nvml" not in e.__class__.__name__.lower():
+            # If the error is not related to NVML, re-raise it.
+            raise e
+
         # CUDA is supported on Jetson, but NVML may not be.
         import os