multi gpu

src/linux/gpu/nvidia.rs

@@ -1,4 +1,4 @@
-use crate::data::{GpuMemory, GpuUsage};
+use crate::data::{GpuMemory, GpuPowerUsage, GpuUsage};
 use nvml_wrapper::enum_wrappers::device::TemperatureSensor;
 use nvml_wrapper::{Device, Nvml};
 use once_cell::sync::Lazy;
@@ -6,32 +6,50 @@ use std::borrow::Cow;
 
 static NVIDIA: Lazy<Option<Nvml>> = Lazy::new(|| Nvml::init().ok());
 
-fn device() -> Option<Device<'static>> {
-    NVIDIA.as_ref()?.device_by_index(0).ok()
+fn devices() -> Option<impl Iterator<Item = Device<'static>>> {
+    let count = NVIDIA.as_ref()?.device_count().unwrap_or_default();
+    Some((0..count).flat_map(device))
 }
 
-pub fn temperature() -> Option<f32> {
-    let temp = device()?.temperature(TemperatureSensor::Gpu).ok()?;
-    Some(temp as f32)
+fn device(index: u32) -> Option<Device<'static>> {
+    NVIDIA.as_ref()?.device_by_index(index).ok()
 }
 
-pub fn power() -> Option<u64> {
-    device()?
-        .total_energy_consumption()
-        .ok()
-        .map(|mj| mj * 1_000)
+pub fn temperature() -> Option<impl Iterator<Item = f32>> {
+    Some(devices()?.flat_map(|device| {
+        device
+            .temperature(TemperatureSensor::Gpu)
+            .ok()
+            .map(|t| t as f32)
+    }))
 }
 
-pub fn memory() -> Option<GpuMemory> {
-    let mem = device()?.memory_info().ok()?;
-    Some(GpuMemory {
-        total: mem.total,
-        free: mem.free,
-    })
+pub fn power() -> Option<impl Iterator<Item = GpuPowerUsage>> {
+    Some(devices()?.flat_map(|device| {
+        let power = device
+            .total_energy_consumption()
+            .ok()
+            .map(|mj| mj * 1_000)?;
+        Some(GpuPowerUsage {
+            card: device.index().unwrap_or_default(),
+            gpu_uj: power,
+        })
+    }))
 }
 
-pub fn utilization() -> impl Iterator<Item = GpuUsage> {
-    let sources = if let Some(device) = device() {
+pub fn memory() -> Option<impl Iterator<Item = GpuMemory>> {
+    Some(devices()?.flat_map(|device| {
+        let mem = device.memory_info().ok()?;
+        Some(GpuMemory {
+            card: device.index().unwrap_or_default(),
+            total: mem.total,
+            free: mem.free,
+        })
+    }))
+}
+
+pub fn utilization() -> Option<impl Iterator<Item = GpuUsage>> {
+    let sources = devices()?.flat_map(|device| {
         let utilization = device.utilization_rates().ok();
         [
             ("compute", utilization.as_ref().map(|u| u.gpu)),
@@ -45,13 +63,12 @@ pub fn utilization() -> impl Iterator<Item = GpuUsage> {
                 device.decoder_utilization().ok().map(|u| u.utilization),
             ),
         ]
-    } else {
-        [("", None); 4]
-    };
-    sources.into_iter().flat_map(|(system, usage)| {
+    });
+    Some(sources.into_iter().flat_map(|(system, usage)| {
         Some(GpuUsage {
+            card: 0,
             system: Cow::Borrowed(system),
             usage: usage?,
         })
-    })
+    }))
 }