experiment mixin

optimum_benchmark/benchmarks/base.py

@@ -3,7 +3,7 @@
 from typing import ClassVar, Generic
 
 from ..backends.base import Backend
-from .report import BenchmarkReport
+from ..report import BenchmarkReport
 from .config import BenchmarkConfigT
 
 LOGGER = getLogger("benchmark")

optimum_benchmark/benchmarks/inference/benchmark.py

@@ -1,15 +1,16 @@
 from logging import getLogger
+from dataclasses import dataclass
 
 from ..base import Benchmark
 from .config import InferenceConfig
 from ...trackers.memory import MemoryTracker
+from ...report import BenchmarkReport, Measurements
 from ...backends.base import Backend, BackendConfigT
 from ...generators.input_generator import InputGenerator
 from ...trackers.energy import EnergyTracker, Efficiency
 from ...trackers.latency import LatencyTracker, Throughput
 from ...import_utils import is_torch_distributed_available
 from ...task_utils import TEXT_GENERATION_TASKS, IMAGE_DIFFUSION_TASKS
-from .report import InferenceReport, TextGenerationReport, ImageDiffusionReport
 
 if is_torch_distributed_available():
     import torch.distributed
@@ -44,6 +45,22 @@
 CALL_EFFICIENCY_UNIT = "images/kWh"
 
 
+@dataclass
+class InferenceReport(BenchmarkReport):
+    forward: Measurements = Measurements()
+
+
+@dataclass
+class ImageDiffusionReport(BenchmarkReport):
+    call: Measurements = Measurements()
+
+
+@dataclass
+class TextGenerationReport(BenchmarkReport):
+    prefill: Measurements = Measurements()
+    decode: Measurements = Measurements()
+
+
 class InferenceBenchmark(Benchmark[InferenceConfig]):
     NAME = "inference"
 
@@ -157,29 +174,29 @@ def run_text_generation_memory_tracking(self, backend: Backend):
         with self.memory_tracker.track():
             _ = backend.forward(self.forward_inputs, self.config.forward_kwargs)
 
-        self.report.prefill.max_memory = self.memory_tracker.get_max_memory()
+        self.report.prefill.memory = self.memory_tracker.get_max_memory()
 
         self.memory_tracker.reset()
         with self.memory_tracker.track():
             _ = backend.generate(self.generate_input, self.config.generate_kwargs)
 
-        self.report.decode.max_memory = self.memory_tracker.get_max_memory()
+        self.report.decode.memory = self.memory_tracker.get_max_memory()
 
     def run_image_diffusion_memory_tracking(self, backend: Backend):
         LOGGER.info("\t+ Running memory tracking")
         self.memory_tracker.reset()
         with self.memory_tracker.track():
             _ = backend.call(self.diffuse_input, self.config.forward_kwargs)
 
-        self.report.call.max_memory = self.memory_tracker.get_max_memory()
+        self.report.call.memory = self.memory_tracker.get_max_memory()
 
     def run_inference_memory_tracking(self, backend: Backend):
         LOGGER.info("\t+ Running memory tracking")
         self.memory_tracker.reset()
         with self.memory_tracker.track():
             _ = backend.forward(self.forward_inputs, self.config.forward_kwargs)
 
-        self.report.forward.max_memory = self.memory_tracker.get_max_memory()
+        self.report.forward.memory = self.memory_tracker.get_max_memory()
 
     ## Latency tracking
     def run_text_generation_latency_tracking(self, backend: Backend):

optimum_benchmark/benchmarks/training/benchmark.py

@@ -1,10 +1,11 @@
 from logging import getLogger
 from contextlib import ExitStack
+from dataclasses import dataclass
 
 from ..base import Benchmark
 from .config import TrainingConfig
-from .report import TrainingReport
 from ...trackers.memory import MemoryTracker
+from ...report import BenchmarkReport, Measurements
 from ...backends.base import Backend, BackendConfigT
 from ...trackers.energy import EnergyTracker, Efficiency
 from ...generators.dataset_generator import DatasetGenerator
@@ -18,6 +19,13 @@
 TRAIN_EFFICIENCY_UNIT = "samples/kWh"
 
 
+@dataclass
+class TrainingReport(BenchmarkReport):
+    overall: Measurements = Measurements()
+    warmup: Measurements = Measurements()
+    train: Measurements = Measurements()
+
+
 class TrainingBenchmark(Benchmark[TrainingConfig]):
     NAME = "training"
 
@@ -69,10 +77,9 @@ def run(self, backend: Backend[BackendConfigT]) -> None:
             )
 
         if self.config.memory:
