From 1db4c785e15099647f2fb8d4af747c6c32832e8d Mon Sep 17 00:00:00 2001 From: Rolv Apneseth Date: Mon, 20 Mar 2023 01:25:37 +0000 Subject: [PATCH] Add functionality to list connected GPU(s) (#140) --- Cargo.toml | 1 + src/android/mod.rs | 4 ++ src/freebsd/mod.rs | 4 ++ src/linux/mod.rs | 29 +++++++++++ src/linux/pci_devices.rs | 109 +++++++++++++++++++++++++++++++++++++++ src/macos/mod.rs | 4 ++ src/netbsd/mod.rs | 4 ++ src/openwrt/mod.rs | 4 ++ src/traits.rs | 8 +++ src/windows/mod.rs | 4 ++ 10 files changed, 171 insertions(+) create mode 100644 src/linux/pci_devices.rs diff --git a/Cargo.toml b/Cargo.toml index d48bc45f..4d6e419c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ build = "build.rs" cfg-if = "1.0.0" libc = "0.2.131" home = "0.5.3" +pciid-parser = "0.6.2" [build-dependencies.vergen] version = "7.3.2" diff --git a/src/android/mod.rs b/src/android/mod.rs index ccb8c4c9..7a4d6090 100644 --- a/src/android/mod.rs +++ b/src/android/mod.rs @@ -289,6 +289,10 @@ impl GeneralReadout for AndroidGeneralReadout { fn disk_space(&self) -> Result<(u128, u128), ReadoutError> { Err(ReadoutError::NotImplemented) } + + fn gpus(&self) -> Result, ReadoutError> { + Err(ReadoutError::NotImplemented) + } } impl MemoryReadout for AndroidMemoryReadout { diff --git a/src/freebsd/mod.rs b/src/freebsd/mod.rs index ee988ec0..724e2ef6 100644 --- a/src/freebsd/mod.rs +++ b/src/freebsd/mod.rs @@ -297,6 +297,10 @@ impl GeneralReadout for FreeBSDGeneralReadout { fn disk_space(&self) -> Result<(u128, u128), ReadoutError> { shared::disk_space(String::from("/")) } + + fn gpus(&self) -> Result, ReadoutError> { + Err(ReadoutError::NotImplemented) + } } impl MemoryReadout for FreeBSDMemoryReadout { diff --git a/src/linux/mod.rs b/src/linux/mod.rs index 90701027..3a195523 100644 --- a/src/linux/mod.rs +++ b/src/linux/mod.rs @@ -1,12 +1,15 @@ #![allow(clippy::unnecessary_cast)] +mod pci_devices; mod sysinfo_ffi; +use self::pci_devices::get_pci_devices; use crate::extra; use crate::extra::get_entries; use crate::extra::path_extension; use crate::shared; use crate::traits::*; use itertools::Itertools; +use pciid_parser::Database; use regex::Regex; use std::fs; use std::fs::read_dir; @@ -543,6 +546,32 @@ impl GeneralReadout for LinuxGeneralReadout { fn disk_space(&self) -> Result<(u128, u128), ReadoutError> { shared::disk_space(String::from("/")) } + + fn gpus(&self) -> Result, ReadoutError> { + let db = match Database::read() { + Ok(db) => db, + _ => panic!("Could not read pci.ids file"), + }; + + let devices = get_pci_devices()?; + let mut gpus = vec![]; + + for device in devices { + if !device.is_gpu(&db) { + continue; + }; + + if let Some(sub_device_name) = device.get_sub_device_name(&db) { + gpus.push(sub_device_name); + }; + } + + if gpus.is_empty() { + Err(ReadoutError::MetricNotAvailable) + } else { + Ok(gpus) + } + } } impl MemoryReadout for LinuxMemoryReadout { diff --git a/src/linux/pci_devices.rs b/src/linux/pci_devices.rs new file mode 100644 index 00000000..788ebfa4 --- /dev/null +++ b/src/linux/pci_devices.rs @@ -0,0 +1,109 @@ +use std::{ + fs::{read_dir, read_to_string}, + io, + path::PathBuf, +}; + +use pciid_parser::{schema::SubDeviceId, Database}; + +use crate::extra::pop_newline; + +fn parse_device_hex(hex_str: &str) -> String { + pop_newline(hex_str).chars().skip(2).collect::() +} + +pub enum PciDeviceReadableValues { + Class, + Vendor, + Device, + SubVendor, + SubDevice, +} + +impl PciDeviceReadableValues { + fn as_str(&self) -> &'static str { + match self { + PciDeviceReadableValues::Class => "class", + PciDeviceReadableValues::Vendor => "vendor", + PciDeviceReadableValues::Device => "device", + PciDeviceReadableValues::SubVendor => "subsystem_vendor", + PciDeviceReadableValues::SubDevice => "subsystem_device", + } + } +} + +#[derive(Debug)] +pub struct PciDevice { + base_path: PathBuf, +} + +impl PciDevice { + fn new(base_path: PathBuf) -> PciDevice { + PciDevice { base_path } + } + + fn _read_value(&self, readable_value: PciDeviceReadableValues) -> String { + let value_path = self.base_path.join(readable_value.as_str()); + + match read_to_string(&value_path) { + Ok(hex_string) => parse_device_hex(&hex_string), + _ => panic!("Could not find value: {:?}", value_path), + } + } + + pub fn is_gpu(&self, db: &Database) -> bool { + let class_value = self._read_value(PciDeviceReadableValues::Class); + let first_pair = class_value.chars().take(2).collect::(); + + match db.classes.get(&first_pair) { + Some(class) => class.name == "Display controller", + _ => false, + } + } + + pub fn get_sub_device_name(&self, db: &Database) -> Option { + let vendor_value = self._read_value(PciDeviceReadableValues::Vendor); + let sub_vendor_value = self._read_value(PciDeviceReadableValues::SubVendor); + let device_value = self._read_value(PciDeviceReadableValues::Device); + let sub_device_value = self._read_value(PciDeviceReadableValues::SubDevice); + + let Some(vendor) = db.vendors.get(&vendor_value) else { + return None; + }; + + let Some(device) = vendor.devices.get(&device_value) else { + return None; + }; + + let sub_device_id = SubDeviceId { + subvendor: sub_vendor_value, + subdevice: sub_device_value, + }; + + if let Some(sub_device) = device.subdevices.get(&sub_device_id) { + let start = match sub_device.find('[') { + Some(i) => i + 1, + _ => panic!( + "Could not find opening square bracket for sub device: {}", + sub_device + ), + }; + let end = sub_device.len() - 1; + + Some(sub_device.chars().take(end).skip(start).collect::()) + } else { + None + } + } +} + +pub fn get_pci_devices() -> Result, io::Error> { + let devices_dir = read_dir("/sys/bus/pci/devices/")?; + + let mut devices = vec![]; + for device_entry in devices_dir.flatten() { + devices.push(PciDevice::new(device_entry.path())); + } + + Ok(devices) +} diff --git a/src/macos/mod.rs b/src/macos/mod.rs index 3d5f1e3c..c9114880 100644 --- a/src/macos/mod.rs +++ b/src/macos/mod.rs @@ -403,6 +403,10 @@ impl GeneralReadout for MacOSGeneralReadout { fn disk_space(&self) -> Result<(u128, u128), ReadoutError> { shared::disk_space(String::from("/")) } + + fn gpus(&self) -> Result, ReadoutError> { + Err(ReadoutError::NotImplemented) + } } impl MacOSGeneralReadout { diff --git a/src/netbsd/mod.rs b/src/netbsd/mod.rs index 44bf2ad9..4d100597 100644 --- a/src/netbsd/mod.rs +++ b/src/netbsd/mod.rs @@ -333,6 +333,10 @@ impl GeneralReadout for NetBSDGeneralReadout { "Error while trying to get statfs structure.", ))) } + + fn gpus(&self) -> Result, ReadoutError> { + Err(ReadoutError::NotImplemented) + } } impl MemoryReadout for NetBSDMemoryReadout { diff --git a/src/openwrt/mod.rs b/src/openwrt/mod.rs index 519c5598..c9183ee8 100644 --- a/src/openwrt/mod.rs +++ b/src/openwrt/mod.rs @@ -214,6 +214,10 @@ impl GeneralReadout for OpenWrtGeneralReadout { fn disk_space(&self) -> Result<(u128, u128), ReadoutError> { shared::disk_space(String::from("/")) } + + fn gpus(&self) -> Result, ReadoutError> { + Err(ReadoutError::NotImplemented) + } } impl MemoryReadout for OpenWrtMemoryReadout { diff --git a/src/traits.rs b/src/traits.rs index 44704890..a39d6382 100644 --- a/src/traits.rs +++ b/src/traits.rs @@ -483,6 +483,11 @@ impl GeneralReadout for MacOSGeneralReadout { fn disk_space(&self) -> Result<(u128, u128), ReadoutError> { Ok((50000000,1000000000)) // Used / Total } + + fn gpus(&self) -> Result, ReadoutError> { + // Get gpu(s) from list of connected pci devices + Ok(vec!(String::from("gpu1"), String::from("gpu2"))) // Return gpu sub-device names + } } ``` @@ -581,6 +586,9 @@ pub trait GeneralReadout { /// /// _e.g._ '1.2TB / 2TB' fn disk_space(&self) -> Result<(u128, u128), ReadoutError>; + + /// This function should return the sub device names of any _GPU(s)_ connected to the host machine. + fn gpus(&self) -> Result, ReadoutError>; } /// Holds the possible variants for battery status. diff --git a/src/windows/mod.rs b/src/windows/mod.rs index b3c44353..74fd1686 100644 --- a/src/windows/mod.rs +++ b/src/windows/mod.rs @@ -341,6 +341,10 @@ impl GeneralReadout for WindowsGeneralReadout { fn disk_space(&self) -> Result<(u128, u128), ReadoutError> { Err(ReadoutError::NotImplemented) } + + fn gpus(&self) -> Result, ReadoutError> { + Err(ReadoutError::NotImplemented) + } } pub struct WindowsProductReadout {