add getting system metrics to rpc

Author: Panaetius

.github/workflows/pr-review.yaml

@@ -66,3 +66,4 @@ jobs:
           VCS__PIPELINE__PULL_NUMBER: ${{ github.event.number}}
           VCS__HTTP_CLIENT__API_URL: "https://api.github.com"
           VCS__HTTP_CLIENT__API_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          REVIEW__IGNORE_CHANGES: '["Cargo.lock"]'

Cargo.lock

@@ -91,6 +91,8 @@ tarpc = { version = "0.37.0", features = [
   "tokio1",
 ] }
 tokio-duplex = "1.0.1"
+sysinfo = "0.38.0"
+nvml-wrapper = "0.11.0"
 
 [build-dependencies]
 anyhow = "1.0.90"

coman/Cargo.toml

@@ -7,7 +7,6 @@ use iroh::{
     endpoint::ConnectionError,
     protocol::{ProtocolHandler, Router},
 };
-use nom::AsBytes;
 use pid1::Pid1Settings;
 use rust_supervisor::{ChildType, Supervisor, SupervisorConfig};
 use tokio::{io::AsyncWriteExt, net::TcpStream};
@@ -127,7 +126,7 @@ async fn port_forward() -> Result<()> {
 
     // add rpc server
     let rpc_handler = RpcHandler;
-    builder = builder.accept(b"/coman/rpc/".as_bytes(), rpc_handler);
+    builder = builder.accept(b"/coman/rpc", rpc_handler);
     let _router = builder.spawn();
     println!("port forwarding started");
 

coman/src/cli/exec.rs

@@ -1,5 +1,8 @@
 use futures::StreamExt;
 use iroh::protocol::ProtocolHandler;
+use nvml_wrapper::Nvml;
+use serde::{Deserialize, Serialize};
+use sysinfo::System;
 use tarpc::{
     serde_transport as transport, server, server::Channel, tokio_serde::formats::Bincode,
     tokio_util::codec::LengthDelimitedCodec,
@@ -8,9 +11,17 @@ use tokio_duplex::Duplex;
 
 use crate::cli::app::COMAN_VERSION;
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ResourceUsage {
+    pub cpu: f32,
+    pub rss: u64,
+    pub gpu: Option<u64>,
+}
+
 #[tarpc::service]
 pub trait ComanRPC {
     async fn version() -> String;
+    async fn resource_usage() -> ResourceUsage;
 }
 #[derive(Debug, Clone)]
 struct RpcServer;
@@ -19,6 +30,40 @@ impl ComanRPC for RpcServer {
     async fn version(self, _: tarpc::context::Context) -> String {
         COMAN_VERSION.to_string()
     }
+
+    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 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
+            }
+        };
+        ResourceUsage {
+            cpu: cpu_usage,
+            rss: sys.used_memory(),
+            gpu: gpu_usage,
+        }
+    }
 }
 
 #[derive(Debug, Default)]

coman/src/cli/rpc.rs

@@ -142,8 +142,15 @@ 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: {}, RSS: {}, GPU: {}",
+        result.cpu,
+        result.rss,
+        result.gpu.map(|g| g.to_string()).unwrap_or("N/A".to_string())
+    );
+
+    Ok(())
 }
 
 #[allow(clippy::too_many_arguments)]

coman/src/cscs/cli.rs

@@ -30,7 +30,10 @@ use tokio_util::codec::LengthDelimitedCodec;
 
 use super::api_client::client::{EdfSpec, ScriptSpec};
 use crate::{
-    cli::{app::COMAN_VERSION, rpc::ComanRPCClient},
+    cli::{
+        app::COMAN_VERSION,
+        rpc::{ComanRPCClient, ResourceUsage},
+    },
     config::{ComputePlatform, Config, get_data_dir},
     cscs::{
         api_client::{
@@ -168,10 +171,10 @@ pub async fn cscs_job_log(
     }
 }
 
-pub async fn cscs_resource_usage(job_id: i64, system: Option<String>) -> Result<()> {
+pub async fn cscs_resource_usage(job_id: i64, system: Option<String>) -> Result<ResourceUsage> {
     let endpoint_id = get_endpoint_id(job_id, system).await?;
 
-    let alpn: Vec<u8> = "/coman/rpc".to_string().into_bytes();
+    let alpn: Vec<u8> = b"/coman/rpc".to_vec();
     let secret_key = SecretKey::generate(&mut rand::rng());
     let endpoint = Endpoint::builder().secret_key(secret_key).bind().await?;
 
@@ -183,10 +186,11 @@ pub async fn cscs_resource_usage(job_id: i64, system: Option<String>) -> Result<
             let framed = codec_builder.new_framed(combined);
             let transport = serde_transport::new(framed, Bincode::default());
             let client = ComanRPCClient::new(client::Config::default(), transport);
-            let result = client.spawn().version(context::current()).await?;
-            let _ = dbg!(result);
-
-            Ok(())
+            client
+                .spawn()
+                .resource_usage(context::current())
+                .await
+                .wrap_err("couldn't get resource usage from remote")
         }
         Err(e) => Err(e).wrap_err("couldn't establish tunnel to remote"),
     }
@@ -517,7 +521,7 @@ async fn inject_coman_squash(
     let config = Config::new().unwrap();
     let local_squash_path = maybe_download_latest_squash(current_system, &config).await?;
     let target = base_path.join("coman.sqsh");
-    let file_meta = std::fs::metadata(local_squash_path.clone())?;
+    let file_meta = std::fs::metadata(local_squash_path.clone()).wrap_err("couldn't load coman squash file")?;
 
     #[cfg(target_family = "unix")]
     let size = file_meta.size() as usize;