formatted with black

Author: benvanwerkhoven

kernel_tuner/observers/nvml.py

@@ -15,9 +15,7 @@
 class nvml:
     """Class that gathers the NVML functionality for one device."""
 
-    def __init__(
-        self, device_id=0, nvidia_smi_fallback="nvidia-smi", use_locked_clocks=False
-    ):
+    def __init__(self, device_id=0, nvidia_smi_fallback="nvidia-smi", use_locked_clocks=False):
         """Create object to control device using NVML."""
         pynvml.nvmlInit()
         self.dev = pynvml.nvmlDeviceGetHandleByIndex(device_id)
@@ -26,9 +24,7 @@ def __init__(
 
         try:
             self.pwr_limit_default = pynvml.nvmlDeviceGetPowerManagementLimit(self.dev)
-            self.pwr_constraints = pynvml.nvmlDeviceGetPowerManagementLimitConstraints(
-                self.dev
-            )
+            self.pwr_constraints = pynvml.nvmlDeviceGetPowerManagementLimitConstraints(self.dev)
         except pynvml.NVMLError_NotSupported:
             self.pwr_limit_default = None
             # inverted range to make all range checks fail
@@ -52,9 +48,7 @@ def __init__(
                 self.gr_clock_default = pynvml.nvmlDeviceGetDefaultApplicationsClock(
                     self.dev, pynvml.NVML_CLOCK_GRAPHICS
                 )
-                self.mem_clock_default = pynvml.nvmlDeviceGetDefaultApplicationsClock(
-                    self.dev, pynvml.NVML_CLOCK_MEM
-                )
+                self.mem_clock_default = pynvml.nvmlDeviceGetDefaultApplicationsClock(self.dev, pynvml.NVML_CLOCK_MEM)
         except pynvml.NVMLError_NotSupported:
             self.gr_clock_default = None
             self.sm_clock_default = None
@@ -67,9 +61,7 @@ def __init__(
         # gather the supported gr clocks for each supported mem clock into a dict
         self.supported_gr_clocks = {}
         for mem_clock in self.supported_mem_clocks:
-            supported_gr_clocks = pynvml.nvmlDeviceGetSupportedGraphicsClocks(
-                self.dev, mem_clock
-            )
+            supported_gr_clocks = pynvml.nvmlDeviceGetSupportedGraphicsClocks(self.dev, mem_clock)
             self.supported_gr_clocks[mem_clock] = supported_gr_clocks
 
         # test whether locked gr clocks and mem clocks are supported
@@ -132,9 +124,7 @@ def persistence_mode(self):
     @persistence_mode.setter
     def persistence_mode(self, new_mode):
         if new_mode not in [0, 1]:
-            raise ValueError(
-                "Illegal value for persistence mode, should be either 0 or 1"
-            )
+            raise ValueError("Illegal value for persistence mode, should be either 0 or 1")
         if self.persistence_mode == new_mode:
             return
         try:
@@ -155,7 +145,9 @@ def set_clocks(self, mem_clock, gr_clock):
         if mem_clock not in self.supported_mem_clocks:
             raise ValueError("Illegal value for memory clock")
         if gr_clock not in self.supported_gr_clocks[mem_clock]:
-            raise ValueError(f"Graphics clock incompatible with memory clock ({mem_clock}), compatible graphics clocks: {self.supported_gr_clocks[mem_clock]}")
+            raise ValueError(
+                f"Graphics clock incompatible with memory clock ({mem_clock}), compatible graphics clocks: {self.supported_gr_clocks[mem_clock]}"
+            )
 
         # Check whether persistence mode is set. Without persistence mode, setting the clocks is not meaningful
         # I deliberately removed the try..except clause here, if setting persistence mode fails, setting the clocks should fail
@@ -185,7 +177,6 @@ def set_clocks(self, mem_clock, gr_clock):
         # Store the fact that we have modified the clocks
         self.modified_clocks = True
 
-
     def reset_clocks(self):
         """Reset the clocks to the default clock if the device uses a non default clock."""
         if self.use_locked_clocks:
@@ -212,16 +203,9 @@ def reset_clocks(self):
                     subprocess.run(args, check=True)
 
         elif self.gr_clock_default is not None:
-            gr_app_clock = pynvml.nvmlDeviceGetApplicationsClock(
-                self.dev, pynvml.NVML_CLOCK_GRAPHICS
-            )
-            mem_app_clock = pynvml.nvmlDeviceGetApplicationsClock(
-                self.dev, pynvml.NVML_CLOCK_MEM
-            )
-            if (
-                gr_app_clock != self.gr_clock_default
-                or mem_app_clock != self.mem_clock_default
-            ):  
+            gr_app_clock = pynvml.nvmlDeviceGetApplicationsClock(self.dev, pynvml.NVML_CLOCK_GRAPHICS)
+            mem_app_clock = pynvml.nvmlDeviceGetApplicationsClock(self.dev, pynvml.NVML_CLOCK_MEM)
+            if gr_app_clock != self.gr_clock_default or mem_app_clock != self.mem_clock_default:
                 self.set_clocks(self.mem_clock_default, self.gr_clock_default)
 
     @property
@@ -246,9 +230,7 @@ def mem_clock(self):
             mem_clock = pynvml.nvmlDeviceGetClockInfo(self.dev, pynvml.NVML_CLOCK_MEM)
             return min(self.supported_mem_clocks, key=lambda x: abs(x - mem_clock))
         else:
-            return pynvml.nvmlDeviceGetApplicationsClock(
-                self.dev, pynvml.NVML_CLOCK_MEM
-            )
+            return pynvml.nvmlDeviceGetApplicationsClock(self.dev, pynvml.NVML_CLOCK_MEM)
 
     @mem_clock.setter
     def mem_clock(self, new_clock):
@@ -269,9 +251,7 @@ def auto_boost(self):
     def auto_boost(self, setting):
         # might need to use pynvml.NVML_FEATURE_DISABLED or pynvml.NVML_FEATURE_ENABLED instead of 0 or 1
         if setting not in [0, 1]:
-            raise ValueError(
-                "Illegal value for auto boost enabled, should be either 0 or 1"
-            )
+            raise ValueError("Illegal value for auto boost enabled, should be either 0 or 1")
         pynvml.nvmlDeviceSetAutoBoostedClocksEnabled(self.dev, setting)
         self._auto_boost = pynvml.nvmlDeviceGetAutoBoostedClocksEnabled(self.dev)[0]
 
@@ -363,9 +343,7 @@ def __init__(
         if any([obs in self.needs_power for obs in observables]):
             self.measure_power = True
             power_observables = [obs for obs in observables if obs in self.needs_power]
-            self.continuous_observer = NVMLPowerObserver(
-                power_observables, self, self.nvml, continous_duration
-            )
+            self.continuous_observer = NVMLPowerObserver(power_observables, self, self.nvml, continous_duration)
 
         # remove power observables
         self.observables = [obs for obs in observables if obs not in self.needs_power]
@@ -380,11 +358,7 @@ def __init__(
         for obs in self.observables:
             self.results[obs + "s"] = []
 
-        self.during_obs = [
-            obs
-            for obs in observables
-            if obs in ["core_freq", "mem_freq", "temperature"]
-        ]
+        self.during_obs = [obs for obs in observables if obs in ["core_freq", "mem_freq", "temperature"]]
         self.iteration = {obs: [] for obs in self.during_obs}
 
     def before_start(self):
@@ -406,15 +380,11 @@ def during(self):
         if "mem_freq" in self.observables:
             self.iteration["mem_freq"].append(self.nvml.mem_clock)
         if self.record_gr_voltage:
-            self.gr_voltage_readings.append(
-                [time.perf_counter() - self.t0, self.nvml.gr_voltage()]
-            )
+            self.gr_voltage_readings.append([time.perf_counter() - self.t0, self.nvml.gr_voltage()])
 
     def after_finish(self):
         if "temperature" in self.observables:
-            self.results["temperatures"].append(
-                np.average(self.iteration["temperature"])
-            )
+            self.results["temperatures"].append(np.average(self.iteration["temperature"]))
         if "core_freq" in self.observables:
             self.results["core_freqs"].append(np.average(self.iteration["core_freq"]))
         if "mem_freq" in self.observables:
@@ -423,12 +393,8 @@ def after_finish(self):
         if "gr_voltage" in self.observables:
             execution_time = time.time() - self.t0
             gr_voltage_readings = self.gr_voltage_readings
-            gr_voltage_readings = [
-                [0.0, gr_voltage_readings[0][1]]
-            ] + gr_voltage_readings
-            gr_voltage_readings = gr_voltage_readings + [
-                [execution_time, gr_voltage_readings[-1][1]]
-            ]
+            gr_voltage_readings = [[0.0, gr_voltage_readings[0][1]]] + gr_voltage_readings
+            gr_voltage_readings = gr_voltage_readings + [[execution_time, gr_voltage_readings[-1][1]]]
             # time in s, graphics voltage in millivolts
             self.results["gr_voltages"].append(np.average(gr_voltage_readings[:][:][1]))
 
@@ -490,8 +456,7 @@ def during(self):
         timestamp = time.perf_counter() - self.t0
         # only store the result if we get a new measurement from NVML
         if len(self.power_readings) == 0 or (
-            self.power_readings[-1][1] != power_usage
-            or timestamp - self.power_readings[-1][0] > 0.01
+            self.power_readings[-1][1] != power_usage or timestamp - self.power_readings[-1][0] > 0.01
         ):
             self.power_readings.append([timestamp, power_usage])
 
@@ -538,9 +503,7 @@ def get_nvml_pwr_limits(device, n=None, quiet=False):
         n = int((power_limit_max - power_limit_min) / power_limit_round) + 1
 
     # Rounded power limit values
-    power_limits = power_limit_round * np.round(
-        (np.linspace(power_limit_min, power_limit_max, n) / power_limit_round)
-    )
+    power_limits = power_limit_round * np.round((np.linspace(power_limit_min, power_limit_max, n) / power_limit_round))
     power_limits = sorted(list(set([int(power_limit) for power_limit in power_limits])))
     tune_params["nvml_pwr_limit"] = power_limits