Implicitly add olive version to saved onnx model proto (microsoft#2183)

## Implicitly add olive version to saved onnx model proto

## 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/common/utils.py

@@ -497,6 +497,10 @@ def hardlink_copy_dir(src_dir, dst_dir, **kwargs):
     copy_dir(src_dir, dst_dir, copy_function=hardlink_copy_file, dirs_exist_ok=True, **kwargs)
 
 
+def is_hardlink(path: Union[str, Path]) -> bool:
+    return Path(path).stat().st_nlink > 1
+
+
 def set_tempdir(tempdir: str = None):
     """Set the root directory for tempfiles.
 

olive/engine/packaging/packaging_generator.py

@@ -18,6 +18,7 @@
 )
 from olive.hardware.accelerator import AcceleratorSpec
 from olive.model import ONNXModelHandler
+from olive.passes.onnx.common import add_version_metadata_to_model_proto
 from olive.resource_path import ResourceType, create_resource_path
 
 logger = logging.getLogger(__name__)
@@ -265,6 +266,8 @@ def _generate_onnx_mlflow_model(model_dir: Path, inference_config: dict):
     # MLFlow will save models with default config save_as_external_data=True
     # https://github.com/mlflow/mlflow/blob/1d6eaaa65dca18688d9d1efa3b8b96e25801b4e9/mlflow/onnx.py#L175
     # There will be an alphanumeric file generated in the same folder as the model file
+    # Add olive version to metadata
+    add_version_metadata_to_model_proto(model_proto)
     mlflow.onnx.save_model(
         model_proto,
         mlflow_model_path,

olive/passes/onnx/common.py

@@ -76,6 +76,27 @@ def get_external_data_config() -> dict[str, PassConfigParam]:
     }
 
 
+def add_version_metadata_to_model_proto(model: onnx.ModelProto) -> onnx.ModelProto:
+    olive_version = None
+    try:
+        import olive
+
+        olive_version = getattr(olive, "__version__", "unknown")
+    except Exception:
+        olive_version = "unknown"
+
+    for md in model.metadata_props:
+        if md.key == "olive_version":
+            md.value = olive_version
+            return model
+
+    md = model.metadata_props.add()
+    md.key = "olive_version"
+    md.value = olive_version
+
+    return model
+
+
 def model_proto_to_file(
     model: onnx.ModelProto,
     output_path: Union[str, Path],
@@ -119,6 +140,8 @@ def model_proto_to_file(
         )
 
     if not save_as_external_data:
+        # Add olive version to metadata
+        add_version_metadata_to_model_proto(model)
         # save model
         onnx.save_model(model, str(output_path))
         return False
@@ -136,6 +159,8 @@ def model_proto_to_file(
         if any(output_dir.iterdir()):
             raise RuntimeError(f"Output directory ({output_dir}) for external data is not empty.")
 
+    # Add olive version to metadata
+    add_version_metadata_to_model_proto(model)
     # save model
     onnx.save_model(
         model,
@@ -392,7 +417,7 @@ def resave_model(
     if not external_file_names:
         if force_external_data:
             # save the model with single external data file
-            model_proto_to_file(onnx.load(model_path), new_model_path, {"save_as_external_data": True})
+            model_proto_to_file(onnx.load(model_path), new_model_path, save_as_external_data=True)
             return True
 
         # no external data, so we can just copy the model
@@ -401,7 +426,7 @@ def resave_model(
 
     if len(external_file_names) > 1:
         # save the model with single external data file
-        model_proto_to_file(onnx.load(model_path), new_model_path, {"save_as_external_data": True})
+        model_proto_to_file(onnx.load(model_path), new_model_path, save_as_external_data=True)
         return True
 
     external_file_path = str(model_path.parent / external_file_names[0])

olive/passes/onnx/static_llm.py

@@ -10,7 +10,12 @@
 from olive.hardware.accelerator import AcceleratorSpec
 from olive.model import CompositeModelHandler, ONNXModelHandler
 from olive.passes import Pass
-from olive.passes.onnx.common import fix_dim_params, process_llm_pipeline, resave_model
+from olive.passes.onnx.common import (
+    add_version_metadata_to_model_proto,
+    fix_dim_params,
+    process_llm_pipeline,
+    resave_model,
+)
 from olive.passes.onnx.onnx_dag import OnnxDAG
 from olive.passes.pass_config import BasePassConfig, PassConfigParam
 
@@ -119,6 +124,8 @@ def process_context_iterator(component_models, llm_pipeline, output_dir):
 
                     # save the model with fixed shapes
                     component_model_path = output_dir / f"{new_component_name}.onnx"
+                    # Add olive version to metadata
+                    add_version_metadata_to_model_proto(component_proto)
                     onnx.save_model(component_proto, component_model_path)
                     new_groups[key][new_component_name] = ONNXModelHandler(
                         model_path=output_dir, onnx_file_name=component_model_path.name

test/passes/onnx/test_common.py

@@ -6,8 +6,13 @@
 import onnx
 import pytest
 
+from olive.common.utils import is_hardlink
 from olive.passes.olive_pass import create_pass_from_dict
-from olive.passes.onnx.common import model_proto_to_olive_model, resave_model
+from olive.passes.onnx.common import (
+    add_version_metadata_to_model_proto,
+    model_proto_to_olive_model,
+    resave_model,
+)
 from olive.passes.onnx.conversion import OnnxConversion
 from test.utils import ONNX_MODEL_PATH, get_hf_model
 
@@ -47,4 +52,11 @@ def test_resave_model(has_external_data, tmp_path):
     assert resave_path.exists()
     if has_external_data:
         assert (resave_path.parent / "resave.onnx.data").exists()
-    assert onnx.load(resave_path) == onnx.load(input_model.model_path)
+
+    input_model = onnx.load(input_model.model_path)
+    resaved_model = onnx.load(resave_path)
+
+    if not is_hardlink(resave_path):
+        input_model = add_version_metadata_to_model_proto(input_model)
+
+    assert resaved_model == input_model