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:
Generated
+1
@@ -2441,6 +2441,7 @@ dependencies = [
|
||||
"futures-util",
|
||||
"gfxinfo",
|
||||
"hostname",
|
||||
"nvml-wrapper",
|
||||
"once_cell",
|
||||
"prost",
|
||||
"prost-build",
|
||||
|
||||
@@ -23,6 +23,10 @@ futures-util = "0.3.31"
|
||||
tracing = { version = "0.1", optional = true }
|
||||
tracing-subscriber = { version = "0.3", features = ["env-filter"], 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"
|
||||
axum-server = { version = "0.7", features = ["tls-rustls"] }
|
||||
rustls = { version = "0.23", features = ["aws-lc-rs"] }
|
||||
@@ -35,7 +39,7 @@ time = { version = "0.3", default-features = false, features = ["formatting", "m
|
||||
|
||||
[features]
|
||||
default = ["gpu"]
|
||||
gpu = ["gfxinfo"]
|
||||
gpu = ["gfxinfo", "nvml-wrapper"]
|
||||
logging = ["tracing", "tracing-subscriber"]
|
||||
|
||||
[build-dependencies]
|
||||
|
||||
+55
-17
@@ -51,30 +51,68 @@ mod worker {
|
||||
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>) {
|
||||
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
|
||||
// stack per request. The agent's negative cache stops asking anyway.
|
||||
let mut probe_failed = false;
|
||||
while let Ok(reply) = rx.recv() {
|
||||
if handle.is_none() && !probe_failed {
|
||||
match gfxinfo::active_gpu() {
|
||||
Ok(g) => handle = Some(g),
|
||||
Err(_) => probe_failed = true,
|
||||
}
|
||||
handle = probe();
|
||||
probe_failed = handle.is_none();
|
||||
}
|
||||
let out = handle.as_ref().map(|gpu| {
|
||||
let info = gpu.info();
|
||||
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(),
|
||||
}]
|
||||
});
|
||||
// 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.
|
||||
let out = handle.as_ref().and_then(collect_from);
|
||||
// A live GPU cannot report 0 total VRAM; zeros mean the session
|
||||
// died (e.g. driver reload). Drop the handle so the next request
|
||||
// re-probes.
|
||||
if let Some(v) = &out
|
||||
&& !v.is_empty()
|
||||
&& v.iter().all(|g| g.mem_total_bytes == 0)
|
||||
|
||||
Reference in New Issue
Block a user