modelshelf 0.1.0

//! Hardware detection for model recommendations.
//!
//! [`Hardware::detect`] answers one question: how much memory can this
//! machine realistically dedicate to a local model? It reads total RAM,
//! enumerates NVIDIA GPUs via `nvidia-smi` (Linux, Windows, WSL2), and flags
//! Apple Silicon unified memory. Detection never fails: anything that cannot
//! be determined degrades to a conservative default.

use serde::{Deserialize, Serialize};

/// Environment variable that overrides detection with a JSON-encoded
/// [`Hardware`] value. Intended for tests and CI
/// (e.g. `MODELSHELF_HW='{"ram_bytes":34359738368}'`).
pub const HW_ENV: &str = "MODELSHELF_HW";

/// RAM assumed when detection fails entirely (8 GiB, deliberately low).
pub const FALLBACK_RAM_BYTES: u64 = 8 * GIB;

const GIB: u64 = 1024 * 1024 * 1024;

/// A discrete GPU visible to the system.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Gpu {
    /// Marketing name, e.g. `NVIDIA GeForce RTX 4090`.
    pub name: String,
    /// Total VRAM in bytes.
    pub vram_bytes: u64,
}

/// What the machine offers for running local models.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Hardware {
    /// Total physical RAM in bytes.
    pub ram_bytes: u64,
    /// True when RAM could not be detected and [`FALLBACK_RAM_BYTES`] was
    /// assumed.
    #[serde(default)]
    pub ram_assumed: bool,
    /// Discrete GPUs (NVIDIA only in this version).
    #[serde(default)]
    pub gpus: Vec<Gpu>,
    /// True on Apple Silicon, where CPU and GPU share one memory pool.
    #[serde(default)]
    pub unified_memory: bool,
}

impl Hardware {
    /// Detect the current machine. Never fails; see the module docs.
    ///
    /// When the [`HW_ENV`] environment variable holds a valid JSON
    /// [`Hardware`] value, it is returned as-is and no probing happens.
    pub fn detect() -> Hardware {
        if let Ok(json) = std::env::var(HW_ENV) {
            if let Ok(hw) = serde_json::from_str::<Hardware>(&json) {
                return hw;
            }
            tracing::warn!("ignoring unparseable {HW_ENV} override");
        }
        let (ram_bytes, ram_assumed) = match detect_ram() {
            Some(bytes) if bytes > 0 => (bytes, false),
            _ => (FALLBACK_RAM_BYTES, true),
        };
        Hardware {
            ram_bytes,
            ram_assumed,
            gpus: detect_nvidia(),
            unified_memory: cfg!(all(target_os = "macos", target_arch = "aarch64")),
        }
    }
}

#[cfg(target_os = "linux")]
fn detect_ram() -> Option<u64> {
    parse_meminfo(&std::fs::read_to_string("/proc/meminfo").ok()?)
}

#[cfg(target_os = "macos")]
fn detect_ram() -> Option<u64> {
    let out = std::process::Command::new("sysctl")
        .args(["-n", "hw.memsize"])
        .output()
        .ok()?;
    if !out.status.success() {
        return None;
    }
    String::from_utf8_lossy(&out.stdout).trim().parse().ok()
}

#[cfg(windows)]
fn detect_ram() -> Option<u64> {
    // MEMORYSTATUSEX + GlobalMemoryStatusEx, declared by hand to avoid a
    // Windows bindings dependency for a single stable API.
    #[repr(C)]
    struct MemoryStatusEx {
        length: u32,
        memory_load: u32,
        total_phys: u64,
        avail_phys: u64,
        total_page_file: u64,
        avail_page_file: u64,
        total_virtual: u64,
        avail_virtual: u64,
        avail_extended_virtual: u64,
    }
    #[link(name = "kernel32")]
    extern "system" {
        fn GlobalMemoryStatusEx(buffer: *mut MemoryStatusEx) -> i32;
    }
    let mut status = MemoryStatusEx {
        length: std::mem::size_of::<MemoryStatusEx>() as u32,
        memory_load: 0,
        total_phys: 0,
        avail_phys: 0,
        total_page_file: 0,
        avail_page_file: 0,
        total_virtual: 0,
        avail_virtual: 0,
        avail_extended_virtual: 0,
    };
    // SAFETY: `status` is a properly initialized MEMORYSTATUSEX with its
    // `length` field set, as the API requires; the call writes only within it.
    let ok = unsafe { GlobalMemoryStatusEx(&mut status) };
    (ok != 0).then_some(status.total_phys)
}

