add getting system metrics to rpc

Author: Panaetius

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

@@ -1,16 +1,27 @@
 use futures::StreamExt;
 use iroh::protocol::ProtocolHandler;
-use tarpc::{
-    serde_transport as transport, server, server::Channel, tokio_serde::formats::Bincode,
-    tokio_util::codec::LengthDelimitedCodec,
-};
+use nvml_wrapper::Nvml;
+use serde::{Deserialize, Serialize};
+use sysinfo::System;
+use tarpc::server::Channel;
+use tarpc::tokio_serde::formats::Bincode;
+use tarpc::tokio_util::codec::LengthDelimitedCodec;
+use tarpc::{serde_transport as transport, server};
 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: u64,
+}
+
 #[tarpc::service]
 pub trait ComanRPC {
     async fn version() -> String;
+    async fn resource_usage() -> ResourceUsage;
 }
 #[derive(Debug, Clone)]
 struct RpcServer;
@@ -19,6 +30,24 @@ 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 nvml = Nvml::init().unwrap();
+        let device = nvml.device_by_index(0).unwrap();
+        let memory_info = device.memory_info().unwrap();
+        ResourceUsage {
+            cpu: cpu_usage,
+            rss: sys.used_memory(),
+            gpu: memory_info.used,
+        }
+    }
 }
 
 #[derive(Debug, Default)]

coman/src/cli/rpc.rs

@@ -142,8 +142,10 @@ 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);
+
+    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"),
     }