gsadmin/rustfs
1
0
Fork 0
mirror of https://github.com/rustfs/rustfs.git synced 2026-07-26 08:18:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(obs): correct gpu collector coverage (#4482)

Browse Source 
Co-authored-by: heihutu <heihutu@gmail.com>
This commit is contained in:
houseme
2026-07-08 21:35:49 +08:00
committed by GitHub
parent 2157470224
commit e7cfc510ec
2 changed files with 194 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/obs/src/metrics/collectors/system_gpu.rs
+174 -26
View File
@@ -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<Self, GpuError> {
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<GpuStats, GpuError> {
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<u64> {
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<u64> {
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::<u64>();
assert_eq!(total, 1280);
}
}
crates/obs/src/metrics/scheduler.rs
+20 -12
View File
@@ -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");
}
}
}