Fix offline_tuning embed overwriting existing tuning results without --force

onnxruntime/python/tools/offline_tuning.py

@@ -34,7 +34,8 @@ def extract(model: onnx.ModelProto):
 
 def embed(model: onnx.ModelProto, tuning_results: list[TuningResults], overwrite=False):
     idx = _find_tuning_results_in_props(model.metadata_props)
-    assert overwrite or idx <= 0, "the supplied onnx file already have tuning results embedded!"
+    if not overwrite and idx >= 0:
+        raise ValueError("the supplied onnx file already have tuning results embedded!")
 
     if idx >= 0:
         model.metadata_props.pop(idx)

onnxruntime/test/python/onnxruntime_test_python_offline_tuning.py

@@ -0,0 +1,54 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import unittest
+
+import onnx
+from onnx import TensorProto, helper
+
+from onnxruntime.tools import offline_tuning
+
+
+def _make_model() -> onnx.ModelProto:
+    graph = helper.make_graph(
+        [helper.make_node("Identity", ["x"], ["y"])],
+        "offline_tuning_test",
+        [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1])],
+        [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1])],
+    )
+    return helper.make_model(graph)
+
+
+def _make_tuning_results(validator_value: str, kernel_id: int) -> list[offline_tuning.TuningResults]:
+    return [
+        {
+            "ep": "ROCMExecutionProvider",
+            "validators": {"version": validator_value},
+            "results": {"op_sig": {"params_sig": kernel_id}},
+        }
+    ]
+
+
+class TestOfflineTuningEmbed(unittest.TestCase):
+    def test_embed_requires_overwrite_when_results_exist(self):
+        # The tuning results are the first (idx 0) and only metadata entry here, which is the usual case.
+        model = _make_model()
+        original = _make_tuning_results("1", 1)
+        offline_tuning.embed(model, original)
+
+        with self.assertRaises(ValueError):
+            offline_tuning.embed(model, _make_tuning_results("2", 2))
+        self.assertEqual(offline_tuning.extract(model), original)
+
+    def test_embed_with_overwrite_replaces_results(self):
+        model = _make_model()
+        offline_tuning.embed(model, _make_tuning_results("1", 1))
+        replacement = _make_tuning_results("2", 2)
+        offline_tuning.embed(model, replacement, overwrite=True)
+
+        self.assertEqual(offline_tuning.extract(model), replacement)
+        self.assertEqual(sum(prop.key == "tuning_results" for prop in model.metadata_props), 1)
+
+
+if __name__ == "__main__":
+    unittest.main()

tools/ci_build/build.py

@@ -1984,6 +1984,8 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs):
             if not args.disable_contrib_ops:
                 run_subprocess([sys.executable, "onnxruntime_test_python_sparse_matmul.py"], cwd=cwd, dll_path=dll_path)
 
+            run_subprocess([sys.executable, "onnxruntime_test_python_offline_tuning.py"], cwd=cwd, dll_path=dll_path)
+
             if args.enable_symbolic_shape_infer_tests:
                 run_subprocess(
                     [sys.executable, "onnxruntime_test_python_symbolic_shape_infer.py"], cwd=cwd, dll_path=dll_path