Skip to content

Commit c66efcb

Browse files
kerthcetkerthcet
authored
Support GPU detect (#27)
* support GPU detect Signed-off-by: kerthcet <kerthcet@gmail.com> * fix lint Signed-off-by: kerthcet <kerthcet@gmail.com> --------- Signed-off-by: kerthcet <kerthcet@gmail.com>
1 parent 7e28012 commit c66efcb

1 file changed

Lines changed: 179 additions & 0 deletions

File tree

‎src/system/system_info.rs‎

Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ use serde::{Deserialize, Serialize};
22
use std::fs;
33
use std::os::unix::fs::MetadataExt;
44
use std::path::PathBuf;
5+
use std::process::Command;
56
use sysinfo::System;
67

78
use crate::registry::model_registry::ModelRegistry;
@@ -15,12 +16,20 @@ pub struct SystemInfo {
1516
pub architecture: String,
1617
pub cpu_cores: usize,
1718
pub total_memory: String,
19+
pub gpu_info: Vec<GpuInfo>,
1820
pub cache_dir: String,
1921
pub cache_size: String,
2022
pub models_count: usize,
2123
pub running_models: usize,
2224
}
2325

26+
#[derive(Debug, Serialize, Deserialize, Clone)]
27+
pub struct GpuInfo {
28+
pub name: String,
29+
pub backend: String, // "CUDA", "Metal", "ROCm", or "Unknown"
30+
pub memory: Option<String>,
31+
}
32+
2433
impl SystemInfo {
2534
pub fn collect() -> Self {
2635
let mut sys = System::new_all();
@@ -32,19 +41,173 @@ impl SystemInfo {
3241
let registry = ModelRegistry::new(None);
3342
let models_count = registry.load_models().unwrap_or_default().len();
3443

44+
let gpu_info = Self::detect_gpus();
45+
3546
SystemInfo {
3647
version: env!("CARGO_PKG_VERSION").to_string(),
3748
os: System::name().unwrap_or_else(|| "Unknown".to_string()),
3849
architecture: System::cpu_arch().unwrap_or_else(|| "Unknown".to_string()),
3950
cpu_cores: sys.cpus().len(),
4051
total_memory: format_size(sys.total_memory()),
52+
gpu_info,
4153
cache_dir: cache_dir.to_string_lossy().to_string(),
4254
cache_size: format_size(cache_size),
4355
models_count,
4456
running_models: 0, // TODO: implement running models tracking
4557
}
4658
}
4759

60+
fn detect_gpus() -> Vec<GpuInfo> {
61+
let mut gpus = Vec::new();
62+
63+
// Try NVIDIA GPUs first (Linux/Windows)
64+
if let Some(nvidia_gpus) = Self::detect_nvidia_gpus() {
65+
gpus.extend(nvidia_gpus);
66+
}
67+
68+
// Try Metal (macOS)
69+
if let Some(metal_gpu) = Self::detect_metal_gpu() {
70+
gpus.push(metal_gpu);
71+
}
72+
73+
// Try AMD ROCm (Linux)
74+
if let Some(amd_gpus) = Self::detect_amd_gpus() {
75+
gpus.extend(amd_gpus);
76+
}
77+
78+
gpus
79+
}
80+
81+
fn detect_nvidia_gpus() -> Option<Vec<GpuInfo>> {
82+
let output = Command::new("nvidia-smi")
83+
.args([
84+
"--query-gpu=name,memory.total",
85+
"--format=csv,noheader,nounits",
86+
])
87+
.output()
88+
.ok()?;
89+
90+
if !output.status.success() {
91+
return None;
92+
}
93+
94+
let output_str = String::from_utf8(output.stdout).ok()?;
95+
let mut gpus = Vec::new();
96+
97+
for line in output_str.lines() {
98+
let parts: Vec<&str> = line.split(',').map(|s| s.trim()).collect();
99+
if parts.len() >= 2 {
100+
gpus.push(GpuInfo {
101+
name: parts[0].to_string(),
102+
backend: "CUDA".to_string(),
103+
memory: Some(format!("{} MB", parts[1])),
104+
});
105+
}
106+
}
107+
108+
if gpus.is_empty() {
109+
None
110+
} else {
111+
Some(gpus)
112+
}
113+
}
114+
115+
fn detect_metal_gpu() -> Option<GpuInfo> {
116+
// Check if running on macOS
117+
if !cfg!(target_os = "macos") {
118+
return None;
119+
}
120+
121+
let output = Command::new("system_profiler")
122+
.arg("SPDisplaysDataType")
123+
.output()
124+
.ok()?;
125+
126+
if !output.status.success() {
127+
return None;
128+
}
129+
130+
let output_str = String::from_utf8(output.stdout).ok()?;
131+
let lines: Vec<&str> = output_str.lines().collect();
132+
133+
// Find GPU name and cores
134+
let mut gpu_name = None;
135+
let mut core_count = None;
136+
137+
for (i, line) in lines.iter().enumerate() {
138+
if line.contains("Chipset Model:") {
139+
let parts: Vec<&str> = line.split("Chipset Model:").collect();
140+
if parts.len() >= 2 {
141+
let name = parts[1].trim();
142+
if !name.is_empty() {
143+
gpu_name = Some(name.to_string());
144+
145+
// Look for core count in the next few lines
146+
for line in lines.iter().skip(i + 1).take(10) {
147+
if line.contains("Total Number of Cores:") {
148+
let core_parts: Vec<&str> =
149+
line.split("Total Number of Cores:").collect();
150+
if core_parts.len() >= 2 {
151+
core_count = Some(core_parts[1].trim().to_string());
152+
}
153+
break;
154+
}
155+
}
156+
break;
157+
}
158+
}
159+
}
160+
}
161+
162+
if let Some(name) = gpu_name {
163+
let memory_str = core_count.map(|cores| format!("{} GPU cores", cores));
164+
165+
return Some(GpuInfo {
166+
name,
167+
backend: "Metal".to_string(),
168+
memory: memory_str,
169+
});
170+
}
171+
172+
None
173+
}
174+
175+
fn detect_amd_gpus() -> Option<Vec<GpuInfo>> {
176+
let output = Command::new("rocm-smi")
177+
.arg("--showproductname")
178+
.output()
179+
.ok()?;
180+
181+
if !output.status.success() {
182+
return None;
183+
}
184+
185+
let output_str = String::from_utf8(output.stdout).ok()?;
186+
let mut gpus = Vec::new();
187+
188+
for line in output_str.lines() {
189+
if line.contains("Card series:") || line.contains("Card model:") {
190+
let parts: Vec<&str> = line.split(':').collect();
191+
if parts.len() >= 2 {
192+
let name = parts[1].trim().to_string();
193+
if !name.is_empty() {
194+
gpus.push(GpuInfo {
195+
name,
196+
backend: "ROCm".to_string(),
197+
memory: None,
198+
});
199+
}
200+
}
201+
}
202+
}
203+
204+
if gpus.is_empty() {
205+
None
206+
} else {
207+
Some(gpus)
208+
}
209+
}
210+
48211
fn calculate_cache_size(cache_dir: &PathBuf) -> u64 {
49212
if !cache_dir.exists() {
50213
return 0;
@@ -93,6 +256,22 @@ impl SystemInfo {
93256
println!(" Architecture: {}", self.architecture);
94257
println!(" CPU Cores: {}", self.cpu_cores);
95258
println!(" Total Memory: {}", self.total_memory);
259+
260+
if !self.gpu_info.is_empty() {
261+
for (i, gpu) in self.gpu_info.iter().enumerate() {
262+
if i == 0 {
263+
print!(" GPU: ");
264+
} else {
265+
print!(" ");
266+
}
267+
print!("{} ({})", gpu.name, gpu.backend);
268+
if let Some(ref memory) = gpu.memory {
269+
print!(" - {}", memory);
270+
}
271+
println!();
272+
}
273+
}
274+
96275
println!();
97276
println!("PUMA Information:");
98277
println!(" PUMA Version: {}", self.version);

0 commit comments

Comments
 (0)