Add pybindings for Model Packaging

cmake/onnxruntime_python.cmake

@@ -76,6 +76,10 @@ endif()
 
 onnxruntime_add_shared_library_module(onnxruntime_pybind11_state ${onnxruntime_pybind_srcs})
 
+if(NOT onnxruntime_MINIMAL_BUILD AND TARGET model_package)
+  target_link_libraries(onnxruntime_pybind11_state PRIVATE model_package)
+endif()
+
 message(STATUS "Python_EXECUTABLE: ${Python_EXECUTABLE}")
 
 # Query Py_GIL_DISABLED (PEP 703)

onnxruntime/__init__.py

@@ -80,6 +80,15 @@
 with contextlib.suppress(ImportError):
     from onnxruntime.capi._pybind_state import OrtEpContextData, OrtEpContextDataBuffer  # noqa: F401
 
+with contextlib.suppress(ImportError):
+    from onnxruntime.capi._pybind_state import (  # noqa: F401
+        ModelPackage,
+        ModelPackageComponentContext,
+        ModelPackageContext,
+        ModelPackageOptions,
+        ModelPackageWriteMode,
+    )
+
 from onnxruntime.capi.onnxruntime_inference_collection import (
     AdapterFormat,  # noqa: F401
     InferenceSession,  # noqa: F401

onnxruntime/python/onnxruntime_inference_collection.py

@@ -1045,6 +1045,18 @@ def clear_ep_context_data_write_func(self) -> None:
         """Clears a previously registered external EPContext data write callback."""
         self._model_compiler.clear_ep_context_data_write_func()
 
+    def set_ep_context_binary_information(self, output_directory: str | os.PathLike, model_name: str) -> None:
+        """
+        Sets the directory and model name that execution providers use to place their EPContext binary files.
+
+        Raises an 'InvalidArgument' exception if either argument is empty or the directory is not valid.
+
+        :param output_directory: Directory for the EP's binary files. A path without a trailing separator and
+            without an extension is rejected as not being a directory.
+        :param model_name: Model name used by the EP when naming its binary files.
+        """
+        self._model_compiler.set_ep_context_binary_information(os.fspath(output_directory), model_name)
+
     def compile_to_bytes(self) -> bytes:
         """
         Compiles to bytes representing the serialized compiled ONNX model.

onnxruntime/python/onnxruntime_pybind_model_compiler.cc

@@ -392,6 +392,14 @@ void PyModelCompiler::ClearEpContextDataWriteFunc() {
   }
 }
 
+onnxruntime::Status PyModelCompiler::SetEpContextBinaryInformation(const std::string& output_directory,
+                                                                   const std::string& model_name) {
+  std::lock_guard<std::mutex> lock{compilation_mutex_};
+  ORT_RETURN_IF(compilation_in_progress_,
+                "Cannot change the EPContext binary information while compilation is in progress.");
+  return model_compile_options_.SetEpContextBinaryInformation(output_directory, model_name);
+}
+
 PyModelCompiler::PyModelCompiler(onnxruntime::Environment& env, const OrtSessionOptions& sess_options,
                                  std::shared_ptr<PyEpSelectionRegistration> py_ep_selection_registration,
                                  std::shared_ptr<PyEpContextDataReadRegistration> py_ep_context_data_read_registration,

onnxruntime/python/onnxruntime_pybind_model_compiler.h

@@ -158,6 +158,15 @@ class PyModelCompiler {
   void SetEpContextDataWriteFunc(PyEpContextDataWriteFunc write_func);
   void ClearEpContextDataWriteFunc();
 
+  /// <summary>
+  /// Sets the directory and model name that EPs use to place their EPContext binary files.
+  /// </summary>
+  /// <param name="output_directory">Directory for EP binary files. Must not be empty.</param>
+  /// <param name="model_name">Model name used by the EP. Must not be empty.</param>
+  /// <returns>A Status indicating error or success.</returns>
+  onnxruntime::Status SetEpContextBinaryInformation(const std::string& output_directory,
+                                                    const std::string& model_name);
+
  private:
   onnxruntime::Status BeginCompilation();
   void EndCompilation() noexcept;

onnxruntime/python/onnxruntime_pybind_state.cc

@@ -3589,8 +3589,19 @@ including arg name, arg type (contains both type and shape).)pbdoc")
             model_compiler->ClearEpContextDataWriteFunc();
           },
           R"pbdoc(Clear the external EPContext data write callback.)pbdoc")
+      .def(
+          "set_ep_context_binary_information",
+          [](PyModelCompiler* model_compiler, const std::string& output_directory, const std::string& model_name) {
+            OrtPybindThrowIfError(model_compiler->SetEpContextBinaryInformation(output_directory, model_name));
+          },
+          py::arg("output_directory"), py::arg("model_name"),
+          R"pbdoc(Set the directory and model name EPs use to place their EPContext binary files.)pbdoc")
 #endif  // !defined(ORT_MINIMAL_BUILD)
       ;
+
+#if !defined(ORT_MINIMAL_BUILD)
+  addModelPackageMethods(m);
+#endif
 }
 
 bool InitArray() {

onnxruntime/python/onnxruntime_pybind_state_common.h

@@ -316,6 +316,14 @@ struct PyInferenceSession {
   }
 #endif
 
+  // Adopts an already initialized session returned by a stable C API.
+  PyInferenceSession(std::unique_ptr<InferenceSession> sess, const PySessionOptionsSnapshot& snapshot)
+      : py_ep_selection_registration_(snapshot.py_ep_selection_registration),
+        py_ep_context_data_read_registration_(snapshot.py_ep_context_data_read_registration),
+        session_options_(snapshot.options),
+        sess_(std::move(sess)) {
+  }
+
   // Returns true if the session options have provider information from either
   // setting explicit providers, setting a provider that supports a OrtEpDevice(s), or
   // setting a selection policy (e.g., prefer gpu).
@@ -518,6 +526,10 @@ void addSparseTensorMethods(pybind11::module& m);
 
 void addAdapterFormatMethods(pybind11::module& m);
 
+#if !defined(ORT_MINIMAL_BUILD)
+void addModelPackageMethods(pybind11::module& m);
+#endif
+
 void addGlobalSchemaFunctions(pybind11::module& m);
 
 void addOpKernelSubmodule(pybind11::module& m);

onnxruntime/test/python/onnxruntime_test_python_compile_api.py

@@ -102,6 +102,16 @@ def __del__(self):
         )
         self.assertEqual(result.returncode, 0, msg=f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}")
 
+    def test_set_ep_context_binary_information(self):
+        compiler = onnxrt.ModelCompiler(onnxrt.SessionOptions(), get_name("mul_1.onnx"))
+        output_dir = os.path.join(self._tmp_dir_path, "")
+        compiler.set_ep_context_binary_information(output_dir, "model")
+
+        with self.assertRaises(InvalidArgument):
+            compiler.set_ep_context_binary_information("", "model")
+        with self.assertRaises(InvalidArgument):
+            compiler.set_ep_context_binary_information(output_dir, "")
+
     def test_compile_with_files_prefer_npu_policy(self):
         """
         Tests compiling a model (to/from files) using an EP selection policy (PREFER_NPU).

onnxruntime/test/python/onnxruntime_test_python_model_package.py

@@ -0,0 +1,150 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import json
+import shutil
+import tempfile
+import unittest
+from pathlib import Path
+
+import numpy as np
+from helper import get_name
+
+import onnxruntime as ort
+from onnxruntime.capi import _pybind_state as C
+
+# mul_1.onnx computes Y = X * W with this constant initializer.
+MUL_1_WEIGHTS = np.array([[1, 2], [3, 4], [5, 6]], dtype=np.float32)
+
+
+@unittest.skipUnless(hasattr(C, "ModelPackage"), "Model package authoring requires a non-minimal build")
+class TestModelPackageAuthoring(unittest.TestCase):
+    def test_author_package_with_shared_asset_and_variant(self):
+        package_name = "my_model"
+        variant_name = f"{package_name}.cpu"
+        with tempfile.TemporaryDirectory() as temp:
+            output_dir = Path(temp) / "package"
+            og_staging = Path(temp) / "og_staging"
+            output_dir.mkdir()
+            og_staging.mkdir()
+            (output_dir / "manifest.json").write_text(
+                json.dumps({"schema_version": "1.0", "components": {}}), encoding="utf-8"
+            )
+            shutil.copyfile(get_name("mul_1.onnx"), og_staging / "mul_1.onnx")
+            variant_dir = output_dir / package_name / variant_name
+            variant_dir.mkdir(parents=True)
+            shutil.copyfile(get_name("mul_1.onnx"), variant_dir / "model.onnx")
+
+            package = ort.ModelPackage.open(
+                str(output_dir), allow_external_paths=False, follow_symlinks=True, strict_unknown_fields=False
+            )
+            package.set_metadata(package_name, "1.0.0")
+            shared_asset_uri = package.add_shared_asset(str(og_staging), copy_in=True)
+            package.set_component_inline(package_name, json.dumps({"variants": {}}))
+            package.set_variant(
+                package_name,
+                variant_name,
+                json.dumps(
+                    {
+                        "variant_directory": f"{package_name}/{variant_name}",
+                        "ep": "CPUExecutionProvider",
+                        "device": "cpu",
+                    }
+                ),
+            )
+            ort_info = {
+                "model_file": "model.onnx",
+                "session_options": {
+                    "session.model_external_initializers_file_folder_path": shared_asset_uri,
+                },
+            }
+            package.set_variant_executor_info_inline(package_name, variant_name, "ort", json.dumps(ort_info))
+            package.commit(mode=ort.ModelPackageWriteMode.PRESERVE)
+            package.close()
+            shutil.rmtree(og_staging)
+
+            with ort.ModelPackage.open(str(output_dir)) as package:
+                info = package.get_info()
+                self.assertEqual(info["package_name"], package_name)
+                self.assertEqual(info["package_version"], "1.0.0")
+                self.assertEqual([a["uri"] for a in info["shared_assets"]], [shared_asset_uri])
+                (variant,) = info["components"][0]["variants"]
+                self.assertEqual(variant["ep"], "CPUExecutionProvider")
+                self.assertEqual(variant["device"], "cpu")
+                body = json.loads(package.get_variant_json(package_name, variant_name))
+                self.assertEqual(body["executor_info"]["ort"], ort_info)
+                asset_path = Path(package.resolve_asset_uri(shared_asset_uri))
+                self.assertTrue((asset_path / "mul_1.onnx").is_file())
+
+            context = ort.ModelPackageContext(str(output_dir))
+            selected = context.select_component(package_name, ort.ModelPackageOptions(ort.SessionOptions()))
+            self.assertEqual(selected.get_selected_variant_name(), variant_name)
+            session = selected.create_session()
+            x = np.arange(6, dtype=np.float32).reshape(3, 2)
+            (output,) = session.run(["Y"], {"X": x}, None)
+            np.testing.assert_array_equal(output, x * MUL_1_WEIGHTS)
+
+    def test_invalid_inputs_report_native_errors(self):
+        with ort.ModelPackage() as package:
+            with self.assertRaises(RuntimeError):
+                package.commit()
+            with self.assertRaises(RuntimeError):
+                package.set_component_inline("model", "not json")
+            with self.assertRaises(RuntimeError):
+                package.set_layout("invalid")
+            with self.assertRaises(RuntimeError):
+                ort.ModelPackage.open("missing-model-package-directory")
+
+
+@unittest.skipUnless(hasattr(C, "ModelPackageContext"), "Model package API requires a non-minimal build")
+class TestStableModelPackageApi(unittest.TestCase):
+    def test_inspect_select_and_create_session(self):
+        with tempfile.TemporaryDirectory() as temp:
+            root = Path(temp) / "package"
+            variant = root / "model" / "cpu"
+            variant.mkdir(parents=True)
+            shutil.copyfile(get_name("mul_1.onnx"), variant / "model.onnx")
+            manifest = {
+                "schema_version": "1.0",
+                "components": {
+                    "model": {
+                        "variants": {
+                            "cpu": {
+                                "variant_directory": "model/cpu",
+                                "ep": "CPUExecutionProvider",
+                                "device": "cpu",
+                                "executor_info": {"ort": {"model_file": "model.onnx"}},
+                            }
+                        }
+                    }
+                },
+            }
+            (root / "manifest.json").write_text(json.dumps(manifest), encoding="utf-8")
+            package = ort.ModelPackageContext(str(root))
+            self.assertEqual(package.get_schema_version(), 1)
+            self.assertEqual(package.get_component_count(), 1)
+            self.assertEqual(package.get_component_names(), ["model"])
+            self.assertEqual(package.get_variant_count("model"), 1)
+            self.assertEqual(package.get_variant_names("model"), ["cpu"])
+            self.assertEqual(package.get_variant_ep_name("model", "cpu"), "CPUExecutionProvider")
+            resolved = package.resolve_string_ref("model/cpu/model.onnx", must_exist=True)
+            self.assertEqual(Path(resolved).resolve(), (variant / "model.onnx").resolve())
+            with self.assertRaises(RuntimeError):
+                package.get_variant_count("missing")
+
+            options = ort.ModelPackageOptions(ort.SessionOptions())
+            component = package.select_component("model", options)
+            del package, options
+            self.assertEqual(component.get_selected_variant_name(), "cpu")
+            self.assertEqual(Path(component.get_selected_variant_folder_path()).resolve(), variant.resolve())
+            sessions = [component.create_session(), component.create_session(ort.SessionOptions())]
+            del component
+            x = np.arange(6, dtype=np.float32).reshape(3, 2)
+            for session in sessions:
+                self.assertIsInstance(session, C.InferenceSession)
+                (output,) = session.run(["Y"], {"X": x}, None)
+                np.testing.assert_array_equal(output, x * MUL_1_WEIGHTS)
+
+
+if __name__ == "__main__":
+    unittest.main()