diff --git a/Cargo.lock b/Cargo.lock index a96b308..7bfa0b8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2441,6 +2441,7 @@ dependencies = [ "futures-util", "gfxinfo", "hostname", + "nvml-wrapper", "once_cell", "prost", "prost-build", diff --git a/socktop_agent/Cargo.toml b/socktop_agent/Cargo.toml index bec254c..ff700c5 100644 --- a/socktop_agent/Cargo.toml +++ b/socktop_agent/Cargo.toml @@ -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] diff --git a/socktop_agent/src/gpu.rs b/socktop_agent/src/gpu.rs index 7e0278a..d32596c 100644 --- a/socktop_agent/src/gpu.rs +++ b/socktop_agent/src/gpu.rs @@ -51,30 +51,68 @@ mod worker { tx } + enum Handle { + /// gfxinfo's own detection (AMD sysfs, NVIDIA via unversioned NVML). + Gfx(Box), + /// 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 { + 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> { + 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) { - let mut handle: Option> = None; + let mut handle: Option = 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)