energy-bench 0.2.5

use std::sync::LazyLock;

use indexmap::IndexMap;

static NVML: LazyLock<Option<nvml_wrapper::Nvml>> =
    LazyLock::new(|| nvml_wrapper::Nvml::init().ok());

pub struct Nvml<'a> {
    devices: Vec<NvmlDevice<'a>>,
}

pub struct NvmlDevice<'a> {
    device: nvml_wrapper::Device<'a>,
    name: String,
    energy: u64,
}

impl<'a> Nvml<'a> {
    pub fn now() -> Option<Self> {
        let nvml = NVML.as_ref()?;
        let count = nvml.device_count().ok()?;
        let devices = (0..count).filter_map(NvmlDevice::new).collect();
        Some(Self { devices })
    }

    pub fn elapsed(&self) -> IndexMap<String, f32> {
        self.devices
            .iter()
            .map(|device| {
                let name = device.name.clone();
                let energy = device.elapsed();
                (name, energy)
            })
            .collect()
    }

    pub fn reset(&mut self) {
        self.devices.iter_mut().for_each(NvmlDevice::reset);
    }
}

impl<'a> NvmlDevice<'a> {
    fn new(index: u32) -> Option<Self> {
        let nvml = NVML.as_ref()?;
        let device = nvml.device_by_index(index).ok()?;
        let name = format!("GPU({}) {}", index, device.name().ok()?);
        let energy = device.total_energy_consumption().ok()?;
        Some(Self {
            device,
            name,
            energy,
        })
    }

    fn elapsed(&self) -> f32 {
        let prev = self.energy;
        let next = self.device.total_energy_consumption().unwrap();
        (next - prev) as f32 / 1000.0
    }

    fn reset(&mut self) {
        self.energy = self.device.total_energy_consumption().unwrap();
    }
}