open-standards/eoc/eoc-rs/crates/eoc-meter/tests/nvml_hardware.rs

//! NVML against the hardware (feature `cuda`): GPU 0's total-energy counter,
//! read by the meter and by an oracle the tests declare here.
//!
//! In a test binary of its own, so the NVML probe behind `nvml_status()`
//! runs in this process, and the tests take one lock. The oracle opens
//! `libnvidia-ml` itself, measures how often the counter is republished,
//! samples it at the probe's own wait (100 ms) by the sustained-advance rule
//! and holds it to the board's power ceiling, which it asks of `nvidia-smi`.
//! None of these tests passes without asserting: where there is no NVML (a
//! CI container has no GPU), the status must say so and agree with the
//! oracle; where the counter is live, both must say so. Either way the GPU
//! counter is never the counter `detect()` hands eoc's consumers, which
//! grade any advance of it `Measured` with no coverage field to say it read
//! only the GPU: by the family's rule a GPU-only figure in such a format is
//! an estimate, never a whole-work hardware reading.

#![cfg(all(target_os = "linux", feature = "cuda"))]

use std::sync::{Mutex, MutexGuard};
use std::time::Duration;

use eoc_meter::{CounterStatus, NvmlCounter};

/// One test at a time: each reads machine-wide counters, and the oracle's
/// own sampling must not overlap the meter's.
static SERIAL: Mutex<()> = Mutex::new(());

fn serial() -> MutexGuard<'static, ()> {
    SERIAL.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
}

/// The NVML probe's wait between samples (`NVML_PROBE_WAIT`).
const TICK: Duration = Duration::from_millis(100);

/// The tests' own reading of NVIDIA's total-energy counter: `libnvidia-ml`
/// opened here with `dlopen` and called through symbols looked up here, not
/// through the crate under test or its bindings, and each GPU's power limit
/// and draw asked of `nvidia-smi`, a separate process. It measures how often
/// the counter is republished, and judges it by the family's rule: every
/// GPU's counter advanced in at least five intervals, at least 80% of them,
/// over at least 300 ms, sampled at the crate's own tick; and it is plausibly
/// counting its board's energy: over the window it ran no faster than 1.5
/// times the board's maximum power limit, the board's own power draw (as
/// `nvidia-smi` reports it just before and just after) was itself within that
/// bound, and the counter ran no faster than twice that draw plus 5 W (the
/// draw's documented accuracy). A counter that runs several times faster than
/// its board draws is not counting that board's energy.
#[allow(dead_code)]
mod nvml_oracle {
    use std::ffi::{c_char, c_int, c_uint, c_void, CString};
    use std::time::{Duration, Instant};

    unsafe extern "C" {
        fn dlopen(filename: *const c_char, flag: c_int) -> *mut c_void;
        fn dlsym(handle: *mut c_void, symbol: *const c_char) -> *mut c_void;
    }

    /// `RTLD_NOW` in `<dlfcn.h>`.
    const RTLD_NOW: c_int = 2;

    type Device = *mut c_void;
    type EnergyFn = unsafe extern "C" fn(Device, *mut u64) -> c_int;
    type LimitsFn = unsafe extern "C" fn(Device, *mut c_uint, *mut c_uint) -> c_int;
    type MilliwattsFn = unsafe extern "C" fn(Device, *mut c_uint) -> c_int;

    /// NVML, opened by the oracle.
    pub struct Nvml {
        energy: EnergyFn,
        limits: Option<LimitsFn>,
        power: Option<MilliwattsFn>,
        /// Every GPU whose total-energy counter reads, in index order.
        pub devices: Vec<Device>,
    }

    fn sym(lib: *mut c_void, name: &str) -> Option<*mut c_void> {
        let name = CString::new(name).ok()?;
        let p = unsafe { dlsym(lib, name.as_ptr()) };
        (!p.is_null()).then_some(p)
    }

