Struct MemoryGuard 

pub struct MemoryGuard { /* private fields */ }

Central memory-safety controller for a single GPU.

Wraps a GpuDevice and provides:

• Optional upfront VRAM reservation (sentinel allocation).
• Budget enforcement — allocations that would exceed the budget are rejected before touching the driver.
• Configurable OOM recovery via OomPolicy.
• An emergency-checkpoint callback.

Construct via MemoryGuardBuilder.

Implementations

impl MemoryGuard

pub fn set_budget(&self, bytes: usize)

Set a hard budget in bytes. Allocations that would push used_bytes past this limit return GpuError::BudgetExceeded without touching the driver.

Pass 0 to remove the budget (unlimited).

pub fn budget(&self) -> usize

Current budget (0 = unlimited).

pub fn on_oom<F: Fn() + Send + Sync + 'static>(&self, f: F)

Register a callback that will be invoked on OOM when the policy is OomPolicy::CheckpointAndFail. Typically used to save a training checkpoint so work is not lost.

pub fn set_oom_policy(&self, policy: OomPolicy)

Change the OOM policy at runtime.

pub fn register_hook(&self, hook: MemoryHook)

Register a pre-OOM hook.

Hooks are called (in priority order, lowest first) when an allocation would exceed the budget. Each hook gets a chance to free memory before the guard falls through to the OomPolicy.

pub fn remove_hook(&self, name: &str) -> bool

Remove a previously registered hook by name.

Returns true if a hook with that name was found and removed.

pub fn pressure_level(&self) -> PressureLevel

Current pressure level based on budget usage.

If no budget is set (budget = 0 / unlimited), always returns PressureLevel::None.

pub fn add_pressure_listener(&self, listener: Box<dyn MemoryPressureListener>)

Register a listener that is notified whenever the pressure level changes (checked after every allocation and free through the guard).

pub fn safe_alloc_with_hooks<T>(&self, count: usize) -> GpuResult<CudaBuffer<T>>

where T: DeviceRepr + ValidAsZeroBits,

Allocate count zero-initialized elements, trying pre-OOM hooks before falling through to the OomPolicy.

The algorithm:

1. Check if the allocation fits within the budget – if so, allocate directly.
2. If not, compute the shortfall.
3. Sort hooks by (priority, estimated_free_bytes descending).
4. Call hooks one at a time, skipping any whose execution_overhead_bytes exceeds current headroom, until enough cumulative memory has been freed.
5. Retry the allocation.
6. If still insufficient after all hooks, fall through to the regular OomPolicy path.

pub fn release_reservation(&self) -> usize

Release the upfront reservation, making its memory available for normal allocations. Returns the number of bytes released, or 0 if there was no active reservation.

pub fn has_reservation(&self) -> bool

Whether an upfront reservation is currently held.

pub fn safe_alloc<T>(&self, count: usize) -> GpuResult<CudaBuffer<T>>

where T: DeviceRepr + ValidAsZeroBits,

Allocate count zero-initialized elements on the device, enforcing the budget and OOM policy.

This is the primary allocation entry point when using the memory guard. Prefer this over raw CudaAllocator::alloc_zeros.

pub fn safe_alloc_copy<T>(&self, data: &[T]) -> GpuResult<CudaBuffer<T>>

where T: DeviceRepr,

Allocate by copying host data to the device, enforcing budget and OOM policy.

pub fn free<T>(&self, buffer: CudaBuffer<T>)

Return a buffer to the guard, freeing GPU memory and updating statistics.

pub fn stats(&self) -> MemoryStats

Snapshot the current memory statistics.

pub fn reset_peak_stats(&self)

Reset the peak-usage counter to the current usage level.

pub fn device(&self) -> &GpuDevice

The underlying device.

pub fn device_arc(&self) -> &Arc<GpuDevice>

The underlying device as an Arc.

Trait Implementations

impl Debug for MemoryGuard

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

Formats the value using the given formatter.

impl Send for MemoryGuard

impl Sync for MemoryGuard