Fixing the GPU tab

narendasan · Dec 13, 2024 · 8c3adda · 8c3adda
1 parent 7f84675
commit 8c3adda
Show file tree

Hide file tree

Showing 15 changed files with 385 additions and 235 deletions.
diff --git a/build.rs b/build.rs
@@ -1,45 +1,102 @@
 extern crate dirs;
 use std::env;
-use std::os::unix::fs;
+use std::fs::exists;
 use std::process::Command;
 use std::str;
 
 //Apologies if you're trying to figure this out, ping me (@narendasan) if you need help :)
 fn main() -> std::io::Result<()> {
     let gpu_monitor = env::var_os("CARGO_FEATURE_GPU_MONITOR").is_some();
     if gpu_monitor {
-        let ldconfig = Command::new("ldconfig").arg("-p")
-                                               .output()
-                                               .expect("Failed to run ldconfig");
+        let ldconfig = Command::new("ldconfig")
+            .arg("-p")
+            .output()
+            .expect("Failed to run ldconfig");
         let ldc_output = String::from_utf8_lossy(&(ldconfig.stdout[..]));
-        let nvml_installed = ldc_output.split('\n').collect::<Vec<&str>>()
-                                      .iter()
-                                      .map(|x| x.split(' ').collect::<Vec<&str>>()[0])
-                                      .map(|x| x.replace('\t', ""))
-                                      .fold(false, |acc, lib| lib == "libnvidia-ml.so" || acc);
+        let nvml_installed = ldc_output
+            .split('\n')
+            .collect::<Vec<&str>>()
+            .iter()
+            .map(|x| x.split(' ').collect::<Vec<&str>>()[0])
+            .map(|x| x.replace('\t', ""))
+            .fold(false, |acc, lib| lib == "libnvidia-ml.so" || acc);
         if !nvml_installed {
-            let nvidia_driver_version = match Command::new("cat").arg("/proc/driver/nvidia/version").output() {
+            let nvidia_driver_version = match Command::new("cat")
+                .arg("/proc/driver/nvidia/version")
+                .output()
+            {
                 Ok(out) => {
                     let nvidia_driver_version_info = match str::from_utf8(&out.stdout) {
                         Ok(out) => out.to_string(),
-                        Err(_) => return Err(std::io::Error::new(std::io::ErrorKind::Other, "Could not parse driver version"))
+                        Err(_) => {
+                            return Err(std::io::Error::new(
+                                std::io::ErrorKind::Other,
+                                "Could not parse driver version",
+                            ))
+                        }
                     };
-                    let first_line = nvidia_driver_version_info.split('.').map(|x| x.to_string()).collect::<Vec<String>>()[0].clone();
-                    match first_line.split("  ").map(|x| x.to_string()).collect::<Vec<String>>().last().cloned() {
+                    let first_line = nvidia_driver_version_info
+                        .split('.')
+                        .map(|x| x.to_string())
+                        .collect::<Vec<String>>()[0]
+                        .clone();
+                    match first_line
+                        .split("  ")
+                        .map(|x| x.to_string())
+                        .collect::<Vec<String>>()
+                        .last()
+                        .cloned()
+                    {
                         Some(ver) => ver,
-                        None => return Err(std::io::Error::new(std::io::ErrorKind::Other, "Could not parse driver version"))
+                        None => {
+                            return Err(std::io::Error::new(
+                                std::io::ErrorKind::Other,
+                                "Could not parse driver version",
+                            ))
+                        }
                     }
                 }
-                Err(_) => return Err(std::io::Error::new(std::io::ErrorKind::Other, "Could not find NVIDIA Driver verion"))
+                Err(_) => {
+                    return Err(std::io::Error::new(
+                        std::io::ErrorKind::Other,
+                        "Could not find NVIDIA Driver verion",
+                    ))
+                }
             };
-            println!("cargo:rustc-env=LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/nvidia-{}/libnvidia-ml.so", nvidia_driver_version);
-            println!("cargo:rustc-env=LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu");
+            let potential_install_locations = vec![
+                "/usr/lib/x86_64-linux-gnu".to_string(),
+                format!("/usr/lib/nvidia-{}", nvidia_driver_version),
+            ];
 
-            let home_dir = match dirs::home_dir() {
-                Some(p) => format!("{}", p.display()),
-                None => return Err(std::io::Error::new(std::io::ErrorKind::Other, "Failed to find path to home directory"))
-            };
-            fs::symlink(format!("/usr/lib/nvidia-{}/libnvidia-ml.so", nvidia_driver_version), format!("{}/.local/lib/libnvidia-ml.so", home_dir))?;
+            let mut install_location: Option<String> = None;
+            for loc in potential_install_locations {
+                match exists(format!("{}/libnvidia-ml.so.1", loc)) {
+                    Ok(e) => {
+                        if e {
+                            install_location = Some(loc.to_string());
+                            break;
+                        }
+                    }
+                    Err(_) => {
+                        return Err(std::io::Error::new(
+                            std::io::ErrorKind::Other,
+                            format!("Could not determine if libnvidia-ml.so.1 exists at {}", loc),
+                        ))
+                    }
+                }
+            }
+
+            match install_location {
+                Some(loc) => {
+                    println!("cargo:rustc-env=LD_LIBRARY_PATH=$LD_LIBRARY_PATH:{}", loc);
+                }
+                None => {
+                    return Err(std::io::Error::new(
+                        std::io::ErrorKind::Other,
+                        "Unable to find libnvidia-ml.so.1",
+                    ))
+                }
+            }
         }
     }
     Ok(())

diff --git a/src/main.rs b/src/main.rs
@@ -115,6 +115,6 @@ fn _main() -> Result<(), Error> {
     }
     terminal.show_cursor().unwrap();
     terminal.clear().unwrap();
-    terminal::disable_raw_mode();
+    terminal::disable_raw_mode()?;
     Ok(())
 }
diff --git a/src/rtop/appdatastreams.rs b/src/rtop/appdatastreams.rs
@@ -5,7 +5,11 @@ use crate::rtop::error::Error;
 #[cfg(feature = "battery-monitor")]
 use battery::Manager;
 #[cfg(feature = "gpu-monitor")]
-use nvml_wrapper::NVML;
+use nvml_wrapper::Nvml;
+#[cfg(feature = "gpu-monitor")]
+use std::ffi::OsStr;
+#[cfg(feature = "gpu-monitor")]
+use std::fs::exists;
 use sysinfo::System as SysInfoSystem;
 
 #[cfg(feature = "battery-monitor")]
@@ -25,7 +29,7 @@ pub struct AppDataStreams {
     pub battery_info: BatteryMonitor,
     pub sys_info_src: SysInfoSystem,
     #[cfg(feature = "gpu-monitor")]
-    pub gpu_info_src: NVML,
+    pub gpu_info_src: Nvml,
     #[cfg(feature = "battery-monitor")]
     pub battery_info_src: Manager,
 }
@@ -46,7 +50,9 @@ impl<'a> AppDataStreams {
             #[cfg(feature = "battery-monitor")]
             battery_info_src: Manager::new()?,
             #[cfg(feature = "gpu-monitor")]
-            gpu_info_src: NVML::init()?,
+            gpu_info_src: Nvml::builder()
+                .lib_path(OsStr::new(&get_nvml_install_path()?))
+                .init()?,
         })
     }
 