-            # it's the same
-            self.report.overall.max_memory = memory_tracker.get_max_memory()
-            self.report.warmup.max_memory = memory_tracker.get_max_memory()
-            self.report.train.max_memory = memory_tracker.get_max_memory()
+            self.report.overall.memory = memory_tracker.get_max_memory()
+            self.report.warmup.memory = memory_tracker.get_max_memory()
+            self.report.train.memory = memory_tracker.get_max_memory()
 
             self.report.log_memory()
 

optimum_benchmark/experiment.py

@@ -1,15 +1,12 @@
 import os
 from logging import getLogger
 from tempfile import TemporaryDirectory
-from dataclasses import dataclass, field
-from typing import Any, Dict, Type, Optional, TYPE_CHECKING
-
-from hydra.utils import get_class
-from transformers.configuration_utils import PushToHubMixin
+from dataclasses import dataclass, field, asdict
+from typing import Any, Dict, Type, Optional, Union, TYPE_CHECKING
 
+from .report import BenchmarkReport
 from .env_utils import get_system_info
 from .import_utils import get_hf_libs_info
-from .benchmarks.report import BenchmarkReport
 from .benchmarks.config import BenchmarkConfig
 from .launchers.config import LauncherConfig
 from .backends.config import BackendConfig
@@ -22,9 +19,15 @@
     from .launchers.base import Launcher
     from .backends.base import Backend
 
+from json import dump
+from flatten_dict import flatten
+from hydra.utils import get_class
+from transformers.configuration_utils import PushToHubMixin
 
 LOGGER = getLogger("experiment")
 
+EXPERIMENT_FILE_NAME = "experiment_config.json"
+
 
 @dataclass
 class ExperimentConfig(PushToHubMixin):
@@ -46,6 +49,58 @@ class ExperimentConfig(PushToHubMixin):
     # ENVIRONMENT CONFIGURATION
     environment: Dict = field(default_factory=lambda: {**get_system_info(), **get_hf_libs_info()})
 
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)
+
+    def to_flat_dict(self) -> Dict[str, Any]:
+        report_dict = self.to_dict()
+        return flatten(report_dict, reducer="dot")
+
+    def to_json(self, path: str, flat: bool = False) -> None:
+        if flat:
+            with open(path, "w") as f:
+                dump(self.to_flat_dict(), f, indent=4)
+        else:
+            with open(path, "w") as f:
+                dump(self.to_dict(), f, indent=4)
+
+    def save_pretrained(
+        self,
+        save_directory: Union[str, os.PathLike],
+        config_file_name: Optional[Union[str, os.PathLike]] = None,
+        push_to_hub: bool = False,
+        **kwargs,
+    ):
+        use_auth_token = kwargs.pop("use_auth_token", None)
+
+        if use_auth_token is not None:
+            kwargs["token"] = use_auth_token
+
+        config_file_name = config_file_name if config_file_name is not None else EXPERIMENT_FILE_NAME
+
+        if os.path.isfile(save_directory):
+            raise AssertionError(f"Provided path ({save_directory}) should be a directory, not a file")
+
+        os.makedirs(save_directory, exist_ok=True)
+
+        if push_to_hub:
+            commit_message = kwargs.pop("commit_message", None)
+            repo_id = kwargs.pop("repo_id", save_directory.split(os.path.sep)[-1])
+            repo_id = self._create_repo(repo_id, **kwargs)
+            files_timestamps = self._get_files_timestamps(save_directory)
+
+        output_config_file = os.path.join(save_directory, config_file_name)
+        self.to_json(output_config_file, flat=False)
+
+        if push_to_hub:
+            self._upload_modified_files(
+                save_directory,
+                repo_id,
+                files_timestamps,
+                commit_message=commit_message,
+                token=kwargs.get("token"),
+            )
+
 
 def run(benchmark_config: BenchmarkConfig, backend_config: BackendConfig) -> BenchmarkReport:
     try:

optimum_benchmark/launchers/torchrun/launcher.py

@@ -6,9 +6,9 @@
 
 from ..base import Launcher
 from .config import TorchrunConfig
+from ...report import BenchmarkReport
 from ...logging_utils import setup_logging
 from ..isolation_utils import device_isolation
-from ...benchmarks.report import BenchmarkReport
 
 import torch.distributed
 from torch.distributed import FileStore
@@ -65,12 +65,15 @@ def launch(self, worker: Callable, *worker_args) -> Dict[str, Any]:
         while not queue.empty():
             outputs.append(queue.get())
 
-        if len(outputs) == 1:
-            report: BenchmarkReport = outputs[0]
-        else:
+        if len(outputs) > 1:
             LOGGER.info(f"\t+ Merging benchmark reports from {len(outputs)} workers")
-            report: BenchmarkReport = outputs[0].aggregate(outputs)
-            report.log_all()
+            report = outputs[0].aggregate(outputs)
+        elif len(outputs) == 1:
+            report = outputs[0]
+        else:
+            raise ValueError("No benchmark report was returned by the workers")
+
+        report.log_all()
 
         return report