Fix issue with CodeCarbon lock (#265)

huggingface · Sep 20, 2024 · 1992de3 · 1992de3
1 parent 6151c5f
commit 1992de3
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 1 deletion.
diff --git a/optimum_benchmark/hub_utils.py b/optimum_benchmark/hub_utils.py
@@ -9,7 +9,7 @@
 import pandas as pd
 from flatten_dict import flatten, unflatten
 from huggingface_hub import create_repo, hf_hub_download, upload_file
-from huggingface_hub.utils._errors import HfHubHTTPError
+from huggingface_hub.utils import HfHubHTTPError
 from typing_extensions import Self
 
 LOGGER = getLogger("hub_utils")

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()