Add size_on_disk API to OliveModelHandler and corresponding metric (microsoft#2262)

## Add size_on_disk API to OliveModelHandler and corresponding metric

## Checklist before requesting a review
- [ ] Add unit tests for this change.
- [x] Make sure all tests can pass.
- [ ] Update documents if necessary.
- [x] Lint and apply fixes to your code by running `lintrunner -a`
- [ ] Is this a user-facing change? If yes, give a description of this
change to be included in the release notes.

## (Optional) Issue link

Author: shaahji

olive/engine/engine.py

@@ -459,21 +459,29 @@ def _create_pareto_frontier_footprint(self, artifacts_dir: Path):
             )
 
     def _dump_run_history(self, run_history, output_path: Path):
-        if not run_history:
-            logger.info("No run history to dump!")
-            return
-        headers = run_history[0]._fields
-        try:
-            from tabulate import tabulate
+        from olive.logging import get_verbosity, set_verbosity
 
-            formatted_rls = tabulate([tuple(rh) for rh in run_history], headers=headers, tablefmt="grid")
-            logger.info("run history:\n%s", formatted_rls)
-        except ImportError:
-            logger.info("Please install tabulate for better run history output")
-            formatted_rls = run_history
-        if not self.skip_saving_artifacts:
-            with Path(output_path).open("w") as f:
-                f.write(f"{formatted_rls}")
+        def _dump_run_history_internal():
+            if not run_history:
+                logger.info("No run history to dump!")
+                return
+            headers = run_history[0]._fields
+            try:
+                from tabulate import tabulate
+
+                formatted_rls = tabulate([tuple(rh) for rh in run_history], headers=headers, tablefmt="grid")
+                logger.info("run history:\n%s", formatted_rls)
+            except ImportError:
+                logger.info("Please install tabulate for better run history output")
+                formatted_rls = run_history
+            if not self.skip_saving_artifacts:
+                with Path(output_path).open("w") as f:
+                    f.write(f"{formatted_rls}")
+
+        verbosity = get_verbosity()
+        set_verbosity(logging.INFO)
+        _dump_run_history_internal()
+        set_verbosity(verbosity)
 
     def resolve_objectives(
         self,

olive/evaluator/metric.py

@@ -10,7 +10,13 @@
 from olive.common.utils import StrEnumBase
 from olive.data.config import DataConfig
 from olive.evaluator.accuracy import AccuracyBase
-from olive.evaluator.metric_config import LatencyMetricConfig, MetricGoal, ThroughputMetricConfig, get_user_config_class
+from olive.evaluator.metric_config import (
+    LatencyMetricConfig,
+    MetricGoal,
+    SizeOnDiskMetricConfig,
+    ThroughputMetricConfig,
+    get_user_config_class,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -20,6 +26,7 @@ class MetricType(StrEnumBase):
     ACCURACY = "accuracy"
     LATENCY = "latency"
     THROUGHPUT = "throughput"
+    SIZE_ON_DISK = "size_on_disk"
     CUSTOM = "custom"
 
 
@@ -58,6 +65,10 @@ class ThroughputSubType(StrEnumBase):
     P999 = "p999"
 
 
+class SizeOnDiskSubType(StrEnumBase):
+    BYTES = "bytes"
+
+
 class SubMetric(ConfigBase):
     name: Union[AccuracySubType, LatencyMetricConfig, str]
     metric_config: ConfigBase = None
@@ -158,6 +169,8 @@ def validate_sub_types(cls, v, values):
                     sub_metric_type_cls = LatencySubType
                 elif values["type"] == MetricType.THROUGHPUT:
                     sub_metric_type_cls = ThroughputSubType
+                elif values["type"] == MetricType.SIZE_ON_DISK:
+                    sub_metric_type_cls = SizeOnDiskSubType
                 # if not exist, will raise ValueError
                 v["name"] = sub_metric_type_cls(v["name"])
             except ValueError:
@@ -182,6 +195,9 @@ def validate_sub_types(cls, v, values):
         elif values["type"] == MetricType.THROUGHPUT:
             v["higher_is_better"] = v.get("higher_is_better", True)
             metric_config_cls = ThroughputMetricConfig
+        elif values["type"] == MetricType.SIZE_ON_DISK:
+            v["higher_is_better"] = False
+            metric_config_cls = SizeOnDiskMetricConfig
         v["metric_config"] = validate_config(v.get("metric_config", {}), metric_config_cls)
 
         return v

olive/evaluator/metric_config.py

@@ -59,6 +59,10 @@ class ThroughputMetricConfig(ConfigBase):
     sleep_num: int = SLEEP_NUM
 
 
+class SizeOnDiskMetricConfig(ConfigBase):
+    pass
+
+
 class MetricGoal(ConfigBase):
     type: str  # threshold , deviation, percent-deviation
     value: float

olive/evaluator/olive_evaluator.py

@@ -25,7 +25,14 @@
 from olive.data.config import DataConfig
 from olive.data.container.dummy_data_container import TRANSFORMER_DUMMY_DATA_CONTAINER
 from olive.data.template import dummy_data_config_template
-from olive.evaluator.metric import LatencySubType, Metric, MetricType, ThroughputSubType, get_latency_config_from_metric
+from olive.evaluator.metric import (
+    LatencySubType,
+    Metric,
+    MetricType,
+    SizeOnDiskSubType,
+    ThroughputSubType,
+    get_latency_config_from_metric,
+)
 from olive.evaluator.metric_backend import MetricBackend
 from olive.evaluator.metric_result import MetricResult, SubMetricResult, flatten_metric_result, joint_metric_key
 from olive.evaluator.registry import Registry
@@ -276,6 +283,19 @@ def _evaluate_throughput(
         latencies = self._evaluate_raw_latency(model, metric, dataloader, post_func, device, execution_providers)
         return OliveEvaluator.compute_throughput(metric, latencies)
 
+    def _evaluate_size_on_disk(
+        self,
+        model: "OliveModelHandler",
+        metric: Metric,
+        dataloader: "DataLoader",
+        post_func=None,
+        device: Device = Device.CPU,
+        execution_providers: Union[str, list[str]] = None,
+    ) -> MetricResult:
+        return MetricResult.parse_obj(
+            {SizeOnDiskSubType.BYTES.value: {"value": model.size_on_disk, "priority": -1, "higher_is_better": False}}
+        )
+
     def _evaluate_custom(
         self,
         model: "OliveModelHandler",
@@ -335,6 +355,10 @@ def evaluate(
                 metrics_res[metric.name] = self._evaluate_throughput(
                     model, metric, dataloader, post_func, device, execution_providers
                 )
+            elif metric.type == MetricType.SIZE_ON_DISK:
+                metrics_res[metric.name] = self._evaluate_size_on_disk(
+                    model, metric, dataloader, post_func, device, execution_providers
+                )
             elif metric.type == MetricType.CUSTOM:
                 metrics_res[metric.name] = self._evaluate_custom(
                     model, metric, dataloader, eval_func, post_func, device, execution_providers
@@ -1056,30 +1080,44 @@ def evaluate(
             self.model_class,
             {k: v for k, v in init_args.items() if k in ["device", "ep", "ep_options"]},
         )
-        lmmodel = get_model(self.model_class)(**init_args, batch_size=self.batch_size, max_length=self.max_length)
-
-        results = simple_evaluate(
-            model=lmmodel,
-            tasks=self.tasks,
-            task_manager=TaskManager(),
-            log_samples=False,
-            batch_size=self.batch_size,
-            device=device,
-            limit=self.limit,
-        )
 
         metrics = {}
-        for task_name in sorted(results["results"].keys()):
-            metric_items = sorted(results["results"][task_name].items())
+        if MetricType.SIZE_ON_DISK.value in self.tasks:
+            self.tasks.remove(MetricType.SIZE_ON_DISK.value)
+            metrics[MetricType.SIZE_ON_DISK.value] = MetricResult.parse_obj(
+                {
+                    SizeOnDiskSubType.BYTES.value: {
+                        "value": model.size_on_disk,
+                        "priority": -1,
+                        "higher_is_better": False,
+                    }
+                }
+            )
+
+        if self.tasks:
+            lmmodel = get_model(self.model_class)(**init_args, batch_size=self.batch_size, max_length=self.max_length)
+
+            results = simple_evaluate(
+                model=lmmodel,
+                tasks=self.tasks,
+                task_manager=TaskManager(),
+                log_samples=False,
+                batch_size=self.batch_size,
+                device=device,
+                limit=self.limit,
+            )
+
+            for task_name in sorted(results["results"].keys()):
+                metric_items = sorted(results["results"][task_name].items())
 
-            task_metrics = {}
-            for mf, v in metric_items:
-                if mf != "alias":
-                    m, _ = mf.split(",", 1)
-                    if not m.endswith("_stderr"):
-                        task_metrics[m] = SubMetricResult(value=v, priority=-1, higher_is_better=True)
+                task_metrics = {}
+                for mf, v in metric_items:
+                    if mf != "alias":
+                        m, _ = mf.split(",", 1)
+                        if not m.endswith("_stderr"):
+                            task_metrics[m] = SubMetricResult(value=v, priority=-1, higher_is_better=True)
 
-            metrics[task_name] = MetricResult.parse_obj(task_metrics)
+                metrics[task_name] = MetricResult.parse_obj(task_metrics)
 
         return flatten_metric_result(metrics)
 

olive/model/handler/base.py

@@ -63,6 +63,12 @@ def model_path(self) -> str:
         """Return local model path."""
         return self.get_resource("model_path")
 
+    @property
+    @abstractmethod
+    def size_on_disk(self) -> int:
+        """Compute size of the model on disk."""
+        raise NotImplementedError
+
     @abstractmethod
     def load_model(self, rank: int = None, cache_model: bool = True) -> object:
         """Load model from disk, return in-memory model object.

olive/model/handler/composite.py

@@ -80,6 +80,11 @@ def get_model_components(self) -> list[tuple[str, OliveModelHandler]]:
     def load_model(self, rank: int = None, cache_model: bool = True):
         raise NotImplementedError
 
+    @property
+    def size_on_disk(self) -> int:
+        """Compute size of the model on disk."""
+        raise NotImplementedError
+
     def prepare_session(
         self,
         inference_settings: Optional[dict[str, Any]] = None,

olive/model/handler/hf.py

@@ -188,6 +188,15 @@ def load_model(self, rank: int = None, cache_model: bool = True) -> HfModelHandl
             model_attributes=self.model_attributes,
         )
 
+    @property
+    def size_on_disk(self) -> int:
+        """Compute size of the model on disk."""
+        nbytes = 0
+        for rank in range(self.num_ranks):
+            model = self.load_model(rank, cache_model=False)
+            nbytes += model.size_on_disk
+        return nbytes
+
     def prepare_session(
         self,
         inference_settings: Optional[dict[str, Any]] = None,

olive/model/handler/onnx.py

@@ -76,6 +76,12 @@ def model_path(self) -> str:
         model_path = super().model_path
         return get_onnx_file_path(model_path, self.onnx_file_name) if model_path else None
 
+    @property
+    def size_on_disk(self) -> int:
+        """Compute size of the model on disk."""
+        model = self.load_model()
+        return model.ByteSize()
+
     @property
     def external_initializers_path(self) -> Optional[str]:
         model_path = super().model_path
@@ -245,6 +251,15 @@ def load_model(self, rank: int = None, cache_model: bool = True) -> ONNXModelHan
             model_attributes=self.model_attributes,
         )
 
+    @property
+    def size_on_disk(self) -> int:
+        """Compute size of the model on disk."""
+        nbytes = 0
+        for rank in range(self.num_ranks):
+            model = self.load_model(rank, cache_model=False)
+            nbytes += model.ByteSize()
+        return nbytes
+
     def prepare_session(
         self,
         inference_settings: Optional[dict[str, Any]] = None,

olive/model/handler/openvino.py

@@ -61,6 +61,11 @@ def load_model(self, rank: int = None, cache_model: bool = True):
         core = ov.Core()
         return core.read_model(self.model_config["model"])
 
+    @property
+    def size_on_disk(self) -> int:
+        """Compute size of the model on disk."""
+        raise NotImplementedError
+
     def prepare_session(
         self,
         inference_settings: Optional[dict[str, Any]] = None,

olive/model/handler/pytorch.py

@@ -40,6 +40,26 @@ def __init__(
             io_config=io_config,
         )
 
+    @property
+    def size_on_disk(self) -> int:
+        """Compute size of the model on disk."""
+        import torch
+
+        class ByteCounter:
+            def __init__(self):
+                self.nbytes = 0
+
+            def write(self, data):
+                self.nbytes += len(data)
+
+            def flush(self):
+                pass
+
+        counter = ByteCounter()
+        model = self.load_model()
+        torch.save(model.state_dict(), counter)
+        return counter.nbytes
+
     def prepare_session(
         self,
         inference_settings: Optional[dict[str, Any]] = None,