Repository navigation
Add EP-agnostic callbacks for persisting EP caches - #34013
Open
Danny On The Air (dannyota) wants to merge 1 commit into
Open
Danny On The Air (dannyota) wants to merge 1 commit into
Danny On The Air (dannyota) wants to merge 1 commit into
Conversation
EPs keep persistent caches (JIT kernels, compiled engines, timing data) under filesystem paths, which sandboxed hosts cannot use. The EPContext read/write callbacks already let applications own model artifacts, but not these optional, performance-only caches. Add SessionOptionsSetEpCacheDataFuncs, reusing OrtReadNamedBufferFunc and OrtWriteNamedBufferFunc. EPs read the callbacks through OrtEpApi::SessionOptionsGetEpCacheDataFuncs and advertise READ/WRITE support through an optional OrtEp::GetEpCacheDataCallbackSupport; applications query it with SessionGetEpCacheDataCallbackSupport. When a callback is set and the EP supports it, the EP uses callback-only storage: read failures are cache misses and write failures are logged. Compilation inherits the callbacks through the copied session options. The example plugin EP persists a versioned Mul plan to show the flow. Signed-off-by: Danny <dannyontheair@outlook.com>
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Contributor
There was a problem hiding this comment.
🟢 Approval recommended
There are no active review issues for the application-managed execution-provider cache changes.
No open issues 🎉
What changed in this PR
Adds application-managed persistent caches for execution providers (EPs), addressing #33748. The APIs let hosts supply storage callbacks instead of filesystem paths while keeping EPContext behavior unchanged.
Changes:
- Adds C APIs and C++ wrappers for cache callbacks and per-session capability queries.
- Stores callbacks in session options and adds version-gated plugin support.
- Demonstrates validated cache reuse in the example EP, with tests for failures, compatibility, compilation inheritance, and concurrent sessions.
| File | Description |
|---|---|
onnxruntime/test/autoep/test_execution.cc |
Adds cache lifecycle, failure, and compatibility tests. |
onnxruntime/test/autoep/test_autoep_utils.h |
Declares cache test-hook pointers. |
onnxruntime/test/autoep/test_autoep_utils.cc |
Loads cache test hooks on Windows and POSIX. |
onnxruntime/test/autoep/library/example_plugin_ep/example_plugin_ep_library.lds |
Exports cache test hooks. |
onnxruntime/test/autoep/library/example_plugin_ep/ep.h |
Declares cache configuration and plan helpers. |
onnxruntime/test/autoep/library/example_plugin_ep/ep.cc |
Implements bounded, validated Mul-plan persistence. |
onnxruntime/test/autoep/library/example_plugin_ep/ep_test_hooks.h |
Declares cache instrumentation hooks. |
onnxruntime/test/autoep/library/example_plugin_ep/ep_test_hooks.cc |
Implements atomic cache counters. |
onnxruntime/test/autoep/library/example_plugin_ep/ep_factory.cc |
Retrieves callbacks and configures compatibility fixtures. |
onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.h |
Declares the plugin capability-query override. |
onnxruntime/core/session/plugin_ep/ep_plugin_provider_interfaces.cc |
Gates capability queries by version and validates flags. |
onnxruntime/core/session/plugin_ep/ep_api.h |
Declares the callback getter implementation. |
onnxruntime/core/session/plugin_ep/ep_api.cc |
Implements and registers the callback getter. |
onnxruntime/core/session/ort_apis.h |
Declares application-facing API implementations. |
onnxruntime/core/session/onnxruntime_c_api.cc |
Implements session capability queries and registers APIs. |
onnxruntime/core/session/abi_session_options.cc |
Implements callback replacement and clearing. |
onnxruntime/core/framework/session_options.h |
Stores callbacks and shared application state. |
include/onnxruntime/core/session/onnxruntime_ep_c_api.h |
Defines and documents EP-facing cache APIs. |
include/onnxruntime/core/session/onnxruntime_cxx_inline.h |
Implements C++ wrappers. |
include/onnxruntime/core/session/onnxruntime_cxx_api.h |
Declares C++ wrappers. |
include/onnxruntime/core/session/onnxruntime_c_api.h |
Defines support flags and public cache contracts. |
include/onnxruntime/core/framework/execution_provider.h |
Adds a default unsupported capability response. |
Review effort: Balanced
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Implements the proposal in #33748: application-supplied callbacks for an EP's persistent caches (JIT kernel caches, compiled-engine caches, timing caches), as an alternative to filesystem paths.
OrtApi::SessionOptionsSetEpCacheDataFuncs(options, read_func, write_func, state)reusesOrtReadNamedBufferFunc/OrtWriteNamedBufferFunc. Each call replaces the tuple; NULL clears a direction.OrtEpApi::SessionOptionsGetEpCacheDataFuncsand advertise READ/WRITE support through an optionalOrtEp::GetEpCacheDataCallbackSupport(version-gated). Applications query support after session creation withOrtApi::SessionGetEpCacheDataCallbackSupport.Follow-ups: preflight/strict capability enforcement, Python/C# bindings, and adoption by hardware EPs.
Testing (macOS arm64 and Linux x64,
--build_shared_lib)onnxruntime_autoep_test --gtest_filter='OrtEpLibrary.*': 79 pass, including newEpCacheOptions,EpCacheRoundTrip,EpCacheFailuresandEpCacheCompatibility(cold store and warm reuse, corruption and size limits, one-sided callbacks, compilation inheritance, concurrent sessions, older-EP version gating)..ldsexports the new test hooks; same 79 pass.Motivation and Context
Fixes #33748. Sandboxed hosts cannot give EPs a writable cache directory; this lets them keep EP caches in their own storage.