#[cfg(not(any(target_os = "linux", target_os = "macos", windows)))]
fn detect_ram() -> Option<u64> {
    None
}

/// Total RAM from `/proc/meminfo` content (`MemTotal:` is in kB).
#[cfg_attr(not(target_os = "linux"), allow(dead_code))]
fn parse_meminfo(contents: &str) -> Option<u64> {
    let line = contents.lines().find_map(|l| l.strip_prefix("MemTotal:"))?;
    let kib: u64 = line.split_whitespace().next()?.parse().ok()?;
    Some(kib * 1024)
}

fn detect_nvidia() -> Vec<Gpu> {
    let out = match std::process::Command::new("nvidia-smi")
        .args([
            "--query-gpu=name,memory.total",
            "--format=csv,noheader,nounits",
        ])
        .output()
    {
        Ok(out) if out.status.success() => out,
        _ => return Vec::new(),
    };
    parse_nvidia_smi(&String::from_utf8_lossy(&out.stdout))
}

/// Parse `nvidia-smi --query-gpu=name,memory.total
/// --format=csv,noheader,nounits` output: one `name, MiB` line per GPU.
/// Lines with unparseable memory (e.g. `[N/A]`) are skipped.
fn parse_nvidia_smi(output: &str) -> Vec<Gpu> {
    let mut gpus = Vec::new();
    for line in output.lines() {
        let line = line.trim();
        if line.is_empty() {
            continue;
        }
        // The name itself never contains a comma; memory is the last field.
        let Some((name, mib)) = line.rsplit_once(',') else {
            continue;
        };
        let Ok(mib) = mib.trim().parse::<u64>() else {
            continue;
        };
        gpus.push(Gpu {
            name: name.trim().to_owned(),
            vram_bytes: mib * 1024 * 1024,
        });
    }
    gpus
}

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

    #[test]
    fn meminfo_memtotal_is_parsed_as_kib() {
        let sample = "MemTotal:       32657844 kB\nMemFree:         1074956 kB\n";
        assert_eq!(parse_meminfo(sample), Some(32_657_844 * 1024));
        assert_eq!(parse_meminfo("MemFree: 1 kB\n"), None);
        assert_eq!(parse_meminfo(""), None);
        assert_eq!(parse_meminfo("MemTotal: garbage kB"), None);
    }

    #[test]
    fn nvidia_smi_single_and_multi_gpu() {
        let one = parse_nvidia_smi("NVIDIA GeForce RTX 4090, 24564\n");
        assert_eq!(
            one,
            vec![Gpu {
                name: "NVIDIA GeForce RTX 4090".into(),
                vram_bytes: 24564 * 1024 * 1024,
            }]
        );

        let two = parse_nvidia_smi("NVIDIA RTX A6000, 49140\nNVIDIA RTX A6000, 49140\n");
        assert_eq!(two.len(), 2);
        assert_eq!(two[1].vram_bytes, 49140 * 1024 * 1024);
    }

    #[test]
    fn nvidia_smi_skips_na_and_junk() {
        assert!(parse_nvidia_smi("").is_empty());
        assert!(parse_nvidia_smi("NVIDIA GeForce GT 710, [N/A]\n").is_empty());
        assert!(parse_nvidia_smi("no comma here\n").is_empty());
        // A healthy GPU on a later line still parses.
        let mixed = parse_nvidia_smi("Broken GPU, [N/A]\nNVIDIA T4, 15360\n");
        assert_eq!(mixed.len(), 1);
        assert_eq!(mixed[0].name, "NVIDIA T4");
    }

    #[test]
    fn hardware_env_json_shape_roundtrips() {
        // The same JSON shape the MODELSHELF_HW override accepts.
        let hw: Hardware = serde_json::from_str(
            r#"{"ram_bytes": 68719476736,
                "gpus": [{"name": "NVIDIA GeForce RTX 4090", "vram_bytes": 25757220864}]}"#,
        )
        .unwrap();
        assert_eq!(hw.ram_bytes, 64 * GIB);
        assert!(!hw.ram_assumed);
        assert!(!hw.unified_memory);
        assert_eq!(hw.gpus.len(), 1);
        let json = serde_json::to_string(&hw).unwrap();
        assert_eq!(serde_json::from_str::<Hardware>(&json).unwrap(), hw);
    }

    #[test]
    fn detect_never_panics_and_reports_positive_ram() {
        // Whatever the host looks like, detect() must produce usable data.
        let hw = Hardware::detect();
        assert!(hw.ram_bytes > 0);
    }
}