use crate::Result;
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct GpuDeviceInfo {
pub index: u32,
pub name: Option<String>,
pub total_bytes: u64,
pub free_bytes: u64,
pub used_bytes: u64,
}
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct ProcessGpuInfo {
pub used_bytes: u64,
pub is_per_process: bool,
pub source: GpuQuerySource,
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GpuQuerySource {
Dxgi,
Nvml,
NvidiaSmi,
}
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct GpuProcessEntry {
pub pid: u32,
pub name: Option<String>,
pub used_bytes: u64,
pub source: GpuQuerySource,
}
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct Snapshot {
pub ram_bytes: u64,
pub gpu: Option<ProcessGpuInfo>,
pub gpu_device: Option<GpuDeviceInfo>,
}
impl Snapshot {
pub fn now(device_index: u32) -> Result<Self> {
let ram_bytes = crate::ram::process_rss()?;
let gpu = crate::gpu::process_gpu_info(device_index).ok();
let gpu_device = crate::gpu::device_info(device_index).ok();
Ok(Self {
ram_bytes,
gpu,
gpu_device,
})
}
pub fn all() -> Result<Vec<Self>> {
let ram_bytes = crate::ram::process_rss()?;
let nvidia_count = crate::gpu::device_count().unwrap_or(0);
#[allow(clippy::as_conversions)]
let mut snapshots: Vec<Self> = Vec::with_capacity(nvidia_count as usize);
for idx in 0..nvidia_count {
snapshots.push(Self {
ram_bytes,
gpu: crate::gpu::process_gpu_info(idx).ok(),
gpu_device: crate::gpu::device_info(idx).ok(),
});
}
#[cfg(all(windows, feature = "dxgi"))]
for (gpu_device, gpu) in crate::gpu::dxgi_non_nvidia_devices(nvidia_count) {
snapshots.push(Self {
ram_bytes,
gpu: Some(gpu),
gpu_device: Some(gpu_device),
});
}
Ok(snapshots)
}
}
#[cfg(feature = "report")]
impl Snapshot {
#[must_use]
pub fn ram_mb(&self) -> f64 {
#[allow(clippy::cast_precision_loss, clippy::as_conversions)]
let mb = self.ram_bytes as f64 / 1_048_576.0;
mb
}
#[must_use]
pub fn vram_mb(&self) -> Option<f64> {
#[allow(clippy::cast_precision_loss, clippy::as_conversions)]
let mb = self.gpu.as_ref().map(|p| p.used_bytes as f64 / 1_048_576.0);
mb
}
}
impl GpuDeviceInfo {
#[must_use]
pub fn name_or_unknown(&self) -> &str {
self.name.as_deref().unwrap_or("unknown GPU")
}
}
#[cfg(feature = "report")]
impl GpuDeviceInfo {
#[must_use]
pub fn format_free(&self) -> String {
#[allow(clippy::cast_precision_loss, clippy::as_conversions)]
let free_mb = self.free_bytes as f64 / 1_048_576.0;
#[allow(clippy::cast_precision_loss, clippy::as_conversions)]
let total_mb = self.total_bytes as f64 / 1_048_576.0;
let name_suffix = self
.name
.as_deref()
.map_or(String::new(), |n| format!(" [{n}]"));
format!(
" GPU {}: free {free_mb:.0} MB / {total_mb:.0} MB{name_suffix}\n",
self.index
)
}
pub fn print_free(&self) {
print!("{}", self.format_free());
}
#[must_use]
pub fn format_total(&self) -> String {
#[allow(clippy::cast_precision_loss, clippy::as_conversions)]
let total_mb = self.total_bytes as f64 / 1_048_576.0;
let name_suffix = self
.name
.as_deref()
.map_or(String::new(), |n| format!(" [{n}]"));
format!(
" GPU {}: total {total_mb:.0} MB{name_suffix}\n",
self.index
)
}
#[must_use]
pub fn format_used(&self) -> String {
#[allow(clippy::cast_precision_loss, clippy::as_conversions)]
let used_mb = self.used_bytes as f64 / 1_048_576.0;
let name_suffix = self
.name
.as_deref()
.map_or(String::new(), |n| format!(" [{n}]"));
format!(" GPU {}: used {used_mb:.0} MB{name_suffix}\n", self.index)
}
}
#[cfg(feature = "test-helpers")]
#[derive(Debug, Clone, Default)]
pub struct GpuDeviceInfoBuilder {
index: u32,
name: Option<String>,
total_bytes: u64,
free_bytes: u64,
used_bytes: u64,
}
#[cfg(feature = "test-helpers")]
impl GpuDeviceInfo {
#[must_use]
pub fn builder() -> GpuDeviceInfoBuilder {
GpuDeviceInfoBuilder::default()
}
}
#[cfg(feature = "test-helpers")]
impl GpuDeviceInfoBuilder {
#[must_use]
pub const fn index(mut self, index: u32) -> Self {
self.index = index;
self
}
#[must_use]
pub fn name(mut self, name: Option<String>) -> Self {
self.name = name;
self
}
#[must_use]
pub const fn total_bytes(mut self, total: u64) -> Self {
self.total_bytes = total;
self
}
#[must_use]
pub const fn free_bytes(mut self, free: u64) -> Self {
self.free_bytes = free;
self
}
#[must_use]
pub const fn used_bytes(mut self, used: u64) -> Self {
self.used_bytes = used;
self
}
#[must_use]
pub fn build(self) -> GpuDeviceInfo {
GpuDeviceInfo {
index: self.index,
name: self.name,
total_bytes: self.total_bytes,
free_bytes: self.free_bytes,
used_bytes: self.used_bytes,
}
}
}
#[cfg(test)]
#[allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::missing_docs_in_private_items
)]
mod tests {
use super::*;
fn make_snapshot(
ram: u64,
vram_used: Option<u64>,
is_per_process: bool,
total: u64,
) -> Snapshot {
Snapshot {
ram_bytes: ram,
gpu: vram_used.map(|used| ProcessGpuInfo {
used_bytes: used,
is_per_process,
source: if is_per_process {
GpuQuerySource::Nvml
} else {
GpuQuerySource::NvidiaSmi
},
}),
gpu_device: if total > 0 {
Some(GpuDeviceInfo {
index: 0,
name: None,
total_bytes: total,
free_bytes: total.saturating_sub(vram_used.unwrap_or(0)),
used_bytes: vram_used.unwrap_or(0),
})
} else {
None
},
}
}
#[test]
fn snapshot_constructs_with_no_gpu() {
let snap = make_snapshot(0, None, false, 0);
assert_eq!(snap.ram_bytes, 0);
assert!(snap.gpu.is_none());
assert!(snap.gpu_device.is_none());
}
#[test]
fn snapshot_constructs_with_full_gpu() {
let snap = make_snapshot(1_048_576, Some(500 * 1_048_576), true, 16_384 * 1_048_576);
assert_eq!(snap.ram_bytes, 1_048_576);
assert_eq!(snap.gpu.as_ref().unwrap().used_bytes, 500 * 1_048_576);
assert!(snap.gpu.as_ref().unwrap().is_per_process);
assert_eq!(
snap.gpu_device.as_ref().unwrap().total_bytes,
16_384 * 1_048_576
);
}
#[cfg(feature = "report")]
#[test]
fn ram_mb_conversion() {
let snap = make_snapshot(1_048_576, None, false, 0);
assert!((snap.ram_mb() - 1.0).abs() < 0.001);
}
#[cfg(feature = "report")]
#[test]
fn ram_mb_zero() {
let snap = make_snapshot(0, None, false, 0);
assert!(snap.ram_mb().abs() < 0.001);
}
#[cfg(feature = "report")]
#[test]
fn vram_mb_none_when_no_gpu() {
let snap = make_snapshot(100, None, false, 0);
assert!(snap.vram_mb().is_none());
}
#[cfg(feature = "report")]
#[test]
fn vram_mb_some_when_gpu_present() {
let snap = make_snapshot(100, Some(2 * 1_048_576), true, 16 * 1_048_576);
assert!((snap.vram_mb().unwrap() - 2.0).abs() < 0.001);
}
#[cfg(feature = "report")]
#[test]
fn format_free_with_name() {
let dev = GpuDeviceInfo {
index: 0,
name: Some("NVIDIA Test GPU".to_owned()),
total_bytes: 16_384 * 1_048_576,
free_bytes: 13_284 * 1_048_576,
used_bytes: 3_100 * 1_048_576,
};
assert_eq!(
dev.format_free(),
" GPU 0: free 13284 MB / 16384 MB [NVIDIA Test GPU]\n"
);
}
#[cfg(feature = "report")]
#[test]
fn format_free_without_name() {
let dev = GpuDeviceInfo {
index: 1,
name: None,
total_bytes: 8_192 * 1_048_576,
free_bytes: 4_096 * 1_048_576,
used_bytes: 4_096 * 1_048_576,
};
assert_eq!(dev.format_free(), " GPU 1: free 4096 MB / 8192 MB\n");
}
#[cfg(feature = "report")]
#[test]
fn format_free_full_device() {
let dev = GpuDeviceInfo {
index: 2,
name: Some("Saturated GPU".to_owned()),
total_bytes: 4_096 * 1_048_576,
free_bytes: 0,
used_bytes: 4_096 * 1_048_576,
};
assert_eq!(
dev.format_free(),
" GPU 2: free 0 MB / 4096 MB [Saturated GPU]\n"
);
}
#[cfg(feature = "report")]
#[test]
fn print_free_does_not_panic() {
let dev = GpuDeviceInfo {
index: 0,
name: None,
total_bytes: 1_000,
free_bytes: 500,
used_bytes: 500,
};
dev.print_free();
}
#[cfg(feature = "report")]
#[test]
fn format_total_with_name() {
let dev = GpuDeviceInfo {
index: 0,
name: Some("NVIDIA Test GPU".to_owned()),
total_bytes: 16_384 * 1_048_576,
free_bytes: 13_284 * 1_048_576,
used_bytes: 3_100 * 1_048_576,
};
assert_eq!(
dev.format_total(),
" GPU 0: total 16384 MB [NVIDIA Test GPU]\n"
);
}
#[cfg(feature = "report")]
#[test]
fn format_total_without_name() {
let dev = GpuDeviceInfo {
index: 1,
name: None,
total_bytes: 8_192 * 1_048_576,
free_bytes: 4_096 * 1_048_576,
used_bytes: 4_096 * 1_048_576,
};
assert_eq!(dev.format_total(), " GPU 1: total 8192 MB\n");
}
#[cfg(feature = "report")]
#[test]
fn format_total_full_device() {
let dev = GpuDeviceInfo {
index: 2,
name: Some("Saturated GPU".to_owned()),
total_bytes: 4_096 * 1_048_576,
free_bytes: 0,
used_bytes: 4_096 * 1_048_576,
};
assert_eq!(
dev.format_total(),
" GPU 2: total 4096 MB [Saturated GPU]\n"
);
}
#[cfg(feature = "report")]
#[test]
fn format_used_with_name() {
let dev = GpuDeviceInfo {
index: 0,
name: Some("NVIDIA Test GPU".to_owned()),
total_bytes: 16_384 * 1_048_576,
free_bytes: 13_284 * 1_048_576,
used_bytes: 3_100 * 1_048_576,
};
assert_eq!(
dev.format_used(),
" GPU 0: used 3100 MB [NVIDIA Test GPU]\n"
);
}
#[cfg(feature = "report")]
#[test]
fn format_used_without_name() {
let dev = GpuDeviceInfo {
index: 1,
name: None,
total_bytes: 8_192 * 1_048_576,
free_bytes: 4_096 * 1_048_576,
used_bytes: 4_096 * 1_048_576,
};
assert_eq!(dev.format_used(), " GPU 1: used 4096 MB\n");
}
#[cfg(feature = "report")]
#[test]
fn format_used_idle_device() {
let dev = GpuDeviceInfo {
index: 2,
name: Some("Idle GPU".to_owned()),
total_bytes: 4_096 * 1_048_576,
free_bytes: 4_096 * 1_048_576,
used_bytes: 0,
};
assert_eq!(dev.format_used(), " GPU 2: used 0 MB [Idle GPU]\n");
}
#[test]
fn name_or_unknown_returns_inner_when_some() {
let dev = GpuDeviceInfo {
index: 0,
name: Some("NVIDIA Test GPU".to_owned()),
total_bytes: 0,
free_bytes: 0,
used_bytes: 0,
};
assert_eq!(dev.name_or_unknown(), "NVIDIA Test GPU");
}
#[test]
fn name_or_unknown_returns_fallback_when_none() {
let dev = GpuDeviceInfo {
index: 0,
name: None,
total_bytes: 0,
free_bytes: 0,
used_bytes: 0,
};
assert_eq!(dev.name_or_unknown(), "unknown GPU");
}
#[cfg(feature = "test-helpers")]
#[test]
fn builder_defaults_produce_zeroed_device() {
let dev = GpuDeviceInfo::builder().build();
assert_eq!(dev.index, 0);
assert!(dev.name.is_none());
assert_eq!(dev.total_bytes, 0);
assert_eq!(dev.free_bytes, 0);
assert_eq!(dev.used_bytes, 0);
}
#[cfg(feature = "test-helpers")]
#[test]
fn builder_index_setter() {
let dev = GpuDeviceInfo::builder().index(7).build();
assert_eq!(dev.index, 7);
}
#[cfg(feature = "test-helpers")]
#[test]
fn builder_name_setter_some() {
let dev = GpuDeviceInfo::builder()
.name(Some("Synthetic GPU".to_owned()))
.build();
assert_eq!(dev.name.as_deref(), Some("Synthetic GPU"));
}
#[cfg(feature = "test-helpers")]
#[test]
fn builder_name_setter_none_is_default() {
let dev = GpuDeviceInfo::builder().name(None).build();
assert!(dev.name.is_none());
}
#[cfg(feature = "test-helpers")]
#[test]
fn builder_byte_setters() {
let dev = GpuDeviceInfo::builder()
.total_bytes(16 * 1_024 * 1_024 * 1_024)
.free_bytes(14 * 1_024 * 1_024 * 1_024)
.used_bytes(2 * 1_024 * 1_024 * 1_024)
.build();
assert_eq!(dev.total_bytes, 16 * 1_024 * 1_024 * 1_024);
assert_eq!(dev.free_bytes, 14 * 1_024 * 1_024 * 1_024);
assert_eq!(dev.used_bytes, 2 * 1_024 * 1_024 * 1_024);
}
#[cfg(feature = "test-helpers")]
#[test]
fn builder_full_round_trip() {
let dev = GpuDeviceInfo::builder()
.index(1)
.name(Some("Round-Trip GPU".to_owned()))
.total_bytes(17_179_869_184)
.free_bytes(15_246_684_160)
.used_bytes(1_933_185_024)
.build();
assert_eq!(dev.index, 1);
assert_eq!(dev.name.as_deref(), Some("Round-Trip GPU"));
assert_eq!(dev.total_bytes, 17_179_869_184);
assert_eq!(dev.free_bytes, 15_246_684_160);
assert_eq!(dev.used_bytes, 1_933_185_024);
}
}