    /// `libnvidia-ml.so.1`, initialised, with every GPU whose counter reads;
    /// `None` where there is no library, no GPU, or no GPU with the counter.
    pub fn open() -> Option<Nvml> {
        let file = CString::new("libnvidia-ml.so.1").ok()?;
        let lib = unsafe { dlopen(file.as_ptr(), RTLD_NOW) };
        if lib.is_null() {
            return None;
        }
        unsafe {
            let init: unsafe extern "C" fn() -> c_int = std::mem::transmute(sym(lib, "nvmlInit_v2")?);
            let count: unsafe extern "C" fn(*mut c_uint) -> c_int = std::mem::transmute(sym(lib, "nvmlDeviceGetCount_v2")?);
            let by_index: unsafe extern "C" fn(c_uint, *mut Device) -> c_int =
                std::mem::transmute(sym(lib, "nvmlDeviceGetHandleByIndex_v2")?);
            let energy: EnergyFn = std::mem::transmute(sym(lib, "nvmlDeviceGetTotalEnergyConsumption")?);
            let limits = sym(lib, "nvmlDeviceGetPowerManagementLimitConstraints").map(|p| std::mem::transmute::<_, LimitsFn>(p));
            let power = sym(lib, "nvmlDeviceGetPowerUsage").map(|p| std::mem::transmute::<_, MilliwattsFn>(p));
            if init() != 0 {
                return None;
            }
            let mut n: c_uint = 0;
            if count(&mut n) != 0 {
                return None;
            }
            let mut devices = Vec::new();
            for i in 0..n {
                let mut d: Device = std::ptr::null_mut();
                let mut mj = 0u64;
                if by_index(i, &mut d) == 0 && !d.is_null() && energy(d, &mut mj) == 0 {
                    devices.push(d);
                }
            }
            (!devices.is_empty()).then_some(Nvml { energy, limits, power, devices })
        }
    }

    impl Nvml {
        /// Every GPU's total-energy counter, in mJ; `None` if one fails.
        pub fn energy_mj(&self) -> Option<Vec<u64>> {
            self.devices
                .iter()
                .map(|&d| {
                    let mut mj = 0u64;
                    (unsafe { (self.energy)(d, &mut mj) } == 0).then_some(mj)
                })
                .collect()
        }

        /// GPU `i`'s maximum power-management limit, in mW, as NVML gives it.
        pub fn max_limit_mw(&self, i: usize) -> Option<u32> {
            let (mut lo, mut hi) = (0, 0);
            (unsafe { (self.limits?)(self.devices[i], &mut lo, &mut hi) } == 0 && hi > 0).then_some(hi)
        }

        /// GPU `i`'s power draw, in mW, as NVML gives it.
        pub fn power_mw(&self, i: usize) -> Option<u32> {
            let mut mw = 0;
            (unsafe { (self.power?)(self.devices[i], &mut mw) } == 0).then_some(mw)
        }
    }

    /// `nvidia-smi`'s view of every GPU: (maximum power limit, enforced limit,
    /// draw), in watts; `None` without `nvidia-smi`.
    pub fn smi() -> Option<Vec<(Option<f64>, Option<f64>, Option<f64>)>> {
        let out = std::process::Command::new("nvidia-smi")
            .args(["--query-gpu=power.max_limit,enforced.power.limit,power.draw", "--format=csv,noheader,nounits"])
            .output()
            .ok()?;
        out.status.success().then(|| {
            String::from_utf8_lossy(&out.stdout)
                .lines()
                .map(|l| {
                    let f: Vec<Option<f64>> = l.split(',').map(|v| v.trim().parse().ok()).collect();
                    (f.first().copied().flatten(), f.get(1).copied().flatten(), f.get(2).copied().flatten())
                })
                .collect()
        })
    }

    /// GPU `i`'s ceiling in watts: `nvidia-smi`'s maximum power limit, else
    /// NVML's.
    pub fn ceiling_w(nvml: &Nvml, smi: &Option<Vec<(Option<f64>, Option<f64>, Option<f64>)>>, i: usize) -> Option<f64> {
        smi.as_ref()
            .and_then(|s| s.get(i))
            .and_then(|(max, _, _)| *max)
            .or_else(|| nvml.max_limit_mw(i).map(|mw| f64::from(mw) / 1000.0))
    }

    /// How often GPU 0's counter is republished: the median gap between its
    /// changes, polled every 2 ms for up to 1.5 s (`None` if it changed
    /// fewer than 6 times).
    pub fn update_period(nvml: &Nvml) -> Option<Duration> {
        let started = Instant::now();
        let mut last = nvml.energy_mj()?[0];
        let mut changes = Vec::new();
        while started.elapsed() < Duration::from_millis(1500) && changes.len() < 16 {
            std::thread::sleep(Duration::from_millis(2));
            let now = nvml.energy_mj()?[0];
            if now != last {
                changes.push(started.elapsed());
                last = now;
            }
        }
        let mut gaps: Vec<Duration> = changes.windows(2).map(|w| w[1] - w[0]).collect();
        if gaps.len() < 5 {
            return None;
        }
        gaps.sort();
        Some(gaps[gaps.len() / 2])
    }

    /// The oracle's verdict on the counters.
    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
    pub enum Verdict {
        /// No NVML, no GPU, or no GPU whose total-energy counter reads.
        Unavailable,
        /// A counter did not keep advancing.
        Frozen,
        /// Every counter kept advancing, but one is not plausibly its board's
        /// energy (see the module doc).
        Implausible,
        /// Every counter kept advancing, plausibly its board's energy.
        Live,
    }

