switch to nvidia-smi

Author: Panaetius

Cargo.lock

@@ -91,7 +91,6 @@ tarpc = { version = "0.37.0", features = [
 ] }
 tokio-duplex = "1.0.1"
 sysinfo = "0.38.0"
-nvml-wrapper = "0.11.0"
 bytesize = "2.3.1"
 
 [build-dependencies]

coman/Cargo.toml

@@ -1,6 +1,6 @@
+use bytesize::ByteSize;
 use futures::StreamExt;
 use iroh::protocol::ProtocolHandler;
-use nvml_wrapper::Nvml;
 use serde::{Deserialize, Serialize};
 use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, System, get_current_pid};
 use tarpc::{
@@ -16,7 +16,7 @@ pub struct ResourceUsage {
     pub cpu: f32,
     pub rss: u64,
     pub vss: u64,
-    pub gpu: Option<u64>,
+    pub gpu: Option<Vec<(u64, u64)>>,
 }
 
 #[tarpc::service]
@@ -43,25 +43,30 @@ impl ComanRPC for RpcServer {
         let Some(process) = sys.process(pid) else {
             return ResourceUsage::default();
         };
-        let gpu_usage = match Nvml::init() {
-            Ok(nvml) => match nvml.device_by_index(0) {
-                Ok(device) => match device.memory_info() {
-                    Ok(memory_info) => Some(memory_info.used),
-                    Err(e) => {
-                        println!("Couldn't get GPU memory info: {e:?}");
-                        None
-                    }
-                },
-                Err(e) => {
-                    println!("couldn't load nvidia device 0: {e:?}");
-                    None
-                }
-            },
-            Err(e) => {
-                println!("Nvidia Device Info not available: {e:?}");
-                None
-            }
+        let gpu_usage = if let Ok(output) = std::process::Command::new("nvidia-smi")
+            .args(vec![
+                "--query-gpu=memory.total,memory.used",
+                "--format=csv,noheader,nounits",
+            ])
+            .output()
+        {
+            let output = String::from_utf8_lossy(&output.stdout);
+            let usage = output
+                .lines()
+                .map(|l| l.split_once(",").unwrap())
+                .map(|(total, used)| {
+                    (
+                        ByteSize::mib(total.trim().parse::<u64>().unwrap()).as_u64(),
+                        ByteSize::mib(used.trim().parse::<u64>().unwrap()).as_u64(),
+                    )
+                })
+                .collect();
+            Some(usage)
+        } else {
+            println!("Failed to execute nvidia-smi, maybe it's not installed");
+            None
         };
+
         ResourceUsage {
             cpu: process.cpu_usage() / sys.cpus().len() as f32,
             rss: process.memory(),

coman/src/cli/rpc.rs

@@ -149,7 +149,18 @@ pub(crate) async fn cli_cscs_job_resource_usage(
         result.cpu,
         ByteSize::b(result.rss).display().iec(),
         ByteSize::b(result.vss).display().iec(),
-        result.gpu.map(|g| g.to_string()).unwrap_or("N/A".to_string())
+        result
+            .gpu
+            .map(|g| g
+                .into_iter()
+                .map(|(total, used)| format!(
+                    "{}/{}({:.1}%)",
+                    ByteSize::b(used).display().iec(),
+                    ByteSize::b(total).display().iec(),
+                    used as f64 / total as f64 * 100.0
+                ))
+                .join(", "))
+            .unwrap_or("N/A".to_string())
     );
 
     Ok(())

coman/src/cscs/cli.rs

@@ -52,7 +52,6 @@ use crate::{
 };
 
 const CSCS_MAX_DIRECT_SIZE: usize = 5242880;
-const DCGM_ENROOT_HOOK: &str = include_str!("./dcgm_enroot_hook.sh");
 
 async fn get_access_token() -> Result<Secret> {
     let client_id = match get_secret(CLIENT_ID_SECRET_NAME).await {
@@ -751,30 +750,6 @@ async fn handle_script(
 
     Ok(script_path)
 }
-async fn setup_dcgm_hook(api_client: &CscsApi, current_system: &str) -> Result<()> {
-    let user_dirs = file_system_roots(Some(FileSystemType::Users)).await?;
-    let user_dir = user_dirs
-        .first()
-        .ok_or(eyre!("couldn't find user root directory on remote"))?;
-    let path = PathBuf::from(user_dir.name.clone())
-        .join(".config")
-        .join("enroot")
-        .join("hooks.d")
-        .join("cscs_jobreport_dcgm_hook.sh");
-
-    let response = api_client.checksum(current_system, path.clone()).await;
-    if let Ok(Some(_)) = response {
-        // file exists
-        return Ok(());
-    }
-    api_client
-        .mkdir(current_system, path.parent().unwrap().to_path_buf())
-        .await?;
-    api_client
-        .upload(current_system, path, DCGM_ENROOT_HOOK.as_bytes().to_vec())
-        .await
-        .wrap_err("couldn't upload dcgm enroot hook ".to_string())
-}
 
 pub async fn cscs_job_start(
     name: Option<String>,
@@ -832,9 +807,6 @@ pub async fn cscs_job_start(
             if coman_squash.is_none() {
                 println!("Warning: coman squash wasn't templated and is needed for ssh through coman to work");
             }
-            if let Err(e) = setup_dcgm_hook(&api_client, current_system).await {
-                println!("Warning: couldn't set up dcgm hook: {e:?}");
-            }
             let environment_path = handle_edf(
                 &api_client,
                 &base_path,