diff --git a/.github/workflows/android.yml b/.github/workflows/android.yml index 252ea7281d981..954c3313faf25 100644 --- a/.github/workflows/android.yml +++ b/.github/workflows/android.yml @@ -78,8 +78,8 @@ jobs: run: | set -e -x BINARY_SIZE_THRESHOLD_ARGS="" - echo "Binary size threshold in bytes: 1436672" - BINARY_SIZE_THRESHOLD_ARGS="--threshold_size_in_bytes 1436672" + echo "Binary size threshold in bytes: 1438720" + BINARY_SIZE_THRESHOLD_ARGS="--threshold_size_in_bytes 1438720" # Ensure ANDROID_NDK_HOME is available and get its real path if [ -z "$ANDROID_NDK_HOME" ]; then diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index c32aa7f4ae75a..bb171b056b400 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -603,6 +603,7 @@ set (onnxruntime_shared_lib_test_SRC ${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/custom_op_utils.cc ${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/test_allocator.cc ${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/test_data_copy.cc + ${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/test_ep_context_data_api.cc ${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/test_experimental_api.cc ${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/test_fixture.h ${ONNXRUNTIME_SHARED_LIB_TEST_SRC_DIR}/test_model_loading.cc @@ -2174,6 +2175,7 @@ if (onnxruntime_BUILD_SHARED_LIB AND # file(GLOB onnxruntime_autoep_test_library_src "${TEST_SRC_DIR}/autoep/library/example_plugin_ep/*.h" "${TEST_SRC_DIR}/autoep/library/example_plugin_ep/*.cc" + "${TEST_SRC_DIR}/autoep/library/ep_context_data_utils.h" "${TEST_SRC_DIR}/autoep/library/plugin_ep_utils.h") onnxruntime_add_shared_library_module(example_plugin_ep ${onnxruntime_autoep_test_library_src}) target_include_directories(example_plugin_ep PRIVATE ${REPO_ROOT}/include/onnxruntime/core/session) diff --git a/include/onnxruntime/core/session/onnxruntime_experimental_c_api.h b/include/onnxruntime/core/session/onnxruntime_experimental_c_api.h index e5b9bd4713a1c..ce76bd385cd85 100644 --- a/include/onnxruntime/core/session/onnxruntime_experimental_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_experimental_c_api.h @@ -46,6 +46,68 @@ ORT_RUNTIME_CLASS(ModelPackageOptions); ORT_RUNTIME_CLASS(ModelPackageContext); ORT_RUNTIME_CLASS(ModelPackageComponentContext); +// Opaque handle holding the EPContext callbacks and opaque state extracted from an OrtSessionOptions instance. Used by +// the experimental OrtEpApi_* EPContext data functions. Create via OrtEpApi_SessionOptions_GetEpContextConfig and +// release with OrtEpApi_ReleaseEpContextConfig. +ORT_RUNTIME_CLASS(EpContextConfig); + +/** \brief Function called to write named binary data. + * + * This callback is currently used for EPContext binary data, but its contract is intentionally generic so future APIs + * can reuse it for other named data payloads. The callback is called synchronously by the component that receives it. + * ORT does not own or retain buffer after the callback returns. ORT does not serialize invocations made by different + * EP instances or worker threads. + * + * Each callback invocation represents one complete write operation for name. The callback signature does not + * provide an offset, sequence number, or final-chunk marker, so the component invoking the callback must define any + * chunked ordering and completion contract with the application. Current EPContext use should prefer a single callback + * invocation per EPContext binary unless chunking semantics are documented by the EP. + * + * The application's implementation can process the data in any way (e.g., encrypt and store, upload to cloud storage, + * or compress) before persisting it. + * + * \param[in] state Opaque pointer holding the user's state. ORT does not own or manage this pointer. The application + * must keep it valid for the duration required by the API that accepted the callback and must provide + * any synchronization required if it can be used concurrently. + * \param[in] name The file name or logical data identifier as a null-terminated UTF-8 string. + * \param[in] buffer The buffer containing data to write. + * \param[in] buffer_num_bytes The size of the buffer in bytes. + * + * \return OrtStatus* Write status. Return nullptr on success. + * On failure, use CreateStatus to provide error info with an appropriate OrtErrorCode + * (e.g., ORT_FAIL); ORT propagates the returned code. ORT will release the OrtStatus* if not null. + */ +typedef OrtStatus*(ORT_API_CALL* OrtWriteNamedBufferFunc)(_In_ void* state, + _In_ const char* name, + _In_ const void* buffer, + _In_ size_t buffer_num_bytes); + +/** \brief Function called to read named binary data. + * + * This callback is currently used for EPContext binary data, but its contract is intentionally generic so future APIs + * can reuse it for other named data payloads. The application reads, processes (e.g., decrypts, decompresses, + * downloads), and returns the requested data. ORT provides an allocator so the application can allocate the output + * buffer directly. The callback is called synchronously by the component that receives it. ORT does not serialize + * invocations made by different EP instances or worker threads. + * + * \param[in] state Opaque pointer holding the user's state. ORT does not own or manage this pointer. The application + * must keep it valid for the duration required by the API that accepted the callback and must provide + * any synchronization required if it can be used concurrently. + * \param[in] name The file name or logical data identifier to read as a null-terminated UTF-8 string. + * \param[in] allocator ORT-provided allocator. The application must use this to allocate the output buffer. + * \param[out] buffer Set by the implementation to the allocated buffer containing the output data. + * \param[out] data_size Set by the implementation to the size of the output data in bytes. + * + * \return OrtStatus* Read status. Return nullptr on success. + * On failure, use CreateStatus to provide error info with an appropriate OrtErrorCode + * (e.g., ORT_FAIL); ORT propagates the returned code. ORT will release the OrtStatus* if not null. + */ +typedef OrtStatus*(ORT_API_CALL* OrtReadNamedBufferFunc)(_In_ void* state, + _In_ const char* name, + _In_ OrtAllocator* allocator, + _Outptr_ void** buffer, + _Out_ size_t* data_size); + // // C function pointer typedefs and name constants // diff --git a/include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc b/include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc index 57a4e472b6f6d..b1140485f7ff1 100644 --- a/include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc +++ b/include/onnxruntime/core/session/onnxruntime_experimental_c_api.inc @@ -282,3 +282,117 @@ ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtModelPackageApi_CreateSession, _In_ OrtModelPackageComponentContext* context, _In_opt_ const OrtSessionOptions* session_options, _Outptr_ OrtSession** session) + +/** \brief Registers a callback to provide EPContext binary data during session load. + * + * When loading a compiled model with external (non-embedded) EPContext binary data, an execution provider can + * retrieve this callback from OrtEpContextConfig and call it instead of reading the binary data from disk. + * + * Reading happens at session load, so this callback is configured on OrtSessionOptions. The corresponding write + * callback runs only at compile time and is configured on OrtModelCompilationOptions via + * OrtCompileApi_ModelCompilationOptions_SetEpContextDataWriteFunc. + * + * The state pointer is stored as-is and is not owned by ORT. It must remain valid while any session or EP created + * from these options may call the callback. If the same state may be used by multiple EPs or threads, the application + * is responsible for synchronization. + * + * \param[in] options The OrtSessionOptions instance. + * \param[in] read_func The OrtReadNamedBufferFunc callback. Pass NULL to clear a previously set callback (any + * previously set state is cleared as well). + * \param[in] state Opaque state passed to read_func. Can be NULL. Ignored when read_func is NULL. + * + * \snippet{doc} snippets.dox OrtStatus Return Value + */ +ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtApi_SessionOptions_SetEpContextDataReadFunc, + _Inout_ OrtSessionOptions* options, _In_opt_ OrtReadNamedBufferFunc read_func, _In_opt_ void* state) + +/** \brief Sets a callback for writing EPContext binary data during compilation. + * + * When EPContext embed mode is disabled, execution providers can retrieve this callback from OrtEpContextConfig and + * call it instead of writing EPContext binary data directly to disk. + * + * This callback may be used together with OrtCompileApi::ModelCompilationOptions_SetEpContextBinaryInformation. The + * binary information still describes the compiled model/output location that EPs may use to generate stable logical + * EPContext data names or as a file-fallback location. If this callback is configured, EPs should call it for + * EPContext binary data instead of writing that data to the fallback file path. + * + * Writing happens only at compile time, so this callback is configured on OrtModelCompilationOptions. The + * corresponding read callback runs at session load and is configured on OrtSessionOptions via + * OrtApi_SessionOptions_SetEpContextDataReadFunc. + * + * The state pointer is stored as-is and is not owned by ORT. It must remain valid for the duration of the compile + * operation that may call the callback. If the same state may be used by multiple EPs or threads, the application is + * responsible for synchronization. + * + * Like OrtApi_SessionOptions_SetEpContextDataReadFunc, passing a NULL write_func clears any previously set callback + * (any previously set state is cleared as well). Calling this multiple times overwrites the previously configured + * callback. + * + * \param[in] model_compile_options The OrtModelCompilationOptions instance. + * \param[in] write_func The OrtWriteNamedBufferFunc callback used to write EPContext bytes. Pass NULL to clear a + * previously set callback (any previously set state is cleared as well). + * \param[in] state Opaque state passed to write_func. Can be NULL. Ignored when write_func is NULL. + * + * \snippet{doc} snippets.dox OrtStatus Return Value + */ +ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtCompileApi_ModelCompilationOptions_SetEpContextDataWriteFunc, + _In_ OrtModelCompilationOptions* model_compile_options, + _In_opt_ OrtWriteNamedBufferFunc write_func, _In_opt_ void* state) + +/** \brief Extracts the EPContext configuration (callbacks and state) from an OrtSessionOptions instance. + * + * The EP should call this during CreateEp() while session_options is still valid, and store the returned handle for + * use during Compile(). On success, `*config` is set to a non-NULL handle that must be released with + * OrtEpApi_ReleaseEpContextConfig. On failure, an error status is returned and `*config` is not modified. + * + * The returned handle owns only ORT's copy of callback function pointers and opaque state pointer values. It does not + * own the application-provided state. The application is responsible for keeping callback state valid and + * synchronized while an EP may call callbacks retrieved from this config. + * + * \param[in] session_options The OrtSessionOptions instance. + * \param[out] config The extracted OrtEpContextConfig. + * + * \snippet{doc} snippets.dox OrtStatus Return Value + */ +ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtEpApi_SessionOptions_GetEpContextConfig, + _In_ const OrtSessionOptions* session_options, _Outptr_ OrtEpContextConfig** config) + +/** \brief Release an OrtEpContextConfig instance. + * + * \param[in] config The OrtEpContextConfig instance to release. May be NULL. + */ +ORT_EXPERIMENTAL_API(28, void, OrtEpApi_ReleaseEpContextConfig, _Frees_ptr_opt_ OrtEpContextConfig* config) + +/** \brief Get the application-provided EPContext data read callback. + * + * Returns the OrtReadNamedBufferFunc and opaque state pointer registered via + * OrtApi_SessionOptions_SetEpContextDataReadFunc. If no callback was registered, *read_func and *state are set to + * NULL. The EP is responsible for calling the callback when present and for using its own normal read path when no + * callback is present. + * + * \param[in] config The OrtEpContextConfig from OrtEpApi_SessionOptions_GetEpContextConfig. + * \param[out] read_func The registered read callback, or NULL if none was registered. + * \param[out] state Opaque state pointer passed to read_func, or NULL if none was registered. + * + * \snippet{doc} snippets.dox OrtStatus Return Value + */ +ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtEpApi_EpContextConfig_GetEpContextDataReadFunc, + _In_ const OrtEpContextConfig* config, _Out_ OrtReadNamedBufferFunc* read_func, + _Out_ void** state) + +/** \brief Get the application-provided EPContext data write callback. + * + * Returns the OrtWriteNamedBufferFunc and opaque state pointer registered via + * OrtCompileApi_ModelCompilationOptions_SetEpContextDataWriteFunc. If no callback was registered, *write_func and + * *state are set to NULL. The EP is responsible for calling the callback when present and for using its own normal + * write path when no callback is present. + * + * \param[in] config The OrtEpContextConfig from OrtEpApi_SessionOptions_GetEpContextConfig. + * \param[out] write_func The registered write callback, or NULL if none was registered. + * \param[out] state Opaque state pointer passed to write_func, or NULL if none was registered. + * + * \snippet{doc} snippets.dox OrtStatus Return Value + */ +ORT_EXPERIMENTAL_API(28, OrtStatusPtr, OrtEpApi_EpContextConfig_GetEpContextDataWriteFunc, + _In_ const OrtEpContextConfig* config, _Out_ OrtWriteNamedBufferFunc* write_func, + _Out_ void** state) diff --git a/include/onnxruntime/core/session/onnxruntime_experimental_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_experimental_cxx_api.h index fbd7ba3659435..d3f3a37eca33d 100644 --- a/include/onnxruntime/core/session/onnxruntime_experimental_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_experimental_cxx_api.h @@ -108,5 +108,79 @@ namespace Experimental { // C++ wrapper types or helpers go here in the `Ort::Experimental` namespace. // +// Move-only RAII owner for an OrtEpContextConfig handle, which carries the EPContext read/write callbacks and opaque +// state extracted from an OrtSessionOptions instance. The handle is released via OrtEpApi_ReleaseEpContextConfig when +// the wrapper is destroyed. +// +// Typical EP usage: construct from the session options during CreateEp(), keep the wrapper for the EP's lifetime, and +// query the callbacks via GetReadFunc() / GetWriteFunc(). +class EpContextConfig { + public: + explicit EpContextConfig(std::nullptr_t) noexcept {} + + explicit EpContextConfig(const SessionOptions& session_options) : EpContextConfig{session_options.GetConst()} {} + + // Extracts the EPContext config from `session_options`. Throws Ort::Exception (ORT_NOT_IMPLEMENTED) if the + // experimental functions are not available in this build, or propagates any error from the extraction. + explicit EpContextConfig(ConstSessionOptions session_options) { + const OrtApi* api = &GetApi(); + // Ensure the release function is available before creating a handle, so the handle can always be freed. + Get_OrtEpApi_ReleaseEpContextConfig_SinceV28_FnOrThrow(api); + auto* get_config = Get_OrtEpApi_SessionOptions_GetEpContextConfig_SinceV28_FnOrThrow(api); + ThrowOnError(get_config(static_cast(session_options), &config_)); + } + + EpContextConfig(EpContextConfig&& other) noexcept : config_{other.config_} { other.config_ = nullptr; } + + EpContextConfig& operator=(EpContextConfig&& other) noexcept { + if (this != &other) { + reset(); + config_ = other.config_; + other.config_ = nullptr; + } + return *this; + } + + EpContextConfig(const EpContextConfig&) = delete; + EpContextConfig& operator=(const EpContextConfig&) = delete; + + ~EpContextConfig() { reset(); } + + OrtEpContextConfig* get() const noexcept { return config_; } + explicit operator bool() const noexcept { return config_ != nullptr; } + + // Relinquishes ownership of the handle without releasing it. + OrtEpContextConfig* release() noexcept { + OrtEpContextConfig* released = config_; + config_ = nullptr; + return released; + } + + // Releases any owned handle and resets to empty. + void reset() noexcept { + if (config_ != nullptr) { + if (auto* release_fn = Get_OrtEpApi_ReleaseEpContextConfig_SinceV28_Fn(&GetApi())) { + release_fn(config_); + } + config_ = nullptr; + } + } + + // Returns the configured read callback and opaque state (both nullptr if none was set). Throws on failure. + void GetReadFunc(OrtReadNamedBufferFunc& read_func, void*& state) const { + auto* get_read_func = Get_OrtEpApi_EpContextConfig_GetEpContextDataReadFunc_SinceV28_FnOrThrow(&GetApi()); + ThrowOnError(get_read_func(config_, &read_func, &state)); + } + + // Returns the configured write callback and opaque state (both nullptr if none was set). Throws on failure. + void GetWriteFunc(OrtWriteNamedBufferFunc& write_func, void*& state) const { + auto* get_write_func = Get_OrtEpApi_EpContextConfig_GetEpContextDataWriteFunc_SinceV28_FnOrThrow(&GetApi()); + ThrowOnError(get_write_func(config_, &write_func, &state)); + } + + private: + OrtEpContextConfig* config_ = nullptr; +}; + } // namespace Experimental } // namespace Ort diff --git a/onnxruntime/core/framework/ep_context_options.cc b/onnxruntime/core/framework/ep_context_options.cc index 99fa21b1e4be8..b53a99084152f 100644 --- a/onnxruntime/core/framework/ep_context_options.cc +++ b/onnxruntime/core/framework/ep_context_options.cc @@ -56,6 +56,10 @@ const BufferWriteFuncHolder* ModelGenOptions::TryGetOutputModelWriteFunc() const return std::get_if(&output_model_location); } +const EpContextDataWriteFuncHolder* ModelGenOptions::TryGetEpContextDataWriteFunc() const { + return ep_context_data_write_func.write_func != nullptr ? &ep_context_data_write_func : nullptr; +} + bool ModelGenOptions::AreInitializersEmbeddedInOutputModel() const { return std::holds_alternative(initializers_location); } diff --git a/onnxruntime/core/framework/ep_context_options.h b/onnxruntime/core/framework/ep_context_options.h index 6643516bfb4c3..344ec5f7f6b58 100644 --- a/onnxruntime/core/framework/ep_context_options.h +++ b/onnxruntime/core/framework/ep_context_options.h @@ -7,6 +7,9 @@ #include #include "core/framework/allocator.h" #include "core/framework/config_options.h" +// Needed for OrtWriteNamedBufferFunc (used by EpContextDataWriteFuncHolder below). This include can be removed +// once the experimental EPContext data callback APIs are promoted to the stable C API. +#include "core/session/onnxruntime_experimental_c_api.h" namespace onnxruntime { namespace epctx { @@ -27,6 +30,14 @@ struct BufferWriteFuncHolder { void* stream_state = nullptr; // Opaque pointer to user's stream state. Passed as first argument to write_func. }; +/// +/// Holds the opaque state and write function that EPs use to write EPContext binary data. +/// +struct EpContextDataWriteFuncHolder { + OrtWriteNamedBufferFunc write_func = nullptr; + void* state = nullptr; +}; + /// /// Holds path and size threshold used to write out initializers to an external file. /// @@ -84,10 +95,13 @@ struct ModelGenOptions { InitializerHandler> // Custom function called for every initializer to determine location. initializers_location = std::monostate{}; + EpContextDataWriteFuncHolder ep_context_data_write_func = {}; + bool HasOutputModelLocation() const; const std::filesystem::path* TryGetOutputModelPath() const; const BufferHolder* TryGetOutputModelBuffer() const; const BufferWriteFuncHolder* TryGetOutputModelWriteFunc() const; + const EpContextDataWriteFuncHolder* TryGetEpContextDataWriteFunc() const; bool AreInitializersEmbeddedInOutputModel() const; const ExternalInitializerFileInfo* TryGetExternalInitializerFileInfo() const; diff --git a/onnxruntime/core/framework/session_options.h b/onnxruntime/core/framework/session_options.h index b328fc916f885..ddd21074afe8a 100644 --- a/onnxruntime/core/framework/session_options.h +++ b/onnxruntime/core/framework/session_options.h @@ -16,6 +16,9 @@ #include "core/framework/ep_context_options.h" #include "core/framework/ort_value.h" #include "core/session/onnxruntime_c_api.h" +// Needed for OrtReadNamedBufferFunc, the type of the EPContext data read callback stored in this struct. This include +// can be removed once the experimental EPContext data callback APIs are promoted to the stable C API. +#include "core/session/onnxruntime_experimental_c_api.h" #include "core/optimizer/graph_transformer_level.h" #include "core/util/thread_utils.h" @@ -226,6 +229,9 @@ struct SessionOptions { bool has_explicit_ep_context_gen_options = false; epctx::ModelGenOptions ep_context_gen_options = {}; epctx::ModelGenOptions GetEpContextGenerationOptions() const; + + OrtReadNamedBufferFunc ep_context_data_read_func = nullptr; + void* ep_context_data_read_state = nullptr; }; inline std::ostream& operator<<(std::ostream& os, const SessionOptions& session_options) { diff --git a/onnxruntime/core/session/experimental_c_api.cc b/onnxruntime/core/session/experimental_c_api.cc index 458c47bcb58cd..5b1a8460d1e1f 100644 --- a/onnxruntime/core/session/experimental_c_api.cc +++ b/onnxruntime/core/session/experimental_c_api.cc @@ -6,12 +6,30 @@ #include #include +#include +#include "core/common/common.h" #include "core/framework/error_code_helper.h" +#include "core/framework/ep_context_options.h" +#include "core/session/abi_session_options_impl.h" #include "core/session/onnxruntime_c_api.h" #include "core/session/onnxruntime_experimental_c_api.h" #include "core/session/ort_apis.h" +#if !defined(ORT_MINIMAL_BUILD) +#include "core/session/model_compilation_options.h" +#endif // !defined(ORT_MINIMAL_BUILD) + +// Backing definition of the OrtEpContextConfig handle used by the experimental OrtEpApi_* EPContext data functions. +// Holds copies of the application's EPContext read/write callbacks and opaque state extracted from an +// OrtSessionOptions instance. +struct OrtEpContextConfig { + OrtWriteNamedBufferFunc write_func = nullptr; + void* write_state = nullptr; + OrtReadNamedBufferFunc read_func = nullptr; + void* read_state = nullptr; +}; + // --------------------------------------------------------------------------- // Experimental function implementations // --------------------------------------------------------------------------- @@ -40,6 +58,94 @@ ORT_API_STATUS_IMPL(OrtApi_ExperimentalApiTest_SinceV28, API_IMPL_END } +ORT_API_STATUS_IMPL(OrtApi_SessionOptions_SetEpContextDataReadFunc_SinceV28, _Inout_ OrtSessionOptions* options, + _In_opt_ OrtReadNamedBufferFunc read_func, _In_opt_ void* state) { + API_IMPL_BEGIN + ORT_API_RETURN_IF(options == nullptr, ORT_INVALID_ARGUMENT, "'options' parameter must not be NULL"); + + // Passing a null read_func clears any previously set callback. Clear the state too so a stale state pointer is + // never paired with a missing callback. + options->value.ep_context_data_read_func = read_func; + options->value.ep_context_data_read_state = read_func != nullptr ? state : nullptr; + return nullptr; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtCompileApi_ModelCompilationOptions_SetEpContextDataWriteFunc_SinceV28, + _In_ OrtModelCompilationOptions* ort_model_compile_options, + _In_opt_ OrtWriteNamedBufferFunc write_func, _In_opt_ void* state) { + API_IMPL_BEGIN +#if !defined(ORT_MINIMAL_BUILD) + ORT_API_RETURN_IF(ort_model_compile_options == nullptr, ORT_INVALID_ARGUMENT, "OrtModelCompilationOptions is NULL"); + + // A null write_func clears any previously set callback (symmetric with OrtApi_SessionOptions_SetEpContextDataReadFunc + // and consistent with calling this multiple times to overwrite the callback). + auto model_compile_options = reinterpret_cast(ort_model_compile_options); + model_compile_options->SetEpContextDataWriteFunc(write_func, state); + return nullptr; +#else + ORT_UNUSED_PARAMETER(ort_model_compile_options); + ORT_UNUSED_PARAMETER(write_func); + ORT_UNUSED_PARAMETER(state); + return OrtApis::CreateStatus(ORT_NOT_IMPLEMENTED, "Compile API is not supported in this build"); +#endif // !defined(ORT_MINIMAL_BUILD) + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtEpApi_SessionOptions_GetEpContextConfig_SinceV28, + _In_ const OrtSessionOptions* session_options, + _Outptr_ OrtEpContextConfig** config) { + API_IMPL_BEGIN + ORT_API_RETURN_IF(session_options == nullptr, ORT_INVALID_ARGUMENT, "OrtSessionOptions is NULL"); + ORT_API_RETURN_IF(config == nullptr, ORT_INVALID_ARGUMENT, "Output OrtEpContextConfig is NULL"); + + auto ep_context_config = std::make_unique(); + if (const auto* write_config = session_options->value.ep_context_gen_options.TryGetEpContextDataWriteFunc()) { + ep_context_config->write_func = write_config->write_func; + ep_context_config->write_state = write_config->state; + } + ep_context_config->read_func = session_options->value.ep_context_data_read_func; + ep_context_config->read_state = session_options->value.ep_context_data_read_state; + + *config = ep_context_config.release(); + return nullptr; + API_IMPL_END +} + +ORT_API(void, OrtEpApi_ReleaseEpContextConfig_SinceV28, _Frees_ptr_opt_ OrtEpContextConfig* config) { + delete config; +} + +ORT_API_STATUS_IMPL(OrtEpApi_EpContextConfig_GetEpContextDataReadFunc_SinceV28, + _In_ const OrtEpContextConfig* config, + _Out_ OrtReadNamedBufferFunc* read_func, + _Out_ void** state) { + API_IMPL_BEGIN + ORT_API_RETURN_IF(config == nullptr, ORT_INVALID_ARGUMENT, "OrtEpContextConfig is NULL"); + ORT_API_RETURN_IF(read_func == nullptr, ORT_INVALID_ARGUMENT, "Output read_func is NULL"); + ORT_API_RETURN_IF(state == nullptr, ORT_INVALID_ARGUMENT, "Output state is NULL"); + + *read_func = config->read_func; + *state = config->read_func != nullptr ? config->read_state : nullptr; + return nullptr; + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtEpApi_EpContextConfig_GetEpContextDataWriteFunc_SinceV28, + _In_ const OrtEpContextConfig* config, + _Out_ OrtWriteNamedBufferFunc* write_func, + _Out_ void** state) { + API_IMPL_BEGIN + ORT_API_RETURN_IF(config == nullptr, ORT_INVALID_ARGUMENT, "OrtEpContextConfig is NULL"); + ORT_API_RETURN_IF(write_func == nullptr, ORT_INVALID_ARGUMENT, "Output write_func is NULL"); + ORT_API_RETURN_IF(state == nullptr, ORT_INVALID_ARGUMENT, "Output state is NULL"); + + *write_func = config->write_func; + *state = config->write_func != nullptr ? config->write_state : nullptr; + return nullptr; + API_IMPL_END +} + } // namespace OrtExperimentalApis // --------------------------------------------------------------------------- diff --git a/onnxruntime/core/session/model_compilation_options.cc b/onnxruntime/core/session/model_compilation_options.cc index a17802bdd7573..9f6d1f9f1a9bc 100644 --- a/onnxruntime/core/session/model_compilation_options.cc +++ b/onnxruntime/core/session/model_compilation_options.cc @@ -133,6 +133,15 @@ void ModelCompilationOptions::SetOutputModelGetInitializerLocationFunc( }; } +void ModelCompilationOptions::SetEpContextDataWriteFunc(OrtWriteNamedBufferFunc write_func, void* state) { + // A null write_func clears any previously set callback. Clear the state too so a stale state pointer is never + // paired with a missing callback. + session_options_.value.ep_context_gen_options.ep_context_data_write_func = epctx::EpContextDataWriteFuncHolder{ + write_func, + write_func != nullptr ? state : nullptr, + }; +} + Status ModelCompilationOptions::SetEpContextBinaryInformation(const std::filesystem::path& output_directory, const std::filesystem::path& model_name) { if (output_directory.empty() || model_name.empty()) { diff --git a/onnxruntime/core/session/model_compilation_options.h b/onnxruntime/core/session/model_compilation_options.h index 47529e794677e..a15af565c4d54 100644 --- a/onnxruntime/core/session/model_compilation_options.h +++ b/onnxruntime/core/session/model_compilation_options.h @@ -13,6 +13,7 @@ #include "core/graph/model_editor_api_types.h" #include "core/session/abi_session_options_impl.h" #include "core/session/onnxruntime_c_api.h" +#include "core/session/onnxruntime_experimental_c_api.h" #include "core/session/onnxruntime_session_options_config_keys.h" namespace onnxruntime { @@ -97,6 +98,13 @@ class ModelCompilationOptions { void SetOutputModelGetInitializerLocationFunc(OrtGetInitializerLocationFunc get_initializer_location_func, void* state); + /// + /// Sets a user-provided function to handle EPContext binary data writes. + /// + /// The user-provided OrtWriteNamedBufferFunc callback used to write EPContext data. + /// The user's state. + void SetEpContextDataWriteFunc(OrtWriteNamedBufferFunc write_func, void* state); + /// /// Sets information relate to EP context binary file. /// EP use this information to decide the location and context binary file name. diff --git a/onnxruntime/test/autoep/ep_context_data_callbacks.h b/onnxruntime/test/autoep/ep_context_data_callbacks.h new file mode 100644 index 0000000000000..d25c7a5c571ee --- /dev/null +++ b/onnxruntime/test/autoep/ep_context_data_callbacks.h @@ -0,0 +1,65 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include + +#include "core/session/onnxruntime_c_api.h" +#include "core/session/onnxruntime_cxx_api.h" + +namespace onnxruntime { +namespace test { + +// Shared EPContext read/write callback test doubles, used by both the EpContextDataUtils unit tests +// (ep_context_data_utils_test.cc) and the PluginEp end-to-end EPContext tests (test_execution.cc). +struct EpContextDataCallbackState { + bool write_called = false; + bool read_called = false; + std::string write_file_name; + std::string read_file_name; + std::vector payload; +}; + +inline OrtStatus* ORT_API_CALL StoreEpContextDataCallback(void* state, const char* file_name, const void* buffer, + size_t buffer_size) { + auto* callback_state = static_cast(state); + callback_state->write_called = true; + callback_state->write_file_name = file_name; + callback_state->payload.clear(); + if (buffer_size != 0) { + if (buffer == nullptr) { + return Ort::GetApi().CreateStatus(ORT_INVALID_ARGUMENT, + "StoreEpContextDataCallback received a null buffer for non-empty data"); + } + callback_state->payload.assign(static_cast(buffer), static_cast(buffer) + buffer_size); + } + return nullptr; +} + +inline OrtStatus* ORT_API_CALL LoadEpContextDataCallback(void* state, const char* file_name, OrtAllocator* allocator, + void** buffer, size_t* data_size) { + auto* callback_state = static_cast(state); + callback_state->read_called = true; + callback_state->read_file_name = file_name; + + *buffer = nullptr; + *data_size = callback_state->payload.size(); + if (callback_state->payload.empty()) { + return nullptr; + } + + OrtStatus* status = Ort::GetApi().AllocatorAlloc(allocator, callback_state->payload.size(), buffer); + if (status != nullptr) { + return status; + } + + std::copy(callback_state->payload.begin(), callback_state->payload.end(), static_cast(*buffer)); + return nullptr; +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/autoep/ep_context_data_utils_test.cc b/onnxruntime/test/autoep/ep_context_data_utils_test.cc new file mode 100644 index 0000000000000..8425280e4845b --- /dev/null +++ b/onnxruntime/test/autoep/ep_context_data_utils_test.cc @@ -0,0 +1,327 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Unit tests for the sample-only EPContext data helpers in +// onnxruntime/test/autoep/library/ep_context_data_utils.h. + +#include +#include +#include +#include + +#include +#include + +#include "core/graph/model_editor_api_types.h" +#include "core/session/onnxruntime_cxx_api.h" +#include "core/session/onnxruntime_experimental_cxx_api.h" + +#include "test/autoep/ep_context_data_callbacks.h" +#include "test/autoep/library/ep_context_data_utils.h" +#include "test/util/include/api_asserts.h" +#include "test/util/include/asserts.h" + +namespace onnxruntime { +namespace test { + +namespace { + +OrtStatus* ORT_API_CALL LoadInvalidEpContextDataCallback(void* state, const char* file_name, + OrtAllocator* /*allocator*/, void** buffer, + size_t* data_size) { + auto* callback_state = static_cast(state); + callback_state->read_called = true; + callback_state->read_file_name = file_name; + + *buffer = nullptr; + *data_size = 1; + return nullptr; +} + +void ExpectOrtStatusError(OrtStatus* status_ptr, OrtErrorCode expected_code, std::string_view expected_message) { + Ort::Status status{status_ptr}; + ASSERT_NE(status_ptr, nullptr) << "Expected a failure status, but the API returned nullptr (OK)."; + ASSERT_FALSE(status.IsOK()); + EXPECT_EQ(status.GetErrorCode(), expected_code); + EXPECT_THAT(std::string{status.GetErrorMessage()}, ::testing::HasSubstr(std::string{expected_message})); +} + +std::filesystem::path PrepareTempTestDir(std::string_view name) { + std::filesystem::path test_dir = std::string{name}; + std::filesystem::remove_all(test_dir); + std::filesystem::create_directories(test_dir); + return test_dir; +} + +} // namespace + +TEST(OrtEpLibrary, EpContextDataUtils_PathHelpersRoundTrip) { + const auto& api = Ort::GetApi(); + const std::string file_name = "context_data.bin"; + +#ifdef _WIN32 + std::wstring wide_file_name; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::Utf8ToWideString(api, file_name, wide_file_name)); + ASSERT_FALSE(wide_file_name.empty()); + std::string round_tripped_file_name; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::WideToUtf8String(api, wide_file_name, round_tripped_file_name)); + EXPECT_EQ(round_tripped_file_name, file_name); + + const std::string invalid_utf8(1, static_cast(0xff)); + std::wstring invalid_wide; + ExpectOrtStatusError(ep_context_data_utils::Utf8ToWideString(api, invalid_utf8, invalid_wide), + ORT_INVALID_ARGUMENT, "not valid UTF-8"); + EXPECT_TRUE(invalid_wide.empty()); +#endif + + std::filesystem::path file_path; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::Utf8Path(api, file_name.c_str(), file_path)); + ASSERT_FALSE(file_path.empty()); + std::string round_tripped_path; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::PathToUtf8String(api, file_path, round_tripped_path)); + EXPECT_EQ(round_tripped_path, file_name); + + std::filesystem::path empty_path; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::Utf8Path(api, nullptr, empty_path)); + EXPECT_TRUE(empty_path.empty()); + ASSERT_ORTSTATUS_OK(ep_context_data_utils::Utf8Path(api, "", empty_path)); + EXPECT_TRUE(empty_path.empty()); +} + +TEST(OrtEpLibrary, EpContextDataUtils_ResolvePathAndInvalidArguments) { + const auto& api = Ort::GetApi(); + std::filesystem::path data_path; + + data_path = "stale.ctx"; + ExpectOrtStatusError(ep_context_data_utils::ResolveEpContextDataPath(api, nullptr, nullptr, data_path), + ORT_INVALID_ARGUMENT, "EPContext data file name must not be empty"); + EXPECT_TRUE(data_path.empty()); + + data_path = "stale.ctx"; + ExpectOrtStatusError(ep_context_data_utils::ResolveEpContextDataPath(api, "", nullptr, data_path), + ORT_INVALID_ARGUMENT, "EPContext data file name must not be empty"); + EXPECT_TRUE(data_path.empty()); + + ASSERT_ORTSTATUS_OK(ep_context_data_utils::ResolveEpContextDataPath(api, "relative.ctx", nullptr, data_path)); + std::string resolved_data_path; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::PathToUtf8String(api, data_path, resolved_data_path)); + EXPECT_EQ(resolved_data_path, "relative.ctx"); + + ExpectOrtStatusError(ep_context_data_utils::WriteEpContextDataToFile(api, "unused.ctx", nullptr, nullptr, 1), + ORT_INVALID_ARGUMENT, "EPContext data buffer must not be null for non-empty data"); + ExpectOrtStatusError(ep_context_data_utils::WriteEpContextDataWithFileFallback(api, nullptr, "unused.ctx", nullptr, + nullptr, 1), + ORT_INVALID_ARGUMENT, "EPContext data buffer must not be null for non-empty data"); + + std::vector data; + ExpectOrtStatusError(ep_context_data_utils::ReadEpContextDataWithFileFallback(api, nullptr, "", nullptr, data), + ORT_INVALID_ARGUMENT, "EPContext data file name must not be empty"); + ExpectOrtStatusError(ep_context_data_utils::WriteEpContextDataWithFileFallback(api, nullptr, "", nullptr, nullptr, 0), + ORT_INVALID_ARGUMENT, "EPContext data file name must not be empty"); + ExpectOrtStatusError(ep_context_data_utils::WriteEpContextDataWithFileFallback( + api, nullptr, "logical_context_data.bin", "", nullptr, nullptr, 0), + ORT_INVALID_ARGUMENT, "EPContext data fallback file name must not be empty"); +} + +TEST(OrtEpLibrary, EpContextDataUtils_ResolvePathRejectsUnsafeNames) { + const auto& api = Ort::GetApi(); + std::filesystem::path data_path; + + ExpectOrtStatusError(ep_context_data_utils::ResolveEpContextDataPath(api, "../escape.ctx", nullptr, data_path), + ORT_INVALID_ARGUMENT, "EPContext data file name must not contain path traversal"); + EXPECT_TRUE(data_path.empty()); + +#ifdef _WIN32 + const char* absolute_file_name = "C:\\temp\\escape.ctx"; + const char* drive_relative_file_name = "C:escape.ctx"; + const char* root_relative_file_name = "\\escape.ctx"; +#else + const char* absolute_file_name = "/tmp/escape.ctx"; +#endif + ASSERT_ORTSTATUS_OK(ep_context_data_utils::ResolveEpContextDataPath(api, absolute_file_name, nullptr, data_path)); + EXPECT_TRUE(data_path.is_absolute()); + +#ifdef _WIN32 + ExpectOrtStatusError(ep_context_data_utils::WriteEpContextDataWithFileFallback( + api, nullptr, drive_relative_file_name, "unused.ctx", nullptr, nullptr, 0), + ORT_INVALID_ARGUMENT, "EPContext data file name must not be absolute or rooted"); + ExpectOrtStatusError(ep_context_data_utils::WriteEpContextDataWithFileFallback( + api, nullptr, root_relative_file_name, "unused.ctx", nullptr, nullptr, 0), + ORT_INVALID_ARGUMENT, "EPContext data file name must not be absolute or rooted"); +#endif + + std::vector data; + ExpectOrtStatusError(ep_context_data_utils::ReadEpContextDataFromFile(api, "../escape.ctx", nullptr, data), + ORT_INVALID_ARGUMENT, "EPContext data file name must not contain path traversal"); + ExpectOrtStatusError(ep_context_data_utils::WriteEpContextDataWithFileFallback( + api, nullptr, absolute_file_name, "unused.ctx", nullptr, nullptr, 0), + ORT_INVALID_ARGUMENT, "EPContext data file name must not be absolute or rooted"); + + ModelEditorGraph empty_model_path_graph; + ExpectOrtStatusError(ep_context_data_utils::ResolveEpContextDataPath(api, "../escape.ctx", + empty_model_path_graph.ToExternal(), data_path), + ORT_INVALID_ARGUMENT, "requires a model path"); + + // A model-derived name that designates a directory ("." or a trailing separator with an empty filename) is + // rejected up front, rather than resolving to a directory and failing later with a confusing I/O error. + ExpectOrtStatusError(ep_context_data_utils::ResolveEpContextDataPath(api, ".", empty_model_path_graph.ToExternal(), + data_path), + ORT_INVALID_ARGUMENT, "must refer to a file"); + ExpectOrtStatusError(ep_context_data_utils::WriteEpContextDataWithFileFallback( + api, nullptr, ".", "unused.ctx", nullptr, nullptr, 0), + ORT_INVALID_ARGUMENT, "must refer to a file"); + ExpectOrtStatusError(ep_context_data_utils::WriteEpContextDataWithFileFallback( + api, nullptr, "sub/", "unused.ctx", nullptr, nullptr, 0), + ORT_INVALID_ARGUMENT, "must refer to a file"); +} + +TEST(OrtEpLibrary, EpContextDataUtils_ResolvePathRejectsSymlinkEscape) { + const auto& api = Ort::GetApi(); + const std::filesystem::path test_dir = PrepareTempTestDir("ort_ep_context_data_utils_symlink_escape_test"); + auto cleanup = gsl::finally([&]() { std::filesystem::remove_all(test_dir); }); + + const std::filesystem::path model_dir = test_dir / "model_dir"; + const std::filesystem::path outside_dir = test_dir / "outside_dir"; + ASSERT_TRUE(std::filesystem::create_directories(model_dir)); + ASSERT_TRUE(std::filesystem::create_directories(outside_dir)); + + const std::filesystem::path symlink_path = model_dir / "linked_outside"; + // Relative symlink targets are resolved by the OS relative to the link's own directory, not the test's working + // directory. Point to the sibling outside_dir using a link-relative target; using the test_dir-relative + // `outside_dir` path here would create a dangling link under model_dir, and weakly_canonical() would not traverse it. + const std::filesystem::path symlink_target = std::filesystem::path{".."} / outside_dir.filename(); + std::error_code symlink_error; + std::filesystem::create_directory_symlink(symlink_target, symlink_path, symlink_error); + if (symlink_error) { + GTEST_SKIP() << "Unable to create directory symlink for containment test: " << symlink_error.message(); + } + + ModelEditorGraph graph; + graph.model_path = model_dir / "model.onnx"; + + std::filesystem::path data_path; + ExpectOrtStatusError(ep_context_data_utils::ResolveEpContextDataPath(api, "linked_outside/escape.ctx", + graph.ToExternal(), data_path), + ORT_INVALID_ARGUMENT, "resolve to a path within the model directory"); + EXPECT_TRUE(data_path.empty()); +} + +TEST(OrtEpLibrary, EpContextDataUtils_FileFallbackReadsAndWrites) { + const auto& api = Ort::GetApi(); + const std::filesystem::path test_dir = PrepareTempTestDir("ort_ep_context_data_utils_file_fallback_test"); + auto cleanup = gsl::finally([&]() { std::filesystem::remove_all(test_dir); }); + + const std::string payload = "file fallback payload"; + const std::filesystem::path data_path = test_dir / "context_data.bin"; + std::string data_file_name; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::PathToUtf8String(api, data_path, data_file_name)); + + ASSERT_ORTSTATUS_OK(ep_context_data_utils::WriteEpContextDataToFile(api, data_file_name.c_str(), nullptr, + payload.data(), payload.size())); + + std::vector data; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::ReadEpContextDataFromFile(api, data_file_name.c_str(), nullptr, data)); + EXPECT_EQ(std::string(data.begin(), data.end()), payload); + + const std::filesystem::path wrapper_data_path = test_dir / "wrapper_context_data.bin"; + std::string wrapper_data_file_name; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::PathToUtf8String(api, wrapper_data_path, wrapper_data_file_name)); + ASSERT_ORTSTATUS_OK(ep_context_data_utils::WriteEpContextDataWithFileFallback( + api, nullptr, wrapper_data_file_name.c_str(), nullptr, payload.data(), payload.size())); + + data.clear(); + ASSERT_ORTSTATUS_OK(ep_context_data_utils::ReadEpContextDataWithFileFallback( + api, nullptr, wrapper_data_file_name.c_str(), nullptr, data)); + EXPECT_EQ(std::string(data.begin(), data.end()), payload); + + const std::filesystem::path fallback_data_path = test_dir / "fallback_context_data.bin"; + std::string fallback_data_file_name; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::PathToUtf8String(api, fallback_data_path, fallback_data_file_name)); + ASSERT_ORTSTATUS_OK(ep_context_data_utils::WriteEpContextDataWithFileFallback( + api, nullptr, "logical_context_data.bin", fallback_data_file_name.c_str(), nullptr, payload.data(), + payload.size())); + + const std::filesystem::path unsafe_logical_fallback_path = test_dir / "unsafe_logical_context_data.bin"; + std::string unsafe_logical_fallback_file_name; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::PathToUtf8String(api, unsafe_logical_fallback_path, + unsafe_logical_fallback_file_name)); + ExpectOrtStatusError(ep_context_data_utils::WriteEpContextDataWithFileFallback( + api, nullptr, "../logical_context_data.bin", + unsafe_logical_fallback_file_name.c_str(), nullptr, payload.data(), payload.size()), + ORT_INVALID_ARGUMENT, "EPContext data file name must not contain path traversal"); + + data.clear(); + ASSERT_ORTSTATUS_OK(ep_context_data_utils::ReadEpContextDataFromFile(api, fallback_data_file_name.c_str(), nullptr, + data)); + EXPECT_EQ(std::string(data.begin(), data.end()), payload); + + const std::filesystem::path empty_data_path = test_dir / "empty_context_data.bin"; + std::string empty_data_file_name; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::PathToUtf8String(api, empty_data_path, empty_data_file_name)); + ASSERT_ORTSTATUS_OK(ep_context_data_utils::WriteEpContextDataWithFileFallback( + api, nullptr, empty_data_file_name.c_str(), nullptr, nullptr, 0)); + + data.assign({'s', 't', 'a', 'l', 'e'}); + ASSERT_ORTSTATUS_OK(ep_context_data_utils::ReadEpContextDataWithFileFallback( + api, nullptr, empty_data_file_name.c_str(), nullptr, data)); + EXPECT_TRUE(data.empty()); + + const std::filesystem::path missing_data_path = test_dir / "missing_context_data.bin"; + std::string missing_data_file_name; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::PathToUtf8String(api, missing_data_path, missing_data_file_name)); + ExpectOrtStatusError(ep_context_data_utils::ReadEpContextDataFromFile(api, missing_data_file_name.c_str(), nullptr, + data), + ORT_FAIL, "Failed to open EPContext data file for read"); +} + +TEST(OrtEpLibrary, EpContextDataUtils_CallbackFallbackUsesCallbacks) { + const auto& api = Ort::GetApi(); + + EpContextDataCallbackState read_callback_state; + read_callback_state.payload = {'c', 'a', 'l', 'l', 'b', 'a', 'c', 'k'}; + EpContextDataCallbackState write_callback_state; + + std::vector data; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::ReadEpContextDataWithFileFallback( + api, LoadEpContextDataCallback, &read_callback_state, "callback_context.bin", nullptr, data)); + ASSERT_TRUE(read_callback_state.read_called); + EXPECT_EQ(read_callback_state.read_file_name, "callback_context.bin"); + EXPECT_EQ(data, read_callback_state.payload); + + const std::string payload = "callback write payload"; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::WriteEpContextDataWithFileFallback( + api, StoreEpContextDataCallback, &write_callback_state, "callback_write_context.bin", + "callback_write_context.bin", nullptr, payload.data(), payload.size())); + ASSERT_TRUE(write_callback_state.write_called); + EXPECT_EQ(write_callback_state.write_file_name, "callback_write_context.bin"); + EXPECT_EQ(std::string(write_callback_state.payload.begin(), write_callback_state.payload.end()), payload); + + write_callback_state = {}; + const std::string payload_with_unused_fallback = "callback write payload with unused fallback"; + // With a callback present the file fallback is never used, so the empty fallback name is accepted (not validated). + ASSERT_ORTSTATUS_OK(ep_context_data_utils::WriteEpContextDataWithFileFallback( + api, StoreEpContextDataCallback, &write_callback_state, + "callback_write_context_unused_fallback.bin", "", nullptr, + payload_with_unused_fallback.data(), payload_with_unused_fallback.size())); + ASSERT_TRUE(write_callback_state.write_called); + EXPECT_EQ(write_callback_state.write_file_name, "callback_write_context_unused_fallback.bin"); + EXPECT_EQ(std::string(write_callback_state.payload.begin(), write_callback_state.payload.end()), + payload_with_unused_fallback); +} + +TEST(OrtEpLibrary, EpContextDataUtils_ReadCallbackRejectsNullBufferForNonEmptyPayload) { + const auto& api = Ort::GetApi(); + + EpContextDataCallbackState read_callback_state; + + std::vector data; + ExpectOrtStatusError(ep_context_data_utils::ReadEpContextDataWithFileFallback( + api, LoadInvalidEpContextDataCallback, &read_callback_state, + "invalid_callback_context.bin", nullptr, data), + ORT_FAIL, "OrtReadNamedBufferFunc returned a null buffer for non-empty EPContext data"); + ASSERT_TRUE(read_callback_state.read_called); + EXPECT_EQ(read_callback_state.read_file_name, "invalid_callback_context.bin"); +} + +} // namespace test +} // namespace onnxruntime diff --git a/onnxruntime/test/autoep/library/ep_context_data_utils.h b/onnxruntime/test/autoep/library/ep_context_data_utils.h new file mode 100644 index 0000000000000..a3f9a377ee92f --- /dev/null +++ b/onnxruntime/test/autoep/library/ep_context_data_utils.h @@ -0,0 +1,501 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef _WIN32 +// Define NOMINMAX (and WIN32_LEAN_AND_MEAN) before so the min/max macros it would otherwise pull in do +// not clobber std::numeric_limits<...>::max() and std::min/std::max used in this header. +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +#endif + +#include "plugin_ep_utils.h" +#include "onnxruntime_experimental_cxx_api.h" + +// Sample-only EPContext data helpers shared by the example plugin EP and its tests. These are intentionally outside +// the ORT C and EP ABI and are provided as a reference for EP authors that need to handle external (non-embedded) +// EPContext binary data. +// +// The intended entry points for EP implementers are the ReadEpContextDataWithFileFallback / +// WriteEpContextDataWithFileFallback overloads: they prefer an application-supplied OrtReadNamedBufferFunc / +// OrtWriteNamedBufferFunc (carried by OrtEpContextConfig) and fall back to file I/O when no callback is configured. +// The other functions are lower-level building blocks. Production EPs should additionally apply their own sandboxing, +// size limits, and path policies; see the per-function notes on how untrusted, model-derived names are treated. +namespace ep_context_data_utils { + +#ifdef _WIN32 +inline std::string WindowsLastErrorMessage(std::string_view message, DWORD error_code) { + return std::string{message} + " GetLastError=" + std::to_string(error_code); +} + +// Converts a UTF-8 string to a wide string. Reports conversion failures (e.g., invalid UTF-8) via OrtStatus* instead +// of silently returning an empty string. An empty input yields an empty output and a success status. +inline OrtStatus* Utf8ToWideString(const OrtApi& api, std::string_view value, std::wstring& wide_value) { + wide_value.clear(); + if (value.empty()) { + return nullptr; + } + if (value.size() > static_cast(std::numeric_limits::max())) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data file name is too long to convert"); + } + + const int wide_length = MultiByteToWideChar(CP_UTF8, MB_ERR_INVALID_CHARS, value.data(), + static_cast(value.size()), nullptr, 0); + if (wide_length <= 0) { + const std::string message = WindowsLastErrorMessage( + "EPContext data file name is not valid UTF-8 or could not be converted to a wide string.", GetLastError()); + return api.CreateStatus(ORT_INVALID_ARGUMENT, message.c_str()); + } + + wide_value.resize(static_cast(wide_length)); + const int converted = MultiByteToWideChar(CP_UTF8, MB_ERR_INVALID_CHARS, value.data(), + static_cast(value.size()), wide_value.data(), wide_length); + if (converted != wide_length) { + wide_value.clear(); + const std::string message = WindowsLastErrorMessage("Failed to convert EPContext data file name to a wide string.", + GetLastError()); + return api.CreateStatus(ORT_FAIL, message.c_str()); + } + return nullptr; +} + +// Converts a wide string to UTF-8. Reports conversion failures via OrtStatus* instead of silently returning an empty +// string. An empty input yields an empty output and a success status. +inline OrtStatus* WideToUtf8String(const OrtApi& api, std::wstring_view value, std::string& utf8_value) { + utf8_value.clear(); + if (value.empty()) { + return nullptr; + } + if (value.size() > static_cast(std::numeric_limits::max())) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data file name is too long to convert"); + } + + const int utf8_length = WideCharToMultiByte(CP_UTF8, 0, value.data(), static_cast(value.size()), + nullptr, 0, nullptr, nullptr); + if (utf8_length <= 0) { + const std::string message = WindowsLastErrorMessage( + "EPContext data file name could not be converted to UTF-8.", GetLastError()); + return api.CreateStatus(ORT_INVALID_ARGUMENT, message.c_str()); + } + + utf8_value.resize(static_cast(utf8_length)); + const int converted = WideCharToMultiByte(CP_UTF8, 0, value.data(), static_cast(value.size()), + utf8_value.data(), utf8_length, nullptr, nullptr); + if (converted != utf8_length) { + utf8_value.clear(); + const std::string message = WindowsLastErrorMessage("Failed to convert EPContext data file name to UTF-8.", + GetLastError()); + return api.CreateStatus(ORT_FAIL, message.c_str()); + } + return nullptr; +} +#endif + +// Converts a UTF-8 path to a std::filesystem::path. A null or empty input yields an empty path and a success status; +// conversion failures are reported via OrtStatus*. +inline OrtStatus* Utf8Path(const OrtApi& api, const char* path, std::filesystem::path& out_path) { + out_path.clear(); + if (path == nullptr || path[0] == '\0') { + return nullptr; + } + +#ifdef _WIN32 + std::wstring wide_path; + RETURN_IF_ERROR(Utf8ToWideString(api, path, wide_path)); + out_path = std::filesystem::path{wide_path}; +#else + (void)api; + out_path = std::filesystem::path{path}; +#endif + return nullptr; +} + +inline OrtStatus* PathToUtf8String(const OrtApi& api, const std::filesystem::path& path, std::string& utf8_path) { + utf8_path.clear(); +#ifdef _WIN32 + RETURN_IF_ERROR(WideToUtf8String(api, path.wstring(), utf8_path)); +#else + (void)api; + utf8_path = path.string(); +#endif + return nullptr; +} + +inline std::string PathToUtf8StringForMessage(const std::filesystem::path& path) { + std::string utf8_path; + Ort::Status status{PathToUtf8String(Ort::GetApi(), path, utf8_path)}; + return status.IsOK() ? utf8_path : std::string{""}; +} + +// Lexical check for a ".." component. This is a coarse guard used when there is no filesystem base directory to +// contain against: logical callback-namespace names and trusted (graph == nullptr) physical paths. It is NOT a +// containment mechanism: it does not resolve symlinks and it rejects benign cases such as "a/b/c/../file.txt". +// Filesystem containment against a model directory is done by IsResolvedPathWithinBase() below, which the untrusted +// (model-relative) resolution path uses. +inline bool ContainsPathTraversal(const std::filesystem::path& path) { + const std::filesystem::path parent_dir{".."}; + for (const auto& component : path) { + if (component == parent_dir) { + return true; + } + } + return false; +} + +inline bool HasAbsoluteOrRootedPath(const std::filesystem::path& path) { + return path.is_absolute() || path.has_root_name() || path.has_root_directory(); +} + +// Returns true if the final component of `path` is empty (e.g., a trailing separator like "sub/") or is the +// current-directory entry ".", i.e. the name designates a directory rather than a file (".." is handled separately by +// ContainsPathTraversal()). Such a name resolves to a directory and would only surface later as a confusing file I/O +// failure, so model-derived names like these are rejected up front. +inline bool IsDirectoryOrEmptyName(const std::filesystem::path& path) { + const std::filesystem::path leaf = path.filename(); + return leaf.empty() || leaf == std::filesystem::path{"."}; +} + +// Returns true if `candidate_full` (a base-relative name already combined with `base`) resolves to a location inside +// `base`. Both are normalized with std::filesystem::weakly_canonical, which resolves "." / ".." and any symlinks in +// the existing portion of the path, so a name that escapes `base` directly or through a symlink is rejected. On +// success the canonical resolved path is written to `resolved`. +inline bool IsResolvedPathWithinBase(const std::filesystem::path& base, const std::filesystem::path& candidate_full, + std::filesystem::path& resolved) { + std::error_code ec; + const std::filesystem::path base_for_canon = base.empty() ? std::filesystem::path{"."} : base; + const std::filesystem::path canonical_base = std::filesystem::weakly_canonical(base_for_canon, ec); + if (ec) { + return false; + } + std::filesystem::path candidate_resolved = std::filesystem::weakly_canonical(candidate_full, ec); + if (ec) { + return false; + } + const std::filesystem::path relative = candidate_resolved.lexically_relative(canonical_base); + if (relative.empty() || *relative.begin() == std::filesystem::path{".."}) { + return false; + } + + resolved = std::move(candidate_resolved); + return true; +} + +inline OrtStatus* ValidateEpContextDataName(const OrtApi& api, const char* file_name, + std::filesystem::path& data_name) { + data_name.clear(); + + if (file_name == nullptr || file_name[0] == '\0') { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data file name must not be empty"); + } + + std::filesystem::path candidate_path; + RETURN_IF_ERROR(Utf8Path(api, file_name, candidate_path)); + if (candidate_path.empty()) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data file name is not a valid path"); + } + + if (HasAbsoluteOrRootedPath(candidate_path)) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data file name must not be absolute or rooted"); + } + + if (ContainsPathTraversal(candidate_path)) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data file name must not contain path traversal"); + } + + if (IsDirectoryOrEmptyName(candidate_path)) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data file name must refer to a file, not a directory"); + } + + data_name = candidate_path; + return nullptr; +} + +// Resolves `file_name` to a filesystem path for reading or writing EPContext data (used by both the read path and +// the write-fallback path). +// +// When `graph` is null the caller is trusted and owns the path: `file_name` is returned as-is and may be absolute (a +// lexical ".." is still rejected as a coarse guard). When `graph` is non-null, `file_name` originates from the +// untrusted EPContext model "ep_cache_context" attribute: the graph must have a model path, the name must be +// relative, and after combining it with the model's directory the result must stay within that directory. Symlinks and +// ".." are resolved (via weakly_canonical), so a name that escapes the model directory - including through a symlink - +// is rejected. +// This helper only decides whether a model-derived file name resolves inside the model directory. Production EPs +// should still choose an application-approved storage root (sandbox), reject special files/locations as appropriate, +// and cap the number of bytes they will read or write for a single EPContext payload. +inline OrtStatus* ResolveEpContextDataPath(const OrtApi& api, const char* file_name, const OrtGraph* graph, + std::filesystem::path& data_path) { + data_path.clear(); + + if (file_name == nullptr || file_name[0] == '\0') { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data file name must not be empty"); + } + + std::filesystem::path candidate_path; + RETURN_IF_ERROR(Utf8Path(api, file_name, candidate_path)); + if (candidate_path.empty()) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data file name is not a valid path"); + } + + // Trusted direct callers (graph == nullptr) own the path and may pass an absolute physical path. + if (graph == nullptr) { + if (ContainsPathTraversal(candidate_path)) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data file name must not contain path traversal"); + } + data_path = candidate_path; + return nullptr; + } + + // Untrusted (model-derived) name: must be relative and must resolve within the model directory. + if (HasAbsoluteOrRootedPath(candidate_path)) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data file name must not be absolute or rooted"); + } + + if (IsDirectoryOrEmptyName(candidate_path)) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data file name must refer to a file, not a directory"); + } + + const ORTCHAR_T* model_path = nullptr; + RETURN_IF_ERROR(api.Graph_GetModelPath(graph, &model_path)); + if (model_path == nullptr || model_path[0] == 0) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, + "EPContext data file fallback requires a model path to resolve relative names"); + } + + const std::filesystem::path base_dir = std::filesystem::path{model_path}.parent_path(); + std::filesystem::path resolved; + if (!IsResolvedPathWithinBase(base_dir, base_dir / candidate_path, resolved)) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, + "EPContext data file name must resolve to a path within the model directory"); + } + + data_path = resolved; + return nullptr; +} + +inline OrtStatus* WriteEpContextDataToResolvedFile(const OrtApi& api, const std::filesystem::path& data_path, + const void* buffer, size_t buffer_size) { + std::ofstream output_stream(data_path, std::ios::binary); + if (!output_stream) { + const std::string message = "Failed to open EPContext data file for write: " + + PathToUtf8StringForMessage(data_path); + return api.CreateStatus(ORT_FAIL, message.c_str()); + } + + if (buffer_size != 0) { + if (buffer_size > static_cast(std::numeric_limits::max())) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data buffer is too large to write"); + } + + output_stream.write(static_cast(buffer), static_cast(buffer_size)); + if (!output_stream) { + const std::string message = "Failed to write EPContext data file: " + + PathToUtf8StringForMessage(data_path); + return api.CreateStatus(ORT_FAIL, message.c_str()); + } + } + + return nullptr; +} + +inline OrtStatus* ReadEpContextDataFromFile(const OrtApi& api, const char* file_name, const OrtGraph* graph, + std::vector& data) { + data.clear(); + + std::filesystem::path data_path; + RETURN_IF_ERROR(ResolveEpContextDataPath(api, file_name, graph, data_path)); + + std::ifstream input_stream(data_path, std::ios::binary); + if (!input_stream) { + const std::string message = "Failed to open EPContext data file for read: " + + PathToUtf8StringForMessage(data_path); + return api.CreateStatus(ORT_FAIL, message.c_str()); + } + + data.assign(std::istreambuf_iterator{input_stream}, std::istreambuf_iterator{}); + if (!input_stream) { + const std::string message = "Failed to read EPContext data file: " + + PathToUtf8StringForMessage(data_path); + return api.CreateStatus(ORT_FAIL, message.c_str()); + } + + return nullptr; +} + +inline OrtStatus* WriteEpContextDataToFile(const OrtApi& api, const char* file_name, const OrtGraph* graph, + const void* buffer, size_t buffer_size) { + if (buffer == nullptr && buffer_size != 0) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data buffer must not be null for non-empty data"); + } + + std::filesystem::path data_path; + RETURN_IF_ERROR(ResolveEpContextDataPath(api, file_name, graph, data_path)); + return WriteEpContextDataToResolvedFile(api, data_path, buffer, buffer_size); +} + +// Low-level overload that takes the read callback and its opaque state directly. Production EPs should use the +// overload below that takes an OrtEpContextConfig; this overload exists so unit tests can inject a callback without +// constructing an OrtEpContextConfig. When `read_func` is null the data is read from a file. +inline OrtStatus* ReadEpContextDataWithFileFallback( + const OrtApi& api, + OrtReadNamedBufferFunc read_func, void* read_state, + const char* file_name, const OrtGraph* graph, + std::vector& data) { + if (file_name == nullptr || file_name[0] == '\0') { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data file name must not be empty"); + } + + if (read_func == nullptr) { + return ReadEpContextDataFromFile(api, file_name, graph, data); + } + + // Use the C allocator API (not Ort::AllocatorWithDefaultOptions, whose constructor throws) so this OrtStatus*-based + // helper stays exception-free. The default allocator is owned by ORT and must not be released here. + OrtAllocator* allocator = nullptr; + RETURN_IF_ERROR(api.GetAllocatorWithDefaultOptions(&allocator)); + void* ep_context_data = nullptr; + size_t ep_context_data_size = 0; + OrtStatus* status = read_func(read_state, file_name, allocator, &ep_context_data, &ep_context_data_size); + auto buffer_deleter = [&api, allocator](void* buffer_to_free) { + if (buffer_to_free != nullptr) { + // Best-effort free during cleanup; release any returned status without throwing. + Ort::Status free_status{api.AllocatorFree(allocator, buffer_to_free)}; + static_cast(free_status); + } + }; + std::unique_ptr ep_context_data_guard(ep_context_data, buffer_deleter); + + if (status != nullptr) { + return status; + } + + if (ep_context_data_size != 0 && ep_context_data == nullptr) { + return api.CreateStatus( + ORT_FAIL, "OrtReadNamedBufferFunc returned a null buffer for non-empty EPContext data"); + } + + data.clear(); + if (ep_context_data != nullptr) { + const char* ep_context_data_begin = static_cast(ep_context_data); + data.assign(ep_context_data_begin, ep_context_data_begin + ep_context_data_size); + } + + return nullptr; +} + +// Reads EPContext binary data named `file_name`. If the session configured an OrtReadNamedBufferFunc (carried by +// `ep_context_config`), it is used; otherwise the data is read from a file. When `graph` is non-null it is the +// EPContext model graph: untrusted absolute/rooted/traversal names are rejected and relative names are resolved +// against the model directory. Pass `graph == nullptr` only for trusted callers supplying a physical path. `data` is +// cleared first and receives the bytes on success. +inline OrtStatus* ReadEpContextDataWithFileFallback( + const OrtApi& api, + const OrtEpContextConfig* ep_context_config, + const char* file_name, const OrtGraph* graph, + std::vector& data) { + OrtReadNamedBufferFunc read_func = nullptr; + void* read_state = nullptr; + if (ep_context_config != nullptr) { + auto get_read_func = + Ort::Experimental::Get_OrtEpApi_EpContextConfig_GetEpContextDataReadFunc_SinceV28_Fn(&api); + if (get_read_func == nullptr) { + return api.CreateStatus(ORT_NOT_IMPLEMENTED, + "OrtEpApi_EpContextConfig_GetEpContextDataReadFunc is not available"); + } + RETURN_IF_ERROR(get_read_func(ep_context_config, &read_func, &read_state)); + } + return ReadEpContextDataWithFileFallback(api, read_func, read_state, file_name, graph, data); +} + +// Low-level overload that takes the write callback and its opaque state directly. Production EPs should use the +// overloads below that take an OrtEpContextConfig; this overload exists so unit tests can inject a callback without +// constructing an OrtEpContextConfig. When `write_func` is null the data is written to the file fallback. +inline OrtStatus* WriteEpContextDataWithFileFallback( + const OrtApi& api, + OrtWriteNamedBufferFunc write_func, void* write_state, + const char* file_name, const char* fallback_file_name, + const OrtGraph* graph, + const void* buffer, size_t buffer_size) { + if (file_name == nullptr || file_name[0] == '\0') { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data file name must not be empty"); + } + + if (buffer == nullptr && buffer_size != 0) { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data buffer must not be null for non-empty data"); + } + + // The app-supplied write callback owns its own logical namespace, so file_name is passed through unmodified. + // Only the file-fallback path below maps a name onto the filesystem, so it validates the logical name there. + if (write_func != nullptr) { + return write_func(write_state, file_name, buffer, buffer_size); + } + + // Even when the physical fallback path is supplied separately, `file_name` is the logical name written into the + // EPContext model's ep_cache_context attribute. Validate it as a safe relative name so a generated model cannot + // contain an unsafe logical reference that later reaches the read-side resolver. + std::filesystem::path logical_path; + RETURN_IF_ERROR(ValidateEpContextDataName(api, file_name, logical_path)); + + if (fallback_file_name == nullptr || fallback_file_name[0] == '\0') { + return api.CreateStatus(ORT_INVALID_ARGUMENT, "EPContext data fallback file name must not be empty"); + } + + std::filesystem::path data_path; + RETURN_IF_ERROR(ResolveEpContextDataPath(api, fallback_file_name, graph, data_path)); + return WriteEpContextDataToResolvedFile(api, data_path, buffer, buffer_size); +} + +// Writes EPContext binary data. If the compilation configured an OrtWriteNamedBufferFunc (carried by +// `ep_context_config`), it is used and `file_name` is passed through unmodified as the logical name. Otherwise the +// data is written to a file at `fallback_file_name`, which is resolved against the model directory when `graph` is +// non-null (and rejected if absolute or rooted in that case). `graph == nullptr` denotes a trusted caller that may +// supply an absolute physical path. `buffer` may be null only when `buffer_size` is 0. +inline OrtStatus* WriteEpContextDataWithFileFallback( + const OrtApi& api, + const OrtEpContextConfig* ep_context_config, + const char* file_name, const char* fallback_file_name, + const OrtGraph* graph, + const void* buffer, size_t buffer_size) { + OrtWriteNamedBufferFunc write_func = nullptr; + void* write_state = nullptr; + if (ep_context_config != nullptr) { + auto get_write_func = + Ort::Experimental::Get_OrtEpApi_EpContextConfig_GetEpContextDataWriteFunc_SinceV28_Fn(&api); + if (get_write_func == nullptr) { + return api.CreateStatus(ORT_NOT_IMPLEMENTED, + "OrtEpApi_EpContextConfig_GetEpContextDataWriteFunc is not available"); + } + RETURN_IF_ERROR(get_write_func(ep_context_config, &write_func, &write_state)); + } + return WriteEpContextDataWithFileFallback(api, write_func, write_state, file_name, fallback_file_name, graph, buffer, + buffer_size); +} + +// Convenience overload that uses `file_name` as both the logical callback name and the file-fallback path. +// Because `file_name` doubles as the fallback path, it must be a safe relative name (this overload validates it and +// rejects absolute/rooted paths and `..` traversal). To write the file fallback to an absolute physical path (a +// trusted caller with `graph == nullptr`), use the overload above that takes a separate `fallback_file_name`. +inline OrtStatus* WriteEpContextDataWithFileFallback( + const OrtApi& api, + const OrtEpContextConfig* ep_context_config, + const char* file_name, const OrtGraph* graph, + const void* buffer, size_t buffer_size) { + return WriteEpContextDataWithFileFallback(api, ep_context_config, file_name, file_name, graph, buffer, buffer_size); +} + +} // namespace ep_context_data_utils diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/ep.cc b/onnxruntime/test/autoep/library/example_plugin_ep/ep.cc index fecf7ac9a4038..90ad4e7976824 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep.cc +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep.cc @@ -4,15 +4,17 @@ #include "ep.h" #include +#include #include #include -#include +#include #include #include #include +#include #include -#include +#include "../ep_context_data_utils.h" #include "ep_factory.h" #include "ep_stream_support.h" @@ -167,13 +169,15 @@ struct EpContextNodeComputeInfo : NodeComputeInfoBase { ExampleEp& ep; }; -ExampleEp::ExampleEp(ExampleEpFactory& factory, const std::string& name, const Config& config, const OrtLogger& logger) +ExampleEp::ExampleEp(ExampleEpFactory& factory, const std::string& name, const Config& config, const OrtLogger& logger, + Ort::Experimental::EpContextConfig ep_context_config) : OrtEp{}, // explicitly call the struct ctor to ensure all optional values are default initialized ApiPtrs{static_cast(factory)}, factory_{factory}, name_{name}, config_{config}, - logger_{logger} { + logger_{logger}, + ep_context_config_{std::move(ep_context_config)} { ort_version_supported = ORT_API_VERSION; // set to the ORT version we were compiled with. // Initialize the execution provider's function table @@ -193,8 +197,6 @@ ExampleEp::ExampleEp(ExampleEpFactory& factory, const std::string& name, const C ORT_FILE, __LINE__, __FUNCTION__)); } -ExampleEp::~ExampleEp() = default; - /*static*/ const char* ORT_API_CALL ExampleEp ::GetNameImpl(const OrtEp* this_ptr) noexcept { const auto* ep = static_cast(this_ptr); @@ -409,6 +411,26 @@ OrtStatus* ORT_API_CALL ExampleEp::CompileImpl(_In_ OrtEp* this_ptr, _In_ const auto fused_node_name = fused_node.GetName(); if (is_ep_context_node) { + Ort::ConstOpAttr embed_mode_attr; + RETURN_IF_ERROR(nodes[0].GetAttributeByName("embed_mode", embed_mode_attr)); + int64_t embed_mode = 1; + RETURN_IF_ERROR(embed_mode_attr.GetValue(embed_mode)); + + if (embed_mode == 0) { + Ort::ConstOpAttr ep_cache_context_attr; + RETURN_IF_ERROR(nodes[0].GetAttributeByName("ep_cache_context", ep_cache_context_attr)); + std::string ep_cache_context; + RETURN_IF_ERROR(ep_cache_context_attr.GetValue(ep_cache_context)); + + // This example only exercises the load-side read flow (callback first, file fallback otherwise) to show how + // an EP retrieves EPContext binary data during compile. A real EP would consume `ep_context_data` (e.g., + // initialize a kernel/engine from it); here it is intentionally read and then discarded. + std::vector ep_context_data; + RETURN_IF_ERROR(ep_context_data_utils::ReadEpContextDataWithFileFallback( + ep->ort_api, ep->ep_context_config_.get(), ep_cache_context.c_str(), ort_graphs[0], + ep_context_data)); + } + // Create EpContextKernel for EPContext nodes - clearly separates from MulKernel ep->ep_context_kernels_.emplace(fused_node_name, std::make_unique(ep->ort_api, ep->logger_)); @@ -449,7 +471,7 @@ OrtStatus* ORT_API_CALL ExampleEp::CompileImpl(_In_ OrtEp* this_ptr, _In_ const // Create EpContext nodes for the fused nodes we compiled (only for Mul, not EPContext). if (ep->config_.enable_ep_context) { assert(ep_context_nodes != nullptr); - RETURN_IF_ERROR(ep->CreateEpContextNodes(gsl::span(fused_nodes, count), + RETURN_IF_ERROR(ep->CreateEpContextNodes(ort_graphs[0], gsl::span(fused_nodes, count), gsl::span(ep_context_nodes, count))); } } @@ -479,7 +501,8 @@ void ORT_API_CALL ExampleEp::ReleaseNodeComputeInfosImpl(OrtEp* this_ptr, // Creates EPContext nodes from the given fused nodes. // This is an example implementation that can be used to generate an EPContext model. However, this example EP // cannot currently run the EPContext model. -OrtStatus* ExampleEp::CreateEpContextNodes(gsl::span fused_nodes, +OrtStatus* ExampleEp::CreateEpContextNodes(const OrtGraph* graph, + gsl::span fused_nodes, /*out*/ gsl::span ep_context_nodes) { try { assert(fused_nodes.size() == ep_context_nodes.size()); @@ -512,11 +535,32 @@ OrtStatus* ExampleEp::CreateEpContextNodes(gsl::span fused_nodes collect_input_output_names(fused_node_outputs, /*out*/ output_names); int64_t is_main_context = (i == 0); - int64_t embed_mode = 1; + int64_t embed_mode = config_.embed_ep_context_in_model ? 1 : 0; // Create node attributes. The CreateNode() function copies the attributes. std::array attributes = {}; - std::string ep_ctx = "binary_data"; + std::string ep_ctx = config_.embed_ep_context_in_model ? "binary_data" : fused_node_name + ".ctx"; + if (!config_.embed_ep_context_in_model) { + const std::string ep_context_data = "binary_data"; + std::string fallback_ep_ctx = ep_ctx; + const OrtGraph* fallback_graph = graph; + if (!config_.ep_context_output_model_path.empty()) { + std::filesystem::path output_model_path; + RETURN_IF_ERROR(ep_context_data_utils::Utf8Path(ort_api, config_.ep_context_output_model_path.c_str(), + output_model_path)); + const std::filesystem::path output_model_dir = output_model_path.parent_path(); + if (!output_model_dir.empty()) { + std::filesystem::path ep_ctx_path; + RETURN_IF_ERROR(ep_context_data_utils::Utf8Path(ort_api, ep_ctx.c_str(), ep_ctx_path)); + RETURN_IF_ERROR(ep_context_data_utils::PathToUtf8String(ort_api, output_model_dir / ep_ctx_path, + fallback_ep_ctx)); + } + fallback_graph = nullptr; + } + RETURN_IF_ERROR(ep_context_data_utils::WriteEpContextDataWithFileFallback( + ort_api, ep_context_config_.get(), ep_ctx.c_str(), fallback_ep_ctx.c_str(), fallback_graph, + ep_context_data.data(), ep_context_data.size())); + } attributes[0] = Ort::OpAttr("ep_cache_context", ep_ctx.data(), static_cast(ep_ctx.size()), ORT_OP_ATTR_STRING); diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/ep.h b/onnxruntime/test/autoep/library/example_plugin_ep/ep.h index 5dcd9f07bef1f..4112abb723d39 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep.h +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep.h @@ -6,6 +6,7 @@ #include #include "../plugin_ep_utils.h" +#include "onnxruntime_experimental_cxx_api.h" class ExampleEpFactory; @@ -61,13 +62,16 @@ class ExampleEp : public OrtEp, public ApiPtrs { public: struct Config { bool enable_ep_context = false; + bool embed_ep_context_in_model = false; bool enable_weightless_ep_context_nodes = false; + std::string ep_context_output_model_path; // Other EP configs (typically extracted from OrtSessionOptions or OrtHardwareDevice(s)) }; - ExampleEp(ExampleEpFactory& factory, const std::string& name, const Config& config, const OrtLogger& logger); + ExampleEp(ExampleEpFactory& factory, const std::string& name, const Config& config, const OrtLogger& logger, + Ort::Experimental::EpContextConfig ep_context_config); - ~ExampleEp(); + ~ExampleEp() = default; std::unordered_map>& MulKernels() { return mul_kernels_; @@ -108,7 +112,8 @@ class ExampleEp : public OrtEp, public ApiPtrs { static OrtStatus* ORT_API_CALL GetDefaultMemoryDeviceImpl(_In_ const OrtEp* this_ptr, _Outptr_ const OrtMemoryDevice** device) noexcept; - OrtStatus* CreateEpContextNodes(gsl::span fused_nodes, + OrtStatus* CreateEpContextNodes(const OrtGraph* graph, + gsl::span fused_nodes, /*out*/ gsl::span ep_context_nodes); // Returns true if the EP should save constant initializers so that they are available during inference. @@ -122,6 +127,7 @@ class ExampleEp : public OrtEp, public ApiPtrs { std::string name_; Config config_{}; const OrtLogger& logger_; + Ort::Experimental::EpContextConfig ep_context_config_; std::unordered_map> mul_kernels_; std::unordered_map> ep_context_kernels_; std::unordered_map float_initializers_; diff --git a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc index 875f70bd29f3c..a323ec0e8c15e 100644 --- a/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc +++ b/onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc @@ -196,6 +196,7 @@ OrtStatus* ORT_API_CALL ExampleEpFactory::CreateEpImpl(OrtEpFactory* this_ptr, const OrtSessionOptions* session_options, const OrtLogger* logger, OrtEp** ep) noexcept { + EXCEPTION_TO_RETURNED_STATUS_BEGIN auto* factory = static_cast(this_ptr); *ep = nullptr; @@ -219,20 +220,33 @@ OrtStatus* ORT_API_CALL ExampleEpFactory::CreateEpImpl(OrtEpFactory* this_ptr, // Create EP configuration from session options, if needed. // Note: should not store a direct reference to the session options object as its lifespan is not guaranteed. std::string ep_context_enable; + std::string ep_context_embed_mode; + std::string ep_context_output_model_path; std::string weightless_ep_context_nodes_enable; RETURN_IF_ERROR(GetSessionConfigEntryOrDefault(*session_options, kOrtSessionOptionEpContextEnable, "0", ep_context_enable)); + RETURN_IF_ERROR(GetSessionConfigEntryOrDefault(*session_options, kOrtSessionOptionEpContextEmbedMode, "0", + ep_context_embed_mode)); + RETURN_IF_ERROR(GetSessionConfigEntryOrDefault(*session_options, kOrtSessionOptionEpContextFilePath, "", + ep_context_output_model_path)); RETURN_IF_ERROR(GetSessionConfigEntryOrDefault(*session_options, kOrtSessionOptionEpEnableWeightlessEpContextNodes, "0", weightless_ep_context_nodes_enable)); ExampleEp::Config config = {}; config.enable_ep_context = ep_context_enable == "1"; + config.embed_ep_context_in_model = ep_context_embed_mode == "1"; + config.ep_context_output_model_path = std::move(ep_context_output_model_path); config.enable_weightless_ep_context_nodes = weightless_ep_context_nodes_enable == "1"; - auto dummy_ep = std::make_unique(*factory, factory->ep_name_, config, *logger); - + // The EpContextConfig wrapper extracts the EPContext callbacks from the session options and owns the handle. It + // throws if the experimental functions are unavailable or extraction fails; EXCEPTION_TO_RETURNED_STATUS_END + // converts that (and any other exception thrown in this function) into an OrtStatus. + auto dummy_ep = std::make_unique( + *factory, factory->ep_name_, config, *logger, + Ort::Experimental::EpContextConfig{Ort::ConstSessionOptions{session_options}}); *ep = dummy_ep.release(); return nullptr; + EXCEPTION_TO_RETURNED_STATUS_END } /*static*/ diff --git a/onnxruntime/test/autoep/test_execution.cc b/onnxruntime/test/autoep/test_execution.cc index 93633f9a375bb..e95918c719324 100644 --- a/onnxruntime/test/autoep/test_execution.cc +++ b/onnxruntime/test/autoep/test_execution.cc @@ -1,9 +1,11 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include #include #include +#include #include // #include #include @@ -12,11 +14,14 @@ #include "core/graph/constants.h" #include "core/graph/onnx_protobuf.h" #include "core/session/onnxruntime_cxx_api.h" +#include "core/session/onnxruntime_experimental_cxx_api.h" #include "core/session/onnxruntime_session_options_config_keys.h" #include "core/session/onnxruntime_ep_device_ep_metadata_keys.h" #include "nlohmann/json.hpp" +#include "test/autoep/ep_context_data_callbacks.h" #include "test/autoep/test_autoep_utils.h" +#include "test/autoep/library/ep_context_data_utils.h" #include "test/autoep/library/example_plugin_ep/ep_test_hooks.h" #include "test/shared_lib/utils.h" #include "test/util/include/api_asserts.h" @@ -29,6 +34,51 @@ namespace test { namespace { +// Invokes the experimental EPContext read setter on the public C API. +void SetEpContextDataReadFunc(Ort::SessionOptions& session_options, OrtReadNamedBufferFunc read_func, void* state) { + auto* set_read_func = + Ort::Experimental::Get_OrtApi_SessionOptions_SetEpContextDataReadFunc_SinceV28_FnOrThrow(&Ort::GetApi()); + ASSERT_ORTSTATUS_OK(set_read_func(session_options, read_func, state)); +} + +// Invokes the experimental EPContext write setter on the public C API. +void SetEpContextDataWriteFunc(Ort::ModelCompilationOptions& compile_options, OrtWriteNamedBufferFunc write_func, + void* state) { + auto* set_write_func = + Ort::Experimental::Get_OrtCompileApi_ModelCompilationOptions_SetEpContextDataWriteFunc_SinceV28_FnOrThrow( + &Ort::GetApi()); + ASSERT_ORTSTATUS_OK(set_write_func(compile_options, write_func, state)); +} + +void LoadModelProtoFromFile(const ORTCHAR_T* model_file, ONNX_NAMESPACE::ModelProto& model_proto) { + std::ifstream model_stream{std::filesystem::path(model_file), std::ios::binary}; + ASSERT_TRUE(model_stream.is_open()); + ASSERT_TRUE(model_proto.ParseFromIstream(&model_stream)); +} + +std::vector GetEpContextNodes(const ONNX_NAMESPACE::ModelProto& model_proto) { + std::vector ep_context_nodes; + + for (const auto& node : model_proto.graph().node()) { + if (node.domain() == kMSDomain && node.op_type() == "EPContext") { + ep_context_nodes.push_back(&node); + } + } + + return ep_context_nodes; +} + +const ONNX_NAMESPACE::AttributeProto* GetNodeAttribute(const ONNX_NAMESPACE::NodeProto& node, + std::string_view attribute_name) { + for (const auto& attribute : node.attribute()) { + if (attribute.name() == attribute_name) { + return &attribute; + } + } + + return nullptr; +} + void RunMulModelWithPluginEp(const ORTCHAR_T* model_path, const Ort::SessionOptions& session_options) { Ort::Session session(*ort_env, model_path, session_options); @@ -521,6 +571,222 @@ TEST(OrtEpLibrary, PluginEp_GenEpContextModel) { } } +TEST(OrtEpLibrary, PluginEp_GenEpContextModel_EmbedModeDoesNotUseCallbacks) { + RegisteredEpDeviceUniquePtr example_ep; + ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); + Ort::ConstEpDevice plugin_ep_device(example_ep.get()); + + const ORTCHAR_T* input_model_file = ORT_TSTR("testdata/mul_1.onnx"); + const ORTCHAR_T* output_model_file = ORT_TSTR("plugin_ep_mul_1_embedded_ctx.onnx"); + std::filesystem::remove(output_model_file); + auto cleanup = gsl::finally([&]() { std::filesystem::remove(output_model_file); }); + + EpContextDataCallbackState write_callback_state; + EpContextDataCallbackState compile_read_callback_state; + { + Ort::SessionOptions session_options; + ASSERT_NO_FATAL_FAILURE( + SetEpContextDataReadFunc(session_options, LoadEpContextDataCallback, &compile_read_callback_state)); + + std::unordered_map ep_options; + session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); + + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + compile_options.SetFlags(OrtCompileApiFlags_ERROR_IF_NO_NODES_COMPILED); + compile_options.SetInputModelPath(input_model_file); + compile_options.SetOutputModelPath(output_model_file); + compile_options.SetEpContextEmbedMode(true); + ASSERT_NO_FATAL_FAILURE( + SetEpContextDataWriteFunc(compile_options, StoreEpContextDataCallback, &write_callback_state)); + + ASSERT_CXX_ORTSTATUS_OK(Ort::CompileModel(*ort_env, compile_options)); + } + + ASSERT_TRUE(std::filesystem::exists(output_model_file)); + EXPECT_FALSE(write_callback_state.write_called); + EXPECT_FALSE(compile_read_callback_state.read_called); + + ONNX_NAMESPACE::ModelProto compiled_model; + ASSERT_NO_FATAL_FAILURE(LoadModelProtoFromFile(output_model_file, compiled_model)); + + auto ep_context_nodes = GetEpContextNodes(compiled_model); + ASSERT_EQ(ep_context_nodes.size(), 1u); + + const ONNX_NAMESPACE::AttributeProto* embed_mode_attr = GetNodeAttribute(*ep_context_nodes[0], "embed_mode"); + ASSERT_NE(embed_mode_attr, nullptr); + EXPECT_EQ(embed_mode_attr->type(), ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + EXPECT_EQ(embed_mode_attr->i(), 1); + + const ONNX_NAMESPACE::AttributeProto* ep_cache_context_attr = GetNodeAttribute(*ep_context_nodes[0], + "ep_cache_context"); + ASSERT_NE(ep_cache_context_attr, nullptr); + EXPECT_EQ(ep_cache_context_attr->type(), ONNX_NAMESPACE::AttributeProto_AttributeType_STRING); + EXPECT_EQ(ep_cache_context_attr->s(), "binary_data"); + + EpContextDataCallbackState load_read_callback_state; + { + Ort::SessionOptions session_options; + ASSERT_NO_FATAL_FAILURE( + SetEpContextDataReadFunc(session_options, LoadEpContextDataCallback, &load_read_callback_state)); + + std::unordered_map ep_options; + session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); + + Ort::Session session(*ort_env, output_model_file, session_options); + } + + EXPECT_FALSE(load_read_callback_state.read_called); +} + +TEST(OrtEpLibrary, PluginEp_GenAndLoadEpContextModel_ExternalDataUsesFileFallback) { + RegisteredEpDeviceUniquePtr example_ep; + ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); + Ort::ConstEpDevice plugin_ep_device(example_ep.get()); + + const ORTCHAR_T* input_model_file = ORT_TSTR("testdata/mul_1.onnx"); + const ORTCHAR_T* output_model_file = ORT_TSTR("plugin_ep_mul_1_file_ctx.onnx"); + std::vector files_to_cleanup{std::filesystem::path{output_model_file}}; + for (const auto& path : files_to_cleanup) { + std::filesystem::remove(path); + } + auto cleanup = gsl::finally([&]() { + for (const auto& path : files_to_cleanup) { + std::filesystem::remove(path); + } + }); + + { + Ort::SessionOptions session_options; + std::unordered_map ep_options; + session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); + + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + compile_options.SetFlags(OrtCompileApiFlags_ERROR_IF_NO_NODES_COMPILED); + compile_options.SetInputModelPath(input_model_file); + compile_options.SetOutputModelPath(output_model_file); + compile_options.SetEpContextEmbedMode(false); + + ASSERT_CXX_ORTSTATUS_OK(Ort::CompileModel(*ort_env, compile_options)); + } + + ASSERT_TRUE(std::filesystem::exists(output_model_file)); + + ONNX_NAMESPACE::ModelProto compiled_model; + ASSERT_NO_FATAL_FAILURE(LoadModelProtoFromFile(output_model_file, compiled_model)); + + auto ep_context_nodes = GetEpContextNodes(compiled_model); + ASSERT_EQ(ep_context_nodes.size(), 1u); + + const ONNX_NAMESPACE::AttributeProto* embed_mode_attr = GetNodeAttribute(*ep_context_nodes[0], "embed_mode"); + ASSERT_NE(embed_mode_attr, nullptr); + EXPECT_EQ(embed_mode_attr->type(), ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + EXPECT_EQ(embed_mode_attr->i(), 0); + + const ONNX_NAMESPACE::AttributeProto* ep_cache_context_attr = GetNodeAttribute(*ep_context_nodes[0], + "ep_cache_context"); + ASSERT_NE(ep_cache_context_attr, nullptr); + EXPECT_EQ(ep_cache_context_attr->type(), ONNX_NAMESPACE::AttributeProto_AttributeType_STRING); + ASSERT_FALSE(ep_cache_context_attr->s().empty()); + + const std::filesystem::path output_model_dir = std::filesystem::path{output_model_file}.parent_path(); + std::filesystem::path ep_cache_context_rel; + ASSERT_ORTSTATUS_OK( + ep_context_data_utils::Utf8Path(Ort::GetApi(), ep_cache_context_attr->s().c_str(), ep_cache_context_rel)); + const std::filesystem::path context_data_path = output_model_dir / ep_cache_context_rel; + files_to_cleanup.push_back(context_data_path); + ASSERT_TRUE(std::filesystem::exists(context_data_path)); + + std::vector context_data; + std::string context_data_file_name; + ASSERT_ORTSTATUS_OK(ep_context_data_utils::PathToUtf8String(Ort::GetApi(), context_data_path, + context_data_file_name)); + ASSERT_ORTSTATUS_OK(ep_context_data_utils::ReadEpContextDataFromFile(Ort::GetApi(), context_data_file_name.c_str(), + nullptr, context_data)); + EXPECT_EQ(std::string(context_data.begin(), context_data.end()), "binary_data"); + + { + Ort::SessionOptions session_options; + std::unordered_map ep_options; + session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); + + Ort::Session session(*ort_env, output_model_file, session_options); + } +} + +TEST(OrtEpLibrary, PluginEp_GenEpContextModel_ExternalDataUsesWriteCallback) { + RegisteredEpDeviceUniquePtr example_ep; + ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); + Ort::ConstEpDevice plugin_ep_device(example_ep.get()); + + const ORTCHAR_T* input_model_file = ORT_TSTR("testdata/mul_1.onnx"); + const ORTCHAR_T* output_model_file = ORT_TSTR("plugin_ep_mul_1_external_ctx.onnx"); + std::filesystem::remove(output_model_file); + auto cleanup = gsl::finally([&]() { std::filesystem::remove(output_model_file); }); + + Ort::SessionOptions session_options; + std::unordered_map ep_options; + session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); + + EpContextDataCallbackState callback_state; + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + compile_options.SetFlags(OrtCompileApiFlags_ERROR_IF_NO_NODES_COMPILED); + compile_options.SetInputModelPath(input_model_file); + compile_options.SetOutputModelPath(output_model_file); + compile_options.SetEpContextEmbedMode(false); + ASSERT_NO_FATAL_FAILURE(SetEpContextDataWriteFunc(compile_options, StoreEpContextDataCallback, &callback_state)); + + ASSERT_CXX_ORTSTATUS_OK(Ort::CompileModel(*ort_env, compile_options)); + ASSERT_TRUE(std::filesystem::exists(output_model_file)); + ASSERT_TRUE(callback_state.write_called); + EXPECT_FALSE(callback_state.write_file_name.empty()); + EXPECT_EQ(std::string(callback_state.payload.begin(), callback_state.payload.end()), "binary_data"); +} + +TEST(OrtEpLibrary, PluginEp_LoadEpContextModel_ExternalDataUsesReadCallback) { + RegisteredEpDeviceUniquePtr example_ep; + ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); + Ort::ConstEpDevice plugin_ep_device(example_ep.get()); + + const ORTCHAR_T* input_model_file = ORT_TSTR("testdata/mul_1.onnx"); + const ORTCHAR_T* compiled_model_file = ORT_TSTR("plugin_ep_mul_1_external_ctx_load.onnx"); + std::filesystem::remove(compiled_model_file); + auto cleanup = gsl::finally([&]() { std::filesystem::remove(compiled_model_file); }); + + EpContextDataCallbackState write_callback_state; + { + Ort::SessionOptions session_options; + std::unordered_map ep_options; + session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); + + Ort::ModelCompilationOptions compile_options(*ort_env, session_options); + compile_options.SetFlags(OrtCompileApiFlags_ERROR_IF_NO_NODES_COMPILED); + compile_options.SetInputModelPath(input_model_file); + compile_options.SetOutputModelPath(compiled_model_file); + compile_options.SetEpContextEmbedMode(false); + ASSERT_NO_FATAL_FAILURE( + SetEpContextDataWriteFunc(compile_options, StoreEpContextDataCallback, &write_callback_state)); + + ASSERT_CXX_ORTSTATUS_OK(Ort::CompileModel(*ort_env, compile_options)); + ASSERT_TRUE(std::filesystem::exists(compiled_model_file)); + ASSERT_TRUE(write_callback_state.write_called); + } + + EpContextDataCallbackState read_callback_state; + read_callback_state.payload = write_callback_state.payload; + { + Ort::SessionOptions session_options; + ASSERT_NO_FATAL_FAILURE(SetEpContextDataReadFunc(session_options, LoadEpContextDataCallback, &read_callback_state)); + + std::unordered_map ep_options; + session_options.AppendExecutionProvider_V2(*ort_env, {plugin_ep_device}, ep_options); + + Ort::Session session(*ort_env, compiled_model_file, session_options); + } + + ASSERT_TRUE(read_callback_state.read_called); + EXPECT_EQ(read_callback_state.read_file_name, write_callback_state.write_file_name); +} + TEST(OrtEpLibrary, PluginEp_GenWeightlessEpContextModel) { RegisteredEpDeviceUniquePtr example_ep; ASSERT_NO_FATAL_FAILURE(Utils::RegisterAndGetExampleEp(*ort_env, Utils::example_ep_info, example_ep)); diff --git a/onnxruntime/test/autoep/test_model_package.cc b/onnxruntime/test/autoep/test_model_package.cc index ee5c8bb567e1e..c9065878da1ce 100644 --- a/onnxruntime/test/autoep/test_model_package.cc +++ b/onnxruntime/test/autoep/test_model_package.cc @@ -621,6 +621,11 @@ TEST(ModelPackageTest, CheckCompiledModelCompatibilityInfo) { compile_options.SetInputModelPath(input_model_file); compile_options.SetOutputModelPath(output_model_file); + // Embed the EPContext binary data inside the compiled model so the model is self-contained. + // This test copies only the compiled .onnx into the model package, so it must not rely on a + // separate sidecar EPContext data file (which non-embedded mode would produce). + compile_options.SetEpContextEmbedMode(true); + ASSERT_CXX_ORTSTATUS_OK(Ort::CompileModel(*ort_env, compile_options)); ASSERT_TRUE(std::filesystem::exists(output_model_file)); } diff --git a/onnxruntime/test/framework/ep_plugin_provider_test.cc b/onnxruntime/test/framework/ep_plugin_provider_test.cc index fef185e20f341..4820f4e5c8898 100644 --- a/onnxruntime/test/framework/ep_plugin_provider_test.cc +++ b/onnxruntime/test/framework/ep_plugin_provider_test.cc @@ -7,12 +7,14 @@ #include #include #include +#include #include #include #include "gsl/gsl" #include "gtest/gtest.h" #include "core/common/logging/sinks/file_sink.h" +#include "core/common/path_string.h" #include "core/framework/config_options.h" #include "core/framework/kernel_def_builder.h" #include "core/framework/op_kernel.h" diff --git a/onnxruntime/test/shared_lib/test_ep_context_data_api.cc b/onnxruntime/test/shared_lib/test_ep_context_data_api.cc new file mode 100644 index 0000000000000..ec8107f92aa7a --- /dev/null +++ b/onnxruntime/test/shared_lib/test_ep_context_data_api.cc @@ -0,0 +1,331 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include + +#include "core/session/onnxruntime_c_api.h" +#include "core/session/onnxruntime_cxx_api.h" +#include "core/session/onnxruntime_experimental_cxx_api.h" + +#include "gmock/gmock.h" +#include "gsl/gsl" +#include "gtest/gtest.h" +#include "test/util/include/api_asserts.h" + +namespace { + +void ExpectFailureOrtStatus(OrtStatus* status_ptr, OrtErrorCode expected_code, const char* expected_message) { + Ort::Status status{status_ptr}; + ASSERT_NE(status_ptr, nullptr) << "Expected a failure status, but the API returned nullptr (OK)."; + ASSERT_FALSE(status.IsOK()); + EXPECT_EQ(status.GetErrorCode(), expected_code); + EXPECT_THAT(status.GetErrorMessage(), ::testing::HasSubstr(expected_message)); +} + +struct EpContextReadCallbackState { + bool called = false; + std::string file_name; + std::vector payload; +}; + +OrtStatus* ORT_API_CALL EpContextReadCallback(void* state, const char* file_name, OrtAllocator* allocator, + void** buffer, size_t* data_size) { + auto* read_state = static_cast(state); + read_state->called = true; + read_state->file_name = file_name; + + *buffer = nullptr; + *data_size = read_state->payload.size(); + + if (read_state->payload.empty()) { + return nullptr; + } + + OrtStatus* status = Ort::GetApi().AllocatorAlloc(allocator, read_state->payload.size(), buffer); + if (status != nullptr) { + return status; + } + + std::memcpy(*buffer, read_state->payload.data(), read_state->payload.size()); + return nullptr; +} + +struct EpContextWriteCallbackState { + bool called = false; + std::string file_name; + std::vector payload; +}; + +OrtStatus* ORT_API_CALL EpContextWriteCallback(void* state, const char* file_name, const void* buffer, + size_t buffer_size) { + auto* write_state = static_cast(state); + write_state->called = true; + write_state->file_name = file_name; + write_state->payload.clear(); + if (buffer_size != 0) { + if (buffer == nullptr) { + return Ort::GetApi().CreateStatus(ORT_INVALID_ARGUMENT, + "EpContextWriteCallback received a null buffer for non-empty data"); + } + + const char* buffer_bytes = static_cast(buffer); + write_state->payload.assign(buffer_bytes, buffer_bytes + buffer_size); + } + + return nullptr; +} + +} // namespace + +TEST(EpContextDataApiTest, ReadFuncIsReturnedByEpApi) { + const auto& ort_api = Ort::GetApi(); + Ort::SessionOptions session_options; + + auto* set_read_func = + Ort::Experimental::Get_OrtApi_SessionOptions_SetEpContextDataReadFunc_SinceV28_FnOrThrow(&ort_api); + + EpContextReadCallbackState callback_state{ + false, + {}, + {'e', 'p', 'c', 't', 'x'}, + }; + ASSERT_ORTSTATUS_OK(set_read_func(session_options, EpContextReadCallback, &callback_state)); + + Ort::Experimental::EpContextConfig ep_context_config{session_options}; + OrtReadNamedBufferFunc read_func = nullptr; + void* callback_state_out = nullptr; + ep_context_config.GetReadFunc(read_func, callback_state_out); + ASSERT_EQ(read_func, EpContextReadCallback); + ASSERT_EQ(callback_state_out, &callback_state); + + Ort::AllocatorWithDefaultOptions allocator; + void* buffer = nullptr; + size_t buffer_size = 0; + ASSERT_ORTSTATUS_OK(read_func(callback_state_out, "context.bin", allocator, &buffer, &buffer_size)); + auto release_buffer = gsl::finally([&]() { + if (buffer != nullptr) { + allocator.Free(buffer); + } + }); + + ASSERT_TRUE(callback_state.called); + EXPECT_EQ(callback_state.file_name, "context.bin"); + ASSERT_EQ(buffer_size, callback_state.payload.size()); + EXPECT_TRUE(std::equal(callback_state.payload.begin(), callback_state.payload.end(), + static_cast(buffer))); +} + +TEST(EpContextDataApiTest, ApiRejectsInvalidArguments) { + const auto& ort_api = Ort::GetApi(); + + auto* get_config = Ort::Experimental::Get_OrtEpApi_SessionOptions_GetEpContextConfig_SinceV28_FnOrThrow(&ort_api); + auto* release_config_func = + Ort::Experimental::Get_OrtEpApi_ReleaseEpContextConfig_SinceV28_FnOrThrow(&ort_api); + auto* get_read_func = + Ort::Experimental::Get_OrtEpApi_EpContextConfig_GetEpContextDataReadFunc_SinceV28_FnOrThrow(&ort_api); + auto* get_write_func = + Ort::Experimental::Get_OrtEpApi_EpContextConfig_GetEpContextDataWriteFunc_SinceV28_FnOrThrow(&ort_api); + auto* set_read_func = + Ort::Experimental::Get_OrtApi_SessionOptions_SetEpContextDataReadFunc_SinceV28_FnOrThrow(&ort_api); + + Ort::SessionOptions session_options; + OrtEpContextConfig* ep_context_config = nullptr; + ExpectFailureOrtStatus(get_config(nullptr, &ep_context_config), ORT_INVALID_ARGUMENT, "OrtSessionOptions is NULL"); + ExpectFailureOrtStatus(get_config(session_options, nullptr), ORT_INVALID_ARGUMENT, + "Output OrtEpContextConfig is NULL"); + + ExpectFailureOrtStatus(set_read_func(nullptr, EpContextReadCallback, nullptr), ORT_INVALID_ARGUMENT, + "'options' parameter must not be NULL"); + + ASSERT_ORTSTATUS_OK(get_config(session_options, &ep_context_config)); + auto release_config = gsl::finally([&]() { release_config_func(ep_context_config); }); + + OrtReadNamedBufferFunc read_func = nullptr; + OrtWriteNamedBufferFunc write_func = nullptr; + void* state = nullptr; + ExpectFailureOrtStatus(get_read_func(nullptr, &read_func, &state), ORT_INVALID_ARGUMENT, + "OrtEpContextConfig is NULL"); + ExpectFailureOrtStatus(get_read_func(ep_context_config, nullptr, &state), ORT_INVALID_ARGUMENT, + "Output read_func is NULL"); + ExpectFailureOrtStatus(get_read_func(ep_context_config, &read_func, nullptr), ORT_INVALID_ARGUMENT, + "Output state is NULL"); + ExpectFailureOrtStatus(get_write_func(nullptr, &write_func, &state), ORT_INVALID_ARGUMENT, + "OrtEpContextConfig is NULL"); + ExpectFailureOrtStatus(get_write_func(ep_context_config, nullptr, &state), ORT_INVALID_ARGUMENT, + "Output write_func is NULL"); + ExpectFailureOrtStatus(get_write_func(ep_context_config, &write_func, nullptr), ORT_INVALID_ARGUMENT, + "Output state is NULL"); + +#if !defined(ORT_MINIMAL_BUILD) + Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "EpContextDataApiRejectsInvalidArguments"}; + Ort::ModelCompilationOptions compilation_options{env, session_options}; + auto* set_write_func = + Ort::Experimental::Get_OrtCompileApi_ModelCompilationOptions_SetEpContextDataWriteFunc_SinceV28_FnOrThrow( + &ort_api); + ExpectFailureOrtStatus(set_write_func(nullptr, EpContextWriteCallback, nullptr), ORT_INVALID_ARGUMENT, + "OrtModelCompilationOptions is NULL"); + // A null write_func is allowed: it clears any previously set callback (covered by WriteFuncCanBeCleared), so it is + // not rejected here. +#endif // !defined(ORT_MINIMAL_BUILD) +} + +TEST(EpContextDataApiTest, AccessorsReturnNullWhenCallbacksUnset) { + Ort::SessionOptions session_options; + Ort::Experimental::EpContextConfig ep_context_config{session_options}; + + OrtReadNamedBufferFunc read_func = EpContextReadCallback; + OrtWriteNamedBufferFunc write_func = EpContextWriteCallback; + void* state = reinterpret_cast(0x1); + + ep_context_config.GetReadFunc(read_func, state); + EXPECT_EQ(read_func, nullptr); + EXPECT_EQ(state, nullptr); + + state = reinterpret_cast(0x1); + ep_context_config.GetWriteFunc(write_func, state); + EXPECT_EQ(write_func, nullptr); + EXPECT_EQ(state, nullptr); +} + +TEST(EpContextDataApiTest, ConfigReturnsConfiguredCallbacks) { + const auto& ort_api = Ort::GetApi(); + Ort::SessionOptions session_options; + + auto* set_read_func = + Ort::Experimental::Get_OrtApi_SessionOptions_SetEpContextDataReadFunc_SinceV28_FnOrThrow(&ort_api); + + EpContextReadCallbackState callback_state{}; + ASSERT_ORTSTATUS_OK(set_read_func(session_options, EpContextReadCallback, &callback_state)); + + Ort::Experimental::EpContextConfig ep_context_config{session_options}; + + OrtReadNamedBufferFunc read_func = nullptr; + void* read_state = nullptr; + ep_context_config.GetReadFunc(read_func, read_state); + EXPECT_EQ(read_func, EpContextReadCallback); + EXPECT_EQ(read_state, &callback_state); + + OrtWriteNamedBufferFunc write_func = nullptr; + void* write_state = nullptr; + ep_context_config.GetWriteFunc(write_func, write_state); + EXPECT_EQ(write_func, nullptr); + EXPECT_EQ(write_state, nullptr); +} + +TEST(EpContextDataApiTest, ReadFuncCanBeCleared) { + const auto& ort_api = Ort::GetApi(); + Ort::SessionOptions session_options; + + auto* set_read_func = + Ort::Experimental::Get_OrtApi_SessionOptions_SetEpContextDataReadFunc_SinceV28_FnOrThrow(&ort_api); + + EpContextReadCallbackState callback_state{}; + ASSERT_ORTSTATUS_OK(set_read_func(session_options, EpContextReadCallback, &callback_state)); + + ASSERT_ORTSTATUS_OK(set_read_func(session_options, nullptr, &callback_state)); + + Ort::Experimental::EpContextConfig ep_context_config{session_options}; + OrtReadNamedBufferFunc read_func = EpContextReadCallback; + void* read_state = reinterpret_cast(0x1); + ep_context_config.GetReadFunc(read_func, read_state); + EXPECT_EQ(read_func, nullptr); + EXPECT_EQ(read_state, nullptr); +} + +#if !defined(ORT_MINIMAL_BUILD) +TEST(EpContextDataApiTest, WriteFuncCanBeSetOnModelCompilationOptions) { + const auto& ort_api = Ort::GetApi(); + Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "EpContextDataWriteFuncCanBeSetOnModelCompilationOptions"}; + Ort::SessionOptions session_options; + Ort::ModelCompilationOptions compilation_options{env, session_options}; + + auto* set_write_func = + Ort::Experimental::Get_OrtCompileApi_ModelCompilationOptions_SetEpContextDataWriteFunc_SinceV28_FnOrThrow( + &ort_api); + + EpContextWriteCallbackState callback_state{}; + ASSERT_ORTSTATUS_OK(set_write_func(compilation_options, EpContextWriteCallback, &callback_state)); + + const std::vector payload{'b', 'i', 'n', 'a', 'r', 'y'}; + ASSERT_ORTSTATUS_OK(EpContextWriteCallback(&callback_state, "engine.bin", payload.data(), payload.size())); + + ASSERT_TRUE(callback_state.called); + EXPECT_EQ(callback_state.file_name, "engine.bin"); + EXPECT_EQ(callback_state.payload, payload); +} + +TEST(EpContextDataApiTest, WriteFuncCanBeCleared) { + const auto& ort_api = Ort::GetApi(); + Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "EpContextDataWriteFuncCanBeCleared"}; + Ort::SessionOptions session_options; + Ort::ModelCompilationOptions compilation_options{env, session_options}; + + auto* set_write_func = + Ort::Experimental::Get_OrtCompileApi_ModelCompilationOptions_SetEpContextDataWriteFunc_SinceV28_FnOrThrow( + &ort_api); + + EpContextWriteCallbackState callback_state{}; + ASSERT_ORTSTATUS_OK(set_write_func(compilation_options, EpContextWriteCallback, &callback_state)); + + // A null write_func clears the previously set callback (symmetric with the read setter) and must be accepted + // rather than rejected with ORT_INVALID_ARGUMENT. + ASSERT_ORTSTATUS_OK(set_write_func(compilation_options, nullptr, &callback_state)); +} + +TEST(EpContextDataApiTest, WriteFuncCanBeUsedWithEpContextBinaryInformation) { + const auto& ort_api = Ort::GetApi(); + Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "EpContextDataWriteFuncCanBeUsedWithEpContextBinaryInformation"}; + Ort::SessionOptions session_options; + Ort::ModelCompilationOptions compilation_options{env, session_options}; + + auto* set_write_func = + Ort::Experimental::Get_OrtCompileApi_ModelCompilationOptions_SetEpContextDataWriteFunc_SinceV28_FnOrThrow( + &ort_api); + + // The EPContext write callback and the EPContext binary information may be configured together; neither call + // rejects the other. + ASSERT_NO_THROW(compilation_options.SetEpContextBinaryInformation(ORT_TSTR("ep_context_dir/"), + ORT_TSTR("compiled_model.onnx"))); + + EpContextWriteCallbackState callback_state{}; + ASSERT_ORTSTATUS_OK(set_write_func(compilation_options, EpContextWriteCallback, &callback_state)); + + const std::vector payload{'c', 't', 'x'}; + ASSERT_ORTSTATUS_OK(EpContextWriteCallback(&callback_state, "logical_context.bin", payload.data(), payload.size())); + + ASSERT_TRUE(callback_state.called); + EXPECT_EQ(callback_state.file_name, "logical_context.bin"); + EXPECT_EQ(callback_state.payload, payload); +} +#endif // !defined(ORT_MINIMAL_BUILD) + +TEST(EpContextDataApiTest, ReturnedReadFuncAllowsEmptyPayloads) { + const auto& ort_api = Ort::GetApi(); + Ort::SessionOptions session_options; + + auto* set_read_func = + Ort::Experimental::Get_OrtApi_SessionOptions_SetEpContextDataReadFunc_SinceV28_FnOrThrow(&ort_api); + + EpContextReadCallbackState callback_state{}; + ASSERT_ORTSTATUS_OK(set_read_func(session_options, EpContextReadCallback, &callback_state)); + + Ort::Experimental::EpContextConfig ep_context_config{session_options}; + OrtReadNamedBufferFunc read_func = nullptr; + void* read_state = nullptr; + ep_context_config.GetReadFunc(read_func, read_state); + ASSERT_EQ(read_func, EpContextReadCallback); + ASSERT_EQ(read_state, &callback_state); + + Ort::AllocatorWithDefaultOptions allocator; + void* buffer = reinterpret_cast(0x1); + size_t buffer_size = 1; + ASSERT_ORTSTATUS_OK(read_func(read_state, "empty.bin", allocator, &buffer, &buffer_size)); + + EXPECT_TRUE(callback_state.called); + EXPECT_EQ(callback_state.file_name, "empty.bin"); + EXPECT_EQ(buffer, nullptr); + EXPECT_EQ(buffer_size, 0U); +}