1  //! GPU detection and information.
 2  
 3  use acdc_core::{Error, Result};
 4  use serde::{Deserialize, Serialize};
 5  use std::process::Command;
 6  
 7  /// GPU information.
 8  #[derive(Debug, Clone, Serialize, Deserialize)]
 9  pub struct GpuInfo {
10      /// GPU name/model
11      pub name: String,
12      /// Whether a compatible GPU is available
13      pub available: bool,
14      /// VRAM in bytes
15      pub vram_bytes: u64,
16      /// CUDA version (if available)
17      pub cuda_version: Option<String>,
18      /// Driver version
19      pub driver_version: Option<String>,
20  }
21  
22  impl GpuInfo {
23      /// Detect GPU information using nvidia-smi.
24      pub fn detect() -> Result<Self> {
25          // Try nvidia-smi first
26          if let Ok(info) = Self::detect_nvidia() {
27              return Ok(info);
28          }
29  
30          // No compatible GPU found
31          Err(Error::SystemCheck("No compatible GPU detected".to_string()))
32      }
33  
34      fn detect_nvidia() -> Result<Self> {
35          let output = Command::new("nvidia-smi")
36              .args([
37                  "--query-gpu=name,memory.total,driver_version",
38                  "--format=csv,noheader,nounits",
39              ])
40              .output()
41              .map_err(|e| Error::SystemCheck(format!("nvidia-smi not found: {}", e)))?;
42  
43          if !output.status.success() {
44              return Err(Error::SystemCheck("nvidia-smi failed to run".to_string()));
45          }
46  
47          let stdout = String::from_utf8_lossy(&output.stdout);
48          let line = stdout
49              .lines()
50              .next()
51              .ok_or_else(|| Error::SystemCheck("No GPU info returned".to_string()))?;
52  
53          let parts: Vec<&str> = line.split(", ").collect();
54          if parts.len() < 3 {
55              return Err(Error::SystemCheck("Invalid nvidia-smi output".to_string()));
56          }
57  
58          let name = parts[0].trim().to_string();
59          let vram_mb: u64 = parts[1].trim().parse().unwrap_or(0);
60          let driver_version = Some(parts[2].trim().to_string());
61  
62          // Get CUDA version
63          let cuda_version = Command::new("nvcc")
64              .arg("--version")
65              .output()
66              .ok()
67              .and_then(|o| {
68                  let stdout = String::from_utf8_lossy(&o.stdout);
69                  stdout
70                      .lines()
71                      .find(|l| l.contains("release"))
72                      .and_then(|l| l.split("release ").nth(1))
73                      .map(|v| v.split(',').next().unwrap_or(v).trim().to_string())
74              });
75  
76          Ok(Self {
77              name,
78              available: true,
79              vram_bytes: vram_mb * 1024 * 1024,
80              cuda_version,
81              driver_version,
82          })
83      }
84  
85      /// Get VRAM in GB.
86      pub fn vram_gb(&self) -> u64 {
87          self.vram_bytes / (1024 * 1024 * 1024)
88      }
89  }