add dcgm support and mount libs

Author: Panaetius

coman/.config/config.toml

@@ -60,6 +60,7 @@ com.hooks.ssh.enabled = "true"
 com.hooks.ssh.authorize_ssh_key = "{{ ssh_public_key }}"
 com.hooks.ssh.port = "15263"
 {% endif %}
+com.hooks.dcgm.enabled = "true"
 """
 
 # set environment variables that should be passed to a job

coman/src/cli/app.rs

@@ -408,7 +408,7 @@ fn remote_path_completer(current: &std::ffi::OsStr) -> Vec<CompletionCandidate>
     let (send, mut recv) = mpsc::unbounded_channel();
     if current.is_empty() || current == "/" {
         tokio::spawn(async move {
-            let roots = file_system_roots().await;
+            let roots = file_system_roots(None).await;
             if let Ok(roots) = roots {
                 for root in roots {
                     send.send(CompletionCandidate::new(root.name.clone())).unwrap();

coman/src/cli/rpc.rs

@@ -2,7 +2,7 @@ use futures::StreamExt;
 use iroh::protocol::ProtocolHandler;
 use nvml_wrapper::Nvml;
 use serde::{Deserialize, Serialize};
-use sysinfo::System;
+use sysinfo::{System, get_current_pid};
 use tarpc::{
     serde_transport as transport, server, server::Channel, tokio_serde::formats::Bincode,
     tokio_util::codec::LengthDelimitedCodec,
@@ -11,11 +11,11 @@ use tokio_duplex::Duplex;
 
 use crate::cli::app::COMAN_VERSION;
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
 pub struct ResourceUsage {
     pub cpu: f32,
-    pub mem_used: u64,
-    pub mem_total: u64,
+    pub rss: u64,
+    pub vss: u64,
     pub gpu: Option<u64>,
 }
 
@@ -35,11 +35,12 @@ impl ComanRPC for RpcServer {
     async fn resource_usage(self, _context: ::tarpc::context::Context) -> ResourceUsage {
         let mut sys = System::new_all();
         sys.refresh_all();
-        let mut cpu_usage = 0.0;
-        for cpu in sys.cpus() {
-            cpu_usage += cpu.cpu_usage();
-        }
-        cpu_usage /= sys.cpus().len() as f32;
+        let Ok(pid) = get_current_pid() else {
+            return ResourceUsage::default();
+        };
+        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() {
@@ -60,9 +61,9 @@ impl ComanRPC for RpcServer {
             }
         };
         ResourceUsage {
-            cpu: cpu_usage,
-            mem_used: sys.used_memory(),
-            mem_total: sys.total_memory(),
+            cpu: process.cpu_usage(),
+            rss: process.memory(),
+            vss: process.virtual_memory(),
             gpu: gpu_usage,
         }
     }

coman/src/cscs/cli.rs

@@ -4,6 +4,7 @@ use std::{
     time::{Duration, Instant},
 };
 
+use bytesize::ByteSize;
 use color_eyre::{Result, eyre::Context};
 use eyre::eyre;
 use futures::StreamExt;
@@ -142,8 +143,16 @@ pub(crate) async fn cli_cscs_job_resource_usage(
     platform: Option<ComputePlatform>,
 ) -> Result<()> {
     let job_id = maybe_job_id_from_name(job, system.clone(), platform.clone()).await?;
-    println!("running port forward for job {job_id}");
-    cscs_resource_usage(job_id, system).await
+    let result = cscs_resource_usage(job_id, system).await?;
+    println!(
+        "CPU: {}, Memory: RSS {} VSS {}, GPU: {}",
+        result.cpu,
+        ByteSize::b(result.rss).display().iec().to_string(),
+        ByteSize::b(result.vss).display().iec().to_string(),
+        result.gpu.map(|g| g.to_string()).unwrap_or("N/A".to_string())
+    );
+
+    Ok(())
 }
 
 #[allow(clippy::too_many_arguments)]

coman/src/cscs/dcgm_enroot_hook.sh

@@ -0,0 +1,60 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+shopt -s lastpipe nullglob
+
+export PATH=\"${PATH}:/usr/sbin:/sbin\"
+
+source \"${ENROOT_LIBRARY_PATH}/common.sh\"
+
+common::checkcmd grep sed ldd ldconfig
+
+if [ \"${OCI_ANNOTATION_com__hooks__dcgm__enabled:-}\" != \"true\" ]; then
+    exit 0
+fi
+
+# Mounting the specified DCGM libraries and directories explicitly
+cat << EOF | enroot-mount --root \"${ENROOT_ROOTFS}\" -
+/usr/local/dcgm /usr/local/dcgm none x-create=dir,bind,ro,nosuid,nodev,private
+/usr/lib64/libnvperf_dcgm_host.so /usr/lib64/libnvperf_dcgm_host.so none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulesysmon.so.3.3.6 /usr/lib64/libdcgmmodulesysmon.so.3.3.6 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulesysmon.so.3 /usr/lib64/libdcgmmodulesysmon.so.3 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulesysmon.so /usr/lib64/libdcgmmodulesysmon.so none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmoduleprofiling.so.3.3.6 /usr/lib64/libdcgmmoduleprofiling.so.3.3.6 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmoduleprofiling.so.3 /usr/lib64/libdcgmmoduleprofiling.so.3 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmoduleprofiling.so /usr/lib64/libdcgmmoduleprofiling.so none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulepolicy.so.3.3.6 /usr/lib64/libdcgmmodulepolicy.so.3.3.6 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulepolicy.so.3 /usr/lib64/libdcgmmodulepolicy.so.3 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulepolicy.so /usr/lib64/libdcgmmodulepolicy.so none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulenvswitch.so.3.3.6 /usr/lib64/libdcgmmodulenvswitch.so.3.3.6 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulenvswitch.so.3 /usr/lib64/libdcgmmodulenvswitch.so.3 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulenvswitch.so /usr/lib64/libdcgmmodulenvswitch.so none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmoduleintrospect.so.3.3.6 /usr/lib64/libdcgmmoduleintrospect.so.3.3.6 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmoduleintrospect.so.3 /usr/lib64/libdcgmmoduleintrospect.so.3 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmoduleintrospect.so /usr/lib64/libdcgmmoduleintrospect.so none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulehealth.so.3.3.6 /usr/lib64/libdcgmmodulehealth.so.3.3.6 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulehealth.so.3 /usr/lib64/libdcgmmodulehealth.so.3 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulehealth.so /usr/lib64/libdcgmmodulehealth.so none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulediag.so.3.3.6 /usr/lib64/libdcgmmodulediag.so.3.3.6 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulediag.so.3 /usr/lib64/libdcgmmodulediag.so.3 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmodulediag.so /usr/lib64/libdcgmmodulediag.so none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmoduleconfig.so.3.3.6 /usr/lib64/libdcgmmoduleconfig.so.3.3.6 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmoduleconfig.so.3 /usr/lib64/libdcgmmoduleconfig.so.3 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgmmoduleconfig.so /usr/lib64/libdcgmmoduleconfig.so none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgm_stub.a /usr/lib64/libdcgm_stub.a none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgm_cublas_proxy12.so /usr/lib64/libdcgm_cublas_proxy12.so none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgm_cublas_proxy11.so /usr/lib64/libdcgm_cublas_proxy11.so none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgm.so.3.3.6 /usr/lib64/libdcgm.so.3.3.6 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgm.so.3 /usr/lib64/libdcgm.so.3 none x-create=file,bind,ro,nosuid,nodev,private
+/usr/lib64/libdcgm.so /usr/lib64/libdcgm.so none x-create=file,bind,ro,nosuid,nodev,private
+EOF
+
+# Refresh the dynamic linker cache to include newly mounted libs
+cat << EOF > \"${ENROOT_ROOTFS}/etc/ld.so.conf.d/enroot-dcgm-hook.conf\"
+/lib64
+/usr/lib64
+EOF
+
+if ! ${ldconfig:-ldconfig} -r \"${ENROOT_ROOTFS}\" >> \"${ENROOT_ROOTFS}/dcgm-hook.log\" 2>&1; then
+    common::err \"Failed to refresh the dynamic linker cache\"
+fi\n"

coman/src/cscs/handlers.rs

@@ -52,6 +52,7 @@ 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 {
@@ -750,6 +751,27 @@ 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
+        .upload(current_system, path, DCGM_ENROOT_HOOK.as_bytes().to_vec())
+        .await
+        .wrap_err(format!("couldn't upload dcgm enroot hook "))
+}
 
 pub async fn cscs_job_start(
     name: Option<String>,
@@ -807,7 +829,9 @@ 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,
@@ -869,7 +893,7 @@ pub async fn cscs_file_list(
         Err(e) => Err(e),
     }
 }
-pub async fn file_system_roots() -> Result<Vec<PathEntry>> {
+pub async fn file_system_roots(type_filter: Option<FileSystemType>) -> Result<Vec<PathEntry>> {
     let config = Config::new().expect("couldn't load config");
     let user_info = cscs_user_info(None, None).await?;
     let systems = cscs_system_list(None).await?;
@@ -878,7 +902,17 @@ pub async fn file_system_roots() -> Result<Vec<PathEntry>> {
         .find(|s| s.name == config.values.cscs.current_system)
         .unwrap_or_else(|| panic!("couldn't get info for system {}", config.values.cscs.current_system));
     let mut subpaths = vec![];
-    for fs in system.file_systems.clone() {
+    let filesystems = if let Some(filter) = type_filter {
+        system
+            .file_systems
+            .clone()
+            .into_iter()
+            .filter(|fs| fs.data_type == filter)
+            .collect()
+    } else {
+        system.file_systems.clone()
+    };
+    for fs in filesystems {
         let entry = match cscs_stat_path(PathBuf::from(fs.path.clone()).join(user_info.name.clone()), None, None).await
         {
             Ok(Some(_)) => PathEntry {

coman/src/cscs/ports.rs

@@ -248,7 +248,7 @@ async fn list_files(id: PathBuf) -> Result<Option<Event<UserEvent>>> {
         .map_err(|_| eyre!("couldn't convert id to string".to_owned()))?;
     if id_str == "/" {
         // load file system roots
-        let subpaths = file_system_roots().await?;
+        let subpaths = file_system_roots(None).await?;
         Ok(Some(Event::User(UserEvent::File(FileEvent::List(id_str, subpaths)))))
     } else {
         let subpaths = cscs_file_list(id, None, None).await?;

firecrest_client/src/filesystem_api.rs

@@ -21,7 +21,7 @@ pub async fn get_filesystem_ops_ls(
     let response = client
         .get(
             format!("filesystem/{system_name}/ops/ls").as_str(),
-            Some(vec![("path", path)]),
+            Some(vec![("path", path), ("showHidden", "false")]),
         )
         .await?;
     let model: GetDirectoryLsResponse = serde_json::from_str(response.as_str())?;