use clawft_types::placement::Provenance;
use super::host::ProbeHost;
use super::probe::{Collected, bytes_attr, cap, str_list};
const CORAL_USB_IDS: &[&str] = &["1a6e:089a", "18d1:9302"];
pub fn probe(host: &dyn ProbeHost, c: &mut Collected) {
cuda(host, c);
hailo(host, c);
coral(host, c);
}
fn cuda(host: &dyn ProbeHost, c: &mut Collected) {
if host.which("nvidia-smi").is_none() {
return;
}
let Some(out) = host.run(
"nvidia-smi",
&[
"--query-gpu=name,memory.total,memory.free",
"--format=csv,noheader,nounits",
],
) else {
c.note(
"accel.gpu.cuda",
"nvidia-smi present but failed; no CUDA GPU advertised",
);
return;
};
for (i, line) in out.lines().filter(|l| !l.trim().is_empty()).enumerate() {
let f: Vec<&str> = line.split(',').map(str::trim).collect();
let mib = |s: Option<&&str>| {
s.and_then(|v| v.parse::<u64>().ok())
.map(|m| m * 1024 * 1024)
};
let (total, free) = (mib(f.get(1)), mib(f.get(2)));
let device = format!("cuda:{i}");
if let Some(mut g) = cap("accel.gpu.cuda", Provenance::Probed) {
g = g
.with_attr("vendor", "nvidia")
.with_attr("device", device.as_str())
.with_attr("unified", false);
if let Some(name) = f.first() {
g = g.with_attr("model", *name);
}
if let Some(t) = total {
g = g.with_attr("mem_bytes", bytes_attr(t));
}
c.push(g);
}
if let (Some(t), Some(v)) = (total, cap("mem.vram", Provenance::Probed)) {
let mut v = v
.with_attr("total", bytes_attr(t))
.with_attr("device", device.as_str());
if let Some(fr) = free {
v = v.with_attr("free", bytes_attr(fr));
}
c.push(v);
}
}
}
fn hailo(host: &dyn ProbeHost, c: &mut Collected) {
if host.which("hailortcli").is_none() {
return;
}
let Some(out) = host.run("hailortcli", &["scan"]) else {
c.note("accel.npu.hailo", "hailortcli present but scan failed");
return;
};
for dev in out
.lines()
.filter_map(|l| l.split("Device:").nth(1))
.map(str::trim)
{
if let Some(h) = cap("accel.npu.hailo", Provenance::Probed) {
c.push(
h.with_attr("device", dev)
.with_attr("formats", str_list(&["hef"]))
.exclusive(),
);
}
}
}
fn coral(host: &dyn ProbeHost, c: &mut Collected) {
let mut devices: Vec<String> = host
.list_dir("/dev")
.unwrap_or_default()
.into_iter()
.filter(|n| n.starts_with("apex_"))
.map(|n| format!("pcie:{n}"))
.collect();
if host.which("lsusb").is_some() {
let usb = host.run("lsusb", &[]).unwrap_or_default();
let n = usb
.lines()
.filter(|l| CORAL_USB_IDS.iter().any(|id| l.contains(id)))
.count();
devices.extend((0..n).map(|i| format!("usb:{i}")));
}
for dev in devices {
if let Some(t) = cap("accel.tpu.coral", Provenance::Probed) {
c.push(
t.with_attr("device", dev.as_str())
.with_attr("formats", str_list(&["tflite"]))
.exclusive(),
);
c.note(
"accel.tpu.coral",
format!("{dev}: Edge TPU device seen; single client, so exclusive"),
);
}
}
}