1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
use crate::model::GpuInfo;
use std::process::Command;
/// Collect GPU information on Linux using NVML for NVIDIA GPUs, fallback to lspci
pub fn collect_gpu_info() -> Vec<GpuInfo> {
// Try NVML first for NVIDIA GPUs
if let Some(nvidia_gpus) = collect_nvidia_gpus()
&& !nvidia_gpus.is_empty()
{
tracing::debug!("Found {} NVIDIA GPU(s) via NVML", nvidia_gpus.len());
return nvidia_gpus;
}
// Fallback to lspci for AMD/Intel/other GPUs
tracing::debug!("NVML not available, falling back to lspci");
collect_gpus_via_lspci()
}
/// Collect NVIDIA GPU information using NVML
#[allow(clippy::cast_precision_loss)]
fn collect_nvidia_gpus() -> Option<Vec<GpuInfo>> {
use nvml_wrapper::Nvml;
// Initialize NVML
let Ok(nvml) = Nvml::init() else {
return None;
};
let Ok(device_count) = nvml.device_count() else {
return None;
};
let mut gpus = Vec::new();
for i in 0..device_count {
match nvml.device_by_index(i) {
Ok(device) => {
let model = device
.name()
.unwrap_or_else(|_| "Unknown NVIDIA GPU".to_owned());
// Get memory info
let memory_info = device.memory_info().ok();
let total_memory_mb = memory_info
.as_ref()
.map(|m| m.total as f64 / 1024.0 / 1024.0);
let used_memory_mb = memory_info
.as_ref()
.map(|m| m.used as f64 / 1024.0 / 1024.0);
// Try to get CUDA cores (not directly available, but we can get compute capability)
let cores = None; // NVML doesn't expose CUDA cores directly
gpus.push(GpuInfo {
model,
cores,
total_memory_mb,
used_memory_mb,
});
tracing::debug!(
"NVIDIA GPU {}: {} (Memory: {:.0} MB / {:.0} MB)",
i,
gpus.last()?.model,
used_memory_mb.unwrap_or(0.0),
total_memory_mb.unwrap_or(0.0)
);
}
Err(e) => {
tracing::warn!("Failed to get NVML device handle for GPU #{}: {}", i, e);
}
}
}
Some(gpus)
}
/// Collect GPU information using lspci (fallback for non-NVIDIA GPUs)
fn collect_gpus_via_lspci() -> Vec<GpuInfo> {
let output = Command::new("lspci").output();
if let Ok(output) = output
&& output.status.success()
{
let output_str = String::from_utf8_lossy(&output.stdout);
let mut gpus = Vec::new();
for line in output_str.lines() {
let line_lower = line.to_lowercase();
if line_lower.contains("vga")
|| line_lower.contains("3d")
|| line_lower.contains("display")
{
// Extract GPU model from lspci output
// Format: "00:02.0 VGA compatible controller: Intel Corporation ..."
if let Some(pos) = line.find(':')
&& let Some(model_start) = line[pos..].find(':')
{
let model = line[pos + model_start + 1..].trim().to_owned();
gpus.push(GpuInfo {
model,
cores: None,
total_memory_mb: None,
used_memory_mb: None,
});
}
}
}
if !gpus.is_empty() {
tracing::debug!("Found {} GPU(s) via lspci", gpus.len());
return gpus;
}
}
// If lspci fails or finds nothing, return empty list
Vec::new()
}