chore: Fixed tests

neptune-ai · Aug 19, 2024 · f62d559 · f62d559
1 parent 2c9cb92
commit f62d559
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/src/neptune/common/hardware/gpu/gpu_monitor.py b/src/neptune/common/hardware/gpu/gpu_monitor.py
@@ -46,7 +46,7 @@ def get_card_used_memory_in_bytes(self, card_index):
         return self.__nvml_get_or_else(lambda: nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(card_index)).used)
 
     def get_card_power_usage(self, card_index):
-        return self.__nvml_get_or_else(lambda: nvmlDeviceGetPowerUsage(nvmlDeviceGetHandleByIndex(card_index)))
+        return self.__nvml_get_or_else(lambda: nvmlDeviceGetPowerUsage(nvmlDeviceGetHandleByIndex(card_index)).used)
 
     def get_card_max_power_rating(self):
         def read_max_power_rating():

diff --git a/tests/unit/neptune/legacy/internal/hardware/gauges/test_gpu_gauges.py b/tests/unit/neptune/legacy/internal/hardware/gauges/test_gpu_gauges.py
@@ -20,7 +20,10 @@
     patch,
 )
 
-from neptune.common.hardware.constants import BYTES_IN_ONE_GB
+from neptune.common.hardware.constants import (
+    BYTES_IN_ONE_GB,
+    MILLIWATTS_IN_ONE_WATT,
+)
 from neptune.common.hardware.gauges.gpu import (
     GpuMemoryGauge,
     GpuPowerGauge,
@@ -79,7 +82,7 @@ def test_gpu_power_gauge(self, nvmlDeviceGetPowerUsage):
         gauge = GpuPowerGauge(card_index=self.card_index)
         # and
         gpu_power_info = MagicMock()
-        gpu_power_info.used = 15.0
+        gpu_power_info.used = 15.0 * MILLIWATTS_IN_ONE_WATT
         nvmlDeviceGetPowerUsage.side_effect = lambda handle: (
             gpu_power_info if handle == self.gpu_card_handle else None
         )