Skip to main content

studio_worker/
sys.rs

1//! Host-system probes: hostname, OS user, VRAM.
2//!
3//! Every probe emits a structured tracing breadcrumb so an operator can
4//! tell from the logs *why* a worker reports the values it does (in
5//! particular, why VRAM came back as `0.0` — was the sysfs tree missing,
6//! present-but-unparseable, or is the worker running on a non-Linux
7//! host?).  Silent `0.0` makes "this worker claims nothing" impossible
8//! to diagnose from logs alone.
9use anyhow::Result;
10use std::path::Path;
11use std::sync::OnceLock;
12
13pub fn machine_name() -> String {
14    let name = hostname::get()
15        .ok()
16        .and_then(|s| s.into_string().ok())
17        .unwrap_or_else(|| "unknown-host".to_string());
18    tracing::debug!(
19        target: "studio_worker::sys",
20        op = "machine_name",
21        value = %name,
22        "resolved host machine name"
23    );
24    name
25}
26
27pub fn username() -> String {
28    username_from_probe(whoami::username())
29}
30
31/// Resolve the OS-user probe into a username, logging the outcome so a
32/// silent fallback can't hide a failing probe.  `whoami::username`
33/// became fallible in whoami 2.x; on the error path we emit a `warn`
34/// breadcrumb naming the underlying error and fall back to
35/// `unknown-user`, mirroring `machine_name`'s `unknown-host` default.
36fn username_from_probe<E: std::fmt::Display>(probe: std::result::Result<String, E>) -> String {
37    let user = match probe {
38        Ok(user) => user,
39        Err(e) => {
40            tracing::warn!(
41                target: "studio_worker::sys",
42                op = "username",
43                error = %e,
44                "failed to resolve OS user; falling back to unknown-user"
45            );
46            "unknown-user".to_string()
47        }
48    };
49    tracing::debug!(
50        target: "studio_worker::sys",
51        op = "username",
52        value = %user,
53        "resolved OS user"
54    );
55    user
56}
57
58/// Cached result of the (relatively expensive) VRAM probe.  Total VRAM
59/// is a static hardware property, so we probe at most once per process —
60/// `build_capabilities` runs on every 5s heartbeat and must not spawn an
61/// `nvidia-smi` subprocess each tick.
62static VRAM_GB: OnceLock<f32> = OnceLock::new();
63
64/// Detect physical VRAM on the host, in GB.  Returns 0.0 when we can't
65/// probe (no NVIDIA GPU, no driver) — the engine still runs in synthetic
66/// mode for low-end / CI machines.
67///
68/// This intentionally avoids a hard dependency on `nvml-wrapper` because
69/// it brings a heavy NVML build dep that we don't want at the CI layer.
70/// On Linux we first try the dependency-free
71/// `/proc/driver/nvidia/gpus/*/information` sysfs probe; current NVIDIA
72/// drivers (5xx) dropped the `Video Memory` line from that file, so we
73/// fall back to `nvidia-smi` (which ships with every driver, on every
74/// platform, and whose `--query-gpu` interface is stable across
75/// versions).  The result is memoised since it can't change while the
76/// process runs.
77pub fn detect_vram_gb() -> Result<f32> {
78    Ok(*VRAM_GB.get_or_init(probe_vram_gb))
79}
80
81fn probe_vram_gb() -> f32 {
82    // Linux exposes cheap, dependency-free sysfs probes; try them first
83    // so the common case never spawns a subprocess.
84    #[cfg(target_os = "linux")]
85    {
86        let from_sysfs = detect_vram_gb_from_sysfs(Path::new("/proc/driver/nvidia/gpus"));
87        if from_sysfs > 0.0 {
88            return from_sysfs;
89        }
90        // AMD (and Intel discrete) expose VRAM via the DRM sysfs tree.
91        let from_amd = detect_vram_gb_from_amd_sysfs(Path::new("/sys/class/drm"));
92        if from_amd > 0.0 {
93            return from_amd;
94        }
95    }
96    // Apple Silicon / Intel Macs: unified memory, sized via sysctl.
97    #[cfg(target_os = "macos")]
98    {
99        if let Some(gb) = detect_vram_gb_via_sysctl() {
100            return gb;
101        }
102    }
103    // NVIDIA on any platform: `nvidia-smi`.  Absent on a non-NVIDIA
104    // host, where the command simply fails to spawn.
105    if let Some(gb) = detect_vram_gb_via_nvidia_smi() {
106        return gb;
107    }
108    // Windows non-NVIDIA (AMD / Intel): CIM `Win32_VideoController`.
109    #[cfg(target_os = "windows")]
110    {
111        if let Some(gb) = detect_vram_gb_via_wmic() {
112            return gb;
113        }
114    }
115    0.0
116}
117
118/// Bytes → GiB (matching the NVIDIA MiB/1024 path, so every vendor
119/// reports the same unit).
120fn bytes_to_gib(bytes: u64) -> f32 {
121    (bytes as f64 / (1024.0 * 1024.0 * 1024.0)) as f32
122}
123
124/// True for a DRM primary-node dir name (`card0`, `card1`, …) but not
125/// the connector (`card0-DP-1`) or render (`renderD128`) siblings.
126fn is_drm_card_dir(name: &str) -> bool {
127    name.strip_prefix("card")
128        .is_some_and(|rest| !rest.is_empty() && rest.chars().all(|c| c.is_ascii_digit()))
129}
130
131/// Sum VRAM across every AMD/Intel-discrete GPU exposed under the DRM
132/// sysfs tree (`<root>/card*/device/mem_info_vram_total`, a byte
133/// count).  Integrated GPUs and drivers that don't publish the file
134/// simply contribute nothing, so an iGPU-only box returns 0 and the
135/// caller keeps the configured threshold as its only signal.
136pub fn detect_vram_gb_from_amd_sysfs(root: &Path) -> f32 {
137    let entries = match std::fs::read_dir(root) {
138        Ok(e) => e,
139        Err(_) => return 0.0,
140    };
141    let mut total_bytes: u64 = 0;
142    let mut cards: u32 = 0;
143    for entry in entries.flatten() {
144        let name = entry.file_name().to_string_lossy().into_owned();
145        if !is_drm_card_dir(&name) {
146            continue;
147        }
148        let vram_file = entry.path().join("device").join("mem_info_vram_total");
149        if let Ok(content) = std::fs::read_to_string(&vram_file) {
150            if let Some(bytes) = parse_amd_vram_total_bytes(&content) {
151                total_bytes = total_bytes.saturating_add(bytes);
152                cards += 1;
153            }
154        }
155    }
156    let gib = bytes_to_gib(total_bytes);
157    if cards > 0 {
158        tracing::info!(
159            target: "studio_worker::sys",
160            op = "probe_vram",
161            source = "amd_drm_sysfs",
162            vram_gb = gib,
163            cards,
164            "detected VRAM via AMD/DRM sysfs"
165        );
166    }
167    gib
168}
169
170/// Parse an AMD `mem_info_vram_total` file: a single decimal byte
171/// count, possibly with trailing whitespace.
172pub fn parse_amd_vram_total_bytes(content: &str) -> Option<u64> {
173    content.trim().parse::<u64>().ok().filter(|b| *b > 0)
174}
175
176/// Parse `sysctl -n hw.memsize` (total RAM in bytes) into a usable VRAM
177/// figure for Apple unified memory.  `fraction` is the share of unified
178/// memory we treat as GPU-addressable (macOS lets the GPU use most of
179/// it; 0.75 is a conservative, widely-cited figure).
180pub fn parse_sysctl_memsize(stdout: &str, fraction: f64) -> Option<f32> {
181    let bytes = stdout.trim().parse::<u64>().ok().filter(|b| *b > 0)?;
182    Some((bytes_to_gib(bytes) as f64 * fraction) as f32)
183}
184
185/// Sum the `AdapterRAM` values from a Windows CIM/WMIC
186/// `Win32_VideoController` dump (one integer per adapter, bytes).
187/// Ignores non-numeric lines (headers, blanks) so it tolerates both
188/// `wmic` and PowerShell `Get-CimInstance` output shapes.
189pub fn parse_wmic_adapter_ram(stdout: &str) -> Option<f32> {
190    let mut total: u64 = 0;
191    let mut found = false;
192    for line in stdout.lines() {
193        if let Some(bytes) = line.trim().parse::<u64>().ok().filter(|b| *b > 0) {
194            total = total.saturating_add(bytes);
195            found = true;
196        }
197    }
198    found.then(|| bytes_to_gib(total))
199}
200
201/// Probe VRAM via `nvidia-smi --query-gpu=memory.total`.  Returns `None`
202/// when the binary is absent (no driver / non-NVIDIA host) or exits
203/// non-zero, in which cases the caller defaults to 0 GB.
204///
205/// Coverage-off: spawning a real `nvidia-smi` is host-dependent (CI has
206/// none), so its success / non-zero-exit arms can't be exercised
207/// deterministically.  The parse + GB conversion + logging it delegates
208/// to ([`vram_gb_from_smi_stdout`], [`parse_nvidia_smi_mib`]) are
209/// unit-tested directly.
210#[cfg_attr(coverage_nightly, coverage(off))]
211fn detect_vram_gb_via_nvidia_smi() -> Option<f32> {
212    let output = std::process::Command::new("nvidia-smi")
213        .args(["--query-gpu=memory.total", "--format=csv,noheader,nounits"])
214        .output();
215    match output {
216        Ok(o) if o.status.success() => vram_gb_from_smi_stdout(&String::from_utf8_lossy(&o.stdout)),
217        Ok(o) => {
218            tracing::warn!(
219                target: "studio_worker::sys",
220                op = "probe_vram",
221                source = "nvidia_smi_failed",
222                code = ?o.status.code(),
223                "nvidia-smi exited non-zero while probing VRAM — defaulting to 0 GB"
224            );
225            None
226        }
227        Err(e) => {
228            tracing::info!(
229                target: "studio_worker::sys",
230                op = "probe_vram",
231                source = "nvidia_smi_absent",
232                error = %e,
233                "nvidia-smi not available — cannot probe VRAM; defaulting to 0 GB"
234            );
235            None
236        }
237    }
238}
239
240/// Live Apple unified-memory probe via `sysctl -n hw.memsize`.
241/// Coverage-off: host-dependent; the parse is unit-tested via
242/// [`parse_sysctl_memsize`].
243#[cfg(target_os = "macos")]
244#[cfg_attr(coverage_nightly, coverage(off))]
245fn detect_vram_gb_via_sysctl() -> Option<f32> {
246    let output = std::process::Command::new("sysctl")
247        .args(["-n", "hw.memsize"])
248        .output()
249        .ok()?;
250    if !output.status.success() {
251        return None;
252    }
253    parse_sysctl_memsize(&String::from_utf8_lossy(&output.stdout), 0.75)
254}
255
256/// Live Windows non-NVIDIA probe via PowerShell CIM
257/// (`Win32_VideoController.AdapterRAM`).  Coverage-off: host-dependent;
258/// the parse is unit-tested via [`parse_wmic_adapter_ram`].  Note
259/// `AdapterRAM` is a UInt32 and saturates at ~4 GiB on larger cards — a
260/// documented Windows limitation, but a conservative floor beats 0.
261#[cfg(target_os = "windows")]
262#[cfg_attr(coverage_nightly, coverage(off))]
263fn detect_vram_gb_via_wmic() -> Option<f32> {
264    let output = std::process::Command::new("powershell")
265        .args([
266            "-NoProfile",
267            "-Command",
268            "(Get-CimInstance Win32_VideoController).AdapterRAM",
269        ])
270        .output()
271        .ok()?;
272    if !output.status.success() {
273        return None;
274    }
275    parse_wmic_adapter_ram(&String::from_utf8_lossy(&output.stdout))
276}
277
278/// Summed VRAM (MiB) from an `nvidia-smi` memory query plus the count of
279/// GPU lines that were dropped from that total.
280///
281/// `dropped` is the number of non-empty lines whose leading token wasn't
282/// a number — nvidia-smi emits `[N/A]` for `memory.total` when a card has
283/// fallen off the bus, hit an ECC fault, or sits in a MIG state with no
284/// resolvable total.  Carrying the count (rather than silently summing
285/// the survivors) means a multi-GPU box that under-reports its VRAM — and
286/// then refuses jobs it could actually run — leaves a breadcrumb instead
287/// of vanishing the card without a trace.
288struct SmiMemTotal {
289    mib: f64,
290    dropped: u32,
291}
292
293/// Convert the stdout of an `nvidia-smi` memory query to GB and emit the
294/// probe breadcrumb.  Split out from the subprocess plumbing so the
295/// parse + conversion + logging are unit-testable without a real
296/// `nvidia-smi` on the box (CI has none).
297fn vram_gb_from_smi_stdout(stdout: &str) -> Option<f32> {
298    let SmiMemTotal { mib, dropped } = parse_nvidia_smi_mib(stdout)?;
299    let vram_gb = (mib / 1024.0) as f32;
300    tracing::info!(
301        target: "studio_worker::sys",
302        op = "probe_vram",
303        source = "nvidia_smi",
304        vram_gb = vram_gb,
305        dropped = dropped,
306        "detected NVIDIA VRAM via nvidia-smi fallback"
307    );
308    Some(vram_gb)
309}
310
311/// Sum the per-GPU MiB totals from
312/// `nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits`.
313/// One line per GPU, each a bare MiB integer (e.g. `24564`).  Tolerates
314/// a trailing unit token (if `nounits` is ever dropped) and ignores
315/// blank lines.  Every non-empty line that fails to parse (e.g. `[N/A]`)
316/// is warn-logged and counted in [`SmiMemTotal::dropped`] before being
317/// left out of the total.  Returns `None` when no line yielded a number.
318fn parse_nvidia_smi_mib(stdout: &str) -> Option<SmiMemTotal> {
319    let mut total: f64 = 0.0;
320    let mut any = false;
321    let mut dropped: u32 = 0;
322    for (idx, line) in stdout.lines().enumerate() {
323        let trimmed = line.trim();
324        if trimmed.is_empty() {
325            continue;
326        }
327        match trimmed
328            .split_whitespace()
329            .next()
330            .and_then(|tok| tok.parse::<f64>().ok())
331        {
332            Some(mib) => {
333                total += mib;
334                any = true;
335            }
336            None => {
337                dropped += 1;
338                tracing::warn!(
339                    target: "studio_worker::sys",
340                    op = "probe_vram",
341                    source = "nvidia_smi",
342                    line = idx,
343                    content = trimmed,
344                    "nvidia-smi VRAM line did not parse as MiB — dropping this GPU from the total"
345                );
346            }
347        }
348    }
349    any.then_some(SmiMemTotal {
350        mib: total,
351        dropped,
352    })
353}
354
355/// VRAM probe driven by a configurable sysfs root.  Public-in-crate so
356/// the integration tests can exercise both the "missing root" and
357/// "populated root" branches without a real `/proc/driver/nvidia` tree.
358///
359/// Emits a summary tracing event per call, plus a `WARN` for every GPU
360/// dropped from the total so a multi-GPU box never under-reports its
361/// VRAM silently:
362///
363/// - `INFO source="no_nvidia_sysfs"` — `root` is not a directory.  This
364///   is the normal case on CI runners / non-GPU hosts.
365/// - `INFO source="nvidia_sysfs"` — at least one GPU's `information`
366///   file was parseable.  `gpu_count` is how many contributed; `dropped`
367///   is how many were present but unreadable / had no parseable `Video
368///   Memory` line (each of those also gets its own `WARN` naming it).
369/// - `WARN source="sysfs_unparseable"` — directories were present but
370///   none parseable (current 5xx drivers dropped the `Video Memory`
371///   line).  The caller then falls back to `nvidia-smi`; the warn is the
372///   breadcrumb that the cheap sysfs path no longer works on this host.
373/// - `WARN source="nvidia_sysfs" reason="no_video_memory_line"|"video_memory_unparseable"|"info_unreadable"`
374///   — a specific GPU was dropped from the total while others survived.
375///   `video_memory_unparseable` means the `Video Memory` line was
376///   present but its value didn't parse (the warn echoes the offending
377///   `content`); `no_video_memory_line` means no such line at all.
378pub fn detect_vram_gb_from_sysfs(root: &Path) -> f32 {
379    let entries = match std::fs::read_dir(root) {
380        Ok(e) => e,
381        Err(_) => {
382            tracing::info!(
383                target: "studio_worker::sys",
384                op = "probe_vram",
385                source = "no_nvidia_sysfs",
386                vram_gb = 0.0,
387                root = %root.display(),
388                "no NVIDIA sysfs tree at probe root — defaulting to 0 GB VRAM"
389            );
390            return 0.0;
391        }
392    };
393
394    let mut total_mib: f64 = 0.0;
395    let mut gpu_count: u32 = 0;
396    let mut parseable: u32 = 0;
397    for entry in entries.flatten() {
398        gpu_count += 1;
399        let gpu_path = entry.path();
400        let info_path = gpu_path.join("information");
401        match std::fs::read_to_string(&info_path) {
402            Ok(content) => {
403                let mut found = false;
404                // A `Video Memory:` line that's present but whose value
405                // can't be parsed (e.g. `N/A` on a driver that stubbed
406                // the field) must be surfaced differently from a GPU
407                // with no such line at all — otherwise the operator is
408                // told the line is missing when it's right there.  Keep
409                // the first offending value to echo in the warn.
410                let mut unparseable: Option<String> = None;
411                for line in content.lines() {
412                    if let Some(rest) = line.trim().strip_prefix("Video Memory:") {
413                        if let Some(mib) = parse_mib(rest) {
414                            total_mib += mib;
415                            found = true;
416                        } else if unparseable.is_none() {
417                            unparseable = Some(rest.trim().to_string());
418                        }
419                    }
420                }
421                if found {
422                    parseable += 1;
423                } else if let Some(content) = unparseable {
424                    tracing::warn!(
425                        target: "studio_worker::sys",
426                        op = "probe_vram",
427                        source = "nvidia_sysfs",
428                        reason = "video_memory_unparseable",
429                        gpu = %gpu_path.display(),
430                        content = content.as_str(),
431                        "sysfs GPU Video Memory line did not parse as MiB — dropping it from the total"
432                    );
433                } else {
434                    tracing::warn!(
435                        target: "studio_worker::sys",
436                        op = "probe_vram",
437                        source = "nvidia_sysfs",
438                        reason = "no_video_memory_line",
439                        gpu = %gpu_path.display(),
440                        "sysfs GPU has no parseable Video Memory line — dropping it from the total"
441                    );
442                }
443            }
444            Err(e) => {
445                tracing::warn!(
446                    target: "studio_worker::sys",
447                    op = "probe_vram",
448                    source = "nvidia_sysfs",
449                    reason = "info_unreadable",
450                    gpu = %gpu_path.display(),
451                    error = %e,
452                    "could not read a sysfs GPU information file — dropping it from the total"
453                );
454            }
455        }
456    }
457
458    let vram_gb = (total_mib / 1024.0) as f32;
459    let dropped = gpu_count.saturating_sub(parseable);
460    if parseable > 0 {
461        tracing::info!(
462            target: "studio_worker::sys",
463            op = "probe_vram",
464            source = "nvidia_sysfs",
465            vram_gb = vram_gb,
466            gpu_count = parseable,
467            dropped = dropped,
468            "detected NVIDIA VRAM via sysfs"
469        );
470    } else {
471        tracing::warn!(
472            target: "studio_worker::sys",
473            op = "probe_vram",
474            source = "sysfs_unparseable",
475            vram_gb = 0.0,
476            gpu_count = gpu_count,
477            root = %root.display(),
478            "NVIDIA sysfs entries present but no Video Memory line (current 5xx drivers dropped it) — falling back to nvidia-smi"
479        );
480    }
481    vram_gb
482}
483
484fn parse_mib(s: &str) -> Option<f64> {
485    // Strings look like " 24576 MiB" or "24576 MB"
486    let trimmed = s.trim();
487    let mut parts = trimmed.split_whitespace();
488    let value = parts.next()?.parse::<f64>().ok()?;
489    let unit = parts.next().unwrap_or("MiB");
490    match unit.to_ascii_lowercase().as_str() {
491        "mib" | "mb" => Some(value),
492        "gib" | "gb" => Some(value * 1024.0),
493        _ => Some(value),
494    }
495}
496
497#[cfg(test)]
498mod tests {
499    use super::*;
500
501    #[test]
502    fn parse_mib_handles_mib() {
503        assert_eq!(parse_mib(" 24576 MiB"), Some(24576.0));
504        assert_eq!(parse_mib("12288 MB"), Some(12288.0));
505        assert_eq!(parse_mib("24 GiB"), Some(24576.0));
506        assert_eq!(parse_mib("8 GB"), Some(8192.0));
507    }
508
509    #[test]
510    fn parse_mib_defaults_to_mib_when_the_unit_is_omitted() {
511        // A bare value (no unit token) is assumed to already be in MiB,
512        // matching nvidia-smi's `--units` output where the suffix is
513        // sometimes stripped.
514        assert_eq!(parse_mib("4096"), Some(4096.0));
515    }
516
517    #[test]
518    fn parse_mib_treats_an_unknown_unit_as_raw_mib() {
519        // An unrecognised suffix must not silently zero the GPU out of
520        // the VRAM total: the worker claims jobs by VRAM, so dropping a
521        // card to 0 would make it refuse work it can actually run. We
522        // keep the numeric value as-is (best-effort MiB) rather than
523        // returning `None`.
524        assert_eq!(parse_mib("2048 KiB"), Some(2048.0));
525        assert_eq!(parse_mib("4 TB"), Some(4.0));
526    }
527
528    #[test]
529    fn parse_mib_rejects_unparseable_or_empty_values() {
530        // A non-numeric leading token (e.g. an `[N/A]` placeholder) or
531        // an empty / whitespace-only line yields `None` so the caller
532        // skips it instead of polluting the total with a bogus number.
533        assert_eq!(parse_mib("N/A MiB"), None);
534        assert_eq!(parse_mib(""), None);
535        assert_eq!(parse_mib("   "), None);
536    }
537
538    #[test]
539    fn machine_name_returns_non_empty() {
540        assert!(!machine_name().is_empty());
541    }
542
543    #[test]
544    fn username_returns_non_empty() {
545        assert!(!username().is_empty());
546    }
547
548    #[test]
549    fn username_from_probe_returns_the_resolved_value() {
550        let user = username_from_probe(Ok::<_, std::io::Error>("alice".to_string()));
551        assert_eq!(user, "alice");
552    }
553
554    #[test]
555    fn username_from_probe_falls_back_to_unknown_user_on_error() {
556        let user =
557            username_from_probe(Err::<String, _>(std::io::Error::other("no entropy source")));
558        assert_eq!(user, "unknown-user");
559    }
560
561    #[test]
562    fn username_from_probe_warns_with_the_error_on_failure() {
563        // whoami 2.x made the probe fallible; a failure must leave an
564        // operator-visible breadcrumb naming the error rather than a
565        // silent fallback that hides why the user came back unknown.
566        let logs = crate::test_support::capture(|| {
567            let _ =
568                username_from_probe(Err::<String, _>(std::io::Error::other("permission denied")));
569        });
570        assert!(logs.contains("WARN"), "expected WARN level, got: {logs}");
571        assert!(
572            logs.contains("op=\"username\""),
573            "expected username op, got: {logs}"
574        );
575        assert!(
576            logs.contains("permission denied"),
577            "expected underlying error, got: {logs}"
578        );
579    }
580
581    #[test]
582    fn username_from_probe_emits_debug_value_on_success() {
583        let logs = crate::test_support::capture(|| {
584            let _ = username_from_probe(Ok::<_, std::io::Error>("bob".to_string()));
585        });
586        assert!(logs.contains("DEBUG"), "expected DEBUG event, got: {logs}");
587        assert!(
588            logs.contains("value=bob"),
589            "expected resolved value, got: {logs}"
590        );
591    }
592
593    #[test]
594    fn detect_vram_gb_from_sysfs_returns_zero_when_root_missing() {
595        let dir = tempfile::tempdir().unwrap();
596        let missing = dir.path().join("nope");
597        assert_eq!(detect_vram_gb_from_sysfs(&missing), 0.0);
598    }
599
600    // The NVIDIA proc-sysfs tree only exists on Linux, and its bus-id
601    // directory names contain colons — an illegal filename character on
602    // Windows — so these fixtures can only be built there.  The pure
603    // parsers (`parse_mib`, `parse_nvidia_smi_mib`) are tested
604    // cross-platform above.
605    #[cfg(target_os = "linux")]
606    #[test]
607    fn detect_vram_gb_from_sysfs_sums_parseable_gpus() {
608        let dir = tempfile::tempdir().unwrap();
609        for (bus, mib) in [("0000:01:00.0", "12288"), ("0000:02:00.0", "24576")] {
610            let gpu = dir.path().join(bus);
611            std::fs::create_dir_all(&gpu).unwrap();
612            std::fs::write(
613                gpu.join("information"),
614                format!("Model: x\nVideo Memory: {mib} MiB\n"),
615            )
616            .unwrap();
617        }
618        // (12288 + 24576) / 1024 = 36 GiB
619        let gb = detect_vram_gb_from_sysfs(dir.path());
620        assert!((gb - 36.0).abs() < 1e-3, "got {gb}");
621    }
622
623    #[cfg(target_os = "linux")]
624    #[test]
625    fn detect_vram_gb_from_sysfs_sums_only_survivors_when_one_gpu_is_unreadable() {
626        // A healthy card next to one whose `information` can't be read
627        // (here a *directory* named `information`, so `read_to_string`
628        // fails on every platform): the survivor still totals, the bad
629        // card is dropped from the sum rather than zeroing the host out.
630        let dir = tempfile::tempdir().unwrap();
631        let good = dir.path().join("0000:01:00.0");
632        std::fs::create_dir_all(&good).unwrap();
633        std::fs::write(good.join("information"), "Video Memory: 12288 MiB\n").unwrap();
634        let bad = dir.path().join("0000:02:00.0");
635        std::fs::create_dir_all(bad.join("information")).unwrap();
636        // Only the healthy card's 12288 MiB / 1024 = 12 GiB counts.
637        let gb = detect_vram_gb_from_sysfs(dir.path());
638        assert!((gb - 12.0).abs() < 1e-3, "got {gb}");
639    }
640
641    // -----------------------------------------------------------------
642    // AMD / Intel-discrete VRAM via the DRM sysfs tree, and the Apple /
643    // Windows non-NVIDIA parsers.  These give the threshold sanity
644    // check + studio matching a real number on non-NVIDIA GPUs, which
645    // all reported 0 before.
646    // -----------------------------------------------------------------
647
648    #[test]
649    fn is_drm_card_dir_matches_only_primary_nodes() {
650        assert!(is_drm_card_dir("card0"));
651        assert!(is_drm_card_dir("card12"));
652        assert!(!is_drm_card_dir("card0-DP-1"));
653        assert!(!is_drm_card_dir("renderD128"));
654        assert!(!is_drm_card_dir("card"));
655        assert!(!is_drm_card_dir("controlD64"));
656    }
657
658    #[test]
659    fn parse_amd_vram_total_bytes_reads_a_byte_count() {
660        assert_eq!(
661            parse_amd_vram_total_bytes("17163091968\n"),
662            Some(17163091968)
663        );
664        assert_eq!(
665            parse_amd_vram_total_bytes("0"),
666            None,
667            "zero = no VRAM file value"
668        );
669        assert_eq!(parse_amd_vram_total_bytes("N/A"), None);
670        assert_eq!(parse_amd_vram_total_bytes(""), None);
671    }
672
673    #[test]
674    fn detect_vram_gb_from_amd_sysfs_sums_cards_and_ignores_siblings() {
675        let dir = tempfile::tempdir().unwrap();
676        // card0 = 16 GiB, card1 = 8 GiB.
677        for (card, bytes) in [
678            ("card0", 16u64 * 1024 * 1024 * 1024),
679            ("card1", 8 * 1024 * 1024 * 1024),
680        ] {
681            let dev = dir.path().join(card).join("device");
682            std::fs::create_dir_all(&dev).unwrap();
683            std::fs::write(dev.join("mem_info_vram_total"), bytes.to_string()).unwrap();
684        }
685        // A connector node + a render node must be ignored.
686        std::fs::create_dir_all(dir.path().join("card0-DP-1")).unwrap();
687        std::fs::create_dir_all(dir.path().join("renderD128")).unwrap();
688        // An iGPU card with no VRAM file contributes nothing.
689        std::fs::create_dir_all(dir.path().join("card2").join("device")).unwrap();
690
691        let gb = detect_vram_gb_from_amd_sysfs(dir.path());
692        assert!((gb - 24.0).abs() < 1e-3, "expected 24 GiB, got {gb}");
693    }
694
695    #[test]
696    fn detect_vram_gb_from_amd_sysfs_returns_zero_without_a_tree() {
697        let missing = std::path::Path::new("/definitely/no/drm/here");
698        assert_eq!(detect_vram_gb_from_amd_sysfs(missing), 0.0);
699    }
700
701    #[test]
702    fn parse_sysctl_memsize_scales_unified_memory() {
703        // 32 GiB unified memory, 75% GPU-addressable = 24 GiB.
704        let bytes = (32u64 * 1024 * 1024 * 1024).to_string();
705        let gb = parse_sysctl_memsize(&bytes, 0.75).unwrap();
706        assert!((gb - 24.0).abs() < 1e-2, "got {gb}");
707        assert_eq!(parse_sysctl_memsize("0", 0.75), None);
708        assert_eq!(parse_sysctl_memsize("garbage", 0.75), None);
709    }
710
711    #[test]
712    fn parse_wmic_adapter_ram_sums_adapters_and_ignores_noise() {
713        // Two adapters: 8 GiB + 4 GiB, with a header + blank lines.
714        let out = format!(
715            "AdapterRAM\n\n{}\n{}\n",
716            8u64 * 1024 * 1024 * 1024,
717            4u64 * 1024 * 1024 * 1024
718        );
719        let gb = parse_wmic_adapter_ram(&out).unwrap();
720        assert!((gb - 12.0).abs() < 1e-2, "got {gb}");
721        assert_eq!(parse_wmic_adapter_ram("AdapterRAM\n\n"), None, "no numbers");
722        assert_eq!(parse_wmic_adapter_ram(""), None);
723    }
724
725    // -----------------------------------------------------------------
726    // nvidia-smi fallback — current NVIDIA drivers (5xx) dropped the
727    // "Video Memory" line from the sysfs `information` file, so the
728    // sysfs probe yields 0 on otherwise-capable hosts.  `nvidia-smi`
729    // ships with every driver and its `--query-gpu` interface is stable
730    // across versions, so it's the layout-proof fallback.
731    // -----------------------------------------------------------------
732
733    #[test]
734    fn parse_nvidia_smi_mib_reads_a_single_bare_value() {
735        let total = parse_nvidia_smi_mib("24564\n").unwrap();
736        assert_eq!(total.mib, 24564.0);
737        assert_eq!(total.dropped, 0);
738    }
739
740    #[test]
741    fn parse_nvidia_smi_mib_sums_multiple_gpus() {
742        let total = parse_nvidia_smi_mib("24564\n24564\n").unwrap();
743        assert_eq!(total.mib, 49128.0);
744        assert_eq!(total.dropped, 0);
745    }
746
747    #[test]
748    fn parse_nvidia_smi_mib_tolerates_units_and_crlf_whitespace() {
749        // If `nounits` is ever dropped the value arrives as "24564 MiB".
750        let total = parse_nvidia_smi_mib("  24564 MiB \r\n").unwrap();
751        assert_eq!(total.mib, 24564.0);
752        assert_eq!(total.dropped, 0);
753    }
754
755    #[test]
756    fn parse_nvidia_smi_mib_returns_none_on_empty_or_na() {
757        assert!(parse_nvidia_smi_mib("").is_none());
758        assert!(parse_nvidia_smi_mib("\n[N/A]\n").is_none());
759    }
760
761    #[test]
762    fn parse_nvidia_smi_mib_sums_survivors_and_counts_a_dropped_gpu() {
763        // A healthy 24 GiB card next to one nvidia-smi reports `[N/A]`
764        // for (fell off the bus / ECC fault): the survivor's VRAM still
765        // totals, but the dropped card is counted, not silently lost.
766        let total = parse_nvidia_smi_mib("24564\n[N/A]\n24564\n").unwrap();
767        assert_eq!(total.mib, 49128.0);
768        assert_eq!(total.dropped, 1);
769    }
770
771    #[test]
772    fn parse_nvidia_smi_mib_warns_on_each_dropped_gpu_line() {
773        // A multi-GPU box that under-reports its VRAM (and then refuses
774        // jobs it can run) must leave a per-line breadcrumb naming the
775        // offending value, not vanish the card without a trace.
776        let logs = crate::test_support::capture(|| {
777            let _ = parse_nvidia_smi_mib("24564\n[N/A]\n");
778        });
779        assert!(logs.contains("WARN"), "expected WARN level, got: {logs}");
780        assert!(
781            logs.contains("op=\"probe_vram\""),
782            "expected probe_vram op, got: {logs}"
783        );
784        assert!(
785            logs.contains("source=\"nvidia_smi\""),
786            "expected source=nvidia_smi, got: {logs}"
787        );
788        assert!(
789            logs.contains("[N/A]"),
790            "the warning must name the unparseable value, got: {logs}"
791        );
792        assert!(
793            logs.contains("dropping this GPU"),
794            "the warning must explain the drop, got: {logs}"
795        );
796    }
797
798    #[test]
799    fn vram_gb_from_smi_stdout_reports_dropped_count_in_breadcrumb() {
800        // The success breadcrumb must surface how many GPUs were dropped
801        // so a truncated VRAM total can't pass for a complete one.
802        let logs = crate::test_support::capture(|| {
803            let gb = vram_gb_from_smi_stdout("24564\n[N/A]\n").unwrap();
804            assert!((gb - 23.99).abs() < 0.05, "survivor still totals: {gb}");
805        });
806        assert!(
807            logs.contains("dropped=1"),
808            "the breadcrumb must report the dropped count, got: {logs}"
809        );
810    }
811
812    #[test]
813    fn vram_gb_from_smi_stdout_converts_mib_to_gb() {
814        // 24564 MiB / 1024 = 23.99 GiB
815        let gb = vram_gb_from_smi_stdout("24564\n").unwrap();
816        assert!((gb - 23.99).abs() < 0.05, "got {gb}");
817    }
818
819    #[test]
820    fn vram_gb_from_smi_stdout_is_none_when_unparseable() {
821        assert_eq!(vram_gb_from_smi_stdout("\n[N/A]\n"), None);
822    }
823
824    #[test]
825    fn vram_gb_from_smi_stdout_emits_info_breadcrumb_on_success() {
826        let logs = crate::test_support::capture(|| {
827            let _ = vram_gb_from_smi_stdout("24564\n");
828        });
829        assert!(logs.contains("INFO"), "expected INFO level, got: {logs}");
830        assert!(
831            logs.contains("op=\"probe_vram\""),
832            "expected probe_vram op, got: {logs}"
833        );
834        assert!(
835            logs.contains("source=\"nvidia_smi\""),
836            "expected source=nvidia_smi, got: {logs}"
837        );
838    }
839}