Repository navigation
memory: stable dynamic-plugin memory ABI (nxmem) — issue #1186 phase 6 - #1440
justinchuby wants to merge 4 commits into
Conversation
…ry-manager' into justinchuby-phase-6-plugin-memory-abi
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## justinchuby-phase-5-process-memory-manager #1440 +/- ##
==============================================================================
- Coverage 80.01% 79.28% -0.73%
==============================================================================
Files 366 373 +7
Lines 164950 165420 +470
Branches 164950 165420 +470
==============================================================================
- Hits 131980 131153 -827
- Misses 28150 29421 +1271
- Partials 4820 4846 +26
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
🔴 Benchmark Regression DetectedComparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).
Visual flags: Host infoWhat this cannot catch
|
> **A100 safety update (head `13c37a7f`):** expected shared-prefix admission failures are now all-or-private: K/V commit is transactional, a failed V share rolls K back before enqueue, and only a successful rollback permits private-KV fallback. Fatal kernel faults remain process-fatal; the new cross-process VMM probe proves a worker fault does not affect an actively executing peer, while the current server is still in-process and does not yet consume physical shared-prefix metadata. See the [final verified conclusion](#1579 (comment)). Collapses phases 1–7 of the #1186 memory architecture rework onto current `main` as a single merge. The stack forked 270 commits ago (176 of them touching `crates/`), so rebasing it layer by layer would mean solving the same 17 conflicts fourteen times, with no reviewer ever looking at the intermediate states. Closes the stack: #1252 #1263 #1279 #1283 #1301 #1341 #1349 #1426 #1440 #1448 #1454 #1462 #1465 #1468 #1533. **A review guide is in the first comment.** It is the part worth reading — this diff is 103 files, but only three decisions in it are ones a compiler cannot check. ## What was already reviewed, and what wasn't Every phase was reviewed and approved by a session that was not its author, under the rejection-lockout rule (a rejected author never writes the next revision). | Phase | PR | Verdict | |---|---|---| | 1–5 | #1252 → #1426 | approved | | 6 | #1440, #1448, #1454 | **rejected** ×3 | | 6 | #1462 | approved | | 7 | #1465 | **rejected** | | 7 | #1468 | approved | | 7 test fixes | #1533 | approved, A100-verified | The four rejected rounds were not replaced — the later rounds are stacked **on top of** them. So the tree here is the approved state, but the history contains the rejected commits. **#1440 / #1448 / #1454 / #1465 should not be reviewed individually**; they close automatically. **Not reviewed anywhere:** the conflict resolutions themselves. That is what this PR is for. ## Verification Run here, on this merged tree: - `cargo check --workspace --all-targets` — clean - `cargo check -p onnx-genai-engine --features cuda,native-backend --all-targets` — clean. Worth calling out: the default feature set does **not** compile the `cfg(cuda)` code, which is where the riskiest edits are. Checking only the default set would have missed a real break (see the guide). - `cargo test` over 7 crates — **1105 passed, 2 failed, 88 ignored** - `cargo clippy --workspace --all-targets` - `cargo fmt` **The 2 failures and the 1 clippy error are pre-existing and proven so, not merge damage:** - `platform_capacity::{disk_capacity_is_measured_for_the_working_directory, an_explicit_byte_limit_is_honored_without_a_device_query}` — `platform_capacity.rs` is byte-identical to `main` (`git diff origin/main HEAD -- ` that file is empty). The cause is a macOS-only FFI layout bug: `fsblkcnt_t` is 4 bytes on macOS (confirmed: `sizeof(struct statvfs)=64`, `sizeof(fsblkcnt_t)=4`) while the Rust struct declares `f_blocks`/`f_bfree`/`f_bavail` as `u64`. CI is Linux, where it is 8 bytes. Left alone: it is a real bug but not this PR's. - `optimizer.rs` clippy `approx_constant` — present on both parents; the merged file is byte-identical to `main`. - Three files had `rustfmt` drift already present on `main` (`matmul_nbits.rs`, `normalization.rs`, `optimizer.rs`). Reverted rather than swept in, to keep the diff readable. **Final A100 revalidation (head `13c37a7f`):** full CUDA-memory GPU suite passed; CUDA EP default-parallel lib suite passed (**488 passed / 17 ignored**); real fp16 GQA covered three requests × two interleaved decode steps with two shared peers and one transactional private fallback, byte-identical to independent GPU and CPU references. A device-started/event-gated worker remained healthy through a peer process `CUDA_ERROR_ILLEGAL_ADDRESS`, owner exit, replacement worker, and process restart. Governor tests, targeted Miri, CI Clippy, fmt, and CUDA honesty also passed. ## Still open after this merges - `memory-plugin-provider-wiring` — the back half of Phase 6 criterion 8 (provider/context pinning), which fell in the gap between phases. **#1186 must not be closed until it lands.** This merge makes it more tractable, not less: see decision 3 in the guide. - `memory-deferred-invariant-asserts` — two surviving mutants found by #1462's reviewer, adjudicated NON-BLOCKING and not defects. - #1533's CUDA honesty guard warns on a legitimately portable anchor. Should be moved out of the CUDA test binary or allow-listed rather than left warning. --------- Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot <copilot@github.com> Copilot-Session: 46c5d75b-8146-489c-b82f-08ee29c27ce4 Copilot-Session: 39ff6824-d35f-4d3b-8f5f-043a7119a100 Copilot-Session: c80f8522-983c-47f7-8241-2155a823aabe
|
Closing: superseded by #1579, which collapsed phases 1–7 into a single merge — landed as This PR was rejected under the lockout rule and superseded by the approved revision that was stacked on top of it (not by a replacement of it). Per #1579's own guidance, it was never meant to be reviewed individually. Containment verified, not assumed: Every one of the 15 stack heads (#1252 → #1533) is a literal ancestor of the merged head, and I separately confirmed the merged head's content reached Why this didn't close itself: #1579's body says #1186 stays open: #1579 explicitly gates it on |
Part of #1186 — Phase 6 only.
Adds nxmem, a versioned C ABI that lets a dynamically loaded plugin supply allocator, virtual-backing, and shared-mapping mechanisms to the memory governance stack built in Phases 1–5. Phase 7 (removing the built-in CUDA eager allocator) is not started here.
Base branch is
justinchuby-phase-5-process-memory-manager, merged atc65cf036.Shape
Three new crates, mirroring the existing
onnx-runtime-ep-nxrt-{abi,host,testplugin}trio:onnx-runtime-memory-abi#[repr(C)]structs, vtables, versioning, status codes. Also shipsinclude/nxmem_memory_abi.handexamples/minimal_plugin.c.onnx-runtime-memory-hostDeviceAllocator/VirtualBacking/SharedMappingtraits fromonnx-runtime-memory-api.onnx-runtime-memory-testplugincdylibpublishing eight mechanisms,dlopened by the ABI tests.The governor never learns it is talking to a plugin: the host adapter presents the same Rust traits Phase 1 defined.
Narrative documentation:
docs/memory/MEMORY_PLUGIN_ABI.md.ABI contract
Structs. Every struct is
#[repr(C)]and begins withuint32_t struct_sizeat offset 0, in every version, forever. Growable structs carryuint32_t abi_minorat offset 4. Records:NxmemStatus,NxmemDeviceId,NxmemAllocation,NxmemAllocRequest/Result,NxmemByteRange,NxmemRangeRequest,NxmemReleaseOutcome,NxmemReleaseCompletion,NxmemReclaimRequest,NxmemUnloadReport,NxmemSharedPrefixHandle/CommitRequest/CommitInfo,NxmemHostCallbacks,NxmemOpenRequest,NxmemVersionRange,NxmemNegotiateRequest/Response.Vtables.
NxmemAllocatorFactoryVtable,NxmemAllocatorVtable,NxmemVirtualBackingVtable,NxmemSharedMappingVtable.Required vs optional. On the allocator vtable,
allocate/deallocate/retain/releaseare required at every level and a vtable missing one is refused. Everything else is a nullable slot; NULL means the capability is absent and surfaces asUnsupportedCapability, never as a silently successful no-op. The host treats a non-NULL capability vtable behind a clear capability flag as a contract violation, not a bonus.Required exports.
NxmemNegotiate,NxmemCreateAllocatorFactories,NxmemQueryUnloadReadiness. All three required — without the third, unload could not be gated at all, so a library missing it is refused at load.Versioning. Major 1 is a hard gate. Minor 1 current, minor 0 baseline; minor only ever appends. Minor 1 appends exactly one slot,
release_allocation, gated byNXMEM_CAP_STRUCTURED_RELEASE.Negotiation and the clamp rule.
NxmemNegotiatefixes a ceiling, not an assignment. Each vtable independently declares in its ownabi_minorwhat it really implements, so one module can ship a current mechanism beside a baseline one.read_prefixis the only supported way to read an untrusted vtable:struct_sizeandabi_minorwithread_unaligned;struct_sizebelow what the level the sender claims requires — a sender contradicting itself is broken, not old;effective_minor = min(declared, negotiated)and rewriteabi_minorto it, so callers need only consult the value they get back;min(struct_size, size_of::<Self>())bytes into a zeroed local;Step 4 is deliberate and is what makes an older host usable with a newer plugin: a newer struct is a strict superset, so reading only the agreed prefix is sound. Rejecting instead would mean no plugin could ever add a slot without breaking every existing host.
Nothing Rust crosses. No trait object, no
Arc, no Rust enum layout, noString/Vec, no allocator ownership. Errors areNxmemStatus— a stableu32code plus a 256-byte inline message buffer, inline precisely so neither side frees the other's heap. Everyextern "C"body on both sides is wrapped incatch_unwind(catch_status_panic/catch_void_panic), so a panic becomesInternalErrorrather than unwinding into foreign frames.Cross-provider misuse. Every allocation, range, and shared-prefix call carries
mechanism_id+deviceand is checked before anything is touched (WrongMechanism/WrongDevice).allocation_idis a host-assigned monotonic counter and is never pointer-derived — the same reasoning that keptAllocationGenerationnon-pointer-derived in Phase 1, so address reuse cannot make one allocation impersonate another.Ownership and lifetimes
NxmemCreateAllocatorFactoriesrelease, exactly onceopen_allocatorrelease;retainadds a referenceallocatedeallocateorrelease_allocationcreate_shared_prefixrelease_shared_prefixEvery in-pointer is borrowed for the call only, with two stated exceptions:
NxmemOpenRequest::callbacksis borrowed for the allocator's whole lifetime, and a factory'snamemust outlive the factory.If the host refuses a vtable that
open_allocatorreturnedOkfor, it still owes the plugin arelease;abandon_allocatorre-reads the prefix defensively and callsreleaseonly when the struct is well-formed enough to locate that slot. Correspondingly, a plugin publishing a malformed vtable must not allocate state first — the test plugin'sshort-structmechanism returns a statelessstatic, which is what a plugin built against a mismatched header would actually do.Release outcomes are three non-interchangeable states:
COMPLETE(creditunmapped_bytes),QUARANTINED(plugin keepsresidual_owned_bytes; address never reissued, residue never refunded),FAILED(nothing mutated; allocation as live as the caller left it). An unrecognised state is treated as quarantine — the only reading that can corrupt neither memory nor accounting.Threading, re-entrancy, callbacks
Every slot may be called concurrently; a plugin does its own locking.
No participant blocks, and no participant holds its own lock, across a call into the other side.
pressuredrops the lock beforeon_pressure;allocate_withholds no lock while running the caller's closure;run_drain_callback_if_readyuses.take()), tightened rather than relaxed, because an ABI call into foreign code is strictly more dangerous than a trait-object call.take_allocationlocks the live map, removes the record, drops the guard, and only then enters the plugin.request_reclaim: the host may re-enter the same plugin on the same thread to satisfy it.drain_releasescallsrelease_completedper retired ticket in enqueue order; take the batch under your lock, drop it, then call the host. The test plugin does exactly this.Unload gating
NxmemQueryUnloadReadinessreports live allocators, allocations, views, capabilities, and queued releases.MemoryPlugin::try_unloadqueries the plugin's report unconditionally, before checking its own counters, so a rejection always carries both sides' tallies and a misreporting plugin is visible rather than fatal. Unload is refused while any count is non-zero.Deferred release keeps everything pinned: an enqueued release counts as live, pins the allocator, which pins the module, and keeps the host's callback table alive because
release_completedwill still be called through it.enqueue_releaseincrements the module's queued counter before the ABI call.Field-ordering decisions (drop order is load-bearing)
PluginModule.libraryis declared last, so thedlclosehappens after every other field has dropped.PluginAllocatordeclares its capability views beforecore, so views drop first; each view holds its ownArc<AllocatorCore>.AllocatorCore.bridgeand.callbacksareBoxes created beforeopen_allocator(stable heap addresses) and are still alive whenDrop for AllocatorCorecalls the plugin'srelease.Arccycle:HostBridgeowns its counters directly rather than pointing back at the allocator.No field was added to
ScopedMemoryBinding,CudaMemoryBinding, orCudaExecutionProvider— nothing outside the three new crates changes.Test plugin and coverage
onnx-runtime-memory-testpluginis acdylibloaded at runtime — out-of-tree in the way that matters (dlopen, no workspace linking) and in-tree in the way that doesn't (built and linted with everything else). It publishes eight named mechanisms rather than switching on env vars or globals, so tests select behaviour by name. It uses host memory only, so the suite is fully portable.eagerlazyshort-structcallback-proberequest_reclaim; fails cleanly when the host refuseslegacy-1-0quarantiningfuture-statestickyRequired scenarios → tests:
a_plugin_outside_the_hosts_major_range_is_refused, plus 11 negotiation unit testsa_short_allocator_vtable_is_refused_before_any_slot_is_read,a_vtable_smaller_than_the_level_it_claims_is_refuseda_mechanism_without_optional_capabilities_reports_none,a_host_with_no_reclaim_hook_reports_the_capability_as_absentallocate_and_release_round_trips_through_the_boundary,a_release_that_misdescribes_the_allocation_is_refused,releasing_an_unknown_address_fails_rather_than_guessing,a_partial_release_is_quarantined_rather_than_refunded,a_release_state_from_the_future_is_quarantined_not_guessedlazy_backing_commits_decommits_and_reports_mapped_bytes,a_shared_prefix_is_reference_counted_and_costed_oncedeferred_releases_retire_in_order_and_pin_the_modulea_refusing_host_callback_fails_the_allocation_cleanlyunload_is_refused_while_an_allocator_is_open/…_an_allocation_is_live/…_a_shared_prefix_is_held/…_a_queued_release_has_not_retired,an_idle_plugin_unloadsa_minor_0_mechanism_works_under_a_minor_1_host(plugin older) andan_older_host_range_still_drives_the_current_plugin(host older)an_allocation_cannot_be_released_by_a_sibling_mechanism,an_object_from_another_mechanism_is_refusedevery_factory_is_released_exactly_onceheader_layout_matches_the_rust_definitions,nxmem_c_example_compiles,a_c_compiler_agrees_with_the_rust_layoutsThree test-design notes, since they affect whether the tests defend their names:
DeviceAllocator::release, not the private outcome-interpreting helper, because the branch worth pinning is the one the production caller actually reaches.nxmem_abi.rstakes a serialising mutex guard declared as its first statement (guard declared first ⇒ dropped last), and the one test that permanently poisons those counters (stickynever retires its queue — that is the behaviour under test) lives in its own integration-test binary,nxmem_abi_unload_gate.rs, because eachtests/*.rsis a separate process..dylibis rebuilt unconditionally. Cargo builds only therlibtarget of a dev-dependency, so the cdylib on disk goes stale silently; the helper runscargo build -p onnx-runtime-memory-testpluginonce per test process and panics loudly on failure. Similarly,every_factory_is_released_exactly_oncereads the count out of the loaded module through a test-only exported symbol — reading the statically linkedrlib's copy of the same static would observe a different variable and pass while proving nothing. (That was a real bug caught during development.)The C header carries machine-readable layout annotations that are checked three ways: against Rust
size_of, against a generated_Static_asserttranslation unit compiled by a real C compiler, and by compilingexamples/minimal_plugin.cwith-Wall -Wextra -Werroragainst nothing but the header. When no C compiler can be found those tests fail loudly; they never skip silently.Validation (exact-base comparison)
Base for comparison:
origin/justinchuby-phase-5-process-memory-manager@c65cf036, checked out into a separate worktree and run with identical commands.Diff vs that base touches only new files plus two purely additive edits:
Cargo.toml(3 members, 2 default-members, 3 workspace deps) andCargo.lock. No existing source file is modified, so no existing package's compilation changes.cargo test -p onnx-runtime-memory-abi -p onnx-runtime-memory-host -p onnx-runtime-memory-testplugin:memory-abilibmemory-abiheader_contractmemory-hostlibmemory-hostnxmem_abimemory-hostnxmem_abi_unload_gatememory-testpluginlibmemory-abidoctestsThe single ignored item is the
export_nxmem_plugin!usage doctest, markedignorebecause it defines#[no_mangle]entry points that cannot be linked into a doctest binary.cargo clippy -p onnx-runtime-memory-abi -p onnx-runtime-memory-host -p onnx-runtime-memory-testplugin --all-targets -- -D warnings— clean.cargo fmt --all -- --check: on this branch the only drift iscrates/onnx-genai-server/src/routes/completions.rs; the base worktree atc65cf036reports the same single file. Inherited, untouched, not fixed here.cargo doc --no-depson all three new crates — no warnings. (The knownonnx-runtime-memory-governorintra-doc-link warnings are inherited and out of scope.)gpu-testsfeature; the whole suite runs on host memory.Other known pre-existing failures on this stack —
matmul_nbits_marlin_numerics, Windows CUDAunnecessary_unwrap, macOS MLAS, governor rustdoc links — are untouched and unrelated to these files.Acceptance criteria not fully met
CUDA-specific paths are not exercised — there are none in this change, and none were added. This host is macOS/arm64 with no CUDA hardware. Nothing under
crates/onnx-runtime-ep-cudaorcrates/onnx-runtime-cuda-memoryis modified by this PR, so there is no CUDA path here that was "compile-checked only": there is no CUDA path here at all. Wiring a CUDA mechanism behind this ABI is Phase 7 work. Stating it plainly rather than implying GPU coverage.Every other acceptance criterion in the Phase 6 section of #1186 is implemented and tested.
Non-goals honoured
No internal policy, holder, victim-selection, or governor type is exposed through the C ABI. No compatibility is promised for anything predating this contract — stated in the crate docs and the header.
Do not merge, do not enable auto-merge, do not enqueue. This PR stays open for human review, like every PR in the #1186 memory stack.