diff --git a/crates/obs/src/metrics/collectors/system_gpu.rs b/crates/obs/src/metrics/collectors/system_gpu.rs index b0da57e07..beb0aca3c 100644 --- a/crates/obs/src/metrics/collectors/system_gpu.rs +++ b/crates/obs/src/metrics/collectors/system_gpu.rs @@ -38,6 +38,7 @@ use crate::metrics::schema::system_gpu::PROCESS_GPU_MEMORY_USAGE_MD; use nvml_wrapper::Nvml; use nvml_wrapper::enums::device::UsedGpuMemory; +use nvml_wrapper::struct_wrappers::device::ProcessInfo; use sysinfo::Pid; @@ -45,7 +46,7 @@ use std::borrow::Cow; use thiserror::Error; -use tracing::warn; +use tracing::{debug, warn}; const LOG_COMPONENT_OBS: &str = "obs"; const LOG_SUBSYSTEM_GPU_METRICS: &str = "gpu_metrics"; @@ -86,6 +87,8 @@ pub struct GpuCollector { nvml: Nvml, /// Process ID to monitor pid: Pid, + /// Cached device count so we only probe NVML topology once at init time. + device_count: u32, } impl GpuCollector { @@ -110,7 +113,13 @@ impl GpuCollector { /// ``` pub fn new(pid: Pid) -> Result { let nvml = Nvml::init().map_err(|e| GpuError::InitError(e.to_string()))?; - Ok(GpuCollector { nvml, pid }) + let device_count = nvml.device_count().map_err(|e| GpuError::DeviceError(e.to_string()))?; + + if device_count == 0 { + return Err(GpuError::DeviceError("No GPU device found".to_string())); + } + + Ok(GpuCollector { nvml, pid, device_count }) } /// Collects GPU metrics for the monitored process. @@ -128,38 +137,125 @@ impl GpuCollector { /// println!("GPU memory usage: {} bytes", stats.memory_usage); /// ``` pub fn collect(&self) -> Result { - if let Ok(device) = self.nvml.device_by_index(0) { - if let Ok(gpu_stats) = device.running_compute_processes() { - for stat in gpu_stats.iter() { - if stat.pid == self.pid.as_u32() { - let memory_used = match stat.used_gpu_memory { - UsedGpuMemory::Used(bytes) => bytes, - UsedGpuMemory::Unavailable => 0, - }; - return Ok(GpuStats { - memory_usage: memory_used, - }); - } + let mut total_memory_usage = 0_u64; + let mut process_found = false; + + for device_index in 0..self.device_count { + let device = match self.nvml.device_by_index(device_index) { + Ok(device) => device, + Err(e) => { + warn!( + event = EVENT_GPU_METRICS_STATE, + component = LOG_COMPONENT_OBS, + subsystem = LOG_SUBSYSTEM_GPU_METRICS, + result = "device_unavailable", + device_index, + error = %e, + "gpu metrics state changed" + ); + continue; } - } else { - warn!( - event = EVENT_GPU_METRICS_STATE, - component = LOG_COMPONENT_OBS, - subsystem = LOG_SUBSYSTEM_GPU_METRICS, - result = "process_stats_unavailable", - fallback_memory_usage = 0, - "gpu metrics state changed" - ); + }; + + let compute_processes = match device.running_compute_processes() { + Ok(processes) => processes, + Err(e) => { + warn!( + event = EVENT_GPU_METRICS_STATE, + component = LOG_COMPONENT_OBS, + subsystem = LOG_SUBSYSTEM_GPU_METRICS, + result = "process_stats_unavailable", + device_index, + process_kind = "compute", + error = %e, + fallback_memory_usage = 0, + "gpu metrics state changed" + ); + Vec::new() + } + }; + + let graphics_processes = match device.running_graphics_processes() { + Ok(processes) => processes, + Err(e) => { + warn!( + event = EVENT_GPU_METRICS_STATE, + component = LOG_COMPONENT_OBS, + subsystem = LOG_SUBSYSTEM_GPU_METRICS, + result = "process_stats_unavailable", + device_index, + process_kind = "graphics", + error = %e, + fallback_memory_usage = 0, + "gpu metrics state changed" + ); + Vec::new() + } + }; + + if let Some(device_memory_usage) = + sum_device_process_memory_usage(&compute_processes, &graphics_processes, self.pid.as_u32()) + { + process_found = true; + total_memory_usage = total_memory_usage.saturating_add(device_memory_usage); } - } else { - return Err(GpuError::DeviceError("No GPU device found".to_string())); } - // Process not found in GPU process list, return 0 usage + if process_found { + return Ok(GpuStats { + memory_usage: total_memory_usage, + }); + } + + debug!( + event = EVENT_GPU_METRICS_STATE, + component = LOG_COMPONENT_OBS, + subsystem = LOG_SUBSYSTEM_GPU_METRICS, + result = "process_not_found", + pid = self.pid.as_u32(), + fallback_memory_usage = 0, + "gpu metrics state changed" + ); + Ok(GpuStats { memory_usage: 0 }) } } +fn sum_device_process_memory_usage( + compute_processes: &[ProcessInfo], + graphics_processes: &[ProcessInfo], + pid: u32, +) -> Option { + let compute_memory = sum_matching_process_memory_usage(compute_processes, pid); + let graphics_memory = sum_matching_process_memory_usage(graphics_processes, pid); + + match (compute_memory, graphics_memory) { + (None, None) => None, + (compute_memory, graphics_memory) => Some(compute_memory.unwrap_or(0).saturating_add(graphics_memory.unwrap_or(0))), + } +} + +fn sum_matching_process_memory_usage(processes: &[ProcessInfo], pid: u32) -> Option { + let mut total_memory_usage = 0_u64; + let mut matched = false; + + for process in processes { + if process.pid == pid { + matched = true; + total_memory_usage = total_memory_usage.saturating_add(process_used_gpu_memory(&process.used_gpu_memory)); + } + } + + matched.then_some(total_memory_usage) +} + +fn process_used_gpu_memory(used_gpu_memory: &UsedGpuMemory) -> u64 { + match used_gpu_memory { + UsedGpuMemory::Used(bytes) => *bytes, + UsedGpuMemory::Unavailable => 0, + } +} + /// Converts GPU stats to Prometheus metrics. /// /// # Arguments @@ -190,6 +286,15 @@ pub fn collect_gpu_metrics(stats: &GpuStats, labels: &[(&'static str, Cow<'stati mod tests { use super::*; + fn process_info(pid: u32, used_gpu_memory: UsedGpuMemory) -> ProcessInfo { + ProcessInfo { + pid, + used_gpu_memory, + gpu_instance_id: None, + compute_instance_id: None, + } + } + #[test] fn test_gpu_stats_default() { let stats = GpuStats::default(); @@ -207,4 +312,47 @@ mod tests { let err = GpuError::ProcessNotFound; assert!(err.to_string().contains("Process not found")); } + + #[test] + fn sums_matching_process_memory_and_treats_unavailable_as_zero() { + let processes = vec![ + process_info(7, UsedGpuMemory::Used(128)), + process_info(7, UsedGpuMemory::Unavailable), + process_info(9, UsedGpuMemory::Used(2048)), + ]; + + assert_eq!(sum_matching_process_memory_usage(&processes, 7), Some(128)); + assert_eq!(sum_matching_process_memory_usage(&processes, 9), Some(2048)); + assert_eq!(sum_matching_process_memory_usage(&processes, 42), None); + } + + #[test] + fn sums_compute_and_graphics_process_memory_for_same_device() { + let compute_processes = vec![process_info(42, UsedGpuMemory::Used(256))]; + let graphics_processes = vec![ + process_info(42, UsedGpuMemory::Used(512)), + process_info(77, UsedGpuMemory::Used(1024)), + ]; + + assert_eq!(sum_device_process_memory_usage(&compute_processes, &graphics_processes, 42), Some(768)); + assert_eq!(sum_device_process_memory_usage(&compute_processes, &graphics_processes, 99), None); + } + + #[test] + fn sums_process_memory_across_multiple_devices() { + let device_zero_compute = vec![process_info(7, UsedGpuMemory::Used(256))]; + let device_zero_graphics = vec![]; + let device_one_compute = vec![]; + let device_one_graphics = vec![process_info(7, UsedGpuMemory::Used(1024))]; + + let total = [ + sum_device_process_memory_usage(&device_zero_compute, &device_zero_graphics, 7), + sum_device_process_memory_usage(&device_one_compute, &device_one_graphics, 7), + ] + .into_iter() + .flatten() + .sum::(); + + assert_eq!(total, 1280); + } } diff --git a/crates/obs/src/metrics/scheduler.rs b/crates/obs/src/metrics/scheduler.rs index 8d96991aa..475627110 100644 --- a/crates/obs/src/metrics/scheduler.rs +++ b/crates/obs/src/metrics/scheduler.rs @@ -1289,6 +1289,19 @@ pub fn init_metrics_runtime(token: CancellationToken) { } }; + #[cfg(feature = "gpu")] + let gpu_collector = current_pid.and_then(|pid| { + use crate::metrics::collectors::GpuCollector; + + match GpuCollector::new(pid) { + Ok(collector) => Some(collector), + Err(e) => { + warn!(event = EVENT_METRICS_RUNTIME_STATE, component = LOG_COMPONENT_OBS, subsystem = LOG_SUBSYSTEM_METRICS_RUNTIME, collector = "gpu_metrics", result = "collector_init_failed", error = %e, "metrics runtime state changed"); + None + } + } + }); + loop { tokio::select! { _ = interval.tick() => { @@ -1312,20 +1325,15 @@ pub fn init_metrics_runtime(token: CancellationToken) { collect_system_monitoring_metrics(&bundle, &labels, &mut host_system, &mut host_networks); #[cfg(feature = "gpu")] - if let Some(pid) = current_pid { - use crate::metrics::collectors::{GpuCollector, collect_gpu_metrics}; + if let Some(collector) = gpu_collector.as_ref() { + use crate::metrics::collectors::collect_gpu_metrics; - match GpuCollector::new(pid) { - Ok(collector) => match collector.collect() { - Ok(gpu_stats) => { - metrics.extend(collect_gpu_metrics(&gpu_stats, &labels)); - } - Err(e) => { - warn!(event = EVENT_METRICS_RUNTIME_STATE, component = LOG_COMPONENT_OBS, subsystem = LOG_SUBSYSTEM_METRICS_RUNTIME, collector = "gpu_metrics", result = "collect_failed", error = %e, "metrics runtime state changed"); - } - }, + match collector.collect() { + Ok(gpu_stats) => { + metrics.extend(collect_gpu_metrics(&gpu_stats, &labels)); + } Err(e) => { - warn!(event = EVENT_METRICS_RUNTIME_STATE, component = LOG_COMPONENT_OBS, subsystem = LOG_SUBSYSTEM_METRICS_RUNTIME, collector = "gpu_metrics", result = "collector_init_failed", error = %e, "metrics runtime state changed"); + warn!(event = EVENT_METRICS_RUNTIME_STATE, component = LOG_COMPONENT_OBS, subsystem = LOG_SUBSYSTEM_METRICS_RUNTIME, collector = "gpu_metrics", result = "collect_failed", error = %e, "metrics runtime state changed"); } } }