use alloc::boxed::Box;
use alloc::string::{String, ToString};
use ruda_core::ir::PciAddress;
use nvml_wrapper::Nvml;
use nvml_wrapper::error::NvmlError;
use crate::runtime::monitoring::{DeviceUtilization, UtilizationUnavailable};
#[derive(Debug)]
pub struct NvmlUtilizationCounter {
nvml: Box<Nvml>,
pci_bus_id: String,
}
impl NvmlUtilizationCounter {
pub fn new(pci_address: PciAddress) -> Result<Self, UtilizationUnavailable> {
let nvml = Box::new(Nvml::init().map_err(unavailability_from_nvml_error)?);
let pci_bus_id = pci_address.to_string();
nvml.device_by_pci_bus_id(pci_bus_id.as_str())
.map_err(unavailability_from_nvml_error)?;
Ok(Self { nvml, pci_bus_id })
}
pub fn read(&self) -> Result<DeviceUtilization, UtilizationUnavailable> {
self.nvml
.device_by_pci_bus_id(self.pci_bus_id.as_str())
.and_then(|device| device.utilization_rates())
.map(|utilization| DeviceUtilization::new(utilization.gpu as f32))
.map_err(unavailability_from_nvml_error)
}
}
fn unavailability_from_nvml_error(error: NvmlError) -> UtilizationUnavailable {
match error {
NvmlError::LibloadingError(_) | NvmlError::LibraryNotFound => {
UtilizationUnavailable::DriverLibraryNotFound(error.to_string())
}
_ => UtilizationUnavailable::QueryFailed(error.to_string()),
}
}