metrale_gpu_runtime/cuda_backend/
memory.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Device free-memory queries, which add host `MemAvailable` only
4//! on an integrated GPU, and the OOM watchdog that polls them.
5//!
6//! Owner: gpu-runtime (CUDA backend).
7//! Invariants:
8//! - At most one OOM watchdog task is spawned per process (`WATCHDOG_RUNNING`).
9
10use super::*;
11
12/// 2026-09-25: Free device memory without a `GpuBackend`, for the OOM
13/// watchdog and the TUI. `None` when `cuMemGetInfo_v2` fails.
14///
15/// The same rule as `free_memory` (`effective_free_bytes`), except that a
16/// failed integrated query counts as discrete (`polled_free_bytes`).
17pub fn cuda_free_memory_bytes() -> Option<usize> {
18    let mut free: usize = 0;
19    let mut total: usize = 0;
20    let status = unsafe { cuMemGetInfo_v2(&mut free, &mut total) };
21    if status != 0 {
22        return None;
23    }
24    Some(polled_free_bytes(
25        free,
26        system_available_memory_bytes(),
27        current_device_is_integrated(),
28    ))
29}
30
31/// 2026-09-25: `CU_DEVICE_ATTRIBUTE_INTEGRATED` of the current context's
32/// device. Requires a current context (`cuCtxGetDevice`); fails, rather than
33/// guessing, when the driver does not answer. Both callers query
34/// `cuMemGetInfo_v2` first, which needs the same context.
35///
36/// Measured 2026-09-04: attribute 18 reads 1 on GB10 and 0 on RTX PRO 6000
37/// Blackwell; `CU_DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS` (99) reads 1 on
38/// both, so it does not tell them apart.
39pub(crate) fn device_is_integrated() -> Result<bool> {
40    const CU_DEVICE_ATTRIBUTE_INTEGRATED: u32 = 18;
41    let mut dev: i32 = 0;
42    let status = unsafe { cuCtxGetDevice(&mut dev) };
43    if status != 0 {
44        bail!("cuCtxGetDevice failed: status {status}");
45    }
46    let mut integrated: i32 = 0;
47    let status =
48        unsafe { cuDeviceGetAttribute(&mut integrated, CU_DEVICE_ATTRIBUTE_INTEGRATED, dev) };
49    if status != 0 {
50        bail!("cuDeviceGetAttribute(INTEGRATED) failed: status {status}");
51    }
52    Ok(integrated != 0)
53}
54
55/// 2026-09-25: [`device_is_integrated`] as an `Option`, for
56/// `cuda_free_memory_bytes`, which returns `Option` and cannot report an
57/// error.
58fn current_device_is_integrated() -> Option<bool> {
59    device_is_integrated().ok()
60}
61
62/// 2026-09-25: Free device memory to report: the larger of `cu_free` and
63/// `mem_available` on an integrated GPU, `cu_free` otherwise or when
64/// `mem_available` is `None`. On a discrete GPU host RAM is a separate pool.
65/// Pure, so the rule is tested without a GPU.
66pub(crate) fn effective_free_bytes(
67    cu_free: usize,
68    mem_available: Option<usize>,
69    integrated: bool,
70) -> usize {
71    match mem_available {
72        Some(avail) if integrated => cu_free.max(avail),
73        _ => cu_free,
74    }
75}
76
77/// 2026-09-25: Whether a watchdog has been spawned. Process-wide, because the
78/// watchdog reads the device's free memory, which does not depend on the
79/// model: a second model load must not start a second watchdog.
80static WATCHDOG_RUNNING: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
81
82/// 2026-09-25: Spawn the OOM watchdog on the tokio runtime, or return `None`
83/// if one was already spawned.
84///
85/// Every `interval` it reads `cuda_free_memory_bytes`. Below `threshold_mb`
86/// MiB it logs an error; at the third consecutive such reading it calls
87/// `std::process::exit(1)`. A reading at or above the threshold resets the
88/// count, and a `None` reading leaves it unchanged. Dropping the returned
89/// handle does not stop the task.
90pub fn spawn_oom_watchdog(
91    threshold_mb: usize,
92    interval: std::time::Duration,
93) -> Option<tokio::task::JoinHandle<()>> {
94    if WATCHDOG_RUNNING.swap(true, std::sync::atomic::Ordering::SeqCst) {
95        return None;
96    }
97    let threshold_bytes = threshold_mb * 1024 * 1024;
98    Some(tokio::spawn(async move {
99        let mut tick = tokio::time::interval(interval);
100        let mut consecutive_low = 0u32;
101        loop {
102            tick.tick().await;
103            if let Some(free) = cuda_free_memory_bytes() {
104                if free < threshold_bytes {
105                    consecutive_low += 1;
106                    let free_mb = free / (1024 * 1024);
107                    tracing::error!(
108                        "OOM watchdog: GPU free memory critically low: {} MB (threshold: {} MB) [{}/3]",
109                        free_mb,
110                        threshold_mb,
111                        consecutive_low,
112                    );
113                    if consecutive_low >= 3 {
114                        tracing::error!(
115                            "OOM watchdog: 3 consecutive readings below threshold. \
116                             Terminating to prevent system freeze."
117                        );
118                        std::process::exit(1);
119                    }
120                } else {
121                    consecutive_low = 0;
122                }
123            }
124        }
125    }))
126}
metrale_gpu_runtime/cuda_backend/memory.rs

metrale_gpu_runtime/cuda_backend/
memory.rs
