Add a stop method to EnergyTracker

huggingface · Sep 18, 2024 · 723d10d · 723d10d
1 parent c2f0afc
commit 723d10d
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 0 deletions.
diff --git a/optimum_benchmark/scenarios/inference/scenario.py b/optimum_benchmark/scenarios/inference/scenario.py
@@ -187,6 +187,7 @@ def run_model_loading_tracking(self, backend: Backend[BackendConfigT]):
             self.report.load.memory = memory_tracker.get_max_memory()
         if self.config.energy:
             self.report.load.energy = energy_tracker.get_energy()
+            energy_tracker.stop()
 
     ## Memory tracking
     def run_text_generation_memory_tracking(self, backend: Backend[BackendConfigT]):
@@ -367,6 +368,7 @@ def run_text_generation_energy_tracking(self, backend: Backend[BackendConfigT]):
         self.report.decode.efficiency = Efficiency.from_energy(
             decode_energy, decode_volume, unit=TEXT_GENERATION_EFFICIENCY_UNIT
         )
+        energy_tracker.stop()
 
     def run_image_diffusion_energy_tracking(self, backend: Backend[BackendConfigT]):
         self.logger.info("\t+ Running Image Diffusion energy tracking")
@@ -391,6 +393,7 @@ def run_image_diffusion_energy_tracking(self, backend: Backend[BackendConfigT]):
         self.report.call.efficiency = Efficiency.from_energy(
             call_energy, call_volume, unit=IMAGE_DIFFUSION_EFFICIENCY_UNIT
         )
+        energy_tracker.stop()
 
     def run_inference_energy_tracking(self, backend: Backend[BackendConfigT]):
         self.logger.info("\t+ Running energy tracking")
@@ -415,6 +418,7 @@ def run_inference_energy_tracking(self, backend: Backend[BackendConfigT]):
         self.report.forward.efficiency = Efficiency.from_energy(
             forward_energy, forward_volume, unit=INFERENCE_EFFICIENCY_UNIT
         )
+        energy_tracker.stop()
 
     @property
     def atomic_forward_volume(self) -> int:  # in samples

diff --git a/optimum_benchmark/scenarios/training/scenario.py b/optimum_benchmark/scenarios/training/scenario.py
@@ -97,6 +97,7 @@ def run(self, backend: Backend[BackendConfigT]) -> BenchmarkReport:
             self.report.overall.efficiency = Efficiency.from_energy(
                 self.report.overall.energy, volume=self.overall_volume, unit=TRAIN_EFFICIENCY_UNIT
             )
+            energy_tracker.stop()
 
         return self.report
 

diff --git a/optimum_benchmark/trackers/energy.py b/optimum_benchmark/trackers/energy.py
@@ -195,3 +195,6 @@ def get_energy(self) -> Energy:
         return Energy(
             unit=ENERGY_UNIT, cpu=self.cpu_energy, gpu=self.gpu_energy, ram=self.ram_energy, total=self.total_energy
         )
+
+    def stop(self):
+        self.emission_tracker.stop()