1use serde::{Deserialize, Serialize};
22
23pub mod collect;
24pub mod energy;
25pub mod energy_sampler;
26pub mod equivalence;
27pub mod ids;
28pub mod limits;
29pub mod parse;
30pub mod policy;
31pub mod report;
32pub mod state;
33pub mod throttle_monitor;
34
35pub use policy::{Decision, Sensitivity, Validity};
36pub use report::HardwareStateReport;
37pub use state::{HardwareState, HardwareStateDelta};
38
39#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
42pub struct Hardware {
43 #[serde(default)]
46 pub gpu: String,
47 #[serde(default)]
49 pub driver: String,
50 #[serde(default, skip_serializing_if = "Option::is_none")]
54 pub sm_clock_mhz: Option<f64>,
55 #[serde(default, skip_serializing_if = "Option::is_none")]
60 pub gpu_count: Option<u32>,
61 #[serde(default)]
64 pub source: String,
65}
66
67impl Hardware {
68 pub fn unknown() -> Self {
69 Self::default()
70 }
71
72 pub fn gate_key(&self) -> String {
82 if self.gpu.is_empty() {
83 return "unknown".to_string();
84 }
85 if let Some(id) = ids::hardware_id_from_gpu_name(&self.gpu) {
86 return id.to_string();
87 }
88 let key: String = self
89 .gpu
90 .to_lowercase()
91 .replace("nvidia", "")
92 .replace("amd", "")
93 .chars()
94 .filter(|c| c.is_alphanumeric())
95 .collect();
96 if key.is_empty() {
97 "unknown".to_string()
98 } else {
99 key
100 }
101 }
102
103 pub fn is_unknown(&self) -> bool {
106 self.gpu.is_empty() && self.driver.is_empty() && self.sm_clock_mhz.is_none()
107 }
108
109 pub fn one_line(&self) -> String {
113 if self.is_unknown() {
114 return "unknown hardware".to_string();
115 }
116 let mut parts = Vec::new();
117 if !self.gpu.is_empty() {
118 parts.push(match self.gpu_count {
119 Some(n) if n > 1 => format!("{} ×{n}", self.gpu),
120 _ => self.gpu.clone(),
121 });
122 }
123 if !self.driver.is_empty() {
124 parts.push(format!("driver {}", self.driver));
125 }
126 if let Some(clock) = self.sm_clock_mhz {
127 parts.push(format!("sm {clock:.0} MHz"));
128 }
129 parts.join(" · ")
130 }
131
132 pub fn probe() -> Self {
135 nvidia_smi()
136 .or_else(rocm_smi)
137 .or_else(sysfs)
138 .unwrap_or_else(Self::unknown)
139 }
140}
141
142fn run(tool: &str, args: &[&str]) -> Option<String> {
145 let out = std::process::Command::new(tool)
146 .args(args)
147 .stdin(std::process::Stdio::null())
148 .output()
149 .ok()?;
150 out.status
151 .success()
152 .then(|| String::from_utf8_lossy(&out.stdout).trim().to_string())
153 .filter(|s| !s.is_empty())
154}
155
156fn nvidia_smi() -> Option<Hardware> {
159 let line = run(
160 "nvidia-smi",
161 &[
162 "--query-gpu=name,driver_version,clocks.sm",
163 "--format=csv,noheader,nounits",
164 ],
165 )?
166 .lines()
167 .next()?
168 .to_string();
169 let mut cells = line.split(',').map(str::trim);
170 let gpu = cells.next().unwrap_or_default().to_string();
171 let driver = cells.next().unwrap_or_default().to_string();
172 let sm_clock_mhz = cells
173 .next()
174 .filter(|v| *v != "[N/A]")
175 .and_then(|v| v.parse().ok());
176 let gpu_count = run("nvidia-smi", &["-L"]).and_then(|out| parse::gpu_count(&out));
177 Some(Hardware {
178 gpu,
179 driver,
180 sm_clock_mhz,
181 gpu_count,
182 source: "nvidia-smi".into(),
183 })
184}
185
186fn rocm_smi() -> Option<Hardware> {
189 let card = run(
190 "rocm-smi",
191 &["--showproductname", "--showdriverversion", "--csv"],
192 )?;
193 let mut gpu = String::new();
194 let mut driver = String::new();
195 for line in card.lines().skip(1) {
196 let cells: Vec<&str> = line.split(',').map(str::trim).collect();
197 if cells.len() >= 3 {
198 gpu = cells[1].to_string();
199 driver = cells[2].to_string();
200 break;
201 }
202 }
203 let sm_clock_mhz = run("rocm-smi", &["--showclocks", "--csv"]).and_then(|out| {
204 out.lines()
205 .find_map(|l| l.split(',').nth(2))
206 .and_then(|v| v.trim().parse().ok())
207 });
208 Some(Hardware {
209 gpu,
210 driver,
211 sm_clock_mhz,
212 gpu_count: None,
213 source: "rocm-smi".into(),
214 })
215}
216
217fn sysfs() -> Option<Hardware> {
221 for entry in std::fs::read_dir("/sys/bus/pci/devices").ok()?.flatten() {
222 let path = entry.path();
223 let Ok(class) = std::fs::read_to_string(path.join("class")) else {
224 continue;
225 };
226 if !class.trim().starts_with("0x03") {
227 continue;
228 }
229 let vendor = std::fs::read_to_string(path.join("vendor")).unwrap_or_default();
230 let device = std::fs::read_to_string(path.join("device")).unwrap_or_default();
231 return Some(Hardware {
232 gpu: format!("pci:{}:{}", vendor.trim(), device.trim()),
233 source: "sysfs".into(),
234 ..Hardware::default()
235 });
236 }
237 None
238}
239
240#[cfg(test)]
241mod tests {
242 use super::*;
243
244 #[test]
245 fn round_trip_keeps_every_field() {
246 let hw = Hardware {
247 gpu: "NVIDIA GB10".into(),
248 driver: "580.126.09".into(),
249 sm_clock_mhz: Some(208.0),
250 gpu_count: Some(1),
251 source: "nvidia-smi".into(),
252 };
253 let back: Hardware = serde_json::from_str(&serde_json::to_string(&hw).unwrap()).unwrap();
254 assert_eq!(hw, back);
255 }
256
257 #[test]
258 fn missing_fields_default_to_unknown() {
259 let hw: Hardware = serde_json::from_str("{}").unwrap();
260 assert_eq!(hw, Hardware::unknown());
261 assert_eq!(hw.one_line(), "unknown hardware");
262 }
263
264 #[test]
265 fn one_line_lists_each_reported_measurement() {
266 let hw = Hardware {
267 gpu: "NVIDIA GB10".into(),
268 driver: "580.126.09".into(),
269 sm_clock_mhz: Some(208.0),
270 gpu_count: None,
271 source: "nvidia-smi".into(),
272 };
273 assert_eq!(
274 hw.one_line(),
275 "NVIDIA GB10 · driver 580.126.09 · sm 208 MHz"
276 );
277 assert_eq!(
278 Hardware {
279 gpu: hw.gpu.clone(),
280 ..Hardware::default()
281 }
282 .one_line(),
283 "NVIDIA GB10"
284 );
285 assert_eq!(
286 Hardware {
287 driver: hw.driver.clone(),
288 ..Hardware::default()
289 }
290 .one_line(),
291 "driver 580.126.09"
292 );
293 assert_eq!(
294 Hardware {
295 sm_clock_mhz: hw.sm_clock_mhz,
296 ..Hardware::default()
297 }
298 .one_line(),
299 "sm 208 MHz"
300 );
301 }
302 #[test]
305 fn a_multi_gpu_box_names_its_width_and_a_single_one_stays_quiet() {
306 let node = |n| Hardware {
307 gpu: "NVIDIA H100 80GB HBM3".into(),
308 driver: "580.126.09".into(),
309 gpu_count: n,
310 ..Hardware::default()
311 };
312 assert_eq!(
313 node(Some(8)).one_line(),
314 "NVIDIA H100 80GB HBM3 ×8 · driver 580.126.09"
315 );
316 assert_eq!(
317 node(Some(1)).one_line(),
318 "NVIDIA H100 80GB HBM3 · driver 580.126.09"
319 );
320 assert_eq!(
321 node(None).one_line(),
322 "NVIDIA H100 80GB HBM3 · driver 580.126.09"
323 );
324 }
325
326 #[test]
329 fn the_gpu_count_round_trips_and_stays_out_of_the_json_when_unmeasured() {
330 let eight = Hardware {
331 gpu: "NVIDIA H200".into(),
332 gpu_count: Some(8),
333 ..Hardware::default()
334 };
335 let json = serde_json::to_string(&eight).unwrap();
336 assert_eq!(
337 serde_json::from_str::<Hardware>(&json).unwrap().gpu_count,
338 Some(8)
339 );
340 let unmeasured = Hardware {
341 gpu: "NVIDIA GB10".into(),
342 ..Hardware::default()
343 };
344 let json = serde_json::to_string(&unmeasured).unwrap();
345 assert!(!json.contains("gpu_count"), "{json}");
346 assert_eq!(
347 serde_json::from_str::<Hardware>(&json).unwrap().gpu_count,
348 None
349 );
350 }
351}
352
353#[cfg(test)]
354#[path = "hardware_gate_key_tests.rs"]
355mod gate_key_tests;