@@ -66,7 +72,45 @@ impl<'a> AppDataStreams {
         #[cfg(feature = "battery-monitor")]
         self.battery_info.poll(&self.battery_info_src)?;
         #[cfg(feature = "gpu-monitor")]
-        self.gpu_info.poll(&self.gpu_info_src)?;
+        self.gpu_info.poll(&self.sys_info_src, &self.gpu_info_src)?;
         Ok(())
     }
 }
+
+#[cfg(feature = "gpu-monitor")]
+fn get_nvml_install_path() -> Result<String, Error> {
+    let potential_install_locations = vec![
+        "/usr/lib/x86_64-linux-gnu".to_string(),
+        "/usr/lib/aarch64-linux-gnu".to_string(),
+    ];
+
+    let mut install_location: Option<String> = None;
+    for loc in potential_install_locations {
+        match exists(format!("{}/libnvidia-ml.so.1", loc)) {
+            Ok(e) => {
+                if e {
+                    install_location = Some(loc.to_string());
+                    break;
+                }
+            }
+            Err(_) => {
+                return Err(Error::from(std::io::Error::new(
+                    std::io::ErrorKind::Other,
+                    format!("Could not determine if libnvidia-ml.so.1 exists at {}", loc),
+                )))
+            }
+        }
+    }
+
+    match install_location {
+        Some(loc) => {
+            return Ok(format!("{}/libnvidia-ml.so.1", loc));
+        }
+        None => {
+            return Err(Error::from(std::io::Error::new(
+                std::io::ErrorKind::Other,
+                "Unable to find libnvidia-ml.so.1",
+            )))
+        }
+    }
+}
diff --git a/src/rtop/datastreams/datastream.rs b/src/rtop/datastreams/datastream.rs
@@ -1,10 +1,10 @@
-use sysinfo::System as SysInfoSystem;
-#[cfg(feature = "gpu-monitor")]
-use nvml_wrapper::NVML;
-#[cfg(feature = "battery-monitor")]
-use battery::Manager;
 #[allow(unused_imports)]
 use crate::rtop::error::Error;
+#[cfg(feature = "battery-monitor")]
+use battery::Manager;
+#[cfg(feature = "gpu-monitor")]
+use nvml_wrapper::Nvml;
+use sysinfo::System as SysInfoSystem;
 
 pub trait SysDataStream {
     fn new(max_hist_len: usize, interpolation_len: u16) -> Self;
@@ -20,7 +20,6 @@ pub trait BatteryDataStream {
 #[cfg(feature = "gpu-monitor")]
 pub trait GPUDataStream {
     fn new(max_hist_len: usize, interpolation_len: u16) -> Self;
-    fn init(&mut self, nvml: &NVML) -> Result<(), Error>;
-    fn poll(&mut self, nvml: &NVML) -> Result<(), Error>;
+    fn init(&mut self, nvml: &Nvml) -> Result<(), Error>;
+    fn poll(&mut self, system_info: &SysInfoSystem, nvml: &Nvml) -> Result<(), Error>;
 }
-