supress warning

Author: Teagan42

benches/monitor.py

@@ -14,16 +14,31 @@
 _NV_OK = False
 pynvml = None
 if not (os.environ.get('LLM_BENCH_SKIP_GPU') == '1'):
+    # prefer newer package name if available
     try:
-        import pynvml
+        import nvidia_ml_py as nvml_pkg  # type: ignore
         try:
-            pynvml.nvmlInit()
+            nvml_pkg.nvmlInit()
+            pynvml = nvml_pkg
             _NV_OK = True
         except Exception:
             _NV_OK = False
     except Exception:
-        pynvml = None  # type: ignore
-        _NV_OK = False
+        try:
+            # suppress deprecation FutureWarning from pynvml package if present
+            import warnings
+            with warnings.catch_warnings():
+                warnings.simplefilter('ignore', FutureWarning)
+                import pynvml as nvml_pkg  # type: ignore
+            try:
+                nvml_pkg.nvmlInit()
+                pynvml = nvml_pkg
+                _NV_OK = True
+            except Exception:
+                _NV_OK = False
+        except Exception:
+            pynvml = None  # type: ignore
+            _NV_OK = False
 
 
 @dataclass