nd / nd-core/src/specs.rs

//! Machine specs, gathered at upload time (best effort, linux-first).
//! Opt-in: the cli asks before including them (`--specs` / `--no-specs`).

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Specs {
    pub os: String,
    pub arch: String,
    pub cpu: String,
    pub cpu_cores: usize,
    pub memory_gb: f64,
    pub gpu: String,
}

impl Specs {
    pub fn is_empty(&self) -> bool {
        self.os.is_empty() && self.cpu.is_empty() && self.gpu.is_empty() && self.memory_gb == 0.0
    }

    /// One-line "os · arch · cpu (n cores) · ram · gpu: x"
    pub fn summary(&self) -> String {
        let mut parts = Vec::new();
        if !self.os.is_empty() {
            parts.push(self.os.clone());
        }
        if !self.arch.is_empty() {
            parts.push(self.arch.clone());
        }
        if !self.cpu.is_empty() {
            if self.cpu_cores > 0 {
                parts.push(format!("{} ({} cores)", self.cpu, self.cpu_cores));
            } else {
                parts.push(self.cpu.clone());
            }
        }
        if self.memory_gb > 0.0 {
            parts.push(format!("{:.1} GB ram", self.memory_gb));
        }
        if !self.gpu.is_empty() {
            parts.push(format!("gpu: {}", self.gpu));
        }
        parts.join(" · ")
    }
}

pub fn gather() -> Specs {
    Specs {
        os: detect_os(),
        arch: std::env::consts::ARCH.to_string(),
        cpu: read_cpu_model(),
        cpu_cores: std::thread::available_parallelism()
            .map(|n| n.get())
            .unwrap_or(0),
        memory_gb: read_mem_gb(),
        gpu: detect_gpu(),
    }
}

fn detect_os() -> String {
    if let Ok(raw) = std::fs::read_to_string("/etc/os-release") {
        for line in raw.lines() {
            if let Some(v) = line.strip_prefix("PRETTY_NAME=") {
                return v.trim_matches('"').trim().to_string();
            }
        }
    }
    std::env::consts::OS.to_string()
}

fn read_cpu_model() -> String {
    std::fs::read_to_string("/proc/cpuinfo")
        .ok()
        .and_then(|raw| {
            raw.lines()
                .find(|l| l.starts_with("model name"))
                .and_then(|l| l.split(':').nth(1))
                .map(|s| s.trim().to_string())
        })
        .unwrap_or_default()
}

fn read_mem_gb() -> f64 {
    std::fs::read_to_string("/proc/meminfo")
        .ok()
        .and_then(|raw| {
            raw.lines()
                .find(|l| l.starts_with("MemTotal:"))
                .and_then(|l| l.split_whitespace().nth(1))
                .and_then(|kb| kb.parse::<f64>().ok())
        })
        .map(|kb| (kb / 1024.0 / 1024.0 * 10.0).round() / 10.0)
        .unwrap_or(0.0)
}

/// nvidia-smi -> rocm-smi -> lspci, first hit wins.
fn detect_gpu() -> String {
    gpu_nvidia()
        .or_else(gpu_amd)
        .or_else(gpu_lspci)
        .unwrap_or_default()
}

fn gpu_nvidia() -> Option<String> {
    let out = std::process::Command::new("nvidia-smi")
        .args(["--query-gpu=name", "--format=csv,noheader"])
        .output()
        .ok()?;
    if !out.status.success() {
        return None;
    }
    first_nonempty_line(&String::from_utf8_lossy(&out.stdout))
}

fn gpu_amd() -> Option<String> {
    let out = std::process::Command::new("rocm-smi")
        .args(["--showproductname"])
        .output()
        .ok()?;
    if !out.status.success() {
        return None;
    }
    let s = String::from_utf8_lossy(&out.stdout);
    // "GPU[0] : Product Name:        AMD Radeon RX 7900 XTX"
    for line in s.lines() {
        if let Some(idx) = line.find("Product Name:") {
            let v = line[idx + "Product Name:".len()..].trim();
            if !v.is_empty() {
                return Some(v.to_string());
            }
        }
    }
    // some rocm versions print "Card series: ..." instead
    for line in s.lines() {
        let l = line.to_lowercase();
        if (l.contains("radeon") || l.contains("instinct")) && !line.trim().is_empty() {
            return Some(line.trim().to_string());
        }
    }
    None
}

fn gpu_lspci() -> Option<String> {
    let out = std::process::Command::new("lspci").output().ok()?;
    if !out.status.success() {
        return None;
    }
    let s = String::from_utf8_lossy(&out.stdout);
    for line in s.lines() {
        let l = line.to_lowercase();
        if l.contains("vga") || l.contains("3d controller") || l.contains("display controller") {
            // "01:00.0 VGA compatible controller: NVIDIA Corporation GA104 ..."
            if let Some(v) = line.splitn(3, ": ").nth(2) {
                let v = v.trim();
                if !v.is_empty() {
                    return Some(v.to_string());
                }
            }
        }
    }
    None
}

fn first_nonempty_line(s: &str) -> Option<String> {
    s.lines().map(|l| l.trim()).find(|l| !l.is_empty()).map(|l| l.to_string())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn summary_and_empty() {
        let s = Specs::default();
        assert!(s.is_empty());
        assert_eq!(s.summary(), "");

        let s = Specs {
            os: "Ubuntu 24.04".into(),
            arch: "x86_64".into(),
            cpu: "AMD EPYC".into(),
            cpu_cores: 8,
            memory_gb: 31.2,
            gpu: "NVIDIA RTX".into(),
        };
        assert!(!s.is_empty());
        assert_eq!(s.summary(), "Ubuntu 24.04 · x86_64 · AMD EPYC (8 cores) · 31.2 GB ram · gpu: NVIDIA RTX");
    }

    #[test]
    fn serde_roundtrip() {
        let s = Specs { os: "t".into(), ..Default::default() };
        let v = serde_json::to_value(&s).unwrap();
        let back: Specs = serde_json::from_value(v).unwrap();
        assert_eq!(back.os, "t");
    }
}