fix(agent): detect NVIDIA GPUs on distros without the unversioned NVML soname

On Debian and derivatives the NVIDIA driver ships only libnvidia-ml.so.1
(the unversioned symlink belongs to the dev package), and nvml-wrapper's
default init dlopens the unversioned name — so gfxinfo reported 'No GPU
found' on a fully functional RTX A2000 host while nvidia-smi worked
fine. Arch-family distros ship the symlink, which is why the desktop
never showed this.

The GPU worker now falls back to initializing NVML directly with the
versioned soname when gfxinfo's probe fails, collecting name/util/vram
through the same handle-caching path. nvml-wrapper was already in the
tree via gfxinfo — same version, no new build cost.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
jasonwitty
2026-08-21 11:16:41 -07:00
parent 8e0effe361
commit 2a7951cf65
3 changed files with 61 additions and 18 deletions
Generated
+1
View File
@@ -2441,6 +2441,7 @@ dependencies = [
"futures-util", "futures-util",
"gfxinfo", "gfxinfo",
"hostname", "hostname",
"nvml-wrapper",
"once_cell", "once_cell",
"prost", "prost",
"prost-build", "prost-build",
+5 -1
View File
@@ -23,6 +23,10 @@ futures-util = "0.3.31"
tracing = { version = "0.1", optional = true } tracing = { version = "0.1", optional = true }
tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = true } tracing-subscriber = { version = "0.3", features = ["env-filter"], optional = true }
gfxinfo = { version = "0.1.2", optional = true } gfxinfo = { version = "0.1.2", optional = true }
# Direct NVML fallback for distros that ship only libnvidia-ml.so.1 (Debian
# and derivatives) — gfxinfo's default init dlopens the unversioned name.
# Same version gfxinfo already pulls in, so this adds no new build cost.
nvml-wrapper = { version = "0.10", optional = true }
once_cell = "1.19" once_cell = "1.19"
axum-server = { version = "0.7", features = ["tls-rustls"] } axum-server = { version = "0.7", features = ["tls-rustls"] }
rustls = { version = "0.23", features = ["aws-lc-rs"] } rustls = { version = "0.23", features = ["aws-lc-rs"] }
@@ -35,7 +39,7 @@ time = { version = "0.3", default-features = false, features = ["formatting", "m
[features] [features]
default = ["gpu"] default = ["gpu"]
gpu = ["gfxinfo"] gpu = ["gfxinfo", "nvml-wrapper"]
logging = ["tracing", "tracing-subscriber"] logging = ["tracing", "tracing-subscriber"]
[build-dependencies] [build-dependencies]
+55 -17
View File
@@ -51,30 +51,68 @@ mod worker {
tx tx
} }
enum Handle {
/// gfxinfo's own detection (AMD sysfs, NVIDIA via unversioned NVML).
Gfx(Box<dyn gfxinfo::Gpu>),
/// Direct NVML with an explicit versioned soname. Debian & friends
/// ship only libnvidia-ml.so.1 (the unversioned symlink lives in the
/// dev package), so gfxinfo's default dlopen fails there even though
/// the driver is fully functional.
Nvml(nvml_wrapper::Nvml),
}
fn probe() -> Option<Handle> {
if let Ok(g) = gfxinfo::active_gpu() {
return Some(Handle::Gfx(g));
}
nvml_wrapper::Nvml::builder()
.lib_path(std::ffi::OsStr::new("libnvidia-ml.so.1"))
.init()
.ok()
.map(Handle::Nvml)
}
fn collect_from(handle: &Handle) -> Option<Vec<GpuMetrics>> {
match handle {
Handle::Gfx(gpu) => {
let info = gpu.info();
Some(vec![GpuMetrics {
name: gpu.model().to_string(),
utilization_gpu_pct: info.load_pct().clamp(0, 100),
mem_used_bytes: info.used_vram(),
mem_total_bytes: info.total_vram(),
}])
}
Handle::Nvml(nvml) => {
let device = nvml.device_by_index(0).ok()?;
let mem = device.memory_info().ok()?;
Some(vec![GpuMetrics {
name: device.name().unwrap_or_else(|_| "NVIDIA GPU".into()),
utilization_gpu_pct: device
.utilization_rates()
.map(|u| u.gpu.clamp(0, 100))
.unwrap_or(0),
mem_used_bytes: mem.used,
mem_total_bytes: mem.total,
}])
}
}
}
fn run(rx: mpsc::Receiver<Reply>) { fn run(rx: mpsc::Receiver<Reply>) {
let mut handle: Option<Box<dyn gfxinfo::Gpu>> = None; let mut handle: Option<Handle> = None;
// Probing failed: remember and answer None without re-initing the GPU // Probing failed: remember and answer None without re-initing the GPU
// stack per request. The agent's negative cache stops asking anyway. // stack per request. The agent's negative cache stops asking anyway.
let mut probe_failed = false; let mut probe_failed = false;
while let Ok(reply) = rx.recv() { while let Ok(reply) = rx.recv() {
if handle.is_none() && !probe_failed { if handle.is_none() && !probe_failed {
match gfxinfo::active_gpu() { handle = probe();
Ok(g) => handle = Some(g), probe_failed = handle.is_none();
Err(_) => probe_failed = true,
} }
} let out = handle.as_ref().and_then(collect_from);
let out = handle.as_ref().map(|gpu| { // A live GPU cannot report 0 total VRAM; zeros mean the session
let info = gpu.info(); // died (e.g. driver reload). Drop the handle so the next request
vec![GpuMetrics { // re-probes.
name: gpu.model().to_string(),
utilization_gpu_pct: info.load_pct().clamp(0, 100),
mem_used_bytes: info.used_vram(),
mem_total_bytes: info.total_vram(),
}]
});
// A live GPU cannot report 0 total VRAM; gfxinfo returns zeros
// when the underlying session died (e.g. driver reload). Drop the
// handle so the next request re-probes.
if let Some(v) = &out if let Some(v) = &out
&& !v.is_empty() && !v.is_empty()
&& v.iter().all(|g| g.mem_total_bytes == 0) && v.iter().all(|g| g.mem_total_bytes == 0)