    /// The family's plausibility rule for one GPU over a window: its counter's
    /// average `rate_w`, the board's draws `draws_w` around the window and
    /// its maximum power limit `ceiling_w`. A bound whose input is missing
    /// is not applied.
    pub fn plausible(rate_w: f64, draws_w: &[f64], ceiling_w: Option<f64>) -> bool {
        if let Some(c) = ceiling_w {
            if rate_w > 1.5 * c || draws_w.iter().any(|d| *d > 1.5 * c) {
                return false;
            }
        }
        if !draws_w.is_empty() {
            let mean = draws_w.iter().sum::<f64>() / draws_w.len() as f64;
            if rate_w > 2.0 * mean + 5.0 {
                return false;
            }
        }
        true
    }

    /// What the oracle saw.
    #[derive(Debug)]
    pub struct Watch {
        pub verdict: Verdict,
        pub intervals: u32,
        pub advances: Vec<u32>,
        pub window: Duration,
        /// Each GPU's average rate over the window, in watts.
        pub rate_w: Vec<f64>,
        /// Each GPU's draw just before and just after the window, in watts.
        pub draws_w: Vec<Vec<f64>>,
        /// Each GPU's ceiling, in watts.
        pub ceiling_w: Vec<Option<f64>>,
    }

    impl Watch {
        /// The verdict, each GPU's advances, its counter's rate, the board's
        /// draw and its maximum limit, for the test's log line.
        pub fn summary(&self) -> String {
            format!(
                "{:?} (advanced {:?} of {} intervals over {} ms; counter rate {:?} W, board draw {:?} W, max power limit {:?} W)",
                self.verdict,
                self.advances,
                self.intervals,
                self.window.as_millis(),
                self.rate_w.iter().map(|w| (w * 10.0).round() / 10.0).collect::<Vec<_>>(),
                self.draws_w,
                self.ceiling_w
            )
        }
    }

    /// The oracle's run over the first `gpus` GPUs (all, if more), sleeping
    /// `tick` between samples, by the sustained-advance rule (stopping when
    /// live, or after a second with at least five intervals, or five
    /// seconds); then each counter's plausibility over the whole window.
    pub fn watch(tick: Duration, gpus: usize) -> Watch {
        let none = Watch {
            verdict: Verdict::Unavailable,
            intervals: 0,
            advances: Vec::new(),
            window: Duration::ZERO,
            rate_w: Vec::new(),
            draws_w: Vec::new(),
            ceiling_w: Vec::new(),
        };
        let Some(nvml) = open() else { return none };
        let n = nvml.devices.len().min(gpus);
        let before = smi();
        let read = || nvml.energy_mj().map(|v| v[..n].to_vec());
        let started = Instant::now();
        let Some(first) = read() else { return none };
        let (mut prev, mut last_at) = (first.clone(), Duration::ZERO);
        let mut advances = vec![0u32; n];
        let mut intervals = 0u32;
        let live = |adv: &[u32], k: u32, w: Duration| {
            w >= Duration::from_millis(300) && adv.iter().all(|&a| a >= 5 && u64::from(a) * 5 >= u64::from(k) * 4)
        };
        loop {
            let el = started.elapsed();
            if live(&advances, intervals, last_at) || el >= Duration::from_secs(5) || (el >= Duration::from_secs(1) && intervals >= 5) {
                break;
            }
            std::thread::sleep(tick);
            if let Some(now) = read() {
                last_at = started.elapsed();
                intervals += 1;
                for ((a, x), p) in advances.iter_mut().zip(&now).zip(&prev) {
                    *a += u32::from(x > p);
                }
                prev = now;
            }
        }
        let after = smi();
        let window = last_at;
        let rate_w: Vec<f64> =
            first.iter().zip(&prev).map(|(a, b)| b.saturating_sub(*a) as f64 / 1000.0 / window.as_secs_f64().max(1e-9)).collect();
        let draw = |s: &Option<Vec<(Option<f64>, Option<f64>, Option<f64>)>>, i: usize| s.as_ref().and_then(|s| s.get(i)).and_then(|t| t.2);
        let draws_w: Vec<Vec<f64>> = (0..n).map(|i| [draw(&before, i), draw(&after, i)].into_iter().flatten().collect()).collect();
        let ceiling_w: Vec<Option<f64>> = (0..n).map(|i| ceiling_w(&nvml, &after, i)).collect();
        let verdict = if !live(&advances, intervals, window) {
            Verdict::Frozen
        } else if (0..n).any(|i| !plausible(rate_w[i], &draws_w[i], ceiling_w[i])) {
            Verdict::Implausible
        } else {
            Verdict::Live
        };
        Watch { verdict, intervals, advances, window, rate_w, draws_w, ceiling_w }
    }

