Repository: aimhubio/aim

Issue:
GPU utilization is not tracked if querying power usage fails
## 🐛 Bug

I am running experiments on a machine with a GPU, yet no GPU metrics are tracked. It seems like this code is responsible:

https://github.com/aimhubio/aim/blob/480e063cde063897283bcd8adb221e9baa861637/aim/ext/resource/stat.py#L152-L186

When any part of the GPU stats collection fails, we just give up entirely and store no information. In my case querying the power usage seems not supported by nvml, it raises `nvml.NVMLError_NotSupported`. Querying utilization and memory usage works just fine though and it would be nice if we could track those stats anyway.

### To reproduce

I'm not sure how to reproduce this, since it depends on a setup where `nvml` fails to determine the GPU power usage.

### Expected behavior

Aim tracks all the information that it can query without exceptions.

### Environment

- Aim Version (e.g., 3.15.1)
- Python version 3.10.9
- pip version 22.0.3
- OS (e.g., Linux) Linux
- Any other relevant information

### Additional context

--



Discussion:
@timokau great point! Thanks for pointing this 🙌 

Resolving patch:
diff --git a/aim/ext/resource/stat.py b/aim/ext/resource/stat.py
--- a/aim/ext/resource/stat.py
+++ b/aim/ext/resource/stat.py
@@ -4,13 +4,7 @@
 
 from aim.ext.resource.utils import round10e5
 
-try:
-    # Import python wrapper for the NVIDIA Management Library
-    # Initialize it or pass if NVIDIA ML is not initialized
-    from py3nvml import py3nvml as nvml
-    nvml.nvmlInit()
-except Exception:
-    pass
+from py3nvml import py3nvml as nvml
 
 
 class StatDict(object):
@@ -151,38 +145,49 @@ def get_stats(self):
         # Collect GPU statistics
         gpus = []
         try:
+            nvml.nvmlInit()
             gpu_device_count = nvml.nvmlDeviceGetCount()
             for i in range(gpu_device_count):
+                gpu_info = dict()
                 handle = nvml.nvmlDeviceGetHandleByIndex(i)
-                nvml_tmp = nvml.NVML_TEMPERATURE_GPU
-
-                # Get device memory and temperature
-                util = nvml.nvmlDeviceGetUtilizationRates(handle)
-                memory = nvml.nvmlDeviceGetMemoryInfo(handle)
-                temp = nvml.nvmlDeviceGetTemperature(handle, nvml_tmp)
-
-                # Compute power usage in watts and percent
-                power_watts = nvml.nvmlDeviceGetPowerUsage(handle) / 1000
-                power_cap = nvml.nvmlDeviceGetEnforcedPowerLimit(handle)
-                power_cap_watts = power_cap / 1000
-                power_watts / power_cap_watts * 100
-
-                gpus.append({
+                try:
+                    util = nvml.nvmlDeviceGetUtilizationRates(handle)
                     # GPU utilization percent
-                    'gpu': round10e5(util.gpu),
-
+                    gpu_info["gpu"] = round10e5(util.gpu)
+                except nvml.NVMLError_NotSupported:
+                    pass
+                try:
+                    # Get device memory
+                    memory = nvml.nvmlDeviceGetMemoryInfo(handle)
                     # Device memory usage
                     # 'memory_used': round10e5(memory.used / 1024 / 1024),
-                    'gpu_memory_percent': round10e5(memory.used * 100 / memory.total),
-
-                    # Power usage in watts and percent
-                    'gpu_power_watts': round10e5(power_watts),
-                    # 'power_percent': round10e5(power_usage),
-
+                    gpu_info["gpu_memory_percent"] = (
+                        round10e5(memory.used * 100 / memory.total),
+                    )
+                except nvml.NVMLError_NotSupported:
+                    pass
+                try:
+                    # Get device temperature
+                    nvml_tmp = nvml.NVML_TEMPERATURE_GPU
+                    temp = nvml.nvmlDeviceGetTemperature(handle, nvml_tmp)
                     # Device temperature
-                    'gpu_temp': round10e5(temp),
-                })
-        except Exception:
+                    gpu_info["gpu_temp"] = round10e5(temp)
+                except nvml.NVMLError_NotSupported:
+                    pass
+                try:
+                    # Compute power usage in watts and percent
+                    power_watts = nvml.nvmlDeviceGetPowerUsage(handle) / 1000
+                    power_cap = nvml.nvmlDeviceGetEnforcedPowerLimit(handle)
+                    power_cap_watts = power_cap / 1000
+                    power_watts / power_cap_watts * 100
+                    # Power usage in watts and percent
+                    gpu_info["gpu_power_watts"]: round10e5(power_watts)
+                    # gpu_info["power_percent"] = round10e5(power_usage)
+                except nvml.NVMLError_NotSupported:
+                    pass
+                gpus.append(gpu_info)
+            nvml.nvmlShutdown()
+        except (nvml.NVMLError_LibraryNotFound, nvml.NVMLError_NotSupported):
             pass
 
         return system, gpus


Test patch: