rigtop_core/
gpu.rs

1//! NVIDIA GPU statistics via NVML.
2//!
3//! `nvml-wrapper` loads `libnvidia-ml` at runtime, so machines without an
4//! NVIDIA driver simply report no GPUs instead of failing to start.
5
6use std::collections::HashMap;
7
8use nvml_wrapper::Nvml;
9use nvml_wrapper::bitmasks::device::ThrottleReasons;
10use nvml_wrapper::enum_wrappers::device::{Clock, EccCounter, MemoryError, TemperatureSensor};
11use nvml_wrapper::enums::device::UsedGpuMemory;
12use nvml_wrapper::{Device, struct_wrappers::device::ProcessInfo as NvmlProcessInfo};
13
14use crate::snapshot::{EccErrors, GpuProcess, GpuStats, PcieLink};
15
16/// Collects per-GPU and per-process statistics from NVML.
17pub struct NvmlCollector {
18    /// `None` when the library or driver is unavailable.
19    nvml: Option<Nvml>,
20    /// Why NVML could not be initialized, for diagnostics.
21    init_error: Option<String>,
22    /// Newest per-process utilization sample timestamp seen, per GPU index.
23    /// Passing it back to NVML returns only samples newer than the last poll.
24    last_sample_ts: HashMap<u32, u64>,
25    /// Whether to query per-process utilization, the costliest NVML call
26    /// (about 2 ms per GPU). Only the full process table shows it.
27    process_utilization: bool,
28    /// Whether ECC is on, per GPU index (changing it takes a reboot).
29    ecc_enabled: HashMap<u32, bool>,
30}
31
32impl NvmlCollector {
33    /// Initializes NVML. Never fails; check [`NvmlCollector::init_error`] for the reason
34    /// when no GPUs are reported.
35    pub fn new() -> Self {
36        let (nvml, init_error) = Self::init();
37        Self {
38            nvml,
39            init_error,
40            last_sample_ts: HashMap::new(),
41            process_utilization: true,
42            ecc_enabled: HashMap::new(),
43        }
44    }
45
46    /// NVIDIA ships no NVML for macOS (nvml-wrapper would look for `nvml.dll`
47    /// there), so it is not tried, and having no GPUs is not worth a note.
48    #[cfg(target_os = "macos")]
49    fn init() -> (Option<Nvml>, Option<String>) {
50        (None, None)
51    }
52
53    #[cfg(not(target_os = "macos"))]
54    fn init() -> (Option<Nvml>, Option<String>) {
55        match Nvml::init() {
56            Ok(nvml) => (Some(nvml), None),
57            Err(err) => (None, Some(err.to_string())),
58        }
59    }
60
61    /// Turns the per-process utilization query on or off (off while only
62    /// the top processes are wanted, as for an overview).
63    pub fn set_process_utilization(&mut self, enabled: bool) {
64        self.process_utilization = enabled;
65    }
66
67    /// The NVML initialization error, if any.
68    pub fn init_error(&self) -> Option<&str> {
69        self.init_error.as_deref()
70    }
71
72    /// Queries every NVIDIA GPU. Individual metrics a device does not support
73    /// are left as `None`.
74    pub fn collect(&mut self) -> Vec<GpuStats> {
75        let Some(nvml) = self.nvml.as_ref() else {
76            return Vec::new();
77        };
78        let count = nvml.device_count().unwrap_or(0);
79        let mut gpus = Vec::with_capacity(count as usize);
80        for index in 0..count {
81            let Ok(device) = nvml.device_by_index(index) else {
82                continue;
83            };
84            let last_ts = self.last_sample_ts.get(&index).copied();
85            let ecc = *self.ecc_enabled.entry(index).or_insert_with(|| {
86                device
87                    .is_ecc_enabled()
88                    .is_ok_and(|mode| mode.currently_enabled)
89            });
90            let (mut stats, newest_ts) =
91                collect_device(index, &device, last_ts, self.process_utilization);
92            if ecc {
93                stats.ecc = ecc_errors(&device);
94            }
95            if let Some(ts) = newest_ts {
96                self.last_sample_ts.insert(index, ts);
97            }
98            gpus.push(stats);
99        }
100        gpus
101    }
102}
103
104impl Default for NvmlCollector {
105    fn default() -> Self {
106        Self::new()
107    }
108}
109
110/// Reads one device. Also returns the newest process sample timestamp.
111fn collect_device(
112    index: u32,
113    device: &Device<'_>,
114    last_ts: Option<u64>,
115    process_utilization: bool,
116) -> (GpuStats, Option<u64>) {
117    let utilization = device.utilization_rates().ok();
118    let memory = device.memory_info().ok();
119    let (processes, newest_ts) = collect_processes(device, last_ts, process_utilization);
120    let stats = GpuStats {
121        index,
122        vendor: "NVIDIA".to_owned(),
123        name: device.name().unwrap_or_else(|_| format!("GPU {index}")),
124        // NVML covers discrete NVIDIA GPUs only.
125        integrated: false,
126        utilization_percent: utilization.as_ref().map(|u| u.gpu as f32),
127        memory_utilization_percent: utilization.as_ref().map(|u| u.memory as f32),
128        memory_used_bytes: memory.as_ref().map(|m| m.used),
129        memory_total_bytes: memory.as_ref().map(|m| m.total),
130        temperature_c: device
131            .temperature(TemperatureSensor::Gpu)
132            .ok()
133            .map(|t| t as f32),
134        // NVML reports milliwatts.
135        power_w: device.power_usage().ok().map(|mw| mw as f32 / 1000.0),
136        power_limit_w: device
137            .enforced_power_limit()
138            .ok()
139            .map(|mw| mw as f32 / 1000.0),
140        fan_percent: device.fan_speed(0).ok().map(|p| p as f32),
141        clock_mhz: device.clock_info(Clock::Graphics).ok(),
142        throttle_reasons: device
143            .current_throttle_reasons()
144            .map(throttle_reasons)
145            .unwrap_or_default(),
146        ecc: None,
147        pcie: pcie_link(device),
148        processes,
149    };
150    (stats, newest_ts)
151}
152
153/// The reasons worth showing: idling, application or display clock
154/// settings and sync boost are not held-down clocks.
155fn throttle_reasons(reasons: ThrottleReasons) -> Vec<String> {
156    let mut names = Vec::new();
157    if reasons.contains(ThrottleReasons::SW_POWER_CAP) {
158        names.push("power cap");
159    }
160    if reasons.contains(ThrottleReasons::SW_THERMAL_SLOWDOWN) {
161        names.push("thermal");
162    }
163    if reasons.contains(ThrottleReasons::HW_THERMAL_SLOWDOWN) {
164        names.push("hw thermal");
165    }
166    if reasons.contains(ThrottleReasons::HW_POWER_BRAKE_SLOWDOWN) {
167        names.push("power brake");
168    }
169    // Set together with the two above; alone, its cause is not known.
170    if reasons.contains(ThrottleReasons::HW_SLOWDOWN)
171        && !reasons.intersects(
172            ThrottleReasons::HW_THERMAL_SLOWDOWN | ThrottleReasons::HW_POWER_BRAKE_SLOWDOWN,
173        )
174    {
175        names.push("hw slowdown");
176    }
177    names.into_iter().map(str::to_owned).collect()
178}
179
180/// ECC errors since the driver loaded.
181fn ecc_errors(device: &Device<'_>) -> Option<EccErrors> {
182    let count = |kind| device.total_ecc_errors(kind, EccCounter::Volatile).ok();
183    Some(EccErrors {
184        corrected: count(MemoryError::Corrected)?,
185        uncorrected: count(MemoryError::Uncorrected)?,
186    })
187}
188
189fn pcie_link(device: &Device<'_>) -> Option<PcieLink> {
190    Some(PcieLink {
191        generation: device.current_pcie_link_gen().ok()?,
192        width: device.current_pcie_link_width().ok()?,
193        max_generation: device.max_pcie_link_gen().ok()?,
194        max_width: device.max_pcie_link_width().ok()?,
195    })
196}
197
198/// Merges compute and graphics process lists with per-process utilization
199/// (when `with_utilization`).
200fn collect_processes(
201    device: &Device<'_>,
202    last_ts: Option<u64>,
203    with_utilization: bool,
204) -> (Vec<GpuProcess>, Option<u64>) {
205    // A process can appear in both lists (e.g. CUDA + OpenGL); keep one entry
206    // per pid with the larger memory figure.
207    let mut by_pid: HashMap<u32, GpuProcess> = HashMap::new();
208    let lists = [
209        device.running_compute_processes(),
210        device.running_graphics_processes(),
211    ];
212    for info in lists.into_iter().flat_map(|r| r.unwrap_or_default()) {
213        let memory = used_memory(&info);
214        let entry = by_pid.entry(info.pid).or_insert_with(|| GpuProcess {
215            pid: info.pid,
216            ..Default::default()
217        });
218        entry.memory_bytes = entry.memory_bytes.max(memory);
219    }
220
221    // Utilization samples: keep the newest sample of each pid. Skipped when
222    // not wanted or when no process uses the GPU (often a second GPU).
223    let mut newest_ts = last_ts;
224    if !with_utilization || by_pid.is_empty() {
225        let mut processes: Vec<GpuProcess> = by_pid.into_values().collect();
226        processes.sort_unstable_by_key(|p| std::cmp::Reverse(p.memory_bytes));
227        return (processes, newest_ts);
228    }
229    // NotFound means "no samples since last_ts", which is normal for idle
230    // GPUs; NotSupported is common on older devices. Both just mean no data.
231    let samples = device
232        .process_utilization_stats(last_ts)
233        .unwrap_or_default();
234    let mut latest: HashMap<u32, (u64, u32)> = HashMap::new();
235    for sample in samples {
236        newest_ts = Some(newest_ts.map_or(sample.timestamp, |t| t.max(sample.timestamp)));
237        let slot = latest.entry(sample.pid).or_insert((0, 0));
238        if sample.timestamp >= slot.0 {
239            *slot = (sample.timestamp, sample.sm_util);
240        }
241    }
242    for (pid, (_, sm_util)) in latest {
243        if let Some(process) = by_pid.get_mut(&pid) {
244            process.sm_percent = Some(sm_util as f32);
245        }
246    }
247
248    let mut processes: Vec<GpuProcess> = by_pid.into_values().collect();
249    processes.sort_unstable_by_key(|p| std::cmp::Reverse(p.memory_bytes));
250    (processes, newest_ts)
251}
252
253fn used_memory(info: &NvmlProcessInfo) -> Option<u64> {
254    match info.used_gpu_memory {
255        UsedGpuMemory::Used(bytes) => Some(bytes),
256        UsedGpuMemory::Unavailable => None,
257    }
258}
259
260#[cfg(test)]
261mod tests {
262    use super::*;
263
264    #[test]
265    fn throttle_reasons_worth_showing() {
266        assert!(throttle_reasons(ThrottleReasons::GPU_IDLE).is_empty());
267        assert!(throttle_reasons(ThrottleReasons::NONE).is_empty());
268        assert_eq!(
269            throttle_reasons(ThrottleReasons::SW_POWER_CAP | ThrottleReasons::GPU_IDLE),
270            ["power cap"]
271        );
272        assert_eq!(
273            throttle_reasons(ThrottleReasons::HW_SLOWDOWN | ThrottleReasons::HW_THERMAL_SLOWDOWN),
274            ["hw thermal"]
275        );
276        assert_eq!(
277            throttle_reasons(ThrottleReasons::HW_SLOWDOWN),
278            ["hw slowdown"]
279        );
280    }
281}
rigtop_core/gpu.rs

rigtop_core/
gpu.rs
