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

@@ -222,8 +222,8 @@ impl CscsApi {
             None => Ok("".to_string()),
         }
     }
-    pub async fn list_path(&self, system_name: &str, path: PathBuf) -> Result<Vec<PathEntry>> {
-        let result = get_filesystem_ops_ls(&self.client, system_name, path)
+    pub async fn list_path(&self, system_name: &str, path: PathBuf, show_hidden: bool) -> Result<Vec<PathEntry>> {
+        let result = get_filesystem_ops_ls(&self.client, system_name, path, show_hidden)
             .await
             .wrap_err("couldn't list path")?;
         match result.output {

coman/src/cscs/api_client/client.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,
@@ -890,7 +862,7 @@ pub async fn cscs_file_list(
             let api_client = CscsApi::new(access_token.0, platform).unwrap();
             let config = Config::new().unwrap();
             api_client
-                .list_path(&system.unwrap_or(config.values.cscs.current_system), path)
+                .list_path(&system.unwrap_or(config.values.cscs.current_system), path, false)
                 .await
         }
         Err(e) => Err(e),
@@ -946,7 +918,7 @@ pub async fn cscs_file_delete(
             let api_client = CscsApi::new(access_token.0, platform).unwrap();
             let config = Config::new().unwrap();
             let current_system = &system.unwrap_or(config.values.cscs.current_system);
-            let paths = api_client.list_path(current_system, remote.clone()).await?;
+            let paths = api_client.list_path(current_system, remote.clone(), false).await?;
             let path = paths.first().ok_or(eyre!("remote path doesn't exist"))?;
             if let PathType::Directory = path.path_type {
                 return Err(eyre!("remote path must be a file, not directory"));
@@ -975,7 +947,7 @@ pub async fn cscs_file_download(
             let api_client = CscsApi::new(access_token.0, platform).unwrap();
             let config = Config::new().unwrap();
             let current_system = &system.unwrap_or(config.values.cscs.current_system);
-            let paths = api_client.list_path(current_system, remote.clone()).await?;
+            let paths = api_client.list_path(current_system, remote.clone(), false).await?;
             let path = paths.first().ok_or(eyre!("remote path doesn't exist"))?;
             if let PathType::Directory = path.path_type {
                 return Err(eyre!("remote path must be a file, not directory"));
@@ -1008,7 +980,7 @@ pub async fn cscs_file_upload(
             let api_client = CscsApi::new(access_token.0, platform).unwrap();
             let config = Config::new().unwrap();
             let current_system = &system.unwrap_or(config.values.cscs.current_system);
-            let existing = api_client.list_path(current_system, remote.clone()).await?;
+            let existing = api_client.list_path(current_system, remote.clone(), false).await?;
             let remote = if !existing.is_empty() {
                 if existing.len() == 1 && existing[0].path_type == PathType::File {
                     return Err(eyre!("remote file already exists"));

coman/src/cscs/dcgm_enroot_hook.sh

@@ -16,12 +16,16 @@ pub async fn get_filesystem_ops_ls(
     client: &FirecrestClient,
     system_name: &str,
     path: PathBuf,
+    show_hidden: bool,
 ) -> Result<GetDirectoryLsResponse> {
     let path = path.as_os_str().to_str().ok_or(eyre!("couldn't cast path to string"))?;
     let response = client
         .get(
             format!("filesystem/{system_name}/ops/ls").as_str(),
-            Some(vec![("path", path), ("showHidden", "true")]),
+            Some(vec![
+                ("path", path),
+                ("showHidden", if show_hidden { "true" } else { "false" }),
+            ]),
         )
         .await?;
     let model: GetDirectoryLsResponse = serde_json::from_str(response.as_str())?;