add (provisional) support for injecting coman into jobs

Author: Panaetius

.github/workflows/release.yaml

@@ -23,7 +23,7 @@ jobs:
             cargo_command: cargo
 
           - os_name: Linux-aarch64
-            os: ubuntu-latest
+            os: ubuntu-24.04-arm
             target: aarch64-unknown-linux-musl
             bin: coman
             name: coman-Linux-aarch64-musl.tar.gz

Cargo.toml

@@ -5,3 +5,8 @@ resolver = "3"
 
 [workspace.dependencies]
 firecrest_client = { path = "./firecrest_client", version = "2.4.0" }
+
+[profile.release]
+strip = true
+opt-level = "z"
+lto = true

coman/.config/config.toml

@@ -24,7 +24,7 @@ sbatch_script_template = """
 #SBATCH --job-name={{name}}
 #SBATCH --ntasks=1
 #SBATCH --time=1:00:00
-srun {% if environment_file %}--environment={{environment_file}}{% endif %} {{command}}
+srun {% if environment_file %}--environment={{environment_file}}{% endif %} {% if coman_squash %}/coman/coman exec {% endif %}{{command}}
 """
 
 # the edf environment toml file template
@@ -37,7 +37,10 @@ srun {% if environment_file %}--environment={{environment_file}}{% endif %} {{co
 #   mount: a dictionary of key/value pairs for folders to mount to the container, with key being the path in the cluster and value being the path in the container
 edf_file_template = """
 {% if edf_image %}image = "{{edf_image}}"{% endif %}
-mounts = [{% for source, target in mount %}"{{source}}:{{target}}",{% endfor %}]
+mounts = [
+  {% for source, target in mount %}"{{source}}:{{target}}",{% endfor %}
+  {% if coman_squash %}"{{coman_squash}}:/coman:sqsh"{% endif%}
+]
 workdir = "{{container_workdir}}"
 
 [env]
@@ -49,7 +52,7 @@ workdir = "{{container_workdir}}"
 {% if ssh_public_key %}
 com.hooks.ssh.enabled = "true"
 com.hooks.ssh.authorize_ssh_key = "{{ ssh_public_key }}"
-com.hooks.ssh.port = 15263
+com.hooks.ssh.port = "15263"
 {% endif %}
 """
 

coman/src/cli.rs

@@ -247,6 +247,8 @@ pub enum CscsJobCommands {
         no_ssh: bool,
         #[clap(short, long, help="ssh public key to use", value_hint=ValueHint::FilePath)]
         ssh_key: Option<PathBuf>,
+        #[clap(long, action, help = "don't upload and inject coman into the container")]
+        no_coman: bool,
         #[clap(trailing_var_arg = true, help = "The command to run in the container", value_hint=ValueHint::Other)]
         command: Option<Vec<String>>,
     },

coman/src/config.rs

@@ -73,6 +73,8 @@ pub struct ComanConfig {
     #[serde(default)]
     pub name: Option<String>,
     #[serde(default)]
+    pub coman_squash_path: Option<PathBuf>,
+    #[serde(default)]
     pub cscs: CscsConfig,
 }
 

coman/src/cscs/api_client/client.rs

@@ -52,6 +52,7 @@ pub struct JobStartOptions {
     pub script_spec: ScriptSpec,
     pub no_ssh: bool,
     pub ssh_key: Option<PathBuf>,
+    pub no_coman: bool,
 }
 
 pub struct CscsApi {
@@ -160,8 +161,8 @@ impl CscsApi {
             .wrap_err("couldn't change directory permission")?;
         Ok(())
     }
-    pub async fn upload(&self, system_name: &str, path: PathBuf, file: Vec<u8>) -> Result<()> {
-        post_filesystem_ops_upload(&self.client, system_name, path, file)
+    pub async fn upload(&self, system_name: &str, target: PathBuf, file: Vec<u8>) -> Result<()> {
+        post_filesystem_ops_upload(&self.client, system_name, target, file)
             .await
             .wrap_err("couldn't upload file")?;
         Ok(())
@@ -170,10 +171,10 @@ impl CscsApi {
         &self,
         system_name: &str,
         account: Option<String>,
-        path: PathBuf,
+        target: PathBuf,
         size: i64,
     ) -> Result<(i64, S3Upload)> {
-        let job = post_filesystem_transfer_upload(&self.client, system_name, account, path, size)
+        let job = post_filesystem_transfer_upload(&self.client, system_name, account, target, size)
             .await
             .wrap_err("couldn't upload file")?;
         if let DownloadFileResponseTransferDirectives::S3(directives) = job.transfer_directives {

coman/src/cscs/cli.rs

@@ -272,7 +272,7 @@ pub(crate) async fn cli_cscs_file_upload(
     }
 }
 
-async fn upload_chunk(path: PathBuf, offset: u64, size: u64, url: Url) -> Result<String> {
+pub(crate) async fn upload_chunk(path: PathBuf, offset: u64, size: u64, url: Url) -> Result<String> {
     let client = reqwest::Client::new();
 
     let source_file = File::open(path).await?;

coman/src/cscs/handlers.rs

@@ -9,6 +9,7 @@ use std::{
 
 use color_eyre::{Result, eyre::eyre};
 use eyre::Context;
+use itertools::Itertools;
 use reqwest::Url;
 
 use super::api_client::client::{EdfSpec, ScriptSpec};
@@ -19,6 +20,7 @@ use crate::{
             client::{CscsApi, JobStartOptions},
             types::{FileStat, FileSystemType, Job, JobDetail, PathEntry, PathType, S3Upload, System, UserInfo},
         },
+        cli::upload_chunk,
         oauth2::{
             CLIENT_ID_SECRET_NAME, CLIENT_SECRET_SECRET_NAME, client_credentials_login, finish_cscs_device_login,
             start_cscs_device_login,
@@ -201,14 +203,74 @@ async fn setup_ssh(
     }
 }
 
+async fn inject_coman_squash(
+    api_client: &CscsApi,
+    base_path: &Path,
+    current_system: &str,
+    options: &JobStartOptions,
+) -> Result<Option<PathBuf>> {
+    if options.no_coman {
+        return Ok(None);
+    }
+    let config = Config::new().unwrap();
+    let local_squash_path = match config.values.coman_squash_path.clone() {
+        Some(path) => path,
+        None => todo!(),
+    };
+    let target = base_path.join("coman.sqsh");
+    let file_meta = std::fs::metadata(local_squash_path.clone())?;
+
+    #[cfg(target_family = "unix")]
+    let size = file_meta.size() as usize;
+
+    #[cfg(target_family = "windows")]
+    let size = file_meta.file_size() as usize;
+
+    //upload squash file
+    let transfer_data = api_client
+        .transfer_upload(current_system, config.values.cscs.account, target.clone(), size as i64)
+        .await?;
+    let mut etags: Vec<String> = Vec::new();
+    let client = reqwest::Client::new();
+    let num_parts = transfer_data.1.num_parts;
+    for (chunk_id, transfer_url) in transfer_data.1.parts_upload_urls.into_iter().enumerate() {
+        println!(
+            "Uploading part {}/{} ({}Mb)",
+            chunk_id + 1,
+            num_parts,
+            transfer_data.1.part_size / 1024 / 1024
+        );
+        let etag = upload_chunk(
+            local_squash_path.clone(),
+            (chunk_id as u64) * transfer_data.1.part_size,
+            transfer_data.1.part_size,
+            transfer_url,
+        )
+        .await?;
+        etags.push(etag);
+    }
+
+    let body = etags
+        .into_iter()
+        .enumerate()
+        .map(|(i, etag)| (i + 1, etag))
+        .map(|(i, etag)| format!("<Part><PartNumber>{}</PartNumber><ETag>{}</ETag></Part>", i, etag))
+        .join("");
+    let body = format!("<CompleteMultipartUpload>{}</CompleteMultipartUpload>", body);
+    let req = client.post(transfer_data.1.complete_upload_url).body(body).build()?;
+    let resp = client.execute(req).await?;
+    resp.error_for_status()?;
+    Ok(Some(target))
+}
+
 async fn handle_edf(
     api_client: &CscsApi,
     base_path: &Path,
     current_system: &str,
     envvars: &HashMap<String, String>,
     workdir: &str,
     options: &JobStartOptions,
-) -> Result<PathBuf> {
+) -> Result<(PathBuf, Option<PathBuf>)> {
     let config = Config::new().unwrap();
     let environment_path = base_path.join("environment.toml");
     match options.edf_spec.clone() {
@@ -248,21 +310,23 @@ async fn handle_edf(
             }
 
             let ssh_path = setup_ssh(api_client, base_path, current_system, options).await?;
+            let coman_squash = inject_coman_squash(api_client, base_path, current_system, options).await?;
 
             let mut context = tera::Context::new();
             context.insert("edf_image", &docker_image.to_edf());
             context.insert("container_workdir", &workdir);
             context.insert("env", &envvars);
             context.insert("mount", &mount);
             context.insert("ssh_public_key", &ssh_path);
+            context.insert("coman_squash", &coman_squash);
 
             let environment_file = tera.render("environment.toml", &context)?;
             api_client.mkdir(current_system, base_path.to_path_buf()).await?;
             api_client.chmod(current_system, base_path.to_path_buf(), "700").await?;
             api_client
                 .upload(current_system, environment_path.clone(), environment_file.into_bytes())
                 .await?;
-            Ok(environment_path)
+            Ok((environment_path, coman_squash))
         }
         EdfSpec::Local(local_path) => {
             let environment_file = std::fs::read_to_string(local_path.clone())?;
@@ -271,17 +335,20 @@ async fn handle_edf(
             api_client
                 .upload(current_system, environment_path.clone(), environment_file.into_bytes())
                 .await?;
-            Ok(environment_path)
+            Ok((environment_path, None))
         }
-        EdfSpec::Remote(path) => Ok(path),
+        EdfSpec::Remote(path) => Ok((path, None)),
     }
 }
+
+#[allow(clippy::too_many_arguments)]
 async fn handle_script(
     api_client: &CscsApi,
     job_name: &str,
     base_path: &Path,
     current_system: &str,
     environment_path: &Path,
+    coman_squash: Option<PathBuf>,
     workdir: &str,
     options: &JobStartOptions,
 ) -> Result<PathBuf> {
@@ -300,6 +367,9 @@ async fn handle_script(
             );
             context.insert("environment_file", &environment_path.to_path_buf());
             context.insert("container_workdir", &workdir);
+            if let Some(path) = coman_squash {
+                context.insert("coman_squash", &path);
+            }
             let script = tera.render("script.sh", &context)?;
             api_client
                 .upload(current_system, script_path.clone(), script.into_bytes())
@@ -360,7 +430,7 @@ pub async fn cscs_job_start(
             let mut envvars = config.values.cscs.env.clone();
             envvars.extend(options.env.clone());
 
-            let environment_path = handle_edf(
+            let (environment_path, coman_squash) = handle_edf(
                 &api_client,
                 &base_path,
                 current_system,
@@ -376,6 +446,7 @@ pub async fn cscs_job_start(
                 &base_path,
                 current_system,
                 &environment_path,
+                coman_squash,
                 &container_workdir,
                 &options,
             )

coman/src/main.rs

@@ -98,6 +98,7 @@ async fn main() -> Result<()> {
                         script_spec,
                         no_ssh,
                         ssh_key,
+                        no_coman,
                     } => {
                         cli_cscs_job_start(
                             name,
@@ -113,6 +114,7 @@ async fn main() -> Result<()> {
                                 script_spec: script_spec.unwrap_or_default().into(),
                                 no_ssh,
                                 ssh_key,
+                                no_coman,
                             },
                             system,
                             platform,