onnx_runtime_memory_api/
lib.rs

1//! # `onnx-runtime-memory-api`
2//!
3//! Low-dependency memory mechanism contracts shared by allocators, governors,
4//! and execution providers.
5//!
6//! This crate is the lowest layer of the runtime memory stack. It owns the
7//! minimum ordinary allocator contract, explicit optional virtual-backing and
8//! shared-mapping capabilities, manager-issued binding identity/lifetime pins,
9//! and the owning/deferred-release contract that says who owns a physical
10//! release and what is true after one partially fails. It does not own
11//! allocation policy, accounting, synchronization, or process-level transaction
12//! management.
13//!
14//! Governor-specific capacity tokens and grants remain in
15//! `onnx-runtime-memory-governor`; they are not methods every allocator or
16//! optional capability must implement.
17
18pub mod allocator;
19pub mod binding;
20pub mod capability;
21pub mod context_pin;
22pub mod deferred;
23
24pub use allocator::{
25    AllocationCommitRange, DeviceAllocator, DeviceKey, HostAllocator, MappedAllocation,
26    SharedDevicePrefix, SharedPrefixCommitInfo,
27};
28pub use binding::{
29    AllocationGeneration, AllocationIdentity, AuthorityIdentity, BindingError, BindingGeneration,
30    BindingId, BindingIdentity, BindingRegistry, BindingResource, BoundAllocation, BoundMemoryView,
31    BoundSharedMapping, BoundSharedPrefix, BoundVirtualBacking, ExplicitReleaseError,
32    MechanismCoherence, MechanismIdentity, MechanismLifecycle, MechanismSnapshot, MemoryBinding,
33    OwnedView, OwningAllocation, OwningReleaseError, ProviderContextIdentity, RegisteredAuthority,
34    RegisteredMechanism, RegisteredProviderContext, ValidatedMemoryView,
35};
36pub use capability::{SharedMapping, VirtualBacking};
37pub use context_pin::{ProviderContextPin, ProviderContextPinError, ProviderContextPinSource};
38pub use deferred::{
39    AllocationReleaseOutcome, AllocationReleaseState, DeferredEnqueueError,
40    DeferredEnqueueRejection, DeferredReleaseDisposition, DeferredReleaseQueue,
41    PreparedAllocationRelease, QuarantineReason, QuarantinedAllocation, ReleaseAccounting,
42    ReleaseFailure, ResidualOwnership,
43};
44
45/// Where the bytes physically live.
46///
47/// Ordered from fastest to slowest, which is also the demotion order.
48#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
49pub enum Tier {
50    /// Accelerator memory (VRAM).
51    Device,
52    /// Host RAM.
53    Host,
54    /// Spill file on disk.
55    Disk,
56}
57
58impl Tier {
59    /// Every tier, fastest first.
60    pub const ALL: [Tier; 3] = [Tier::Device, Tier::Host, Tier::Disk];
61
62    pub const fn index(self) -> usize {
63        match self {
64            Tier::Device => 0,
65            Tier::Host => 1,
66            Tier::Disk => 2,
67        }
68    }
69
70    /// Human-facing name used in error messages.
71    pub const fn name(self) -> &'static str {
72        match self {
73            Tier::Device => "device",
74            Tier::Host => "host",
75            Tier::Disk => "disk",
76        }
77    }
78}
79
80/// What a reservation is for.
81#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
82pub enum MemoryRole {
83    KvCache,
84    Workspace { step_scoped: bool },
85    Weights,
86    Activation,
87}
88
89/// Shared error vocabulary for mechanism and governance operations.
90///
91/// Not `Clone`/`PartialEq`: a refusal that carries the cause underneath it
92/// cannot be meaningfully duplicated or compared, and keeping the cause is
93/// worth more than either. Match on the variant instead.
94#[derive(Debug, thiserror::Error)]
95pub enum MemoryError {
96    /// The tier does not have room, and no holder released enough.
97    #[error(
98        "cannot reserve {requested} bytes of {tier} memory for {role:?}: {used} of {limit} bytes \
99         are already leased, leaving {available}; free memory by closing sessions, lower the \
100         demand, or raise the {tier} limit"
101    )]
102    TierExhausted {
103        /// Which tier ran out.
104        tier: &'static str,
105        /// What the caller asked for.
106        requested: u64,
107        /// Bytes leased before this request.
108        used: u64,
109        /// The tier ceiling.
110        limit: u64,
111        /// `limit - used`.
112        available: u64,
113        /// The role that was refused.
114        role: MemoryRole,
115    },
116    /// The request itself is not representable.
117    #[error("cannot reserve {requested} bytes of {tier} memory: {reason}")]
118    InvalidRequest {
119        /// Which tier was addressed.
120        tier: &'static str,
121        /// What the caller asked for.
122        requested: u64,
123        /// What is wrong with it.
124        reason: &'static str,
125    },
126    /// The request was well formed and within budget, but the allocator behind
127    /// the tier refused it for a reason of its own.
128    ///
129    /// Distinct from [`MemoryError::TierExhausted`], which means *we* declined,
130    /// and from [`MemoryError::InvalidRequest`], which means the caller asked
131    /// for something impossible. This one carries the backing allocator's own
132    /// account of the failure, which is usually the only thing that identifies
133    /// it: a driver that is out of memory and a driver that has no context both
134    /// fail an allocation, and calling them both "out of memory" sends the next
135    /// person to read the log in the wrong direction.
136    #[error("cannot allocate {requested} bytes of {tier} memory: {reason}")]
137    AllocationFailed {
138        /// Which tier was addressed.
139        tier: &'static str,
140        /// What the caller asked for.
141        requested: u64,
142        /// What the backing allocator said.
143        reason: String,
144    },
145    /// A well-formed capacity transfer or backing claim could not make enough
146    /// governed bytes available.
147    #[error(
148        "cannot make {requested} bytes of {tier} capacity available for {role:?}: only \
149         {available} bytes became available; {detail}"
150    )]
151    CapacityUnavailable {
152        tier: &'static str,
153        requested: u64,
154        available: u64,
155        role: MemoryRole,
156        /// What this layer can say about the shortfall on its own.
157        ///
158        /// Names the operation that came up short; the refusal underneath it,
159        /// when there was one, belongs in `source` rather than being folded in
160        /// here, so that a caller can still match on it and a reader is not
161        /// shown the same sentence twice.
162        detail: String,
163        /// The refusal this one is reporting, when it is reporting one.
164        ///
165        /// `None` when this layer decided on its own, as when a reclaim target
166        /// simply was not reached. Typed as `dyn Error` rather than a boxed
167        /// [`MemoryError`] both because the layer underneath is not always a
168        /// governor and because `#[source]` on a `Box<ConcreteError>` hands
169        /// callers a chain node whose concrete type is the *box*, so
170        /// `downcast_ref::<MemoryError>()` would miss it -- which is the whole
171        /// thing this field exists to make possible.
172        #[source]
173        source: Option<Box<dyn std::error::Error + Send + Sync>>,
174    },
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180
181    #[test]
182    fn tier_order_and_names_are_stable() {
183        assert_eq!(Tier::ALL, [Tier::Device, Tier::Host, Tier::Disk]);
184        assert_eq!(Tier::Device.name(), "device");
185        assert_eq!(Tier::Host.name(), "host");
186        assert_eq!(Tier::Disk.name(), "disk");
187    }
188}
onnx_runtime_memory_api/lib.rs

onnx_runtime_memory_api/
lib.rs
