diff --git a/crates/onnx-runtime-ep-cuda/src/deferred_release.rs b/crates/onnx-runtime-ep-cuda/src/deferred_release.rs index 98d50817bb..0d7012d84c 100644 --- a/crates/onnx-runtime-ep-cuda/src/deferred_release.rs +++ b/crates/onnx-runtime-ep-cuda/src/deferred_release.rs @@ -170,9 +170,40 @@ impl DeferredActionOutcome { /// /// `execute` consumes the action, so the ownership it carries is released (or /// retained) exactly once. It runs with **no** queue lock held. -pub trait DeferredReleaseAction: Send + Debug { +/// +/// Actions are `'static` because the queue may outlive every caller that +/// handed it one, and because device-loss settlement stores what an action +/// could not give back as `Box`. +pub trait DeferredReleaseAction: Send + Debug + 'static { fn execute(self: Box) -> DeferredActionOutcome; + /// Settle this action after the device or provider context was lost. + /// + /// This is a *terminal* path, not a deferral: it consumes the action + /// exactly like `execute` does, but it must never call an allocator, never + /// touch the device, and never refund. It runs with no queue lock held. + /// + /// The return value is the ownership the **queue** must keep alive, and it + /// is deliberately narrower than the action itself. An action that carries + /// a settlement of its own — a prepared request that owes a mechanism a + /// terminal state — is expected to perform that settlement here and hand + /// back only the residual physical ownership. Handing back the whole + /// unexecuted action instead would leave that settlement permanently owed, + /// and would keep alive whatever provider context the action pins, which + /// for a context-owned queue is the queue itself. + /// + /// The default retains the whole action, which is correct for an action + /// whose ownership is purely physical (a page's allocator/allowance, a + /// reservation ticket) and owes no mechanism anything. + fn settle_device_lost(self: Box, detail: &str) -> Option { + let bytes = self.bytes(); + Some(RetainedOwnership { + bytes, + detail: detail.to_owned(), + keep_alive: Box::new(self), + }) + } + /// Stable short label for stats and shutdown diagnostics. fn label(&self) -> &'static str; @@ -454,9 +485,16 @@ impl CudaDeferredReleaseQueue { /// Record that the device or provider context was lost. /// /// From this point the queue never queries a fence, never calls an - /// allocator, and never refunds: every pending and future request is - /// retained with its ownership and accounting intact, and its events are - /// kept rather than destroyed through a context that is known bad. + /// allocator, and never refunds. Every pending request is *settled* rather + /// than abandoned: its action records its own terminal device-loss state — + /// for a binding-prepared release that is device-lost quarantine at the + /// exact allocation identity — and hands back only the ownership that must + /// stay alive, which the queue then holds. Its events are kept rather than + /// destroyed through a context that is known bad. + /// + /// Settling rather than freezing the whole action is what lets the + /// mechanism reach quiescence, so the provider context can confirm + /// termination and discharge the quarantine it was handed. pub fn mark_device_lost(&self, reason: impl Into) { let reason = reason.into(); let gate = self @@ -508,26 +546,10 @@ impl CudaDeferredReleaseQueue { .unwrap_or_else(|poisoned| poisoned.into_inner()); if self.is_device_lost() { drop(execution); - let PendingRelease { fences, action } = entry; - for fence in fences { - fence.retain_after_device_loss(); - } - let bytes = action.bytes(); - retained.push(RetainedRelease { - label: action.label(), - state: AllocationReleaseState::DeviceLost, - bytes, - detail: String::from("device lost while a poller owned the deferred release"), - ownership: Some(RetainedOwnership { - bytes, - detail: String::from( - "device lost while a poller owned the deferred release", - ), - keep_alive: Box::new(action), - }), - }); - self.counters.quarantined.fetch_add(1, Ordering::Relaxed); - self.outstanding.fetch_sub(1, Ordering::AcqRel); + retained.push(self.settle_lost_entry( + entry, + "device lost while a poller owned the deferred release", + )); continue; } if !entry.fences.iter().all(|fence| fence.is_complete()) { @@ -552,26 +574,11 @@ impl CudaDeferredReleaseQueue { .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()); if self.is_device_lost() { - while let Some(PendingRelease { fences, action }) = carry.pop_front() { - for fence in fences { - fence.retain_after_device_loss(); - } - let bytes = action.bytes(); - let detail = - String::from("device lost after a poller observed an incomplete release fence"); - retained.push(RetainedRelease { - label: action.label(), - state: AllocationReleaseState::DeviceLost, - bytes, - detail: detail.clone(), - ownership: Some(RetainedOwnership { - bytes, - detail, - keep_alive: Box::new(action), - }), - }); - self.counters.quarantined.fetch_add(1, Ordering::Relaxed); - self.outstanding.fetch_sub(1, Ordering::AcqRel); + while let Some(entry) = carry.pop_front() { + retained.push(self.settle_lost_entry( + entry, + "device lost after a poller observed an incomplete release fence", + )); } } if !carry.is_empty() || !retained.is_empty() { @@ -627,8 +634,39 @@ impl CudaDeferredReleaseQueue { }); } - /// Move every pending request into retained ownership without touching the - /// device. Used on device loss. + /// Settle one pending entry that can never reach the device again. + /// + /// The fences are retained rather than destroyed, and the action settles + /// its own device-loss state and hands back only the ownership the queue + /// must keep. The queue deliberately does **not** keep the unexecuted + /// action: an action that owes a mechanism a terminal state would owe it + /// forever, its binding/provider-context pins would never be released, and + /// for a context-owned queue those pins reach back to this queue, so the + /// retained record would keep itself alive. + /// + /// Callers must hold neither the queue lock nor the execution gate. + fn settle_lost_entry(&self, entry: PendingRelease, detail: &str) -> RetainedRelease { + let PendingRelease { fences, action } = entry; + for fence in fences { + // The context is gone; destroying the event would be a driver call + // on state that no longer exists. + fence.retain_after_device_loss(); + } + let label = action.label(); + let bytes = action.bytes(); + let ownership = action.settle_device_lost(detail); + self.counters.quarantined.fetch_add(1, Ordering::Relaxed); + self.outstanding.fetch_sub(1, Ordering::AcqRel); + RetainedRelease { + label, + state: AllocationReleaseState::DeviceLost, + bytes, + detail: detail.to_owned(), + ownership, + } + } + + /// Settle every pending request as device-lost without touching the device. fn retain_all_pending(&self, detail: &str) { let drained: Vec = { let mut state = self.lock(); @@ -638,27 +676,8 @@ impl CudaDeferredReleaseQueue { return; } let mut records = Vec::with_capacity(drained.len()); - for PendingRelease { fences, action } in drained { - for fence in fences { - // The context is gone; destroying the event would be a driver - // call on state that no longer exists. - fence.retain_after_device_loss(); - } - self.counters.quarantined.fetch_add(1, Ordering::Relaxed); - records.push(RetainedRelease { - label: action.label(), - state: AllocationReleaseState::DeviceLost, - bytes: action.bytes(), - detail: detail.to_string(), - // The action is never executed, so every allocator, pointer, - // allowance and pin it holds stays alive and unusable. - ownership: Some(RetainedOwnership { - bytes: action.bytes(), - detail: detail.to_string(), - keep_alive: Box::new(action), - }), - }); - self.outstanding.fetch_sub(1, Ordering::AcqRel); + for entry in drained { + records.push(self.settle_lost_entry(entry, detail)); } let mut state = self.lock(); state.retained.append(&mut records); @@ -973,6 +992,35 @@ impl DeferredReleaseAction for PreparedReleaseAction { } } + /// Settle the prepared request as device-lost quarantine. + /// + /// No allocator is called and nothing is refunded: the binding records the + /// exact allocation identity as device-lost, the mechanism's queued-release + /// count settles, and the active-operation pin is released, so the + /// mechanism can reach quiescence and the provider context can confirm + /// termination and discharge the quarantine. + /// + /// Only the pinned allocator is handed back — the same residual the normal + /// quarantine path retains. It is what keeps the physical ownership from + /// being handed out again, and unlike the request it does not pin the + /// provider context, so retaining it cannot keep this queue alive. + fn settle_device_lost(mut self: Box, detail: &str) -> Option { + let request = self.request.take()?; + let bytes = request.len() as u64; + let allocator = Arc::clone(request.allocator()); + let outcome = request.quarantine_device_lost(); + // The observer sees the real terminal outcome. It refunds only what the + // allocator reported unmapped, which is zero here, and counts no free. + if let Some(observer) = self.observer.as_ref() { + observer.released(&outcome); + } + Some(RetainedOwnership { + bytes, + detail: detail.to_owned(), + keep_alive: Box::new(allocator), + }) + } + fn label(&self) -> &'static str { "provider allocation" } diff --git a/crates/onnx-runtime-ep-cuda/tests/deferred_release_queue.rs b/crates/onnx-runtime-ep-cuda/tests/deferred_release_queue.rs index 4070c8a3a3..33be623202 100644 --- a/crates/onnx-runtime-ep-cuda/tests/deferred_release_queue.rs +++ b/crates/onnx-runtime-ep-cuda/tests/deferred_release_queue.rs @@ -8,6 +8,7 @@ //! Owner/view/ABA semantics are covered by the memory-api tests and are not //! duplicated here. +use std::ptr::NonNull; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, Barrier}; use std::time::Duration; @@ -16,7 +17,12 @@ use onnx_runtime_ep_cuda::deferred_release::{ CudaDeferredReleaseQueue, DeferredActionOutcome, DeferredReleaseAction, ReleaseFence, ReleaseFenceSource, RetainedOwnership, }; -use onnx_runtime_memory_governor::{AllocationReleaseState, DeferredEnqueueRejection}; +use onnx_runtime_memory_governor::{ + AllocationIdentity, AllocationReleaseState, BindingRegistry, BindingResource, + DeferredEnqueueRejection, DeviceAllocator, DeviceKey, HostAllocator, MemoryBinding, + MemoryError, QuarantineReason, QuarantinedAllocation, RegisteredMechanism, + RegisteredProviderContext, +}; /// A fence whose completion the test controls. #[derive(Debug)] @@ -681,3 +687,351 @@ fn dropping_a_queue_with_pending_work_neither_frees_nor_blocks() { "the action is dropped without being executed, which is a retain" ); } + +// --------------------------------------------------------------------------- +// Production-path device-loss settlement +// +// The tests above drive the scheduler through a toy action. The ones below use +// the *production* entry point — `enqueue_prepared` with a real +// `PreparedAllocationRelease` taken from a real `MemoryBinding` — because the +// thing that has to be true after device loss is a property of the binding and +// of the provider context, not of the queue's own counters. +// --------------------------------------------------------------------------- + +/// A host-backed stand-in for a device allocator that counts physical releases. +/// +/// Device loss must never reach it, so `release_calls` staying at zero is the +/// load-bearing assertion; the bytes it handed out stay owned, exactly as +/// quarantined device ownership does. +#[derive(Debug, Default)] +struct CountingReleaseAllocator { + release_calls: AtomicUsize, +} + +impl CountingReleaseAllocator { + fn release_calls(&self) -> usize { + self.release_calls.load(Ordering::Acquire) + } +} + +impl DeviceAllocator for CountingReleaseAllocator { + fn allocate(&self, bytes: usize, align: usize) -> Result, MemoryError> { + HostAllocator.allocate(bytes, align) + } + + unsafe fn deallocate(&self, ptr: NonNull, bytes: usize, align: usize) { + self.release_calls.fetch_add(1, Ordering::AcqRel); + // SAFETY: forwarded unchanged from this method's contract. + unsafe { HostAllocator.deallocate(ptr, bytes, align) }; + } + + fn device(&self) -> DeviceKey { + DeviceKey::HOST + } +} + +/// The provider-context resource, shaped exactly like the CUDA provider's. +/// +/// It pins the deferred queue, which is what makes an unsettled retained +/// request a *cycle* rather than a leak: request -> binding -> mechanism -> +/// provider context -> this pin -> queue -> retained request. +#[derive(Debug)] +struct ContextPin { + #[allow(dead_code)] + queue: Arc, +} + +/// A real binding registry over a host allocator, wired to `queue` the way the +/// CUDA provider wires its own. +struct BoundFixture { + registry: BindingRegistry, + context: RegisteredProviderContext, + mechanism: RegisteredMechanism, + binding: MemoryBinding, + allocator: Arc, +} + +impl BoundFixture { + fn new(queue: &Arc) -> Self { + let registry = BindingRegistry::new().expect("registry"); + let context = registry + .register_provider_context( + DeviceKey::HOST, + Arc::new(ContextPin { + queue: Arc::clone(queue), + }) as Arc, + ) + .expect("provider context"); + let authority = registry + .register_authority(DeviceKey::HOST, Arc::new(()) as Arc) + .expect("authority"); + let allocator = Arc::new(CountingReleaseAllocator::default()); + let mechanism = registry + .register_allocator( + context, + authority, + Arc::clone(&allocator) as Arc, + ) + .expect("mechanism"); + let binding = registry.bind(DeviceKey::HOST).expect("binding"); + Self { + registry, + context, + mechanism, + binding, + allocator, + } + } + + fn snapshot(&self) -> onnx_runtime_memory_governor::MechanismSnapshot { + self.registry.snapshot(self.mechanism).expect("snapshot") + } + + fn quarantined(&self) -> Vec { + self.registry + .quarantined(self.mechanism) + .expect("quarantined ownership") + } + + /// Hand one real allocation's final ownership to `queue`. + fn enqueue_one( + &self, + queue: &Arc, + bytes: usize, + ) -> AllocationIdentity { + let allocation = self.binding.allocate(bytes, 256).expect("allocation"); + let request = self + .binding + .prepare_release(allocation) + .expect("prepared release"); + let identity = request.identity(); + queue.enqueue_prepared(request, None).expect("accepted"); + identity + } + + /// The documented device-loss teardown: invalidate, confirm termination + /// (which discharges quarantine), then remove the registration and the + /// provider context. + fn tear_down_after_device_loss(self) { + self.registry + .invalidate_device(DeviceKey::HOST, "test device loss") + .expect("invalidate"); + self.registry + .confirm_context_terminated(self.context) + .expect("a settled context can confirm termination"); + self.registry + .remove(self.mechanism) + .expect("a discharged mechanism can be removed"); + drop(self.binding); + drop(self.allocator); + self.registry + .remove_provider_context(self.context) + .expect("an unused provider context can be removed"); + drop(self.registry); + } +} + +/// Assert everything that must be true of one settled device-lost request. +fn assert_settled_device_lost( + fixture: &BoundFixture, + identity: AllocationIdentity, + bytes: usize, + queue: &Arc, +) { + assert_eq!( + fixture.allocator.release_calls(), + 0, + "the allocator is never called after device loss" + ); + + let quarantined = fixture.quarantined(); + assert_eq!( + quarantined.len(), + 1, + "the binding records exactly one retained allocation" + ); + let record = quarantined[0]; + assert_eq!(record.identity, identity, "the exact allocation identity"); + assert_eq!(record.state, AllocationReleaseState::DeviceLost); + assert_eq!(record.reason, QuarantineReason::DeviceLost); + assert_eq!(record.bytes, bytes); + assert_eq!(record.retained_bytes, bytes as u64); + + let snapshot = fixture.snapshot(); + assert_eq!( + snapshot.queued_releases, 0, + "the queued release settled instead of staying owed forever" + ); + assert_eq!( + snapshot.active_operations, 0, + "the operation pin is released, so the context can reach quiescence" + ); + assert_eq!(snapshot.live_allocations, 0); + assert_eq!(snapshot.quarantined_allocations, 1); + + let stats = queue.stats(); + assert_eq!(stats.pending, 0, "nothing is owed to the device any more"); + assert_eq!(stats.completed, 0, "a retain is not a free"); + assert_eq!(stats.quarantined, 1); + assert_eq!(stats.retained, 1); + assert_eq!(stats.mapped_refunded_bytes, 0, "no refund on assumption"); + + let retained = queue.retained(); + assert_eq!(retained.len(), 1); + assert_eq!(retained[0].state, AllocationReleaseState::DeviceLost); + assert_eq!(retained[0].bytes, bytes as u64); + assert_eq!(retained[0].label, "provider allocation"); +} + +#[test] +fn device_loss_settles_a_pending_prepared_release_and_frees_the_context() { + let h = harness(8); + let queue = Arc::clone(&h.queue); + let fixture = BoundFixture::new(&queue); + let before_enqueue = Arc::strong_count(&queue); + + let identity = fixture.enqueue_one(&queue, 4096); + assert_eq!(queue.pending(), 1); + assert_eq!(fixture.snapshot().queued_releases, 1); + assert_eq!(fixture.snapshot().active_operations, 1); + + // The reported path: the request is still pending when loss is published. + queue.mark_device_lost("Xid 13"); + + assert_eq!( + h.destroyed.load(Ordering::Acquire), + 0, + "events are retained, not destroyed through a lost context" + ); + assert_settled_device_lost(&fixture, identity, 4096, &queue); + assert_eq!( + Arc::strong_count(&queue), + before_enqueue, + "settlement must not leave the queue holding a new reference to itself" + ); + + fixture.tear_down_after_device_loss(); + // The harness's own handle is the only other one; drop it so the count + // below is exactly "what the settled release left behind". + drop(h); + assert_eq!( + Arc::strong_count(&queue), + 1, + "the retained residual must not keep the provider context — and therefore \ + this queue — alive" + ); + // The residual is still held, so the physical ownership is not reusable. + assert_eq!(queue.retained().len(), 1); +} + +#[test] +fn device_loss_settles_a_prepared_release_a_poller_was_holding() { + let entered = Arc::new(Barrier::new(2)); + let resume = Arc::new(Barrier::new(2)); + let queue = CudaDeferredReleaseQueue::manual( + Box::new(BlockingStreams { + entered: Arc::clone(&entered), + resume: Arc::clone(&resume), + }), + 8, + ); + let fixture = BoundFixture::new(&queue); + let before_enqueue = Arc::strong_count(&queue); + let identity = fixture.enqueue_one(&queue, 8192); + + let poll_queue = Arc::clone(&queue); + let poller = std::thread::spawn(move || poll_queue.poll()); + entered.wait(); + let loss_queue = Arc::clone(&queue); + let loss = std::thread::spawn(move || loss_queue.mark_device_lost("concurrent Xid")); + resume.wait(); + + assert_eq!(poller.join().expect("poller"), 0); + loss.join().expect("loss marker"); + + // The carry path settles the same way the pending path does. + assert_settled_device_lost(&fixture, identity, 8192, &queue); + assert_eq!( + Arc::strong_count(&queue), + before_enqueue, + "the carry path must not retain a reference to the queue either" + ); + + fixture.tear_down_after_device_loss(); + assert_eq!( + Arc::strong_count(&queue), + 1, + "a carried request must not keep the provider context alive" + ); +} + +#[test] +fn concurrent_loss_settles_every_prepared_release_whichever_path_takes_it() { + // Pollers and device loss race, so the entries are settled by a mix of the + // already-lost, carry, and retain-all paths. Whichever site wins, every + // request must reach the same terminal state and the queue must end up + // holding nothing that points back at it. + let h = harness(64); + let queue = Arc::clone(&h.queue); + let fixture = BoundFixture::new(&queue); + let before_enqueue = Arc::strong_count(&queue); + + let mut identities = Vec::new(); + for _ in 0..32 { + identities.push(fixture.enqueue_one(&queue, 1024)); + } + + let start = Arc::new(Barrier::new(3)); + let pollers: Vec<_> = (0..2) + .map(|_| { + let queue = Arc::clone(&queue); + let start = Arc::clone(&start); + std::thread::spawn(move || { + start.wait(); + for _ in 0..64 { + queue.poll(); + } + }) + }) + .collect(); + start.wait(); + queue.mark_device_lost("racing Xid"); + for poller in pollers { + poller.join().expect("poller"); + } + // Anything a poller carried back after loss is settled by a final poll. + queue.poll(); + + assert_eq!( + fixture.allocator.release_calls(), + 0, + "no release may reach the allocator once loss is published" + ); + assert_eq!(queue.pending(), 0, "nothing stays owed"); + let quarantined = fixture.quarantined(); + assert_eq!(quarantined.len(), identities.len()); + for record in &quarantined { + assert_eq!(record.state, AllocationReleaseState::DeviceLost); + assert_eq!(record.reason, QuarantineReason::DeviceLost); + } + let recorded: std::collections::HashSet<_> = + quarantined.iter().map(|record| record.identity).collect(); + for identity in &identities { + assert!( + recorded.contains(identity), + "allocation {identity:?} was never settled" + ); + } + let snapshot = fixture.snapshot(); + assert_eq!(snapshot.queued_releases, 0); + assert_eq!(snapshot.active_operations, 0); + assert_eq!( + Arc::strong_count(&queue), + before_enqueue, + "no settled request may leave a queue reference behind" + ); + + fixture.tear_down_after_device_loss(); + drop(h); + assert_eq!(Arc::strong_count(&queue), 1); +} diff --git a/crates/onnx-runtime-memory-api/src/deferred.rs b/crates/onnx-runtime-memory-api/src/deferred.rs index ccc0d81186..c8308c7654 100644 --- a/crates/onnx-runtime-memory-api/src/deferred.rs +++ b/crates/onnx-runtime-memory-api/src/deferred.rs @@ -697,6 +697,38 @@ impl PreparedAllocationRelease { } } + /// Settle this request as device-lost quarantine, without calling the + /// allocator and without refunding anything. + /// + /// This is the same settlement [`execute`](Self::execute) performs when it + /// observes a device-lost release gate, exposed for the case where the + /// *provider's* queue learns the context is unusable first: the mechanism + /// lifecycle may not have been invalidated yet, so `execute` would still be + /// allowed to call the allocator, which is exactly what must not happen. + /// + /// It is deliberately distinct from + /// [`quarantine(QuarantineReason::DeviceLost)`](Self::quarantine), which + /// records the generic [`AllocationReleaseState::Quarantined`]. Device loss + /// is its own terminal state because it is discharged by confirmed + /// context/process termination rather than by anything the runtime can do + /// to the device. + /// + /// Consuming the request is the point: the binding records the exact + /// allocation identity as device-lost, the queued-release count settles, + /// and the active-operation pin — with the mechanism, provider-context, and + /// binding references behind it — is released, so a queue that holds the + /// residual does not keep its own provider context alive. + pub fn quarantine_device_lost(mut self) -> AllocationReleaseOutcome { + self.armed = false; + let bytes = self.bytes as u64; + self.settle_quarantine( + ReleaseAccounting::new(bytes, 0), + AllocationReleaseState::DeviceLost, + QuarantineReason::DeviceLost, + bytes, + ) + } + /// Retain ownership deliberately without calling the allocator. /// /// This is the explicit form of what `Drop` does implicitly.