resource_sampler/
gpu.rs

1// NVIDIA/AMD GPU sampling.
2//
3// Logic moved from `peek-core::proc::gpu`. This crate owns the low-level GPU
4// metrics and struct; peek-core re-exports `GpuInfo` so callers keep using
5// `peek_core::GpuInfo` as before.
6
7use serde::{Deserialize, Serialize};
8
9/// GPU utilisation snapshot for a single device.
10#[derive(Debug, Serialize, Deserialize, Clone)]
11pub struct GpuInfo {
12    pub index: usize,
13    pub name: String,
14    pub utilization_percent: Option<f64>,
15    pub memory_used_mb: Option<f64>,
16    pub memory_total_mb: Option<f64>,
17    /// Memory used by the inspected PID on this GPU (NVIDIA, when available).
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub process_used_mb: Option<f64>,
20    /// "nvml", "sysfs", or "nvidia-smi"
21    pub source: String,
22}
23
24/// Attempt to collect GPU utilisation for this process.
25///
26/// Strategy (in order):
27/// 1. Run `nvidia-smi` (NVIDIA) and parse CSV output.
28/// 2. Walk `/sys/class/drm/card*/device/` for AMD via ROCm sysfs.
29/// 3. Return empty Vec if nothing is found.
30///
31/// Per-process memory (process_used_mb) is filled when nvidia-smi supports
32/// `--query-compute-apps` and this PID has compute processes on a GPU.
33pub fn collect_gpu(pid: i32) -> Vec<GpuInfo> {
34    // Try nvidia-smi first (includes per-PID attribution when available)
35    let mut gpus = try_nvidia_smi(pid);
36    if !gpus.is_empty() {
37        return gpus;
38    }
39
40    // Try AMD sysfs
41    gpus = try_amd_sysfs();
42    gpus
43}
44
45// ─── NVIDIA via nvidia-smi ────────────────────────────────────────────────────
46
47/// Map from GPU uuid (from nvidia-smi) to memory used by the given PID (MiB).
48fn query_nvidia_compute_apps(pid: i32) -> std::collections::HashMap<String, f64> {
49    let output = match std::process::Command::new("nvidia-smi")
50        .args([
51            "--query-compute-apps=pid,gpu_uuid,used_memory",
52            "--format=csv,noheader,nounits",
53        ])
54        .output()
55    {
56        Ok(o) if o.status.success() => o,
57        _ => return std::collections::HashMap::new(),
58    };
59    let stdout = String::from_utf8_lossy(&output.stdout);
60    let pid_str = pid.to_string();
61    let mut map = std::collections::HashMap::new();
62    for line in stdout.lines() {
63        let parts: Vec<&str> = line.split(',').map(str::trim).collect();
64        if parts.len() < 3 {
65            continue;
66        }
67        if parts[0] != pid_str {
68            continue;
69        }
70        let uuid = parts[1].to_string();
71        if let Ok(used) = parts[2].parse::<f64>() {
72            *map.entry(uuid).or_insert(0.0) += used;
73        }
74    }
75    map
76}
77
78fn try_nvidia_smi(pid: i32) -> Vec<GpuInfo> {
79    // Query GPUs with uuid so we can match compute-apps to GPU index
80    let output = match std::process::Command::new("nvidia-smi")
81        .args([
82            "--query-gpu=index,uuid,name,utilization.gpu,memory.used,memory.total",
83            "--format=csv,noheader,nounits",
84        ])
85        .output()
86    {
87        Ok(o) if o.status.success() => o,
88        _ => return Vec::new(),
89    };
90
91    let process_mem = query_nvidia_compute_apps(pid);
92
93    let stdout = String::from_utf8_lossy(&output.stdout);
94    let mut gpus = Vec::new();
95
96    for line in stdout.lines() {
97        let parts: Vec<&str> = line.split(',').map(str::trim).collect();
98        // index,uuid,name,util,mem_used,mem_total (6 columns; uuid may contain dashes only)
99        if parts.len() < 6 {
100            continue;
101        }
102        let index: usize = parts[0].parse().unwrap_or(0);
103        let uuid = parts[1].to_string();
104        let name = parts[2].to_string();
105        let util: Option<f64> = parts[3].parse().ok();
106        let mem_used: Option<f64> = parts[4].parse().ok();
107        let mem_total: Option<f64> = parts[5].parse().ok();
108        let process_used_mb = process_mem.get(&uuid).copied();
109
110        gpus.push(GpuInfo {
111            index,
112            name,
113            utilization_percent: util,
114            memory_used_mb: mem_used,
115            memory_total_mb: mem_total,
116            process_used_mb,
117            source: "nvidia-smi".to_string(),
118        });
119    }
120
121    gpus
122}
123
124// ─── AMD via sysfs ───────────────────────────────────────────────────────────
125
126fn try_amd_sysfs() -> Vec<GpuInfo> {
127    let mut gpus = Vec::new();
128
129    let drm = match std::fs::read_dir("/sys/class/drm") {
130        Ok(d) => d,
131        Err(_) => return Vec::new(),
132    };
133
134    let mut index = 0usize;
135    let mut entries: Vec<_> = drm.flatten().collect();
136    entries.sort_by_key(|e| e.file_name());
137
138    for entry in entries {
139        let name = entry.file_name();
140        let s = name.to_string_lossy();
141        // Only top-level card entries, not renderD
142        if !s.starts_with("card") || s.contains('-') {
143            continue;
144        }
145
146        let base = entry.path().join("device");
147        if !base.exists() {
148            continue;
149        }
150
151        // GPU busy percent (AMD)
152        let util: Option<f64> = std::fs::read_to_string(base.join("gpu_busy_percent"))
153            .ok()
154            .and_then(|s| s.trim().parse().ok());
155
156        if util.is_none() {
157            // Not an AMD GPU with known sysfs interface
158            continue;
159        }
160
161        // VRAM used / total (bytes)
162        let vram_used_mb: Option<f64> = std::fs::read_to_string(base.join("mem_info_vram_used"))
163            .ok()
164            .and_then(|s| s.trim().parse::<u64>().ok())
165            .map(|b| b as f64 / 1_048_576.0);
166
167        let vram_total_mb: Option<f64> = std::fs::read_to_string(base.join("mem_info_vram_total"))
168            .ok()
169            .and_then(|s| s.trim().parse::<u64>().ok())
170            .map(|b| b as f64 / 1_048_576.0);
171
172        // Try to get a friendly name from uevent
173        let gpu_name = std::fs::read_to_string(base.join("uevent"))
174            .ok()
175            .and_then(|s| {
176                s.lines()
177                    .find(|l| l.starts_with("DRIVER="))
178                    .map(|l| l.to_string())
179            })
180            .unwrap_or_else(|| format!("AMD GPU ({})", s));
181
182        gpus.push(GpuInfo {
183            index,
184            name: gpu_name,
185            utilization_percent: util,
186            memory_used_mb: vram_used_mb,
187            memory_total_mb: vram_total_mb,
188            process_used_mb: None,
189            source: "sysfs/amdgpu".to_string(),
190        });
191        index += 1;
192    }
193
194    gpus
195}
196
197#[cfg(test)]
198mod tests {
199    use super::*;
200
201    #[test]
202    fn collect_gpu_returns_vec() {
203        // Should not panic regardless of whether a GPU is present.
204        let result = collect_gpu(1);
205        // Each entry should have non-empty name and source
206        for g in &result {
207            assert!(!g.name.is_empty());
208            assert!(!g.source.is_empty());
209        }
210    }
211}
resource_sampler/gpu.rs

resource_sampler/
gpu.rs
