Repository navigation
Add Python bindings for EPContext data callbacks - #32271
Merged
Tianlei Wu (tianleiwu) merged 26 commits intoOct 6, 2026
Merged
Conversation
added 2 commits
August 25, 2026 14:47
Promote named-buffer EPContext callbacks to the stable C and C++ APIs with a required finite read-size policy. Add immutable EP config snapshots, strict participating-EP support negotiation, V28 compatibility shims, hardened sample helpers, example plugin adoption, documentation, and focused API/auto-EP tests.
Mark the strict-read flag unused when ONNX partitioning is excluded, and preserve the V28 EPContext config shims directly in minimal builds where the plugin EP API implementation is intentionally not linked.
Gopalakrishnan Nallasamy (GopalakrishnanN)
force-pushed
the
GopalakrishnanN/python-epcontext-data-io
branch
from
August 26, 2026 07:11
cd0a987 to
c4412a7
Compare
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Gopalakrishnan Nallasamy (GopalakrishnanN)
requested a review
from Edward Chen (edgchen1)
August 28, 2026 02:52
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Gopalakrishnan Nallasamy (GopalakrishnanN)
force-pushed
the
GopalakrishnanN/python-epcontext-data-io
branch
from
August 29, 2026 07:25
c4412a7 to
c59803f
Compare
Use opaque read options consistently, clarify application/EP size enforcement, update version-31 callback support negotiation, and cover ownership and older-plugin boundaries. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…untime/contrib_ops/cuda/moe/moe_quantization.cc. (#32932) ### Description <!-- Describe your changes. --> Replace old use of `routing_record` with `routing_snapshot_` in onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc. ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> Fix build error about undeclared identifier `routing_record`. Likely an unlucky merge artifact from two recent PRs #32761 and #32738. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Increase the minimal Android section-size budget by 2 KiB. CI measured 1,598,685 bytes versus the prior 1,597,440-byte budget; the new 1,599,488-byte limit retains 803 bytes of headroom. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Implement review option 3: remove the proposed read-options APIs and runtime size-limit storage, keep callback/state snapshots and strict support negotiation, and leave allocation/deserialization limits to applications and EPs. Retain an EP-selected bound in the sample helper with boundary coverage. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Clarify the external-data lifecycle query, simplify callback flag documentation, keep allocation-policy guidance in the design document, and require explicit C++ callback state arguments. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Remove native size-policy fields and verify the retained buffer limit against the finalized API 31 transport. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Check final current-provider assignments as well as newly accepted capabilities, including resource-accounted singleton survivors omitted from the final capability list. Keep placement and accounting unchanged and cover ONNX/ORT loads, supported callbacks, and excluded claims. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Resolve the Android workflow and MoE expert-counting conflicts using the current main versions. Preserve the stable EPContext APIs and retained-NHWC READ support fix. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Use effective provider settings to negotiate callback support before every capability-discovery path, including AOT inlining, NHWC passes, and ORT-format loading. TensorRT external output cannot bypass WRITE negotiation through ORT's embed setting, and QNN capability-time context reads require READ support before backend setup. Preserve post-assignment validation and no-callback behavior. Add real-I/O regressions and control cases. Ten baseline failures reproduce the gaps; all 162 targeted framework tests pass after the fix. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The compiler-side read callback path is retained but never invoked by the tests, leaving its state forwarding unverified.
Review effort: Balanced
Findings: 1
Open (1)
What changed in this PR
Adds Python bindings for bounded EPContext read/write callbacks with safe lifetimes, status propagation, and reentrancy handling.
Changes:
- Adds allocator-backed callback data views and size enforcement.
- Retains callback registrations across session/compiler snapshots.
- Adds callback failure, lifetime, concurrency, and ownership tests.
| File | Description |
|---|---|
onnxruntime/__init__.py |
Exports EPContext data types. |
onnxruntime/python/onnxruntime_inference_collection.py |
Adds compiler callback methods. |
onnxruntime/python/onnxruntime_pybind_model_compiler.cc |
Implements callbacks, buffers, GIL handling, and compilation guards. |
onnxruntime/python/onnxruntime_pybind_model_compiler.h |
Declares callback types and retained state. |
onnxruntime/python/onnxruntime_pybind_state.cc |
Exposes bindings and snapshots callback registrations. |
onnxruntime/python/onnxruntime_pybind_state_common.h |
Adds retained session callback state and initialization guards. |
onnxruntime/test/python/onnxruntime_test_python_compile_api.py |
Tests callback behavior and lifetime. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Gopalakrishnan Nallasamy (GopalakrishnanN)
requested a balanced review from Copilot
October 5, 2026 17:32
Copilot started reviewing on behalf of
Gopalakrishnan Nallasamy (GopalakrishnanN)
October 5, 2026 17:34
Contributor
There was a problem hiding this comment.
Copilot review overview
🔵 Needs a closer look
Cross-language callback lifetimes, allocator ownership, and GIL/concurrency behavior warrant final human review despite focused coverage.
Review effort: Balanced
Findings: None
Resolved since last review (1)
Tianlei Wu (tianleiwu)
left a comment
Tianlei Wu (tianleiwu)
left a comment
Contributor
There was a problem hiding this comment.
Full static review of fb6c867 against this PR's stacked base: no actionable findings. Reviewed native callback/allocator lifetimes, retained-view invalidation, registration snapshots, GIL and concurrency/reentrancy, exception propagation, minimal-build/platform guards, and the added tests. The previously reported compiler-read coverage gap is fixed; I read the author reply and did not duplicate that resolved concern.
CI: Optional Lint, Generate Python API docs, Python format, Optional Lint C++, lintrunner and license/cla succeed; Linux Android Emulator QNN CI Pipeline is neutral. No failing or pending checks are reported, but the rollup contains no Python/native build-and-test job, so this is not a CI-backed approval.
Integration: GitHub reports MERGEABLE/CLEAN against GopalakrishnanN/stable-epcontext-data-api; merge-tree against the fetched base is clean. Separately, merge-tree against current origin/main (0ea42b2) reports conflicts in graph_partitioner.cc, graph_partitioner.h and inference_session.cc. Those are parent-stack integration conflicts outside this PR's changed files, not base-branch conflicts.
No local builds, tests, lint or validation were run, as requested.
Tianlei Wu (tianleiwu)
previously approved these changes
Oct 5, 2026
Gopalakrishnan Nallasamy (GopalakrishnanN)
added a commit
that referenced
this pull request
Oct 5, 2026
### Description <!-- Describe your changes. --> Promotes the experimental EPContext named-buffer callback transport to the stable ONNX Runtime 1.31 C and C++ APIs. - Adds stable read/write callback registration and immutable EP-facing snapshots of the callback and application-state pointers. - Implements review option 3: removes `max_data_size` and the proposed read-options type, constructor, getter, setter, and release APIs. Registration and retrieval carry only callback/state pointers. - Leaves allocation limits to the application callback and data validation/deserialization limits to the EP. ORT does not invoke the callback or inspect its result. The sample helper can enforce an EP-selected limit without putting that policy into the stable ABI. - Adds append-only `GetEpContextDataCallbackSupport` READ/WRITE advertisement, gated on plugin API version 31. Unsupported participating EPs are rejected before `Compile()`, including the ORT-format read path. Registered callbacks cannot silently fall back to filesystem I/O. - Removes the experimental EPContext lookup entries and C++ wrapper; experimental APIs have no stability guarantee. - Updates the C++ wrappers, provider bridge, example EP, tests, and design document. - Retains the early support query for providers that can produce external context data without `Compile()`, with VitisAI's direct-assignment lifecycle rationale documented. - Includes the upstream GPU build fix from #32932 and the previously approved 2 KiB Android size-budget adjustment. Provider/binding adoption follows in QNN (#32266), OpenVINO (#32267), C# (#32270), and Python (#32271). Those follow-ups need to use the finalized callback/state-only signatures. ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> This graduates the transport introduced experimentally by #28624 and refined by #29294 and #31195. Applications can keep external EPContext data in encrypted or application-managed storage without silently falling back to plaintext filesystem I/O. Following the payload-policy review, the stable contract is limited to transport, ownership, and support negotiation. Centralized result-size validation would not itself prevent allocation inside the callback; passing a limit to the callback would still depend on callback/EP enforcement. Removing the policy avoids permanent API surface for a guarantee ORT does not provide. ### Testing Windows VS18 builds: - `cmake --build build\WindowsEpContextStable\RelWithDebInfo --config RelWithDebInfo --parallel 8 --target onnxruntime_shared_lib_test onnxruntime_autoep_test onnxruntime_test_all` — passed. - `cmake --build build\WindowsEpContextMinimal\MinSizeRel --config MinSizeRel --parallel 4 --target onnxruntime_shared_lib_test` — passed; minimal configuration disables exceptions, RTTI, and ML ops. Tests run from `build\WindowsEpContextStable\RelWithDebInfo\RelWithDebInfo`: - `.\onnxruntime_shared_lib_test.exe --gtest_filter=EpContextDataApiTest.*` — 14 passed, including immutable snapshots, cloning/lifetime, unchanged error outputs, and callback clearing. - `.\onnxruntime_autoep_test.exe --gtest_filter=*EpContext*` — 25 passed, 2 skipped for unavailable Windows symlink privileges. Covers strict support negotiation, callback/file paths, and EP-selected helper size-limit boundaries. - `.\onnxruntime_test_all.exe --gtest_filter=InternalTestingEP.*EpContext*:InternalTestingEP.Nhwc*` — 19 passed. Tests run from `build\WindowsEpContextMinimal\MinSizeRel\MinSizeRel`: - `.\onnxruntime_shared_lib_test.exe --gtest_filter=EpContextDataApiTest.*` — 2 passed. Formatting and `git diff --check` passed. No references to the removed read-options API or runtime size fields remain in tracked files. Not validated locally: GPU/Android builds, VitisAI/VAIP hardware behavior, other hardware EP integrations, or the full suite. The VitisAI lifecycle discussion remains open for reviewer confirmation. --------- Co-authored-by: Gopalakrishnan Nallasamy <gopalakrishnan.nallasamy@microsoft.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Co-authored-by: Edward Chen <18449977+edgchen1@users.noreply.github.com>
Base automatically changed from
GopalakrishnanN/stable-epcontext-data-api
to
main
October 5, 2026 21:31
Gopalakrishnan Nallasamy (GopalakrishnanN)
dismissed
Tianlei Wu (tianleiwu)’s stale review
October 5, 2026 21:31
The base branch was changed.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Gopalakrishnan Nallasamy (GopalakrishnanN)
requested a review
from Tianlei Wu (tianleiwu)
October 6, 2026 01:48
Tianlei Wu (tianleiwu)
left a comment
Tianlei Wu (tianleiwu)
left a comment
Contributor
There was a problem hiding this comment.
Reviewed the complete diff at abed89c, including callback/allocator ownership, retained session/compiler registrations, GIL handling, reentrancy, build guards, and the example EP fixture. No blocking correctness findings; one non-blocking test-coverage suggestion is attached inline.
The previously reported compiler-read coverage gap and test-worker exception-handling concern are addressed in the current code and were not duplicated.
Validation: static review and git diff --check passed. No local runtime tests were run because the available Python/example-EP binaries predate this API. The current GitHub check rollup reports 103 successful checks and one neutral check, with no failures or pending checks.
Tianlei Wu (tianleiwu)
approved these changes
Oct 6, 2026
Tianlei Wu (tianleiwu)
deleted the
GopalakrishnanN/python-epcontext-data-io
branch
October 6, 2026 04:35
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
Exposes stable named EPContext data read/write callbacks through Python.
max_data_sizeis retained and enforced by Python-owned callback state; no removed native read-options API or native size-policy field is used.e0cc266cdefef4e27104b14e4091a9d61f75a214, including its main-branch update, retained-NHWC READ-support validation, and provider-specific callback preflight before capability-time I/O.Motivation and Context
This independent binding PR remains stacked on #32265. It enables bounded, application-managed EPContext I/O without reintroducing payload policy into the stable native transport ABI.
The latest core checks provider-specific callback requirements before capability discovery, including AOT, NHWC, and ORT-format paths, rather than relying solely on checks immediately before
Compile(). Existing Python callback failure and lifetime behavior is preserved.Retarget to
mainafter #32265 lands.Testing
Revalidated on Windows using the existing isolated Python environment:
cmake --build .\build\WindowsBindings\RelWithDebInfo --config RelWithDebInfo --parallel 8 --target onnxruntime_pybind11_state example_plugin_ep— passed after merginge0cc266cde, including regeneration for the updated core dependencies.e0cc266cdeand that Python imports the freshly rebuilt 1.31.0 module with the callback API present.build\WindowsBindings\RelWithDebInfo\RelWithDebInfo,python -m pytest C:\Users\gnallasamy\Codes\ort-bindings-471c3cdf\python\onnxruntime\test\python\onnxruntime_test_python_compile_api.py -q— 14 passed, 2 hardware-dependent skips.ruff check,ruff format --check, andgit diff --check e0cc266cdefef4e27104b14e4091a9d61f75a214 HEAD— passed.The source-tree package emitted its expected missing build/package-metadata warning. GPU/hardware-provider-specific preflight paths and non-Windows runs were not performed.