Repository navigation
feat(memory): add process memory manager (Phase 5) - #1426
Closed
justinchuby wants to merge 3 commits into
Closed
justinchuby wants to merge 3 commits into
justinchuby wants to merge 3 commits into
Conversation
Centralize memory registration and scoped allocation transactions while preserving governor, holder, and execution-provider ownership boundaries. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Wait only for the exact managed workspace allocation, keep timeout retries recoverable, and quarantine compatibility charges on failed deallocation. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Keep quarantined charges visible with active waiters and exercise the production workspace barrier and deallocation branches directly. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## justinchuby-fix-phase4-device-loss-settlement #1426 +/- ##
=================================================================================
+ Coverage 79.52% 80.01% +0.48%
=================================================================================
Files 365 366 +1
Lines 162327 164950 +2623
Branches 162327 164950 +2623
=================================================================================
+ Hits 129096 131980 +2884
+ Misses 28513 28150 -363
- Partials 4718 4820 +102
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
|
Contributor
There was a problem hiding this comment.
Pull request overview
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Adds a process-level ProcessMemoryManager and migrates key allocation flows (notably session workspaces and native CUDA components) to manager-issued scoped bindings/transactions, improving identity, settlement, and device-loss handling without centralizing policy.
Changes:
- Introduces
ProcessMemoryManagerpublic API surface and extensive governor-level tests for process transactions, settlement, device loss, and reentrancy. - Migrates session workspace allocation/deallocation to
WorkspaceAllocation+replace_workspace/deallocate_workspaceAPIs. - Integrates CUDA EP with the process manager, including context operation gating, managed settlement tokens, and post-drain cleanup hooks.
Reviewed changes
Copilot reviewed 31 out of 33 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| docs/memory/MEMORY_ARCHITECTURE.md | Documents Phase 5 process manager responsibilities, transaction model, and accounting axes. |
| crates/onnx-runtime-session/src/executor/state.rs | Switches workspace representation to WorkspaceAllocation. |
| crates/onnx-runtime-session/src/executor/mod.rs | Updates imports and ensures workspace deallocation uses deallocate_workspace. |
| crates/onnx-runtime-session/src/executor/dispatch.rs | Replaces manual reserve/allocate/lease flow with replace_workspace. |
| crates/onnx-runtime-session/src/executor/bindings.rs | Updates persistent/step workspace replacement + release to new workspace APIs. |
| crates/onnx-runtime-memory-governor/tests/process_memory_manager.rs | Adds comprehensive tests for manager behavior (transactions, settlement, device loss, concurrency). |
| crates/onnx-runtime-memory-governor/src/lib.rs | Exposes manager module and adds mapped-attribution context-termination + retained-growth settlement. |
| crates/onnx-runtime-memory-governor/Cargo.toml | Adds thiserror dependency for new manager/error types. |
| crates/onnx-runtime-memory-api/src/binding.rs | Adds register_allocator_with_device and snapshots() for process-manager integration. |
| crates/onnx-runtime-memory-api/src/allocator.rs | Extends MappedAllocation with additional_owned_bytes. |
| crates/onnx-runtime-ep-plugin/src/device.rs | Handles both binding-owned and manager-owned buffer releases. |
| crates/onnx-runtime-ep-cuda/src/weight_paging.rs | Adds context-scope gating + termination reconciliation to residency/page-in paths. |
| crates/onnx-runtime-ep-cuda/src/provider.rs | Integrates CUDA EP allocation/release with process manager and workspace settlement barrier. |
| crates/onnx-runtime-ep-cuda/src/deferred_release.rs | Adds re-entrant execution gate and a drain callback hook for cleanup after queue drain. |
| crates/onnx-runtime-ep-api/src/provider.rs | Introduces managed buffer ownership, WorkspaceAllocation, and workspace allocation/deallocation APIs. |
| crates/onnx-runtime-ep-api/src/lib.rs | Re-exports new workspace/ownership types. |
| crates/onnx-runtime-cuda-memory/src/vmm_allocator.rs | Plumbs additional_owned_bytes out of CUDA VMM allocation commits. |
| crates/onnx-runtime-cuda-memory/src/virtual_memory.rs | Adds context-identity tracking + external context-termination reconciliation for physical pools. |
| crates/onnx-genai-server/src/state.rs | Adds a shared ProcessMemoryManager to server memory authority provider. |
| crates/onnx-genai-engine/src/pipeline/routing.rs | Loads native CUDA pipeline components using shared authority + process manager. |
| crates/onnx-genai-engine/src/pipeline/mod.rs | Threads CUDA manager/governor/policy through component session construction. |
| crates/onnx-genai-engine/src/pipeline/flat_autoregressive.rs | Threads process manager into native CUDA component loads. |
| crates/onnx-genai-engine/src/pipeline/decoder_component.rs | Threads process manager into native decoder construction. |
| crates/onnx-genai-engine/src/native_decode/proposer.rs | Adds load_with_cuda_memory that uses a shared manager. |
| crates/onnx-genai-engine/src/native_decode/load.rs | Threads process manager through native decode CUDA EP initialization. |
| crates/onnx-genai-engine/src/native_component.rs | Adds governed CUDA component loader using shared manager. |
| crates/onnx-genai-engine/src/memory_authority.rs | Extends authority provider trait to supply a process manager. |
| crates/onnx-genai-engine/src/lib.rs | Re-exports ProcessMemoryManager. |
| crates/onnx-genai-engine/src/engine/speculative_load.rs | Loads speculative proposer with shared manager under CUDA. |
| crates/onnx-genai-engine/src/engine/load.rs | Threads process manager into native decode and speculative proposer loading. |
| crates/onnx-genai-engine/src/engine/governor.rs | Stores a process manager in governor and threads reservation bytes separately. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+333
to
+348
| impl Drop for ExecutionGateGuard<'_> { | ||
| fn drop(&mut self) { | ||
| let current = std::thread::current().id(); | ||
| let mut state = self | ||
| .gate | ||
| .state | ||
| .lock() | ||
| .unwrap_or_else(|poisoned| poisoned.into_inner()); | ||
| debug_assert_eq!(state.owner.as_ref(), Some(¤t)); | ||
| state.depth = state.depth.saturating_sub(1); | ||
| if state.depth == 0 { | ||
| state.owner = None; | ||
| self.gate.wake.notify_all(); | ||
| } | ||
| } | ||
| } |
Comment on lines
+2270
to
+2278
| assert!( | ||
| lease.is_none(), | ||
| "CUDA manager-backed workspace must keep accounting in ManagedAllocation, not an outer \ | ||
| compatibility lease" | ||
| ); | ||
| assert!( | ||
| captured, | ||
| "CUDA workspace must carry allocation-specific manager settlement" | ||
| ); |
Comment on lines
+450
to
+465
| static QUARANTINED_WORKSPACE_LEASES: std::sync::OnceLock<std::sync::Mutex<Vec<MemoryLease>>> = | ||
| std::sync::OnceLock::new(); | ||
|
|
||
| fn quarantine_failed_workspace_lease(lease: MemoryLease) { | ||
| eprintln!( | ||
| "execution provider workspace deallocation failed before physical release was proven; \ | ||
| retaining its {} byte {:?} lease in compatibility quarantine", | ||
| lease.bytes(), | ||
| lease.tier() | ||
| ); | ||
| QUARANTINED_WORKSPACE_LEASES | ||
| .get_or_init(|| std::sync::Mutex::new(Vec::new())) | ||
| .lock() | ||
| .unwrap_or_else(|poisoned| poisoned.into_inner()) | ||
| .push(lease); | ||
| } |
Comment on lines
+403
to
420
| pub fn into_raw_with_owner(self) -> (*mut c_void, Option<BoundBufferOwnership>) { | ||
| let ptr = self.ptr.as_ptr(); | ||
| match self.owner { | ||
| BufferOwner::Bound(owner) => (ptr, Some(BoundBufferOwnership::Binding(*owner))), | ||
| BufferOwner::Managed(owner) => (ptr, Some(BoundBufferOwnership::Managed(*owner))), | ||
| _ => (ptr, None), | ||
| } | ||
| } | ||
|
|
||
| /// Consume the buffer into its raw address and complete binding ownership. | ||
| pub fn into_raw_with_bound_ownership(self) -> (*mut c_void, Option<BoundBufferOwnership>) { | ||
| let ptr = self.ptr.as_ptr(); | ||
| match self.owner { | ||
| BufferOwner::Bound(owner) => (ptr, Some(*owner)), | ||
| BufferOwner::Bound(owner) => (ptr, Some(BoundBufferOwnership::Binding(*owner))), | ||
| BufferOwner::Managed(owner) => (ptr, Some(BoundBufferOwnership::Managed(*owner))), | ||
| _ => (ptr, None), | ||
| } | ||
| } |
justinchuby
added a commit
that referenced
this pull request
Aug 22, 2026
> **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
Owner
Author
|
Closing: superseded by #1579, which collapsed phases 1–7 into a single merge — landed as This PR was approved by a session that was not its author, and its work is in 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 |
justinchuby
added a commit
that referenced
this pull request
Aug 22, 2026
…ees were destroyed (#1717) Two decision records existed **only on this machine** — written into `.squad/decisions/inbox/` but never committed. Recovered here. The inbox README already names this exact failure mode as the reason the directory was made tracked on 2026-07-29: > **Drops survive worktree deletion.** Previously the inbox was gitignored, so drops existed only on the machine that wrote them. Removing a worktree before Scribe ran destroyed unmerged drops — this cost real records more than once. The policy fixed the *gitignore* half of the problem. It cannot fix the half where a drop is written and simply never `git add`ed, which is what happened to both of these. I found them while cleaning up 11 stale worktrees; one was inside a worktree I was about to delete. ## What is being recovered **1. `Copilot-phase-5-process-memory-manager-ownership-boundarie.md`** Written in the `justinchuby-phase-5-process-memory-manager` worktree (PR #1426, now closed as superseded by #1579). Copied out immediately before that worktree was removed. Verified it is not a duplicate before adding it: ``` $ git grep 'Phase 5 process memory manager ownership boundaries' origin/main -- .squad (no match) ``` So the Scribe never merged it — it was not sitting redundantly next to an existing entry in `.squad/decisions.md`. Content-wise this one is not merely historical. It records the Phase 5 ownership split: `ProcessMemoryManager` owns process quota, canonical authority/view mapping, holder identity, context transaction gates, loss broadcast and settlement coordination, while **governors keep budget policy, holders keep victim choice, and EP contexts keep stream/copy/commit/decommit/fence/deferred-release ordering**. It also records two invariants that are easy to get wrong later — quarantine never infers a refund from allocation length, and confirmed context termination is the sole device-loss discharge boundary. That boundary is still load-bearing after #1579, and **#1186 is still open precisely on the provider/context-pinning half of it**, so losing this note would have removed context from work that has not happened yet. **2. `coordinator-criterion-failure-is-structural-not-probabilistic.md`** Written in this worktree. References #1579, #1586, #1533, #1600 and three sessions that converged on it independently. ## Scope No code changes; two files under `.squad/decisions/inbox/`. Per the inbox contract these are transient — the Scribe merges them into `.squad/decisions.md` and deletes the drops. ## Worth considering separately Nothing detects an uncommitted drop. Both of these survived by coincidence, because a human happened to be cleaning worktrees and looked at `git status` before running `git worktree remove --force`. A pre-removal check for untracked files under `.squad/decisions/inbox/` would close that gap; I have not written one here. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Copilot-Session: c80f8522-983c-47f7-8241-2155a823aabe
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.
Summary
Introduces the process-level
ProcessMemoryManagerand manager-issued scoped memory bindings for Phase 5 of #1186.The manager coordinates registration, mechanism selection, canonical authority and holder identity, parent-quota delegation, allocation transactions, device-loss publication, and accounting settlement. It does not absorb policy or device ordering:
BindingRegistryremains the only registration/selection/allocation-generation ledger.Native CUDA decoder, pipeline decoder, prompt/every-step components, and proposer sessions now share the engine/server process manager. Plain construction still creates a local unlimited parent manager, preserving current defaults.
Part of #1186 — Phase 5 only
Stacked on the final reviewed Phase-4 corrective branch / PR #1349 (
justinchuby-fix-phase4-device-loss-settlement).Architecture and ownership
flowchart TD PMM[ProcessMemoryManager] REG[BindingRegistry<br/>single registration + lifecycle ledger] GOV[MemoryGovernor / authority<br/>budget and lease policy] HOLDER[Holder / PressureResponder<br/>victim choice + reclaim] BIND[ScopedMemoryBinding<br/>device + mechanism + context + authority] TX[Allocation transaction<br/>reserve → allocate → commit/map → publish] OWN[ManagedAllocation<br/>owning allocation + settlement token] EP[Execution provider context<br/>streams + copies + fences + deferred queue] Q[Phase-4 structured release / quarantine] PMM --> REG PMM --> GOV PMM -. pressure request .-> HOLDER REG --> BIND GOV --> TX BIND --> TX TX --> OWN OWN --> EP EP --> Q Q -. actual outcome .-> PMM HOLDER -. released bytes only .-> GOVSource of truth
ProcessMemoryManagerBindingRegistryThe manager keeps policy metadata keyed by opaque registry identities, but no second context/mechanism membership ledger. Mechanism snapshots come directly from
BindingRegistry::snapshots().Scoped identity and lifecycle
ProcessAuthorityIdand one governor.active → retiring/lost → terminated. Allocation commit/publication and weight page-in hold context-operation tokens; termination waits for those tokens with no manager/registry/governor lock held.Lock order
No strong manager cycle is introduced: queue cleanup captures a weak manager, registry/context pins do not point back to the manager, and allocator teardown completion is signaled only after queue-producing destructors return.
Transaction and rollback matrix
unmapped_bytesFinite process limits either reserve the full per-allocation authority maximum or use a fixed full-capacity authority delegation. Lowering a parent below live slices is rejected, and a local authority whose extent outgrows its slice cannot grant another transaction.
Accounting axes
CUDA VMM now carries
additional_owned_bytesseparately fromnewly_mapped_bytes. Pool-backed allocations do not claim allocation-span bytes as physical residency. Quarantine also never uses allocation length to refund rounded authority/process charges.Migration and adapters
onnx-runtime-sessionworkspace preparation usesallocate_workspace/replace_workspacetransactions.ManagedAllocationinsideDeviceBuffer; charges settle only after the deferred Phase-4 outcome.DeviceBufferextraction returnsBoundBufferOwnership, preserving manager settlement at compile time instead of panicking or stripping it.CudaExecutionProvider::with_memory(Arc<dyn DeviceAllocator>)remains supported. Replacement registers/selects through the manager; retired mechanisms remain tracked until their old allocations and teardown complete.DeviceAllocator; the manager wraps registration and transaction identity.Validation
deallocate_workspacesuccess/error coverage, and waiter-safe quarantined snapshot retention including the context-termination race.-D warnings; CUDA library passes-D warningsafter allowing only exact inherited lint classes.-D warningsafter allowing only inherited broken/private/redundant link warnings in unchanged files.onnx-runtime-ep-cuda --features cuda --all-targetstest binaries compile; hardware tests were compile-only / ignored, not executed on this non-CUDA host.git diff --checkpasses.Exact base comparison
The exact #1349 head (
9ac2b41b) was tested before edits:onnx-genai-server/src/routes/completions.rs.onnx-runtime-ep-cpulints; additional CUDA strict failures are likewise in unchanged kernel tests.deferred_release_queue,matmul_nbits_marlin_numerics,vmm_release_quarantine).Scope guard
This PR does not choose weight/KV victims, inspect model-specific contents, reclaim unrelated external processes, add a dynamic plugin ABI, remove eager CUDA, flip defaults, or change scheduler/model behavior. It does not weaken or bypass the Phase-4 ownership, deferred-release, quarantine, exact-once, or device-loss state machines.
This PR must remain open for human review. Do not merge, auto-merge, enqueue, or stack-merge it.