From 860d50c6d0f05ed1941a24fd08b7f6ec5c33f525 Mon Sep 17 00:00:00 2001 From: Preetha Veeramalai Date: Tue, 27 Oct 2026 02:50:42 -0700 Subject: [PATCH] Add pybindings for Model Packaging --- cmake/onnxruntime_python.cmake | 4 + onnxruntime/__init__.py | 9 + .../onnxruntime_inference_collection.py | 12 + .../onnxruntime_pybind_model_compiler.cc | 8 + .../onnxruntime_pybind_model_compiler.h | 9 + .../onnxruntime_pybind_model_package.cc | 296 ++++++++++++++++++ .../python/onnxruntime_pybind_state.cc | 11 + .../python/onnxruntime_pybind_state_common.h | 12 + .../onnxruntime_test_python_compile_api.py | 10 + .../onnxruntime_test_python_model_package.py | 150 +++++++++ 10 files changed, 521 insertions(+) create mode 100644 onnxruntime/python/onnxruntime_pybind_model_package.cc create mode 100644 onnxruntime/test/python/onnxruntime_test_python_model_package.py diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index 603ad2564e7b0..cc52f050141bd 100644 --- a/cmake/onnxruntime_python.cmake +++ b/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) diff --git a/onnxruntime/__init__.py b/onnxruntime/__init__.py index 46594ea89dbcb..0128e0e868497 100644 --- a/onnxruntime/__init__.py +++ b/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 diff --git a/onnxruntime/python/onnxruntime_inference_collection.py b/onnxruntime/python/onnxruntime_inference_collection.py index e3597b6cce158..d974ee8125ac9 100644 --- a/onnxruntime/python/onnxruntime_inference_collection.py +++ b/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. diff --git a/onnxruntime/python/onnxruntime_pybind_model_compiler.cc b/onnxruntime/python/onnxruntime_pybind_model_compiler.cc index 4e238c4b3966c..8c7d747641ebf 100644 --- a/onnxruntime/python/onnxruntime_pybind_model_compiler.cc +++ b/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 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 py_ep_selection_registration, std::shared_ptr py_ep_context_data_read_registration, diff --git a/onnxruntime/python/onnxruntime_pybind_model_compiler.h b/onnxruntime/python/onnxruntime_pybind_model_compiler.h index 325af62ce8350..935843f4da105 100644 --- a/onnxruntime/python/onnxruntime_pybind_model_compiler.h +++ b/onnxruntime/python/onnxruntime_pybind_model_compiler.h @@ -158,6 +158,15 @@ class PyModelCompiler { void SetEpContextDataWriteFunc(PyEpContextDataWriteFunc write_func); void ClearEpContextDataWriteFunc(); + /// + /// Sets the directory and model name that EPs use to place their EPContext binary files. + /// + /// Directory for EP binary files. Must not be empty. + /// Model name used by the EP. Must not be empty. + /// A Status indicating error or success. + onnxruntime::Status SetEpContextBinaryInformation(const std::string& output_directory, + const std::string& model_name); + private: onnxruntime::Status BeginCompilation(); void EndCompilation() noexcept; diff --git a/onnxruntime/python/onnxruntime_pybind_model_package.cc b/onnxruntime/python/onnxruntime_pybind_model_package.cc new file mode 100644 index 0000000000000..f6084da69584d --- /dev/null +++ b/onnxruntime/python/onnxruntime_pybind_model_package.cc @@ -0,0 +1,296 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#if !defined(ORT_MINIMAL_BUILD) + +#include "python/onnxruntime_pybind_state_common.h" +#include "python/onnxruntime_pybind_exceptions.h" + +#include +#include +#include +#include +#include +#include + +#include "core/common/common.h" +#include "core/common/path_string.h" +#include "core/framework/error_code_helper.h" +#include "model_package.h" + +namespace onnxruntime::python { +namespace py = pybind11; + +namespace { + +void CheckPackageStatus(ModelPackageStatus* status) { + if (status != nullptr) { + std::unique_ptr owned_status{ + status, ModelPackageStatus_Release}; + throw std::runtime_error(ModelPackageStatus_Message(owned_status.get())); + } +} + +const char* OptionalString(const std::optional& value) { + return value ? value->c_str() : nullptr; +} + +py::object OptionalText(const char* value) { + return value ? py::object(py::str(value)) : py::none(); +} + +// Copies the C info tree, whose pointers are invalidated by the next package mutation. +py::dict PackageInfoToDict(const ModelPackageInfo& info) { + py::list components; + for (size_t c = 0; c < info.num_components; ++c) { + const ModelComponentInfo& component = info.components[c]; + py::list variants; + for (size_t v = 0; v < component.num_variants; ++v) { + const ModelVariantInfo& variant = component.variants[v]; + py::dict executor_info; + for (size_t e = 0; e < variant.num_executor_infos; ++e) { + executor_info[py::str(variant.executor_infos[e].namespace_key)] = py::str(variant.executor_infos[e].json); + } + py::dict variant_dict; + variant_dict["name"] = OptionalText(variant.name); + variant_dict["variant_directory"] = OptionalText(variant.variant_directory); + variant_dict["ep"] = OptionalText(variant.ep); + variant_dict["device"] = OptionalText(variant.device); + variant_dict["compatibility_string"] = OptionalText(variant.compatibility_string); + variant_dict["additional_metadata_json"] = OptionalText(variant.additional_metadata_json); + variant_dict["executor_info"] = executor_info; + variants.append(variant_dict); + } + py::dict component_dict; + component_dict["name"] = OptionalText(component.name); + component_dict["additional_metadata_json"] = OptionalText(component.additional_metadata_json); + component_dict["variants"] = variants; + components.append(component_dict); + } + + py::list shared_assets; + for (size_t a = 0; a < info.num_shared_assets; ++a) { + py::dict asset; + asset["uri"] = OptionalText(info.shared_assets[a].uri); + asset["resolved_path"] = OptionalText(info.shared_assets[a].resolved_path); + shared_assets.append(asset); + } + + py::dict result; + result["schema_version_major"] = info.schema_version_major; + result["schema_version_minor"] = info.schema_version_minor; + result["package_name"] = OptionalText(info.package_name); + result["package_version"] = OptionalText(info.package_version); + result["description"] = OptionalText(info.description); + result["layout"] = OptionalText(info.layout); + result["additional_metadata_json"] = OptionalText(info.additional_metadata_json); + result["components"] = components; + result["shared_assets"] = shared_assets; + return result; +} + +class PyModelPackage { + public: + using Handle = std::unique_ptr; + + explicit PyModelPackage(Handle package) : package_(std::move(package)) {} + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(PyModelPackage); + + ModelPackage* Get() const { + if (!package_) { + throw py::value_error("Model package is closed."); + } + return package_.get(); + } + + void Close() { package_.reset(); } + + private: + Handle package_; +}; + +} // namespace + +void addModelPackageMethods(py::module& m) { + py::enum_(m, "ModelPackageWriteMode") + .value("PRESERVE", MODEL_PACKAGE_WRITE_PRESERVE) + .value("DENSE", MODEL_PACKAGE_WRITE_DENSE); + + py::class_(m, "ModelPackage", + "Author model packages using the native model package library. JSON arguments are strings.") + .def(py::init([]() { + ModelPackage* raw = nullptr; + CheckPackageStatus(ModelPackage_New(&raw)); + return std::make_unique(PyModelPackage::Handle{raw, ModelPackage_Close}); + })) + .def_static("open", [](const std::string& root, bool allow_external_paths, + bool follow_symlinks, bool strict_unknown_fields) { + ModelPackageOpenOptions options{allow_external_paths, follow_symlinks, strict_unknown_fields}; + ModelPackage* raw = nullptr; + CheckPackageStatus(ModelPackage_Open(root.c_str(), &options, &raw)); + return std::make_unique(PyModelPackage::Handle{raw, ModelPackage_Close}); + }, py::arg("package_root"), py::arg("allow_external_paths") = false, + py::arg("follow_symlinks") = true, py::arg("strict_unknown_fields") = true) + .def("close", &PyModelPackage::Close) + .def("__enter__", [](PyModelPackage& self) -> PyModelPackage& { + self.Get(); + return self; + }, py::return_value_policy::reference_internal) + .def("__exit__", [](PyModelPackage& self, const py::object&, const py::object&, const py::object&) { + self.Close(); + }) + .def("set_metadata", [](PyModelPackage& self, const std::optional& name, + const std::optional& version, + const std::optional& description) { + CheckPackageStatus(ModelPackage_SetMetadata(self.Get(), OptionalString(name), OptionalString(version), + OptionalString(description))); + }, py::arg("name") = std::nullopt, py::arg("version") = std::nullopt, + py::arg("description") = std::nullopt) + .def("set_layout", [](PyModelPackage& self, const std::string& layout) { + CheckPackageStatus(ModelPackage_SetLayout(self.Get(), layout.c_str())); + }, py::arg("layout")) + .def("set_component_inline", [](PyModelPackage& self, const std::string& name, const std::string& json) { + CheckPackageStatus(ModelPackage_SetComponentInline(self.Get(), name.c_str(), json.c_str())); + }, py::arg("name"), py::arg("component_json")) + .def("set_component_external", [](PyModelPackage& self, const std::string& name, const std::string& path) { + CheckPackageStatus(ModelPackage_SetComponentExternal(self.Get(), name.c_str(), path.c_str())); + }, py::arg("name"), py::arg("path")) + .def("remove_component", [](PyModelPackage& self, const std::string& name) { + CheckPackageStatus(ModelPackage_RemoveComponent(self.Get(), name.c_str())); + }, py::arg("name")) + .def("set_variant", [](PyModelPackage& self, const std::string& component, const std::string& variant, + const std::string& json) { + CheckPackageStatus(ModelPackage_SetVariant(self.Get(), component.c_str(), variant.c_str(), json.c_str())); + }, py::arg("component_name"), py::arg("variant_name"), py::arg("variant_json")) + .def("remove_variant", [](PyModelPackage& self, const std::string& component, const std::string& variant) { + CheckPackageStatus(ModelPackage_RemoveVariant(self.Get(), component.c_str(), variant.c_str())); + }, py::arg("component_name"), py::arg("variant_name")) + .def("set_variant_executor_info_inline", [](PyModelPackage& self, const std::string& component, + const std::string& variant, const std::string& ns, + const std::string& json) { + CheckPackageStatus(ModelPackage_SetVariantExecutorInfoInline( + self.Get(), component.c_str(), variant.c_str(), ns.c_str(), json.c_str())); + }, py::arg("component_name"), py::arg("variant_name"), py::arg("namespace_key"), py::arg("info_json")) + .def("set_variant_executor_info_external", [](PyModelPackage& self, const std::string& component, + const std::string& variant, const std::string& ns, + const std::string& path) { + CheckPackageStatus(ModelPackage_SetVariantExecutorInfoExternal( + self.Get(), component.c_str(), variant.c_str(), ns.c_str(), path.c_str())); + }, py::arg("component_name"), py::arg("variant_name"), py::arg("namespace_key"), py::arg("path")) + .def("remove_variant_executor_info", [](PyModelPackage& self, const std::string& component, + const std::string& variant, const std::string& ns) { + CheckPackageStatus(ModelPackage_RemoveVariantExecutorInfo( + self.Get(), component.c_str(), variant.c_str(), ns.c_str())); + }, py::arg("component_name"), py::arg("variant_name"), py::arg("namespace_key")) + .def("set_additional_metadata_json", [](PyModelPackage& self, const std::string& scope, + const std::optional& component, + const std::optional& variant, + const std::optional& json) { + CheckPackageStatus(ModelPackage_SetAdditionalMetadataJson( + self.Get(), scope.c_str(), OptionalString(component), OptionalString(variant), OptionalString(json))); + }, py::arg("scope"), py::arg("component_name") = std::nullopt, + py::arg("variant_name") = std::nullopt, py::arg("json") = std::nullopt) + .def("add_shared_asset", [](PyModelPackage& self, const std::string& source, + const std::optional& expected_uri, bool copy_in) { + const char* uri = nullptr; + CheckPackageStatus(ModelPackage_AddSharedAsset( + self.Get(), source.c_str(), OptionalString(expected_uri), copy_in, &uri)); + return std::string{uri}; + }, py::arg("source_dir"), py::arg("expected_uri") = std::nullopt, py::arg("copy_in") = true) + .def("remove_shared_asset", [](PyModelPackage& self, const std::string& uri) { + CheckPackageStatus(ModelPackage_RemoveSharedAsset(self.Get(), uri.c_str())); + }, py::arg("uri")) + .def("get_info", [](const PyModelPackage& self) { + return PackageInfoToDict(*ModelPackage_Info(self.Get())); + }, "Return package metadata, components, variants and shared assets as a dict.") + .def("get_component_json", [](const PyModelPackage& self, const std::string& component) { + const char* json = nullptr; + CheckPackageStatus(ModelPackage_GetComponentJson(self.Get(), component.c_str(), &json)); + return std::string{json}; + }, py::arg("component_name")) + .def("get_variant_json", [](const PyModelPackage& self, const std::string& component, + const std::string& variant) { + const char* json = nullptr; + CheckPackageStatus(ModelPackage_GetVariantJson(self.Get(), component.c_str(), variant.c_str(), &json)); + return std::string{json}; + }, py::arg("component_name"), py::arg("variant_name")) + .def("resolve_asset_uri", [](const PyModelPackage& self, const std::string& uri) { + const char* path = nullptr; + CheckPackageStatus(ModelPackage_ResolveAssetUri(self.Get(), uri.c_str(), &path)); + return std::string{path}; + }, py::arg("uri")) + .def("resolve_string_ref", [](const PyModelPackage& self, const std::string& ref, + const std::optional& base, bool must_exist) { + const char* path = nullptr; + CheckPackageStatus(ModelPackage_ResolveStringRef( + self.Get(), OptionalString(base), ref.c_str(), must_exist, &path)); + return std::string{path}; + }, py::arg("string_ref"), py::arg("base_dir") = std::nullopt, py::arg("must_exist") = false) + .def_static("compute_directory_hash", [](const std::string& source) { + const char* uri = nullptr; + CheckPackageStatus(ModelPackage_ComputeDirectoryHash(source.c_str(), &uri)); + return std::string{uri}; + }, py::arg("source_dir")) + .def("commit", [](PyModelPackage& self, const std::optional& destination, + ModelPackageWriteMode mode) { + CheckPackageStatus(ModelPackage_Commit(self.Get(), OptionalString(destination), mode)); + }, py::arg("dest_root") = std::nullopt, py::arg("mode") = MODEL_PACKAGE_WRITE_PRESERVE, + "Commit in place, or to an empty destination. Shared asset sources must remain readable until commit.") + .def("prune", [](PyModelPackage& self) { + CheckPackageStatus(ModelPackage_Prune(self.Get())); + }); + + py::class_(m, "ModelPackageOptions", + "Execution provider configuration used for package variant selection.") + .def(py::init([](const PySessionOptions& options) { + auto snapshot = CreatePySessionOptionsSnapshot(options); + OrtModelPackageOptions* result = nullptr; + OrtPybindThrowIfError(ToStatusAndRelease( + Ort::GetModelPackageApi().CreateModelPackageOptionsFromSessionOptions( + GetOrtEnv(), &snapshot.options, &result))); + return Ort::ModelPackageOptions{result}; + }), py::arg("sess_options")); + + py::class_(m, "ModelPackageComponentContext", + "A selected component variant from an existing model package.") + .def("get_selected_variant_name", &Ort::ModelPackageComponentContext::GetSelectedVariantName) + .def("get_selected_variant_folder_path", [](const Ort::ModelPackageComponentContext& self) { + return PathToUTF8String(self.GetSelectedVariantFolderPath()); + }) + .def("create_session", [](Ort::ModelPackageComponentContext& self, const PySessionOptions* options) { + auto snapshot = options ? CreatePySessionOptionsSnapshot(*options) : PySessionOptionsSnapshot{}; + OrtSession* result = nullptr; + OrtPybindThrowIfError(ToStatusAndRelease(Ort::GetModelPackageApi().CreateSession( + GetOrtEnv(), self, options ? &snapshot.options : nullptr, &result))); + // The stable C API's OrtSession handle owns an internal InferenceSession. + std::unique_ptr session{reinterpret_cast(result)}; + return std::make_unique(std::move(session), snapshot); + }, py::arg("sess_options") = static_cast(nullptr), + "Create an initialized InferenceSession. None applies the package's session/provider options; " + "explicit options override them, while missing package path options are still added."); + + py::class_(m, "ModelPackageContext", "Open and inspect an existing model package.") + .def(py::init([](const std::string& root) { + return Ort::ModelPackageContext{ToPathString(root).c_str()}; + }), py::arg("package_root")) + .def("get_schema_version", &Ort::ModelPackageContext::GetSchemaVersion) + .def("get_component_count", &Ort::ModelPackageContext::GetComponentCount) + .def("get_component_names", &Ort::ModelPackageContext::GetComponentNames) + .def("get_variant_count", &Ort::ModelPackageContext::GetVariantCount, py::arg("component_name")) + .def("get_variant_names", &Ort::ModelPackageContext::GetVariantNames, py::arg("component_name")) + .def("get_variant_ep_name", [](const Ort::ModelPackageContext& self, const std::string& component, + const std::string& variant) -> std::optional { + const char* ep = self.GetVariantEpName(component.c_str(), variant.c_str()); + return ep ? std::optional{ep} : std::nullopt; + }, py::arg("component_name"), py::arg("variant_name")) + .def("resolve_string_ref", [](const Ort::ModelPackageContext& self, const std::string& ref, + const std::optional& base, bool must_exist) { + return self.ResolveStringRef(base ? base->c_str() : nullptr, ref.c_str(), must_exist); + }, py::arg("string_ref"), py::arg("base_dir") = std::nullopt, py::arg("must_exist") = false) + .def("select_component", &Ort::ModelPackageContext::SelectComponent, + py::arg("component_name"), py::arg("options")); +} + +} // namespace onnxruntime::python + +#endif // !defined(ORT_MINIMAL_BUILD) diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index 1211899ceb103..406dce4e35bf8 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/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() { diff --git a/onnxruntime/python/onnxruntime_pybind_state_common.h b/onnxruntime/python/onnxruntime_pybind_state_common.h index 71607fefbfdb3..8e1c79430a346 100644 --- a/onnxruntime/python/onnxruntime_pybind_state_common.h +++ b/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 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); diff --git a/onnxruntime/test/python/onnxruntime_test_python_compile_api.py b/onnxruntime/test/python/onnxruntime_test_python_compile_api.py index 0e383922bfb6c..32dc775d05c1e 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_compile_api.py +++ b/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). diff --git a/onnxruntime/test/python/onnxruntime_test_python_model_package.py b/onnxruntime/test/python/onnxruntime_test_python_model_package.py new file mode 100644 index 0000000000000..e54a421655b90 --- /dev/null +++ b/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()