[misc][cuda] add warning for pynvml user (vllm-project#7675)

ShangmingCai · Aug 20, 2024 · ad28a74 · ad28a74
1 parent e6d811d
commit ad28a74
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py
@@ -18,6 +18,12 @@
 _P = ParamSpec("_P")
 _R = TypeVar("_R")
 
+if pynvml.__file__.endswith("__init__.py"):
+    logger.warning(
+        "You are using a deprecated `pynvml` package. Please install"
+        " `nvidia-ml-py` instead. See https://pypi.org/project/pynvml "
+        "for more information.")
+
 # NVML utils
 # Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
 # all the related functions work on real physical device ids.