use crate::device::readers::common_cache::{DetailBuilder, DeviceStaticInfo};
use crate::device::types::{GpuInfo, ProcessInfo};
use crate::device::GpuReader;
use crate::utils::get_hostname;
use chrono::Local;
use libamdgpu_top::stat::{self, FdInfoStat, ProcInfo};
use libamdgpu_top::AMDGPU::{DeviceHandle, GpuMetrics, MetricsInfo, GPU_INFO};
use libamdgpu_top::{AppDeviceInfo, DevicePath, VramUsage};
use std::collections::HashMap;
use std::sync::{Mutex, OnceLock};
const MAX_GPU_UTILIZATION: f64 = 100.0; const MAX_GPU_POWER_WATTS: f64 = 1000.0; const MAX_GPU_TEMP_CELSIUS: u32 = 125; const MAX_GPU_FREQ_MHZ: u32 = 5000; const MAX_GPU_MEMORY_BYTES: u64 = 512 * 1024 * 1024 * 1024;
const MAX_VERSION_COMPONENT: i32 = 999;
struct AmdGpuDevice {
device_path: DevicePath,
device_handle: DeviceHandle,
vram_usage: Mutex<VramUsage>, static_info: OnceLock<DeviceStaticInfo>, }
pub struct AmdGpuReader {
devices: Vec<AmdGpuDevice>,
rocm_version: OnceLock<Option<String>>,
}
impl Default for AmdGpuReader {
fn default() -> Self {
Self::new()
}
}
impl AmdGpuReader {
pub fn new() -> Self {
if !Self::check_amd_gpu_permissions() {
return Self {
devices: Vec::new(),
rocm_version: OnceLock::new(),
};
}
let device_path_list = DevicePath::get_device_path_list();
let mut devices = Vec::new();
const MAX_DEVICES: usize = 256;
let device_paths_to_process: Vec<_> =
device_path_list.into_iter().take(MAX_DEVICES).collect();
for device_path in device_paths_to_process {
match device_path.init() {
Ok(amdgpu_dev) => {
match amdgpu_dev.memory_info() {
Ok(memory_info) => {
let vram_usage = VramUsage::new(&memory_info);
devices.push(AmdGpuDevice {
device_path: device_path.clone(),
device_handle: amdgpu_dev,
vram_usage: Mutex::new(vram_usage),
static_info: OnceLock::new(),
});
}
Err(e) => {
eprintln!(
"Warning: Failed to get memory info for AMD GPU {}: {e}",
device_path.pci
);
}
}
}
Err(e) => {
eprintln!(
"Warning: Failed to initialize AMD GPU {}: {e}",
device_path.pci
);
}
}
}
Self {
devices,
rocm_version: OnceLock::new(),
}
}
fn get_rocm_version(&self) -> Option<String> {
self.rocm_version
.get_or_init(libamdgpu_top::get_rocm_version)
.clone()
}
fn get_device_static_info<'a>(&self, device: &'a AmdGpuDevice) -> &'a DeviceStaticInfo {
device
.static_info
.get_or_init(|| {
let ext_info = device.device_handle.device_info().ok();
let memory_info = device.device_handle.memory_info().ok();
let (device_name, mut detail) = if let (Some(ext), Some(mem)) =
(ext_info.as_ref(), memory_info.as_ref())
{
let sensors = libamdgpu_top::stat::Sensors::new(
&device.device_handle,
&device.device_path.pci,
ext,
);
let app_device_info = AppDeviceInfo::new(
&device.device_handle,
ext,
mem,
&sensors,
&device.device_path,
);
let mut builder = DetailBuilder::new()
.insert("Device Name", &app_device_info.marketing_name)
.insert("PCI Bus", app_device_info.pci_bus.to_string());
if let Some(ref ver) = self.get_rocm_version() {
builder = builder
.insert("ROCm Version", ver)
.insert("lib_name", "ROCm")
.insert("lib_version", ver);
}
let mut detail = builder.build();
detail.insert(
"Device ID".to_string(),
format!("{:#06x}", ext.device_id()),
);
detail.insert(
"Revision ID".to_string(),
format!("{:#04x}", ext.pci_rev_id()),
);
detail.insert(
"ASIC Name".to_string(),
app_device_info.asic_name.to_string(),
);
if let Some(ref vbios) = app_device_info.vbios {
detail.insert("VBIOS Version".to_string(), vbios.ver.clone());
detail.insert("VBIOS Date".to_string(), vbios.date.clone());
}
if let Some(ref cap) = app_device_info.power_cap {
detail.insert("Power Cap".to_string(), format!("{} W", cap.current));
detail.insert("Power Cap (Min)".to_string(), format!("{} W", cap.min));
detail.insert("Power Cap (Max)".to_string(), format!("{} W", cap.max));
}
if let Some(link) = app_device_info.max_gpu_link {
detail.insert(
"Max GPU Link".to_string(),
format!("Gen{} x{}", link.gen, link.width),
);
}
if let Some(link) = app_device_info.max_system_link {
detail.insert(
"Max System Link".to_string(),
format!("Gen{} x{}", link.gen, link.width),
);
}
if let Some(min_dpm_link) = app_device_info.min_dpm_link {
detail.insert(
"Min DPM Link".to_string(),
format!("Gen{} x{}", min_dpm_link.gen, min_dpm_link.width),
);
}
if let Some(max_dpm_link) = app_device_info.max_dpm_link {
detail.insert(
"Max DPM Link".to_string(),
format!("Gen{} x{}", max_dpm_link.gen, max_dpm_link.width),
);
}
(app_device_info.marketing_name, detail)
} else {
(String::from("Unknown GPU"), HashMap::new())
};
match device.device_handle.get_drm_version_struct() {
Ok(drm) => {
if drm.version_major >= 0
&& drm.version_major <= MAX_VERSION_COMPONENT
&& drm.version_minor >= 0
&& drm.version_minor <= MAX_VERSION_COMPONENT
&& drm.version_patchlevel >= 0
&& drm.version_patchlevel <= MAX_VERSION_COMPONENT
{
let ver = format!(
"{}.{}.{}",
drm.version_major, drm.version_minor, drm.version_patchlevel
);
detail.insert("Driver Version".to_string(), ver);
} else {
eprintln!(
"Warning: Invalid driver version components detected: {}.{}.{} for device {}",
drm.version_major, drm.version_minor, drm.version_patchlevel,
device.device_path.pci
);
}
}
Err(e) => {
eprintln!(
"Warning: Failed to get driver version for device {}: {e}",
device.device_path.pci
);
}
};
DeviceStaticInfo::with_details(device_name, None, detail)
})
}
fn check_amd_gpu_permissions() -> bool {
use std::fs;
let dri_path = std::path::Path::new("/dev/dri");
if !dri_path.exists() {
return false;
}
match fs::read_dir(dri_path) {
Ok(entries) => {
for entry in entries.flatten() {
let path = entry.path();
let file_name = path.file_name().and_then(|n| n.to_str()).unwrap_or("");
if file_name.starts_with("card") || file_name.starts_with("render") {
if unsafe { libc::geteuid() } == 0 {
return true; }
if let Ok(_file) = fs::OpenOptions::new().read(true).write(true).open(&path)
{
return true; }
}
}
false }
Err(_) => false, }
}
}
impl GpuReader for AmdGpuReader {
fn get_gpu_info(&self) -> Vec<GpuInfo> {
let mut gpu_info = Vec::new();
for device in &self.devices {
let static_info = self.get_device_static_info(device);
let mut detail = static_info.detail.clone();
let device_name = static_info.name.clone();
let ext_info = match device.device_handle.device_info() {
Ok(info) => info,
Err(e) => {
eprintln!(
"Warning: Failed to get device info for AMD GPU {}: {e}",
device.device_path.pci
);
continue; }
};
let memory_info = {
let vram_usage_result = device.vram_usage.lock();
match vram_usage_result {
Ok(mut vram_usage) => {
vram_usage.update_usage(&device.device_handle);
vram_usage.update_usable_heap_size(&device.device_handle);
vram_usage.0 }
Err(poisoned) => {
eprintln!(
"Warning: VramUsage mutex was poisoned for device {}, recovering...",
device.device_path.pci
);
match device.device_handle.memory_info() {
Ok(fresh_memory_info) => {
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
poisoned.into_inner()
})) {
Ok(mut guard) => {
*guard = VramUsage::new(&fresh_memory_info);
guard.update_usage(&device.device_handle);
guard.update_usable_heap_size(&device.device_handle);
guard.0
}
Err(_) => {
eprintln!(
"Critical: Failed to recover poisoned mutex for device {}, skipping",
device.device_path.pci
);
continue;
}
}
}
Err(e) => {
eprintln!("Failed to get fresh memory info during recovery: {e}");
continue; }
}
}
}
};
let sensors = libamdgpu_top::stat::Sensors::new(
&device.device_handle,
&device.device_path.pci,
&ext_info,
);
if let Some(ref sensors) = sensors {
if let Some(link) = sensors.current_link {
detail.insert(
"Current Link".to_string(),
format!("Gen{} x{}", link.gen, link.width),
);
}
if let Some(fan) = sensors.fan_rpm {
detail.insert("Fan Speed".to_string(), format!("{fan} RPM"));
}
if let Some(mclk) = sensors.mclk {
detail.insert("Memory Clock".to_string(), format!("{mclk} MHz"));
}
}
let mut utilization = 0.0;
let mut power_consumption = 0.0;
let mut temperature: u32 = 0;
let mut frequency: u32 = 0;
if let Ok(metrics) = GpuMetrics::get_from_sysfs_path(&device.device_path.sysfs_path) {
if let Some(gfx_activity) = metrics.get_average_gfx_activity() {
utilization = (gfx_activity as f64).clamp(0.0, MAX_GPU_UTILIZATION);
}
if let Some(power) = metrics.get_average_socket_power() {
let watts = power as f64 / 1000.0; power_consumption = watts.clamp(0.0, MAX_GPU_POWER_WATTS);
}
if let Some(temp) = metrics.get_temperature_edge() {
temperature = (temp as u32).min(MAX_GPU_TEMP_CELSIUS);
}
if let Some(freq) = metrics.get_current_gfxclk() {
frequency = (freq as u32).min(MAX_GPU_FREQ_MHZ);
}
}
if let Some(ref s) = sensors {
if utilization == 0.0 {
}
if power_consumption == 0.0 {
if let Some(ref p) = s.average_power {
let watts = p.value as f64 / 1000.0; power_consumption = watts.clamp(0.0, MAX_GPU_POWER_WATTS);
} else if let Some(ref p) = s.input_power {
let watts = p.value as f64 / 1000.0; power_consumption = watts.clamp(0.0, MAX_GPU_POWER_WATTS);
}
}
if temperature == 0 {
if let Some(ref t) = s.edge_temp {
temperature = (t.current as u32).min(MAX_GPU_TEMP_CELSIUS);
}
}
if frequency == 0 {
if let Some(clk) = s.sclk {
frequency = clk.min(MAX_GPU_FREQ_MHZ);
}
}
}
let total_memory = if memory_info.vram.total_heap_size > 0 {
memory_info.vram.total_heap_size.min(MAX_GPU_MEMORY_BYTES)
} else if memory_info.vram.usable_heap_size > 0 {
memory_info.vram.usable_heap_size.min(MAX_GPU_MEMORY_BYTES)
} else {
0
};
let used_memory = memory_info.vram.heap_usage.min(total_memory);
let info = GpuInfo {
uuid: format!("GPU-{}", device.device_path.pci), time: Local::now().format("%Y-%m-%d %H:%M:%S").to_string(),
name: device_name, device_type: "GPU".to_string(),
host_id: get_hostname(),
hostname: get_hostname(),
instance: get_hostname(),
utilization,
ane_utilization: 0.0,
dla_utilization: None,
tensorcore_utilization: None,
temperature,
used_memory,
total_memory,
frequency,
power_consumption,
gpu_core_count: None,
detail,
};
gpu_info.push(info);
}
gpu_info
}
fn get_process_info(&self) -> Vec<ProcessInfo> {
use std::collections::{HashMap, HashSet};
let mut process_info_list = Vec::new();
let proc_list = stat::get_process_list();
struct GpuProcessData {
device_id: usize,
device_uuid: String,
pid: u32,
name: String,
vram_usage_kib: u64,
gtt_usage_kib: u64,
}
let mut gpu_processes = Vec::new();
let mut gpu_pids = HashSet::new();
for (device_idx, device) in self.devices.iter().enumerate() {
let mut proc_index: Vec<ProcInfo> = Vec::new();
stat::update_index_by_all_proc(
&mut proc_index,
&[&device.device_path.render, &device.device_path.card],
&proc_list,
);
let mut fdinfo = FdInfoStat::default();
fdinfo.get_all_proc_usage(&proc_index);
for proc_usage in fdinfo.proc_usage {
let vram_usage_kib = proc_usage.usage.vram_usage;
let gtt_usage_kib = proc_usage.usage.gtt_usage;
if vram_usage_kib > 0 || gtt_usage_kib > 0 {
let pid = proc_usage.pid as u32;
gpu_pids.insert(pid);
gpu_processes.push(GpuProcessData {
device_id: device_idx,
device_uuid: format!("GPU-{}", device.device_path.pci),
pid,
name: proc_usage.name,
vram_usage_kib,
gtt_usage_kib,
});
}
}
}
use crate::utils::with_global_system;
use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, UpdateKind};
let system_processes = with_global_system(|system| {
let refresh_kind = ProcessRefreshKind::nothing()
.with_cpu()
.with_memory()
.with_user(UpdateKind::OnlyIfNotSet);
system.refresh_processes_specifics(ProcessesToUpdate::All, true, refresh_kind);
crate::device::process_list::get_all_processes(system, &gpu_pids)
});
let process_map: HashMap<u32, _> = system_processes.iter().map(|p| (p.pid, p)).collect();
for gpu_proc in gpu_processes {
let gpu_memory_bytes = if gpu_proc.vram_usage_kib > 0 {
gpu_proc.vram_usage_kib * 1024
} else {
gpu_proc.gtt_usage_kib * 1024
};
let sys_proc = process_map.get(&gpu_proc.pid);
let process_info = ProcessInfo {
device_id: gpu_proc.device_id,
device_uuid: gpu_proc.device_uuid,
pid: gpu_proc.pid,
process_name: gpu_proc.name,
used_memory: gpu_memory_bytes,
cpu_percent: sys_proc.map(|p| p.cpu_percent).unwrap_or(0.0),
memory_percent: sys_proc.map(|p| p.memory_percent).unwrap_or(0.0),
memory_rss: sys_proc.map(|p| p.memory_rss).unwrap_or(0),
memory_vms: sys_proc.map(|p| p.memory_vms).unwrap_or(0),
user: sys_proc.map(|p| p.user.clone()).unwrap_or_default(),
state: sys_proc.map(|p| p.state.clone()).unwrap_or_default(),
start_time: sys_proc.map(|p| p.start_time.clone()).unwrap_or_default(),
cpu_time: sys_proc.map(|p| p.cpu_time).unwrap_or(0),
command: sys_proc.map(|p| p.command.clone()).unwrap_or_default(),
ppid: sys_proc.map(|p| p.ppid).unwrap_or(0),
threads: sys_proc.map(|p| p.threads).unwrap_or(0),
uses_gpu: true,
priority: sys_proc.map(|p| p.priority).unwrap_or(0),
nice_value: sys_proc.map(|p| p.nice_value).unwrap_or(0),
gpu_utilization: 0.0, };
process_info_list.push(process_info);
}
process_info_list
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_max_version_component_validation() {
assert!(
MAX_VERSION_COMPONENT >= 99,
"Should support two-digit version components"
);
assert!(
MAX_VERSION_COMPONENT <= 9999,
"Should not be excessively large"
);
let common_versions = vec![
(6, 12, 0), (5, 15, 0), (30, 10, 1), (999, 999, 999), ];
for (major, minor, patch) in common_versions {
assert!(
major <= MAX_VERSION_COMPONENT,
"Major version {major} should be valid"
);
assert!(
minor <= MAX_VERSION_COMPONENT,
"Minor version {minor} should be valid"
);
assert!(
patch <= MAX_VERSION_COMPONENT,
"Patch version {patch} should be valid"
);
}
}
#[test]
fn test_version_validation_rejects_invalid() {
let invalid_versions = vec![
(1000, 0, 0), (0, 1000, 0), (0, 0, 1000), (-1, 0, 0), (0, -1, 0), (0, 0, -1), ];
for (major, minor, patch) in invalid_versions {
let major_valid = major >= 0 && major <= MAX_VERSION_COMPONENT;
let minor_valid = minor >= 0 && minor <= MAX_VERSION_COMPONENT;
let patch_valid = patch >= 0 && patch <= MAX_VERSION_COMPONENT;
assert!(
!(major_valid && minor_valid && patch_valid),
"Version {major}.{minor}.{patch} should be invalid"
);
}
}
#[test]
fn test_memory_validation_constants() {
assert_eq!(
MAX_GPU_MEMORY_BYTES,
512 * 1024 * 1024 * 1024,
"Max GPU memory should be 512GB"
);
let mi325x_memory: u64 = 288 * 1024 * 1024 * 1024;
assert!(
mi325x_memory < MAX_GPU_MEMORY_BYTES,
"Should support MI325X 288GB memory"
);
let future_memory: u64 = 400 * 1024 * 1024 * 1024;
assert!(
future_memory < MAX_GPU_MEMORY_BYTES,
"Should have headroom for future GPUs"
);
}
#[test]
fn test_gpu_metric_validation_constants() {
assert_eq!(MAX_GPU_UTILIZATION, 100.0, "Max utilization should be 100%");
assert_eq!(
MAX_GPU_POWER_WATTS, 1000.0,
"Max power should support high-end GPUs"
);
assert_eq!(
MAX_GPU_TEMP_CELSIUS, 125,
"Max temp should be above thermal limits"
);
assert_eq!(
MAX_GPU_FREQ_MHZ, 5000,
"Max frequency should support boost clocks"
);
let mi300x_power = 750.0; assert!(
mi300x_power < MAX_GPU_POWER_WATTS,
"Should support MI300X power draw"
);
let typical_boost_freq = 2500; assert!(
typical_boost_freq < MAX_GPU_FREQ_MHZ,
"Should support typical boost frequencies"
);
}
}