portablesource_rs/
gpu.rs

1// portablesource
2// Copyright (C) 2025  PortableSource / NeuroDonu
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU Affero General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8//
9// This program is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12// GNU Affero General Public License for more details.
13//
14// You should have received a copy of the GNU Affero General Public License
15// along with this program.  If not, see <https://www.gnu.org/licenses/>.
16
17//! GPU detection and management
18
19use crate::{Result, PortableSourceError};
20use std::process::Command;
21#[cfg(windows)]
22use serde::Deserialize;
23#[cfg(windows)]
24use wmi::{COMLibrary, WMIConnection};
25
26#[derive(Debug, Clone, PartialEq)]
27pub enum GpuType {
28    Nvidia,
29    Amd,
30    Intel,
31    Unknown,
32}
33
34#[derive(Debug, Clone)]
35pub struct GpuInfo {
36    pub name: String,
37    pub gpu_type: GpuType,
38    pub memory_mb: u32,
39    pub driver_version: Option<String>,
40}
41
42pub struct GpuDetector;
43
44impl GpuDetector {
45    pub fn new() -> Self {
46        Self
47    }
48    
49    /// Detect NVIDIA GPU using nvidia-smi
50    pub fn detect_nvidia_gpu(&self) -> Result<Option<GpuInfo>> {
51        let mut cmd = Command::new("nvidia-smi");
52        cmd.args(&["--query-gpu=name,memory.total,driver_version", "--format=csv,noheader,nounits"]);
53
54        #[cfg(target_os = "windows")]
55        {
56            use std::os::windows::process::CommandExt;
57            cmd.creation_flags(0x08000000); // CREATE_NO_WINDOW
58        }
59
60        let output = cmd.output();
61
62        match output {
63            Ok(output) if output.status.success() => {
64                let stdout = String::from_utf8_lossy(&output.stdout);
65                if let Some(line) = stdout.lines().next() {
66                    self.parse_nvidia_smi_output(line)
67                } else {
68                    Ok(None)
69                }
70            }
71            _ => {
72                log::debug!("nvidia-smi not available or failed");
73                Ok(None)
74            }
75        }
76    }
77    
78    fn parse_nvidia_smi_output(&self, line: &str) -> Result<Option<GpuInfo>> {
79        let parts: Vec<&str> = line.split(',').map(|s| s.trim()).collect();
80        
81        if parts.len() >= 3 {
82            let name = parts[0].to_string();
83            let memory_mb = parts[1].parse::<u32>()
84                .map_err(|_| PortableSourceError::gpu_detection("Failed to parse GPU memory"))?;
85            let driver_version = Some(parts[2].to_string());
86            
87            Ok(Some(GpuInfo {
88                name,
89                gpu_type: GpuType::Nvidia,
90                memory_mb,
91                driver_version,
92            }))
93        } else {
94            Err(PortableSourceError::gpu_detection("Invalid nvidia-smi output format"))
95        }
96    }
97    
98    /// Detect GPU using Windows WMI (via wmi crate), fallback to WMIC on Windows only
99    pub fn detect_gpu_wmi(&self) -> Result<Vec<GpuInfo>> {
100        #[cfg(windows)]
101        {
102            if let Ok(com) = COMLibrary::new() {
103                if let Ok(wmi_con) = WMIConnection::new(com.into()) {
104                    #[derive(Deserialize)]
105                    #[allow(non_snake_case)]
106                    struct Win32VideoController {
107                        #[serde(rename = "Name")] Name: Option<String>,
108                        #[serde(rename = "AdapterRAM")] AdapterRAM: Option<u64>,
109                        #[serde(rename = "DriverVersion")] DriverVersion: Option<String>,
110                    }
111                    if let Ok(results) = wmi_con.query::<Win32VideoController>() {
112                        let mut gpus = Vec::new();
113                        for r in results {
114                            let name = r.Name.unwrap_or_default();
115                            if name.is_empty() { continue; }
116                            let adapter_ram = r.AdapterRAM.unwrap_or(0);
117                            let memory_mb = (adapter_ram / (1024 * 1024)) as u32;
118                            let driver_version = r.DriverVersion;
119                            let gpu_type = self.determine_gpu_type(&name);
120                            gpus.push(GpuInfo { name, gpu_type, memory_mb, driver_version });
121                        }
122                        if !gpus.is_empty() { return Ok(gpus); }
123                    }
124                }
125            }
126
127            // Fallback: WMIC CLI
128            let mut cmd = Command::new("wmic");
129            cmd.args(&["path", "win32_VideoController", "get", "name,AdapterRAM,DriverVersion", "/format:csv"]);
130            {
131                use std::os::windows::process::CommandExt;
132                cmd.creation_flags(0x08000000);
133            }
134            let output = cmd.output();
135            match output {
136                Ok(output) if output.status.success() => {
137                    let stdout = String::from_utf8_lossy(&output.stdout);
138                    let mut gpus = Vec::new();
139                    for line in stdout.lines().skip(1) {
140                        if line.trim().is_empty() { continue; }
141                        let parts: Vec<&str> = line.split(',').collect();
142                        if parts.len() >= 4 {
143                            let name = parts[3].trim().to_string();
144                            if name.is_empty() || name == "Name" { continue; }
145                            let memory_bytes = parts[2].trim().parse::<u64>().unwrap_or(0);
146                            let memory_mb = (memory_bytes / (1024 * 1024)) as u32;
147                            let driver_version = {
148                                let dv = parts.get(1).map(|s| s.trim()).unwrap_or("");
149                                if dv.is_empty() || dv == "DriverVersion" { None } else { Some(dv.to_string()) }
150                            };
151                            let gpu_type = self.determine_gpu_type(&name);
152                            gpus.push(GpuInfo { name, gpu_type, memory_mb, driver_version });
153                        }
154                    }
155                    Ok(gpus)
156                }
157                _ => Ok(Vec::new()),
158            }
159        }
160        #[cfg(not(windows))]
161        {
162            Ok(Vec::new())
163        }
164    }
165
166    #[cfg(unix)]
167    fn detect_gpu_linux_lspci(&self) -> Vec<GpuInfo> {
168        let mut gpus = Vec::new();
169        let output = Command::new("sh")
170            .arg("-c")
171            .arg("lspci -mm | egrep -i 'VGA|3D|Display'")
172            .output();
173        if let Ok(out) = output {
174            if out.status.success() {
175                let text = String::from_utf8_lossy(&out.stdout);
176                for line in text.lines() {
177                    let l = line.to_string();
178                    let up = l.to_uppercase();
179                    let gpu_type = if up.contains("NVIDIA") { GpuType::Nvidia } else if up.contains("AMD") || up.contains("ATI") || up.contains("RADEON") { GpuType::Amd } else if up.contains("INTEL") { GpuType::Intel } else { GpuType::Unknown };
180                    if gpu_type != GpuType::Unknown {
181                        // Try to extract model name between quotes if present
182                        let name = if let Some(start) = l.find('"') { if let Some(end) = l[start+1..].find('"') { l[start+1..start+1+end].to_string() } else { l.clone() } } else { l.clone() };
183                        gpus.push(GpuInfo { name, gpu_type, memory_mb: 0, driver_version: None });
184                    }
185                }
186            }
187        }
188        gpus
189    }
190
191    #[cfg(unix)]
192    fn detect_gpu_linux_glxinfo(&self) -> Option<GpuInfo> {
193        let out = Command::new("sh").arg("-c").arg("glxinfo -B 2>/dev/null | grep 'renderer string' || true").output().ok()?;
194        if !out.status.success() { return None; }
195        let text = String::from_utf8_lossy(&out.stdout);
196        let line = text.lines().next()?.to_string();
197        let lower = line.to_lowercase();
198        let gpu_type = if lower.contains("nvidia") { GpuType::Nvidia } else if lower.contains("amd") || lower.contains("radeon") { GpuType::Amd } else if lower.contains("intel") { GpuType::Intel } else { GpuType::Unknown };
199        Some(GpuInfo { name: line, gpu_type, memory_mb: 0, driver_version: None })
200    }
201    
202    fn determine_gpu_type(&self, name: &str) -> GpuType {
203        let name_upper = name.to_uppercase();
204        
205        if name_upper.contains("NVIDIA") || name_upper.contains("GEFORCE") || name_upper.contains("QUADRO") || name_upper.contains("TESLA") {
206            GpuType::Nvidia
207        } else if name_upper.contains("AMD") || name_upper.contains("RADEON") {
208            GpuType::Amd
209        } else if name_upper.contains("INTEL") {
210            GpuType::Intel
211        } else {
212            GpuType::Unknown
213        }
214    }
215    
216    /// Get the best available GPU (prioritize NVIDIA)
217    pub fn get_best_gpu(&self) -> Result<Option<GpuInfo>> {
218        // First try nvidia-smi for accurate NVIDIA detection
219        if let Some(nvidia_gpu) = self.detect_nvidia_gpu()? {
220            return Ok(Some(nvidia_gpu));
221        }
222        
223        #[cfg(windows)]
224        {
225            // Fall back to WMI/WMIC on Windows
226            let gpus = self.detect_gpu_wmi()?;
227            for gpu in &gpus { if gpu.gpu_type == GpuType::Nvidia { return Ok(Some(gpu.clone())); } }
228            return Ok(gpus.into_iter().next());
229        }
230        #[cfg(unix)]
231        {
232            // Linux: try lspci then glxinfo as best-effort
233            let mut gpus = self.detect_gpu_linux_lspci();
234            if gpus.is_empty() {
235                if let Some(glx) = self.detect_gpu_linux_glxinfo() { gpus.push(glx); }
236            }
237            for gpu in &gpus { if gpu.gpu_type == GpuType::Nvidia { return Ok(Some(gpu.clone())); } }
238            return Ok(gpus.into_iter().next());
239        }
240    }
241    
242    /// Check if NVIDIA GPU is available
243    pub fn has_nvidia_gpu(&self) -> bool {
244        self.detect_nvidia_gpu().unwrap_or(None).is_some()
245    }
246
247}
248
249// removed raw COM helpers; using wmi crate instead
250
251impl Default for GpuDetector {
252    fn default() -> Self {
253        Self::new()
254    }
255}
portablesource_rs/gpu.rs

portablesource_rs/
gpu.rs
