foxhunt/crates/ml/src/gpu/capabilities.rs

//! GPU hardware capability detection.
//!
//! Detects GPU VRAM and device name via nvidia-smi.
//! Use `GpuCapabilities::detect()` at startup and cache the result.

use serde::{Deserialize, Serialize};
use std::process::Command;
use std::sync::OnceLock;
use tracing::{info, warn};

/// Detected GPU hardware capabilities.
///
/// Construct via `detect()` for real hardware or `cpu_fallback()` / `with_vram()` for testing.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct GpuCapabilities {
    /// GPU device name (e.g., "NVIDIA GeForce RTX 3050 Ti Laptop GPU")
    pub device_name: String,
    /// Total VRAM in MB (0.0 for CPU)
    pub total_vram_mb: f64,
    /// Free VRAM in MB at detection time (0.0 for CPU)
    pub free_vram_mb: f64,
    /// Whether a CUDA GPU was detected
    pub is_cuda: bool,
}

impl GpuCapabilities {
    /// Detect GPU capabilities by querying nvidia-smi.
    ///
    /// Falls back to CPU if nvidia-smi is unavailable or fails.
    /// Call this once at startup and cache the result.
    pub fn detect() -> Self {
        match Self::query_nvidia_smi() {
            Ok(caps) => {
                info!(
                    "GPU detected: {} (Total: {:.0} MB, Free: {:.0} MB)",
                    caps.device_name, caps.total_vram_mb, caps.free_vram_mb
                );
                caps
            }
            Err(reason) => {
                warn!("GPU detection failed ({}), using CPU fallback", reason);
                Self::cpu_fallback()
            }
        }
    }

    /// Explicit CPU-only capabilities (no GPU).
    pub fn cpu_fallback() -> Self {
        Self {
            device_name: "CPU".to_owned(),
            total_vram_mb: 0.0,
            free_vram_mb: 0.0,
            is_cuda: false,
        }
    }

    /// Construct with known VRAM values (for testing or manual override).
    pub fn with_vram<S: Into<String>>(device_name: S, total_mb: f64, free_mb: f64) -> Self {
        Self {
            device_name: device_name.into(),
            total_vram_mb: total_mb,
            free_vram_mb: free_mb,
            is_cuda: total_mb > 0.0,
        }
    }

    /// Query nvidia-smi for GPU memory and device name.
    fn query_nvidia_smi() -> Result<Self, String> {
        let output = Command::new("nvidia-smi")
            .args([
                "--query-gpu=memory.total,memory.free,name",
                "--format=csv,noheader,nounits",
            ])
            .output()
            .map_err(|e| format!("nvidia-smi not found: {}", e))?;

        if !output.status.success() {
            let stderr = String::from_utf8_lossy(&output.stderr);
            return Err(format!("nvidia-smi failed: {}", stderr));
        }

        let stdout = String::from_utf8_lossy(&output.stdout);
        let parts: Vec<&str> = stdout.trim().split(',').collect();

        if parts.len() < 3 {
            return Err(format!("unexpected nvidia-smi format: {}", stdout));
        }

        let total_mb = parts
            .first()
            .ok_or("missing total")?
            .trim()
            .parse::<f64>()
            .map_err(|e| format!("parse total: {}", e))?;
        let free_mb = parts
            .get(1)
            .ok_or("missing free")?
            .trim()
            .parse::<f64>()
            .map_err(|e| format!("parse free: {}", e))?;
        let device_name = parts
            .get(2)
            .ok_or("missing name")?
            .trim()
            .to_string();

        Ok(Self {
            device_name,
            total_vram_mb: total_mb,
            free_vram_mb: free_mb,
            is_cuda: true,
        })
    }
}

/// Cached GPU capabilities -- detected once, reused across all callers.
///
/// Use this instead of `GpuCapabilities::detect()` when you don't need fresh data.
/// The first call runs nvidia-smi; subsequent calls return the cached result.
pub fn cached_capabilities() -> &'static GpuCapabilities {
    static CAPS: OnceLock<GpuCapabilities> = OnceLock::new();
    CAPS.get_or_init(GpuCapabilities::detect)
}

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

    #[test]
    fn test_cpu_fallback() {
        let caps = GpuCapabilities::cpu_fallback();
        assert_eq!(caps.device_name, "CPU");
        assert_eq!(caps.total_vram_mb, 0.0);
        assert_eq!(caps.free_vram_mb, 0.0);
        assert!(!caps.is_cuda);
    }

    #[test]
    fn test_with_vram_rtx_3050_ti() {
        let caps = GpuCapabilities::with_vram("RTX 3050 Ti", 4096.0, 3700.0);
        assert_eq!(caps.device_name, "RTX 3050 Ti");
        assert_eq!(caps.total_vram_mb, 4096.0);
        assert_eq!(caps.free_vram_mb, 3700.0);
        assert!(caps.is_cuda);
    }

    #[test]
    fn test_with_vram_a100() {
        let caps = GpuCapabilities::with_vram("A100-SXM4-80GB", 81920.0, 80000.0);
        assert_eq!(caps.total_vram_mb, 81920.0);
        assert!(caps.is_cuda);
    }

    #[test]
    fn test_with_vram_zero_is_cpu() {
        let caps = GpuCapabilities::with_vram("none", 0.0, 0.0);
        assert!(!caps.is_cuda);
    }

    #[test]
    fn test_detect_does_not_panic() {
        // detect() should never panic -- falls back to CPU gracefully
        let caps = GpuCapabilities::detect();
        assert!(!caps.device_name.is_empty());
    }

    #[test]
    fn test_serde_roundtrip() {
        let caps = GpuCapabilities::with_vram("RTX 4090", 24576.0, 23000.0);
        let json = serde_json::to_string(&caps).unwrap();
        let deserialized: GpuCapabilities = serde_json::from_str(&json).unwrap();
        assert_eq!(caps, deserialized);
    }

    #[test]
    fn test_cached_capabilities_returns_same_instance() {
        let caps1 = cached_capabilities();
        let caps2 = cached_capabilities();
        assert_eq!(caps1, caps2);
    }
}