hematite/ui/
gpu_monitor.rs

1//! Background GPU VRAM monitor.
2//!
3//! Spawns a Tokio task that polls `nvidia-smi` every few seconds and stores
4//! the result in lock-free atomics so the TUI render loop can read it cheaply.
5
6use std::sync::atomic::{AtomicU32, Ordering};
7use std::sync::Arc;
8
9/// Shared GPU state — read by the TUI, written by the background poller.
10#[derive(Debug)]
11pub struct GpuState {
12    /// VRAM used in MiB.
13    pub used_mib: AtomicU32,
14    /// VRAM total in MiB.
15    pub total_mib: AtomicU32,
16    /// GPU name (set once on first successful poll).
17    pub name: std::sync::Mutex<String>,
18}
19
20impl GpuState {
21    pub fn new() -> Self {
22        Self {
23            used_mib: AtomicU32::new(0),
24            total_mib: AtomicU32::new(0),
25            name: std::sync::Mutex::new("GPU".into()),
26        }
27    }
28
29    /// Returns (used_mib, total_mib).
30    pub fn read(&self) -> (u32, u32) {
31        (
32            self.used_mib.load(Ordering::Relaxed),
33            self.total_mib.load(Ordering::Relaxed),
34        )
35    }
36
37    /// Returns the ratio used/total, clamped to [0.0, 1.0].
38    pub fn ratio(&self) -> f64 {
39        let (used, total) = self.read();
40        if total == 0 {
41            return 0.0;
42        }
43        (used as f64 / total as f64).clamp(0.0, 1.0)
44    }
45
46    /// Returns a human-readable label like "7.5 GB / 12.0 GB".
47    pub fn label(&self) -> String {
48        let (used, total) = self.read();
49        if total == 0 {
50            return "N/A".into();
51        }
52        format!(
53            "{:.1} GB / {:.1} GB",
54            used as f64 / 1024.0,
55            total as f64 / 1024.0
56        )
57    }
58
59    /// Returns the GPU name (e.g. "NVIDIA GeForce RTX 4070").
60    pub fn gpu_name(&self) -> String {
61        self.name.lock().unwrap().clone()
62    }
63}
64
65/// Spawn the background polling task. Returns the shared state handle.
66pub fn spawn_gpu_monitor() -> Arc<GpuState> {
67    let state = Arc::new(GpuState::new());
68    let bg = state.clone();
69
70    tokio::spawn(async move {
71        loop {
72            if let Some((used, total, name)) = poll_nvidia_smi().await {
73                bg.used_mib.store(used, Ordering::Relaxed);
74                bg.total_mib.store(total, Ordering::Relaxed);
75                if !name.is_empty() {
76                    *bg.name.lock().unwrap() = name;
77                }
78            }
79            tokio::time::sleep(std::time::Duration::from_secs(2)).await;
80        }
81    });
82
83    state
84}
85
86/// Call nvidia-smi and parse the CSV output.
87async fn poll_nvidia_smi() -> Option<(u32, u32, String)> {
88    let output = tokio::process::Command::new("nvidia-smi")
89        .args([
90            "--query-gpu=memory.used,memory.total,name",
91            "--format=csv,noheader,nounits",
92        ])
93        .output()
94        .await
95        .ok()?;
96
97    if !output.status.success() {
98        return None;
99    }
100
101    let stdout = String::from_utf8_lossy(&output.stdout);
102    let line = stdout.trim();
103    let parts: Vec<&str> = line.splitn(3, ',').collect();
104    if parts.len() < 2 {
105        return None;
106    }
107
108    let used: u32 = parts[0].trim().parse().ok()?;
109    let total: u32 = parts[1].trim().parse().ok()?;
110    let name = parts
111        .get(2)
112        .map(|s| s.trim().to_string())
113        .unwrap_or_default();
114
115    Some((used, total, name))
116}
hematite/ui/gpu_monitor.rs

hematite/ui/
gpu_monitor.rs
