metrale_gpu_sys/
nvml.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: NVML (the NVIDIA management library), resolved with `dlopen`.
4//!
5//! There is no link-time dependency: on a host without the driver,
6//! `NvmlLibrary::open` returns `NvmlUnavailable` instead of the process
7//! failing to load. Only `nvmlInit_v2`, `nvmlShutdown` and
8//! `nvmlDeviceGetHandleByIndex_v2` are required; a missing query symbol leaves
9//! its `NvmlFns` slot `None`.
10//!
11//! `NvmlFns` is a table of function pointers. `NvmlLibrary` fills it from the
12//! shared object; tests, here and in the telemetry crate, fill it with
13//! functions of their own through `Nvml::with_fns`. `Nvml` and `NvmlDevice`
14//! run the same code over either.
15//!
16//! Values are in NVML's units, which the accessor names carry: `power_mw`
17//! (milliwatts), `total_energy_mj` (millijoules), `clock_mhz`, `temperature_c`,
18//! `pcie_kbps` (KB/s); `memory` is in bytes.
19//!
20//! Owner: metrale-gpu-sys.
21//! Invariants:
22//! - An `Nvml` exists only after its table's `init` succeeded, and its drop
23//!   calls the table's `shutdown` once.
24
25#![allow(non_camel_case_types)]
26
27use std::ffi::{c_uint, c_ulonglong, c_void};
28
29use libloading::Library;
30
31/// 2026-09-26: `nvmlReturn_t`.
32pub type NvmlReturn = std::ffi::c_int;
33/// 2026-09-26: `nvmlDevice_t`, an opaque handle owned by the library.
34pub type NvmlDeviceHandle = *mut c_void;
35
36pub const NVML_SUCCESS: NvmlReturn = 0;
37pub const NVML_ERROR_NOT_SUPPORTED: NvmlReturn = 3;
38pub const NVML_ERROR_NO_PERMISSION: NvmlReturn = 4;
39
40/// 2026-09-26: `nvmlClockType_t` values.
41pub const NVML_CLOCK_GRAPHICS: c_uint = 0;
42pub const NVML_CLOCK_SM: c_uint = 1;
43pub const NVML_CLOCK_MEM: c_uint = 2;
44/// 2026-09-26: `nvmlTemperatureSensors_t::NVML_TEMPERATURE_GPU`.
45pub const NVML_TEMPERATURE_GPU: c_uint = 0;
46/// 2026-09-26: `nvmlPcieUtilCounter_t` values.
47pub const NVML_PCIE_UTIL_TX_BYTES: c_uint = 0;
48pub const NVML_PCIE_UTIL_RX_BYTES: c_uint = 1;
49/// 2026-09-26: `NVML_NVLINK_MAX_LINKS` in the CUDA toolkit's nvml.h.
50pub const NVML_NVLINK_MAX_LINKS: c_uint = 18;
51
52/// 2026-09-26: `nvmlMemory_t` (v1), in bytes.
53#[repr(C)]
54#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
55pub struct NvmlMemory {
56    pub total: c_ulonglong,
57    pub free: c_ulonglong,
58    pub used: c_ulonglong,
59}
60
61pub type FnVoid = unsafe extern "C" fn() -> NvmlReturn;
62pub type FnHandleByIndex = unsafe extern "C" fn(c_uint, *mut NvmlDeviceHandle) -> NvmlReturn;
63pub type FnDevU32 = unsafe extern "C" fn(NvmlDeviceHandle, *mut c_uint) -> NvmlReturn;
64pub type FnDevU64 = unsafe extern "C" fn(NvmlDeviceHandle, *mut c_ulonglong) -> NvmlReturn;
65pub type FnDevArgU32 = unsafe extern "C" fn(NvmlDeviceHandle, c_uint, *mut c_uint) -> NvmlReturn;
66pub type FnDevMemory = unsafe extern "C" fn(NvmlDeviceHandle, *mut NvmlMemory) -> NvmlReturn;
67
68/// 2026-09-26: The entry points this crate binds. The first three are required;
69/// each query is `None` when the library does not export it.
70#[derive(Clone, Copy, Debug)]
71pub struct NvmlFns {
72    pub init: FnVoid,
73    pub shutdown: FnVoid,
74    pub device_by_index: FnHandleByIndex,
75    /// 2026-09-26: `nvmlDeviceGetPowerUsage`, milliwatts.
76    pub power_usage: Option<FnDevU32>,
77    /// 2026-09-26: `nvmlDeviceGetTotalEnergyConsumption`, millijoules since the
78    /// driver was last reloaded.
79    pub total_energy: Option<FnDevU64>,
80    /// 2026-09-26: `nvmlDeviceGetClockInfo(type)`, MHz.
81    pub clock_info: Option<FnDevArgU32>,
82    /// 2026-09-26: `nvmlDeviceGetTemperature(sensor)`, degrees C.
83    pub temperature: Option<FnDevArgU32>,
84    /// 2026-09-26: `nvmlDeviceGetCurrentClocksEventReasons`, or
85    /// `nvmlDeviceGetCurrentClocksThrottleReasons` when the library lacks it.
86    pub clocks_event_reasons: Option<FnDevU64>,
87    /// 2026-09-26: `nvmlDeviceGetMemoryInfo`.
88    pub memory_info: Option<FnDevMemory>,
89    /// 2026-09-26: `nvmlDeviceGetPcieThroughput(counter)`, KB/s.
90    pub pcie_throughput: Option<FnDevArgU32>,
91    /// 2026-09-26: `nvmlDeviceGetNvLinkState(link)`, 1 = active.
92    pub nvlink_state: Option<FnDevArgU32>,
93}
94
95/// 2026-09-26: Why NVML cannot be used. The telemetry sampler records it as
96/// the device being unavailable and serving continues.
97#[derive(Clone, Debug, PartialEq, Eq)]
98pub enum NvmlUnavailable {
99    /// 2026-09-26: No candidate shared object could be opened.
100    NotInstalled(String),
101    /// 2026-09-26: The object opened but lacks a required entry point.
102    MissingSymbol(&'static str),
103    /// 2026-09-26: `nvmlInit_v2` returned this code.
104    InitFailed(NvmlReturn),
105    /// 2026-09-26: `nvmlDeviceGetHandleByIndex_v2(index)` returned this code.
106    NoDevice { index: u32, code: NvmlReturn },
107}
108
109impl std::fmt::Display for NvmlUnavailable {
110    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
111        match self {
112            Self::NotInstalled(why) => write!(f, "NVML library not loadable: {why}"),
113            Self::MissingSymbol(s) => write!(f, "NVML library lacks required symbol {s}"),
114            Self::InitFailed(c) => write!(f, "nvmlInit_v2 failed with code {c}"),
115            Self::NoDevice { index, code } => {
116                write!(f, "NVML has no device {index} (code {code})")
117            }
118        }
119    }
120}
121
122impl std::error::Error for NvmlUnavailable {}
123
124/// 2026-09-26: The shared object names tried, in order.
125pub const LIBRARY_NAMES: [&str; 2] = ["libnvidia-ml.so.1", "libnvidia-ml.so"];
126
127/// 2026-09-26: An opened `libnvidia-ml` and the table resolved from it. The
128/// library stays loaded as long as this value lives.
129pub struct NvmlLibrary {
130    fns: NvmlFns,
131    _lib: Library,
132}
133
134impl NvmlLibrary {
135    /// 2026-09-26: `open_from(&LIBRARY_NAMES)`.
136    pub fn open() -> Result<Self, NvmlUnavailable> {
137        Self::open_from(&LIBRARY_NAMES)
138    }
139
140    /// 2026-09-26: Open the first of `candidates` that loads. When that one
141    /// lacks a required symbol, the error is returned and later candidates are
142    /// not tried.
143    pub fn open_from(candidates: &[&str]) -> Result<Self, NvmlUnavailable> {
144        let mut last = String::from("no candidate names");
145        for name in candidates {
146            // 2026-09-26: SAFETY: loading runs the object's initialisers; NVML
147            // itself is initialised only by the `nvmlInit_v2` call in
148            // `Nvml::init`.
149            match unsafe { Library::new(name) } {
150                Ok(lib) => {
151                    let fns = resolve(&lib)?;
152                    return Ok(Self { fns, _lib: lib });
153                }
154                Err(e) => last = format!("{name}: {e}"),
155            }
156        }
157        Err(NvmlUnavailable::NotInstalled(last))
158    }
159
160    pub fn fns(&self) -> NvmlFns {
161        self.fns
162    }
163}
164
165fn required<T: Copy>(lib: &Library, name: &'static str) -> Result<T, NvmlUnavailable> {
166    optional(lib, name).ok_or(NvmlUnavailable::MissingSymbol(name))
167}
168
169fn optional<T: Copy>(lib: &Library, name: &str) -> Option<T> {
170    let mut sym = Vec::with_capacity(name.len() + 1);
171    sym.extend_from_slice(name.as_bytes());
172    sym.push(0);
173    // 2026-09-26: SAFETY: `resolve` calls this with each `name`'s signature
174    // from nvml.h.
175    unsafe { lib.get::<T>(&sym).ok().map(|s| *s) }
176}
177
178fn resolve(lib: &Library) -> Result<NvmlFns, NvmlUnavailable> {
179    Ok(NvmlFns {
180        init: required(lib, "nvmlInit_v2")?,
181        shutdown: required(lib, "nvmlShutdown")?,
182        device_by_index: required(lib, "nvmlDeviceGetHandleByIndex_v2")?,
183        power_usage: optional(lib, "nvmlDeviceGetPowerUsage"),
184        total_energy: optional(lib, "nvmlDeviceGetTotalEnergyConsumption"),
185        clock_info: optional(lib, "nvmlDeviceGetClockInfo"),
186        temperature: optional(lib, "nvmlDeviceGetTemperature"),
187        clocks_event_reasons: optional(lib, "nvmlDeviceGetCurrentClocksEventReasons")
188            .or_else(|| optional(lib, "nvmlDeviceGetCurrentClocksThrottleReasons")),
189        memory_info: optional(lib, "nvmlDeviceGetMemoryInfo"),
190        pcie_throughput: optional(lib, "nvmlDeviceGetPcieThroughput"),
191        nvlink_state: optional(lib, "nvmlDeviceGetNvLinkState"),
192    })
193}
194
195/// 2026-09-26: An initialised NVML session over a function table.
196/// `nvmlShutdown` runs on drop, balancing the successful init.
197pub struct Nvml {
198    fns: NvmlFns,
199    /// 2026-09-26: The opened library that keeps `fns` valid; `None` for a
200    /// table of functions linked into the binary.
201    _keep: Option<NvmlLibrary>,
202}
203
204impl Nvml {
205    /// 2026-09-26: Open and initialise the system library.
206    pub fn open() -> Result<Self, NvmlUnavailable> {
207        let lib = NvmlLibrary::open()?;
208        let fns = lib.fns();
209        Self::init(fns, Some(lib))
210    }
211
212    /// 2026-09-26: Initialise over a table whose functions stay valid for the
213    /// life of the process. Tests use it.
214    pub fn with_fns(fns: NvmlFns) -> Result<Self, NvmlUnavailable> {
215        Self::init(fns, None)
216    }
217
218    fn init(fns: NvmlFns, keep: Option<NvmlLibrary>) -> Result<Self, NvmlUnavailable> {
219        // 2026-09-26: SAFETY: `init` is `nvmlInit_v2` or a test function of
220        // that type.
221        let rc = unsafe { (fns.init)() };
222        if rc != NVML_SUCCESS {
223            return Err(NvmlUnavailable::InitFailed(rc));
224        }
225        Ok(Self { fns, _keep: keep })
226    }
227
228    /// 2026-09-26: The device at NVML index `index`.
229    pub fn device(&self, index: u32) -> Result<NvmlDevice<'_>, NvmlUnavailable> {
230        let mut handle: NvmlDeviceHandle = std::ptr::null_mut();
231        // 2026-09-26: SAFETY: `handle` is a valid out-pointer.
232        let rc = unsafe { (self.fns.device_by_index)(index, &mut handle) };
233        if rc != NVML_SUCCESS {
234            return Err(NvmlUnavailable::NoDevice { index, code: rc });
235        }
236        Ok(NvmlDevice {
237            fns: &self.fns,
238            handle,
239        })
240    }
241}
242
243impl Drop for Nvml {
244    fn drop(&mut self) {
245        // 2026-09-26: SAFETY: balances the successful init in `Nvml::init`. A
246        // failed shutdown is ignored: this side holds nothing to release.
247        let _ = unsafe { (self.fns.shutdown)() };
248    }
249}
250
251/// 2026-09-26: One device's queries. `Ok(None)` when the entry point is absent
252/// or answers `NVML_ERROR_NOT_SUPPORTED` or `NVML_ERROR_NO_PERMISSION`;
253/// `Err(code)` for any other failure.
254pub struct NvmlDevice<'a> {
255    fns: &'a NvmlFns,
256    handle: NvmlDeviceHandle,
257}
258
259// 2026-09-26: SAFETY: nvml.h states that NVML is thread-safe; the table is
260// plain function pointers.
261unsafe impl Send for NvmlDevice<'_> {}
262
263fn answer<T>(rc: NvmlReturn, v: T) -> Result<Option<T>, NvmlReturn> {
264    match rc {
265        NVML_SUCCESS => Ok(Some(v)),
266        NVML_ERROR_NOT_SUPPORTED | NVML_ERROR_NO_PERMISSION => Ok(None),
267        other => Err(other),
268    }
269}
270
271impl NvmlDevice<'_> {
272    fn u32_query(&self, f: Option<FnDevU32>) -> Result<Option<u32>, NvmlReturn> {
273        let Some(f) = f else { return Ok(None) };
274        let mut v: c_uint = 0;
275        // 2026-09-26: SAFETY: `f` has the `(device, unsigned*)` signature; `v`
276        // is a valid out-pointer.
277        answer(unsafe { f(self.handle, &mut v) }, v)
278    }
279
280    fn u64_query(&self, f: Option<FnDevU64>) -> Result<Option<u64>, NvmlReturn> {
281        let Some(f) = f else { return Ok(None) };
282        let mut v: c_ulonglong = 0;
283        // 2026-09-26: SAFETY: `f` has the `(device, unsigned long long*)`
284        // signature.
285        answer(unsafe { f(self.handle, &mut v) }, v)
286    }
287
288    fn arg_query(&self, f: Option<FnDevArgU32>, arg: c_uint) -> Result<Option<u32>, NvmlReturn> {
289        let Some(f) = f else { return Ok(None) };
290        let mut v: c_uint = 0;
291        // 2026-09-26: SAFETY: `f` has the `(device, unsigned, unsigned*)`
292        // signature.
293        answer(unsafe { f(self.handle, arg, &mut v) }, v)
294    }
295
296    pub fn power_mw(&self) -> Result<Option<u32>, NvmlReturn> {
297        self.u32_query(self.fns.power_usage)
298    }
299
300    pub fn total_energy_mj(&self) -> Result<Option<u64>, NvmlReturn> {
301        self.u64_query(self.fns.total_energy)
302    }
303
304    pub fn clock_mhz(&self, clock: c_uint) -> Result<Option<u32>, NvmlReturn> {
305        self.arg_query(self.fns.clock_info, clock)
306    }
307
308    pub fn temperature_c(&self) -> Result<Option<u32>, NvmlReturn> {
309        self.arg_query(self.fns.temperature, NVML_TEMPERATURE_GPU)
310    }
311
312    pub fn clocks_event_reasons(&self) -> Result<Option<u64>, NvmlReturn> {
313        self.u64_query(self.fns.clocks_event_reasons)
314    }
315
316    pub fn memory(&self) -> Result<Option<NvmlMemory>, NvmlReturn> {
317        let Some(f) = self.fns.memory_info else {
318            return Ok(None);
319        };
320        let mut m = NvmlMemory::default();
321        // 2026-09-26: SAFETY: `m` is a valid `nvmlMemory_t` out-pointer.
322        answer(unsafe { f(self.handle, &mut m) }, m)
323    }
324
325    pub fn pcie_kbps(&self, counter: c_uint) -> Result<Option<u32>, NvmlReturn> {
326        self.arg_query(self.fns.pcie_throughput, counter)
327    }
328
329    /// 2026-09-26: The number of links in `0..NVML_NVLINK_MAX_LINKS` whose
330    /// state query succeeds with state 1; `None` when no link query succeeds.
331    /// A failed per-link query counts as no such link, not as an error.
332    pub fn nvlink_active_links(&self) -> Option<u32> {
333        let f = self.fns.nvlink_state?;
334        let mut active = 0u32;
335        let mut any = false;
336        for link in 0..NVML_NVLINK_MAX_LINKS {
337            let mut state: c_uint = 0;
338            // 2026-09-26: SAFETY: `f` has the `(device, unsigned, unsigned*)`
339            // signature.
340            if unsafe { f(self.handle, link, &mut state) } == NVML_SUCCESS {
341                any = true;
342                active += u32::from(state == 1);
343            }
344        }
345        any.then_some(active)
346    }
347}
348
349#[cfg(test)]
350#[path = "nvml_tests.rs"]
351mod tests;
metrale_gpu_sys/nvml.rs

metrale_gpu_sys/
nvml.rs
