mermaid_cli/utils/
host_memory.rs

1//! Best-effort host memory detection for Ollama auto-sizing.
2//!
3//! Mermaid sizes an Ollama model's `num_ctx` to fit the host's memory so the
4//! model stays on the GPU (CPU/RAM offload is 5–20× slower, so it's off by
5//! default). That needs two numbers: total system RAM (the budget when the user
6//! opts into RAM offload) and total VRAM (the default budget).
7//!
8//! Both are **best-effort and cached** — VRAM detection is GPU-vendor specific
9//! and there is no Ollama API for it, so we shell out to `nvidia-smi` and fall
10//! back gracefully. When VRAM can't be detected the sizing logic uses a
11//! conservative fixed cap (see [`crate::models::adapters::ollama_sizing`]), so a
12//! `None` here never breaks anything.
13
14use crate::constants::NVIDIA_SMI_TIMEOUT_SECS;
15use std::sync::OnceLock;
16
17/// Total physical system RAM in bytes, or `None` if it can't be read. Cached for
18/// the process (RAM doesn't change during a session).
19pub fn system_ram_bytes() -> Option<u64> {
20    static CACHE: OnceLock<Option<u64>> = OnceLock::new();
21    *CACHE.get_or_init(detect_system_ram)
22}
23
24fn detect_system_ram() -> Option<u64> {
25    let mut sys = sysinfo::System::new();
26    sys.refresh_memory();
27    let total = sys.total_memory(); // bytes (sysinfo ≥ 0.30)
28    (total > 0).then_some(total)
29}
30
31/// Total GPU VRAM in bytes (best-effort), or `None` if no GPU is detected.
32///
33/// On unified-memory Macs there is no separate, slow "RAM offload" target, so
34/// system RAM *is* the GPU budget. On NVIDIA we read `nvidia-smi`. Everything
35/// else (AMD, integrated, headless) returns `None` and the caller uses the
36/// conservative fallback cap. Cached for the process.
37pub async fn gpu_vram_bytes() -> Option<u64> {
38    static CACHE: tokio::sync::OnceCell<Option<u64>> = tokio::sync::OnceCell::const_new();
39    *CACHE.get_or_init(detect_vram).await
40}
41
42async fn detect_vram() -> Option<u64> {
43    if let Some(v) = nvidia_vram_bytes().await {
44        return Some(v);
45    }
46    // Apple Silicon / Intel Macs: unified memory — the GPU draws from system RAM,
47    // and there's no slow separate pool to "offload" to, so treat RAM as VRAM.
48    if cfg!(target_os = "macos") {
49        return system_ram_bytes();
50    }
51    None
52}
53
54/// Largest single NVIDIA GPU's total VRAM via `nvidia-smi`, in bytes. A model
55/// loads onto one GPU unless it spans, so we take the max single GPU rather than
56/// assume the pool. Best-effort: absent binary, non-zero exit, or timeout → `None`.
57async fn nvidia_vram_bytes() -> Option<u64> {
58    let mut cmd = tokio::process::Command::new("nvidia-smi");
59    cmd.args(["--query-gpu=memory.total", "--format=csv,noheader,nounits"]);
60    cmd.kill_on_drop(true);
61
62    let out = tokio::time::timeout(
63        std::time::Duration::from_secs(NVIDIA_SMI_TIMEOUT_SECS),
64        cmd.output(),
65    )
66    .await
67    .ok()? // timed out
68    .ok()?; // failed to spawn (nvidia-smi not installed)
69
70    if !out.status.success() {
71        return None;
72    }
73
74    // One line per GPU, in MiB (`nounits`).
75    let max_mib = String::from_utf8_lossy(&out.stdout)
76        .lines()
77        .filter_map(|l| l.trim().parse::<u64>().ok())
78        .max()?;
79    max_mib.checked_mul(1024 * 1024)
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85
86    #[test]
87    fn system_ram_is_detected_and_plausible() {
88        // Any real machine (including CI runners) reports > 0 RAM, and a sane
89        // total is at least ~256 MiB — this also catches a units regression
90        // (KiB vs bytes) in the underlying crate.
91        let ram = system_ram_bytes().expect("system RAM should be detectable");
92        assert!(
93            ram > 256 * 1024 * 1024,
94            "implausibly small RAM ({ram} bytes) — units regression?"
95        );
96    }
97
98    #[tokio::test]
99    async fn gpu_vram_probe_is_best_effort() {
100        // Some on a GPU box, None on CI — must not panic or hang either way.
101        let _ = gpu_vram_bytes().await;
102    }
103}
mermaid_cli/utils/host_memory.rs

mermaid_cli/utils/
host_memory.rs
