Struct MemoryTracker 

pub struct MemoryTracker { /* private fields */ }

Memory usage tracker for GPU operations.

Tracks allocated and peak memory usage across operations. Thread-safe via atomics.

Why This Struct

Unlike CPU memory which is managed by the OS, GPU memory requires explicit tracking because:

1. No swap: When VRAM runs out, allocations fail immediately
2. Fragmentation: Repeated allocations can fragment the heap
3. Debugging: Memory leaks on GPU are harder to diagnose than CPU leaks

Usage Pattern

use rust_ai_core::MemoryTracker;

let tracker = MemoryTracker::new();

// Before allocation
tracker.allocate(1024 * 1024).expect("allocation should succeed"); // 1 MB

// After freeing
tracker.deallocate(1024 * 1024);

println!("Peak usage: {} bytes", tracker.peak_bytes());

Implementations

impl MemoryTracker

pub fn new() -> Self

Create a new memory tracker with no limit.

Why No Default Limit

Different GPUs have vastly different VRAM capacities (4GB to 80GB+). Setting a default limit would either be too restrictive for high-end cards or meaningless for consumer cards. Users should set limits explicitly.

pub fn with_limit(limit_bytes: usize) -> Self

Create a tracker with a memory limit.

Arguments

• limit_bytes - Maximum allowed allocation in bytes

Why Memory Limits

Setting a limit below actual VRAM capacity reserves space for:

• CUDA context overhead (~200-500 MB)
• Framework allocations (PyTorch/Candle tensor cache)
• Other processes sharing the GPU

pub fn with_overhead_factor(self, factor: f64) -> Self

Set a custom overhead factor for estimates.

Arguments

• factor - Multiplier applied to estimates (default: 1.1)

pub fn allocate(&self, bytes: usize) -> Result<()>

Record a memory allocation.

Arguments

• bytes - Number of bytes allocated

Returns

Ok(()) if allocation is within limits.

Errors

Returns CoreError::OutOfMemory if allocation would exceed the limit.

Why Track Allocations

Explicit tracking catches memory issues early. Without tracking, OOM errors occur deep in CUDA kernels with unhelpful error messages.

pub fn deallocate(&self, bytes: usize)

Record a memory deallocation.

Arguments

• bytes - Number of bytes freed

pub fn allocated_bytes(&self) -> usize

Get currently allocated bytes.

pub fn peak_bytes(&self) -> usize

Get peak allocation during tracker lifetime.

Why Track Peak

Peak usage is more useful than current usage for capacity planning. It shows the high-water mark needed to complete a workload.

pub fn limit_bytes(&self) -> usize

Get configured memory limit (0 = unlimited).

pub fn estimate_with_overhead(&self, shape: &[usize], dtype: DType) -> usize

Estimate required memory with overhead factor applied.

Arguments

• shape - Tensor dimensions
• dtype - Data type

Returns

Estimated bytes including overhead buffer.

pub fn would_fit(&self, bytes: usize) -> bool

Check if an allocation would fit within limits.

Arguments

• bytes - Proposed allocation size

Returns

true if allocation would succeed.

pub fn reset(&self)

Reset the tracker to initial state.

Why Reset

Between training epochs or inference batches, resetting allows tracking per-phase memory usage without creating new tracker instances.

Trait Implementations

impl Debug for MemoryTracker

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for MemoryTracker

fn default() -> Self

Returns the “default value” for a type.