failing test_power_limit_optimizer

ml-energy · Mar 19, 2024 · f357411 · f357411
1 parent d91a2c5
commit f357411
Show file tree

Hide file tree

Showing 4 changed files with 571 additions and 2 deletions.
diff --git a/tests/optimizer/test_power_limit_optimizer.py b/tests/optimizer/test_power_limit_optimizer.py
@@ -126,7 +126,8 @@ def test_power_limit_optimizer(
     tmp_path: Path,
 ):
     # Mock PyNVML.
-    pynvml_mock = mocker.patch("zeus.optimizer.power_limit.pynvml", autospec=True)
+    # pynvml_mock = mocker.patch("zeus.optimizer.power_limit.pynvml", autospec=True)
+    pynvml_mock = mocker.patch("zeus.device.gpu.pynvml", autospec=True)
     pynvml_mock.nvmlDeviceGetHandleByIndex.side_effect = lambda i: f"handle{i}"
     pynvml_mock.nvmlDeviceGetPowerManagementLimitConstraints.side_effect = \
         lambda _: (min(replay_log.power_limits) * 1000, max(replay_log.power_limits) * 1000)

diff --git a/zeus/device/gpu.py b/zeus/device/gpu.py
@@ -26,13 +26,15 @@
 """
 
 from __future__ import annotations
+from typing import TYPE_CHECKING
 
 import abc
 import pynvml
 import os
 
 # from exception import ZeusBaseGPUError
-import contextlib
+if TYPE_CHECKING:
+    import pynvml
 
 
 # import amdsmi