dynamo_memory/
lib.rs

1// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Clean, minimal storage API for v2 block manager.
5//!
6//! This module provides a simplified storage abstraction with:
7//! - Single trait for type erasure (`MemoryDescriptor`)
8//! - Concrete storage types (no trait implementations required)
9//! - Composition-based NIXL registration via `NixlRegistered<T>` wrapper
10//! - RAII with proper drop ordering (registration handle drops before memory)
11
12#![deny(missing_docs)]
13
14pub mod actions;
15pub mod arena;
16pub mod nixl;
17#[cfg(target_os = "linux")]
18pub mod numa;
19
20/// Offset-based buffer views into underlying storage.
21pub mod offset;
22
23/// CUDA memory pool utilities.
24pub mod pool;
25
26/// Common imports for working with memory types.
27pub mod prelude;
28
29mod device;
30mod disk;
31mod external;
32mod pinned;
33mod system;
34mod tensor;
35
36#[cfg(test)]
37mod tests;
38
39pub use arena::{ArenaAllocator, ArenaBuffer, ArenaError};
40pub use device::DeviceStorage;
41pub use disk::DiskStorage;
42pub use external::ExternalDeviceMemory;
43#[cfg(target_os = "linux")]
44pub use numa::{NumaNode, is_numa_disabled, is_numa_enabled};
45pub use offset::OffsetBuffer;
46pub use pinned::PinnedStorage;
47pub use pool::{CudaMemPool, CudaMemPoolBuilder};
48pub use system::SystemStorage;
49pub use tensor::{TensorDescriptor, TensorDescriptorExt};
50
51use serde::{Deserialize, Serialize};
52use std::any::Any;
53use std::fmt;
54use std::sync::Arc;
55use thiserror::Error;
56
57/// Result type for storage operations.
58pub type Result<T> = std::result::Result<T, StorageError>;
59
60/// Core trait for memory regions that can be type-erased.
61///
62/// This is the only trait in the storage API. Concrete storage types
63/// implement this trait to enable type erasure via `Arc<dyn MemoryDescriptor>`.
64pub trait MemoryDescriptor: Send + Sync + fmt::Debug {
65    /// Base address of the memory region.
66    fn addr(&self) -> usize;
67
68    /// Size of the memory region in bytes.
69    fn size(&self) -> usize;
70
71    /// Type of storage backing this region.
72    fn storage_kind(&self) -> StorageKind;
73
74    /// Enable downcasting to concrete type.
75    fn as_any(&self) -> &dyn Any;
76
77    /// Get the NIXL descriptor for this memory region.
78    fn nixl_descriptor(&self) -> Option<nixl::NixlDescriptor>;
79}
80
81/// Errors that can occur during storage operations.
82#[derive(Debug, Error)]
83#[allow(missing_docs)]
84pub enum StorageError {
85    #[error("allocation failed: {0}")]
86    AllocationFailed(String),
87
88    #[error("registration failed: {0}")]
89    RegistrationFailed(String),
90
91    #[error("operation failed: {0}")]
92    OperationFailed(String),
93
94    #[error("unsupported operation: {0}")]
95    Unsupported(String),
96
97    #[error("I/O error: {0}")]
98    Io(#[from] std::io::Error),
99
100    // #[cfg(feature = "cuda")]
101    #[error("CUDA error: {0}")]
102    Cuda(#[from] cudarc::driver::DriverError),
103
104    #[error("NIXL error: {0}")]
105    Nixl(#[from] nixl_sys::NixlError),
106}
107
108/// Storage type classification.
109#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
110pub enum StorageKind {
111    /// System memory (malloc)
112    System,
113
114    /// CUDA pinned host memory
115    // #[cfg(feature = "cuda")]
116    Pinned,
117
118    /// CUDA device memory with device ID
119    // #[cfg(feature = "cuda")]
120    Device(u32),
121
122    /// Disk-backed memory (mmap)
123    Disk(u64),
124}
125
126impl StorageKind {
127    /// Returns the CUDA device index if this is device memory.
128    pub fn cuda_device_index(&self) -> Option<u32> {
129        match self {
130            StorageKind::Device(idx) => Some(*idx),
131            _ => None,
132        }
133    }
134
135    /// Returns true if this is CUDA device memory.
136    pub fn is_cuda(&self) -> bool {
137        matches!(self, StorageKind::Device(_))
138    }
139
140    /// Returns true if this is system memory (malloc).
141    pub fn is_system(&self) -> bool {
142        matches!(self, StorageKind::System)
143    }
144
145    /// Returns true if this is CUDA pinned host memory.
146    pub fn is_pinned(&self) -> bool {
147        matches!(self, StorageKind::Pinned)
148    }
149
150    /// Returns true if this is disk-backed memory.
151    pub fn is_disk(&self) -> bool {
152        matches!(self, StorageKind::Disk(_))
153    }
154}
155
156/// Type-erased memory region for use in layouts.
157#[derive(Clone)]
158pub struct Buffer(Arc<dyn MemoryDescriptor>);
159
160impl Buffer {
161    /// Wraps a concrete storage type into a type-erased [`Buffer`].
162    ///
163    /// This is the primary way to create a `Buffer` from any type that
164    /// implements [`MemoryDescriptor`].
165    pub fn new<S: MemoryDescriptor + 'static>(memory: S) -> Self {
166        Buffer(Arc::new(memory))
167    }
168}
169
170impl MemoryDescriptor for Buffer {
171    fn addr(&self) -> usize {
172        self.0.addr()
173    }
174    fn size(&self) -> usize {
175        self.0.size()
176    }
177    fn storage_kind(&self) -> StorageKind {
178        self.0.storage_kind()
179    }
180    fn as_any(&self) -> &dyn Any {
181        self.0.as_any()
182    }
183    fn nixl_descriptor(&self) -> Option<nixl::NixlDescriptor> {
184        self.0.nixl_descriptor()
185    }
186}
187
188impl std::ops::Deref for Buffer {
189    type Target = dyn MemoryDescriptor;
190
191    fn deref(&self) -> &Self::Target {
192        self.0.as_ref()
193    }
194}
195
196impl std::fmt::Debug for Buffer {
197    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
198        f.debug_struct("Buffer")
199            .field("addr", &self.addr())
200            .field("size", &self.size())
201            .field("kind", &self.storage_kind())
202            .finish()
203    }
204}
205
206/// Helper function to convert concrete storage to type-erased form.
207pub fn create_buffer<S: MemoryDescriptor + 'static>(memory: S) -> Buffer {
208    Buffer(Arc::new(memory))
209}
210
211impl Buffer {
212    /// Create a Buffer from an existing Arc<dyn MemoryDescriptor>.
213    pub fn from_arc(arc: Arc<dyn MemoryDescriptor>) -> Self {
214        Buffer(arc)
215    }
216}
217
218// From implementations for ergonomic Buffer creation
219impl From<Arc<dyn MemoryDescriptor>> for Buffer {
220    fn from(arc: Arc<dyn MemoryDescriptor>) -> Self {
221        Buffer::from_arc(arc)
222    }
223}
224
225impl From<Arc<dyn nixl::NixlMemory + Send + Sync>> for Buffer {
226    fn from(arc: Arc<dyn nixl::NixlMemory + Send + Sync>) -> Self {
227        // Arc<dyn NixlMemory> implements MemoryDescriptor, so we can wrap it
228        Buffer::new(arc)
229    }
230}
231
232/// An unowned contiguous chunk of memory, not storage specific.
233#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
234pub struct MemoryRegion {
235    /// Start address of the memory region.
236    pub addr: usize,
237
238    /// Size of the memory region in bytes.
239    pub size: usize,
240}
241
242impl MemoryRegion {
243    /// Creates a new memory region with the given base address and size.
244    pub fn new(addr: usize, size: usize) -> Self {
245        Self { addr, size }
246    }
247
248    /// Returns the base address of this memory region.
249    #[inline]
250    pub fn addr(&self) -> usize {
251        self.addr
252    }
253
254    /// Returns the size of this memory region in bytes.
255    #[inline]
256    pub fn size(&self) -> usize {
257        self.size
258    }
259
260    /// Get a slice view of this memory region.
261    ///
262    /// # Safety
263    /// This is unsafe because:
264    /// - The caller must ensure the memory region is valid and properly initialized
265    /// - The caller must ensure no mutable references exist to this memory
266    /// - The caller must ensure the memory remains valid for the lifetime of the slice
267    #[cfg(feature = "unsafe-slices")]
268    pub unsafe fn as_slice(&self) -> Result<&[u8]> {
269        if self.size == 0 {
270            return Ok(&[]);
271        }
272        // SAFETY: Caller guarantees memory is valid
273        unsafe {
274            Ok(std::slice::from_raw_parts(
275                self.addr as *const u8,
276                self.size,
277            ))
278        }
279    }
280
281    /// Get a mutable slice view of this memory region.
282    ///
283    /// # Safety
284    /// This is unsafe because:
285    /// - The caller must ensure the memory region is valid and properly initialized
286    /// - The caller must ensure no other references (mutable or immutable) exist to this memory
287    /// - The caller must ensure the memory remains valid for the lifetime of the slice
288    #[cfg(feature = "unsafe-slices")]
289    pub unsafe fn as_slice_mut(&mut self) -> Result<&mut [u8]> {
290        if self.size == 0 {
291            return Ok(&mut []);
292        }
293        // SAFETY: Caller guarantees memory is valid and exclusively accessible
294        unsafe {
295            Ok(std::slice::from_raw_parts_mut(
296                self.addr as *mut u8,
297                self.size,
298            ))
299        }
300    }
301}
302
303// Canonical truthy/bool parsing, re-exported from the shared `dynamo-truthy`
304// crate (this crate cannot depend on `dynamo-runtime`, whose `config` module
305// re-exports the same helpers).
306pub use dynamo_truthy::{env_is_truthy, parse_bool, parse_bool_opt};
dynamo_memory/lib.rs

dynamo_memory/
lib.rs