    /// One line on the GPUs for the log: count, counters, ceilings, draw, and
    /// the counter's measured republication `period`.
    pub fn listing(period: Option<Duration>) -> String {
        let Some(nvml) = open() else { return "no NVML device with a total-energy counter".to_string() };
        format!(
            "{} GPU(s) with a total-energy counter, counters {:?} mJ; NVML max power limit {:?} mW, power draw {:?} mW; \
             nvidia-smi (max limit, enforced limit, draw) W {:?}; counter republished every {:?}",
            nvml.devices.len(),
            nvml.energy_mj(),
            (0..nvml.devices.len()).map(|i| nvml.max_limit_mw(i)).collect::<Vec<_>>(),
            (0..nvml.devices.len()).map(|i| nvml.power_mw(i)).collect::<Vec<_>>(),
            smi(),
            period
        )
    }

    /// GPU 0's counter's republication period, where there is one.
    pub fn period() -> Option<Duration> {
        open().and_then(|n| update_period(&n))
    }
}

use nvml_oracle::Verdict;

/// The status agrees with the oracle, judged at the probe's own tick, and the
/// tick is no faster than the counter is republished; only a live status
/// constructs the counter, and `detect()` never selects it, live or not.
#[test]
fn the_nvml_status_agrees_with_the_oracle_at_its_own_tick() {
    let _one = serial();
    let period = nvml_oracle::period();
    let status = eoc_meter::nvml_status();
    let watch = nvml_oracle::watch(TICK, 1);
    let constructed = NvmlCounter::try_detect().is_ok();
    let selected = eoc_meter::detect().name();
    eprintln!(
        "nvml hardware: {}; oracle at a {TICK:?} tick {}; status={status:?} constructed={constructed} detect()={selected}",
        nvml_oracle::listing(period),
        watch.summary()
    );
    if let Some(p) = period {
        assert!(TICK >= p * 9 / 10, "the {TICK:?} tick is faster than the counter's {p:?} republication");
    }
    assert_eq!(status == CounterStatus::Live, watch.verdict == Verdict::Live, "the status disagrees with the oracle");
    match watch.verdict {
        Verdict::Unavailable => assert_eq!(status, CounterStatus::Unavailable),
        Verdict::Frozen => assert_eq!(status, CounterStatus::Frozen),
        Verdict::Implausible => assert_eq!(status, CounterStatus::Implausible),
        Verdict::Live => {}
    }
    assert_eq!(constructed, status == CounterStatus::Live, "construction disagrees with the status");
    assert_ne!(selected, "nvml", "detect() handed eoc's consumers the GPU-only counter, which they grade Measured");
}

/// Over a one-second span a live GPU counter advances, inside the oracle's
/// own reads around the span, and no faster than the board can draw.
#[test]
fn a_live_gpu_counter_reads_inside_the_oracles_bracket() {
    let _one = serial();
    let Ok(counter) = NvmlCounter::try_detect() else {
        eprintln!("nvml span: status {:?}; no counter", eoc_meter::nvml_status());
        assert_ne!(eoc_meter::nvml_status(), CounterStatus::Live);
        return;
    };
    let oracle = nvml_oracle::open().expect("the meter reads NVML, yet the oracle cannot open it");
    let o0 = oracle.energy_mj().expect("oracle read")[0];
    let started = std::time::Instant::now();
    let m0 = counter.read_gpu_microjoules().expect("a live counter reads");
    std::thread::sleep(Duration::from_secs(1));
    let m1 = counter.read_gpu_microjoules().expect("a live counter reads");
    let wall = started.elapsed();
    let o1 = oracle.energy_mj().expect("oracle read")[0];
    let ceiling = nvml_oracle::ceiling_w(&oracle, &nvml_oracle::smi(), 0).unwrap_or(f64::INFINITY);
    let uj = m1 - m0;
    eprintln!("nvml span: {wall:?}; meter {uj} µJ; oracle {o0} -> {o1} mJ; ceiling {ceiling} W");
    assert!(uj > 0 && uj <= (o1 - o0) * 1_000, "{uj} µJ outside the oracle's (0, {}] µJ", (o1 - o0) * 1_000);
    let watts = uj as f64 / 1e6 / wall.as_secs_f64();
    assert!(watts <= 1.5 * ceiling, "the GPU counter reads {watts:.1} W, past the board's {ceiling} W ceiling");
}