Improve linux nvidia gpu caching (#2230)

Changes it from a 10s full cache to a 60s conditional cache; we still check the power state now even when using the cached list.
This commit is contained in:
Clement Tsang
2026-09-03 03:18:51 -04:00
committed by GitHub
parent 30fb02c187
commit f9e1703727
2 changed files with 28 additions and 12 deletions
+3 -1
View File
@@ -192,8 +192,10 @@ pub struct DataCollector {
gpu_pids: Option<Vec<IntHashMap<Pid, (u64, u32)>>>, gpu_pids: Option<Vec<IntHashMap<Pid, (u64, u32)>>>,
#[cfg(feature = "gpu")] #[cfg(feature = "gpu")]
gpus_total_mem: Option<u64>, gpus_total_mem: Option<u64>,
#[cfg(all(target_os = "linux", feature = "gpu", feature = "nvidia"))] #[cfg(all(target_os = "linux", feature = "gpu", feature = "nvidia"))]
nvidia_gpu_list_cache: Option<(Vec<String>, Instant)>, /// A vector of GPU names and their corresponding paths, alongside the last update time.
nvidia_gpu_list_cache: Option<(Vec<(String, std::path::PathBuf)>, Instant)>,
#[cfg(feature = "zfs")] #[cfg(feature = "zfs")]
free_arc_mem: bool, free_arc_mem: bool,
+25 -11
View File
@@ -77,7 +77,7 @@ fn is_gpu_class(class_code: &str) -> bool {
/// - <https://us.download.nvidia.com/XFree86/Linux-x86_64/525.89.02/README/dynamicpowermanagement.html> /// - <https://us.download.nvidia.com/XFree86/Linux-x86_64/525.89.02/README/dynamicpowermanagement.html>
/// - <https://www.kernel.org/doc/Documentation/ABI/testing/sysfs-devices-power_state> /// - <https://www.kernel.org/doc/Documentation/ABI/testing/sysfs-devices-power_state>
#[cfg(target_os = "linux")] #[cfg(target_os = "linux")]
fn get_active_pci_bus_ids() -> Vec<String> { fn get_active_pci_bus_ids() -> Vec<(String, std::path::PathBuf)> {
use std::fs; use std::fs;
use crate::collection::linux::utils::is_device_awake; use crate::collection::linux::utils::is_device_awake;
@@ -86,7 +86,7 @@ fn get_active_pci_bus_ids() -> Vec<String> {
return Vec::new(); return Vec::new();
}; };
let mut result: Vec<String> = entries let mut result: Vec<(String, std::path::PathBuf)> = entries
.flatten() .flatten()
.filter_map(|entry| { .filter_map(|entry| {
let path = entry.path(); let path = entry.path();
@@ -128,7 +128,7 @@ fn get_active_pci_bus_ids() -> Vec<String> {
.file_name() .file_name()
.into_string() .into_string()
.ok() .ok()
.map(|name| concat_string::concat_string!("0000", name)) .map(|name| (concat_string::concat_string!("0000", name), path))
} else { } else {
None None
} }
@@ -155,28 +155,42 @@ pub fn get_nvidia_gpu_data(collector: &mut DataCollector) -> Option<GpusData> {
target_os = "linux" => { target_os = "linux" => {
use itertools::Either; use itertools::Either;
// Refresh every ~10 seconds. // We cache for a minute, but still check whether the list of devices is sleeping. This way,
// TODO: IS it possible that our caching keeps stuff awake...? Hm. // solves the problem of waking sleeping devices, but also means we don't check as much AND we
// still support hotplugged devices (in theory).
if let Some((cached_list, cached_time)) = &collector.nvidia_gpu_list_cache if let Some((cached_list, cached_time)) = &collector.nvidia_gpu_list_cache
&& cached_time.elapsed().as_secs() < 10 && cached_time.elapsed().as_secs() < 60
{ {
let devices = Either::Left( let devices = Either::Left(
cached_list cached_list
.iter() .iter()
.filter_map(|id| nvml.device_by_pci_bus_id(id.as_str()).ok()), .filter_map(|(id, device)| {
use crate::collection::linux::utils::is_device_awake;
if is_device_awake(device) {
nvml.device_by_pci_bus_id(id.as_str()).ok()
} else {
None
}
}),
); );
(devices, cached_list.len()) (devices, cached_list.len())
} else { } else {
let pci_bus_ids = get_active_pci_bus_ids(); let pci_bus_ids = get_active_pci_bus_ids();
let num_gpus = pci_bus_ids.len(); let num_gpus = pci_bus_ids.len();
collector.nvidia_gpu_list_cache = collector.nvidia_gpu_list_cache =
Some((pci_bus_ids.clone(), std::time::Instant::now())); Some((pci_bus_ids, std::time::Instant::now()));
let devices = Either::Right( let devices = Either::Right(
pci_bus_ids collector.nvidia_gpu_list_cache
.into_iter() .as_ref()
.filter_map(|id| nvml.device_by_pci_bus_id(id).ok()), .expect("we just inserted the cache entry")
.0
.iter()
.filter_map(|(id, _path)| nvml.device_by_pci_bus_id(id.as_str()).ok())
); );
(devices, num_gpus) (devices, num_gpus)
} }
} }