flodl_hw/
nvidia.rs

1//! The NVIDIA backend: `nvidia-smi`.
2
3use std::path::Path;
4use std::process::Command;
5
6use crate::GpuInfo;
7use crate::report::{GpuSurvey, NoteKind};
8use crate::vendor::{GpuArch, GpuVendor};
9
10/// Cheap, subprocess-free check for "is there an NVIDIA driver here".
11///
12/// Gates the `nvidia-smi` spawn so a pure-AMD or CPU-only box pays a
13/// couple of `stat` calls instead of a process launch. `/dev/nvidiactl`
14/// is the control node every NVIDIA driver creates;
15/// `/proc/driver/nvidia` covers containers that map the proc entry but
16/// not the device nodes.
17fn driver_present() -> bool {
18    Path::new("/dev/nvidiactl").exists() || Path::new("/proc/driver/nvidia").exists()
19}
20
21/// Probe NVIDIA devices, appending to `out`.
22///
23/// Records a note rather than returning an error: a survey is a report
24/// on a whole machine, and one vendor's absence is not a failure of the
25/// sweep.
26pub(crate) fn probe(out: &mut GpuSurvey) {
27    // On non-Linux there are no device nodes to check, so fall through
28    // to the subprocess and let its absence answer.
29    if cfg!(target_os = "linux") && !driver_present() {
30        return;
31    }
32
33    let output = match Command::new("nvidia-smi")
34        // `name` is queried LAST: it is the only field that can contain
35        // the `", "` separator, so with it last `splitn(4, ", ")` keeps
36        // the whole name (commas and all) in the final cell. No CSV
37        // parser needed, since the first three cells are comma-free by
38        // construction.
39        .args([
40            "--query-gpu=index,compute_cap,memory.total,name",
41            "--format=csv,noheader,nounits",
42        ])
43        .output()
44    {
45        Ok(o) if o.status.success() => o,
46        // nvidia-smi present but errored: a driver or permission
47        // problem, or an old bundle that doesn't support the
48        // `compute_cap` query field. Distinct from "not installed",
49        // and the difference is what the user needs told.
50        Ok(o) => {
51            out.note(
52                GpuVendor::Nvidia,
53                NoteKind::ToolFailed,
54                format!(
55                    "`nvidia-smi` exited {}: {}. An NVIDIA driver is loaded but the \
56                     tool cannot enumerate devices, so this box reports 0 NVIDIA GPUs.",
57                    o.status,
58                    String::from_utf8_lossy(&o.stderr).trim(),
59                ),
60            );
61            return;
62        }
63        Err(_) => {
64            // Driver nodes exist (we got past the gate) but the CLI is
65            // missing. Real on minimal containers that map /dev/nvidia*
66            // without installing the utilities.
67            if cfg!(target_os = "linux") {
68                out.note(
69                    GpuVendor::Nvidia,
70                    NoteKind::HardwareUnusable,
71                    "an NVIDIA driver is present but `nvidia-smi` is not on PATH, so \
72                     GPUs cannot be enumerated. Install the NVIDIA utilities package \
73                     (or, in Docker, use a CUDA base image / the NVIDIA container \
74                     toolkit)."
75                        .to_string(),
76                );
77            }
78            return;
79        }
80    };
81
82    let stdout = String::from_utf8_lossy(&output.stdout);
83    for line in stdout.lines().filter(|l| !l.trim().is_empty()) {
84        match parse_csv_row(line) {
85            Some(g) => out.devices.push(g),
86            // A malformed row is a GPU silently dropped. Leave a trace.
87            None => out.note(
88                GpuVendor::Nvidia,
89                NoteKind::Unparsable,
90                format!("could not parse an `nvidia-smi` row, GPU skipped: {line:?}"),
91            ),
92        }
93    }
94}
95
96/// Parse one `index, compute_cap, memory.total, name` CSV row
97/// (`--format=csv,noheader,nounits`). `None` on any bad field.
98fn parse_csv_row(line: &str) -> Option<GpuInfo> {
99    let parts: Vec<&str> = line.splitn(4, ", ").collect();
100    if parts.len() < 4 {
101        return None;
102    }
103    Some(GpuInfo {
104        index: parts[0].trim().parse().ok()?,
105        vendor: GpuVendor::Nvidia,
106        arch: GpuArch::parse(GpuVendor::Nvidia, parts[1])?,
107        total_memory_mb: parts[2].trim().parse().ok()?,
108        name: parts[3].trim().to_string(),
109    })
110}
111
112/// NVIDIA driver version string, or `None` when `nvidia-smi` is absent
113/// or errors.
114pub fn nvidia_driver_version() -> Option<String> {
115    let output = Command::new("nvidia-smi")
116        .args(["--query-gpu=driver_version", "--format=csv,noheader"])
117        .output()
118        .ok()?;
119    if !output.status.success() {
120        return None;
121    }
122    let s = String::from_utf8_lossy(&output.stdout);
123    Some(s.lines().next()?.trim().to_string())
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[test]
131    fn parses_a_well_formed_row() {
132        // Column order: index, compute_cap, memory.total, name.
133        let g = parse_csv_row("1, 6.1, 6078, NVIDIA GeForce GTX 1060 6GB").unwrap();
134        assert_eq!(g.index, 1);
135        assert_eq!(g.name, "NVIDIA GeForce GTX 1060 6GB");
136        assert_eq!(g.arch, GpuArch::Sm { major: 6, minor: 1 });
137        assert_eq!(g.vendor, GpuVendor::Nvidia);
138        assert_eq!(g.total_memory_mb, 6078);
139    }
140
141    #[test]
142    fn keeps_a_comma_inside_the_name() {
143        // `name` last means an embedded ", " stays in the final cell
144        // (splitn(4) stops after 3 separators) rather than truncating.
145        let g = parse_csv_row("0, 8.0, 81920, NVIDIA A100, 80GB").unwrap();
146        assert_eq!(g.name, "NVIDIA A100, 80GB");
147        assert_eq!(g.total_memory_mb, 81920);
148    }
149
150    #[test]
151    fn rejects_malformed_rows() {
152        assert!(parse_csv_row("0, 8.9, three").is_none());
153        assert!(parse_csv_row("x, 8.9, 24564, name").is_none());
154        assert!(parse_csv_row("0, notacap, 24564, name").is_none());
155    }
156}
flodl_hw/nvidia.rs

flodl_hw/
nvidia.rs
