[core] port pynvml into vllm codebase (#12963)

Signed-off-by: youkaichao <youkaichao@gmail.com>
2025-02-09 15:00:00 +08:00
parent 24700c346b
commit cf797aa856
6 changed files with 6169 additions and 35 deletions
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -46,8 +46,9 @@ if current_platform.is_rocm():
        finally:
            amdsmi_shut_down()
 elif current_platform.is_cuda():
-    from pynvml import (nvmlDeviceGetHandleByIndex, nvmlDeviceGetMemoryInfo,
-                        nvmlInit, nvmlShutdown)
+    from vllm.third_party.pynvml import (nvmlDeviceGetHandleByIndex,
+                                         nvmlDeviceGetMemoryInfo, nvmlInit,
+                                         nvmlShutdown)

    @contextmanager
    def _nvml():