allow specifying script source

Author: Panaetius

coman/src/cli.rs

@@ -5,7 +5,7 @@ use strum::VariantNames;
 
 use crate::{
     config::{ComputePlatform, get_config_dir, get_data_dir, get_project_local_config_file},
-    cscs::api_client::client::EdfSpec as EdfSpecEnum,
+    cscs::api_client::client::{EdfSpec as EdfSpecEnum, ScriptSpec as ScriptSpecEnum},
     util::types::DockerImageUrl,
 };
 
@@ -57,12 +57,47 @@ pub enum CscsCommands {
     },
 }
 
+#[derive(Args, Clone, Debug)]
+#[group(multiple = false)]
+pub struct ScriptSpec {
+    #[arg(
+        long,
+        help = "generate and upload script file based on template (on by default unless `--local` or `--remote` are passed)"
+    )]
+    generate_script: bool,
+    #[arg(long, value_name = "PATH", help = "upload local script file")]
+    local_script: Option<PathBuf>,
+    #[arg(long, value_name = "PATH", help = "use script file already present on remote")]
+    remote_script: Option<PathBuf>,
+}
+impl Default for ScriptSpec {
+    fn default() -> Self {
+        Self {
+            generate_script: true,
+            local_script: Default::default(),
+            remote_script: Default::default(),
+        }
+    }
+}
+
+impl From<ScriptSpec> for ScriptSpecEnum {
+    fn from(val: ScriptSpec) -> Self {
+        if let Some(local_script) = val.local_script {
+            ScriptSpecEnum::Local(local_script)
+        } else if let Some(remote_script) = val.remote_script {
+            ScriptSpecEnum::Remote(remote_script)
+        } else {
+            ScriptSpecEnum::Generate
+        }
+    }
+}
+
 #[derive(Args, Clone, Debug)]
 #[group(multiple = false)]
 pub struct EdfSpec {
     #[arg(
         long,
-        help = "generate and upload edf file based on template (on by default unless `--local`, `--remote` or `--no-edf` are passed)"
+        help = "generate and upload edf file based on template (on by default unless `--local` or `--remote` are passed)"
     )]
     generate_edf: bool,
     #[arg(long, value_name = "PATH", help = "upload local edf file")]
@@ -111,8 +146,6 @@ pub enum CscsJobCommands {
     Submit {
         #[clap(short, long, help = "name of the job")]
         name: Option<String>,
-        #[clap(short, long, help = "the path to the srun script file to use")]
-        script_file: Option<PathBuf>,
         #[clap(
             short,
             long,
@@ -131,6 +164,8 @@ pub enum CscsJobCommands {
         stderr: Option<PathBuf>,
         #[command(flatten)]
         edf_spec: Option<EdfSpec>,
+        #[command(flatten)]
+        script_spec: Option<ScriptSpec>,
         #[clap(trailing_var_arg = true, help = "The command to run in the container")]
         command: Option<Vec<String>>,
     },

coman/src/cscs/api_client/client.rs

@@ -25,6 +25,13 @@ use crate::{
     util::types::DockerImageUrl,
 };
 #[derive(Debug, Clone, Default)]
+pub enum ScriptSpec {
+    #[default]
+    Generate,
+    Local(PathBuf),
+    Remote(PathBuf),
+}
+#[derive(Debug, Clone, Default)]
 pub enum EdfSpec {
     #[default]
     Generate,
@@ -34,7 +41,6 @@ pub enum EdfSpec {
 
 #[derive(Debug, Clone, Default)]
 pub struct JobStartOptions {
-    pub script_file: Option<PathBuf>,
     pub image: Option<DockerImageUrl>,
     pub command: Option<Vec<String>>,
     pub stdout: Option<PathBuf>,
@@ -43,6 +49,7 @@ pub struct JobStartOptions {
     pub env: Vec<(String, String)>,
     pub mount: Vec<(String, String)>,
     pub edf_spec: EdfSpec,
+    pub script_spec: ScriptSpec,
 }
 
 pub struct CscsApi {

coman/src/cscs/handlers.rs

@@ -10,7 +10,7 @@ use std::{
 use color_eyre::{Result, eyre::eyre};
 use reqwest::Url;
 
-use super::api_client::client::EdfSpec;
+use super::api_client::client::{EdfSpec, ScriptSpec};
 use crate::{
     config::{ComputePlatform, Config},
     cscs::{
@@ -236,8 +236,49 @@ async fn handle_edf(
         EdfSpec::Remote(path) => Ok(path),
     }
 }
+async fn handle_script(
+    api_client: &CscsApi,
+    job_name: &str,
+    base_path: &Path,
+    current_system: &str,
+    environment_path: &Path,
+    workdir: &str,
+    options: &JobStartOptions,
+) -> Result<PathBuf> {
+    let config = Config::new().unwrap();
+    let script_path = base_path.join("script.sh");
+    match options.script_spec.clone() {
+        ScriptSpec::Generate => {
+            let script_template = config.cscs.sbatch_script_template;
+            let mut tera = tera::Tera::default();
+            tera.add_raw_template("script.sh", &script_template)?;
+            let mut context = tera::Context::new();
+            context.insert("name", &job_name);
+            context.insert(
+                "command",
+                &options.command.clone().unwrap_or(config.cscs.command).join(" "),
+            );
+            context.insert("environment_file", &environment_path.to_path_buf());
+            context.insert("container_workdir", &workdir);
+            let script = tera.render("script.sh", &context)?;
+            api_client
+                .upload(current_system, script_path.clone(), script.into_bytes())
+                .await?;
+
+            Ok(script_path)
+        }
+        ScriptSpec::Local(local_path) => {
+            let script = std::fs::read_to_string(local_path)?;
+            api_client
+                .upload(current_system, script_path.clone(), script.into_bytes())
+                .await?;
+
+            Ok(script_path)
+        }
+        ScriptSpec::Remote(script_path) => Ok(script_path),
+    }
+}
 
-#[allow(clippy::too_many_arguments)]
 pub async fn cscs_job_start(
     name: Option<String>,
     options: JobStartOptions,
@@ -276,6 +317,7 @@ pub async fn cscs_job_start(
 
             let mut envvars = config.cscs.env.clone();
             envvars.extend(options.env.clone());
+
             let environment_path = handle_edf(
                 &api_client,
                 &base_path,
@@ -286,27 +328,16 @@ pub async fn cscs_job_start(
             )
             .await?;
 
-            // upload script
-            let script_path = base_path.join("script.sh");
-            let script_template = options
-                .script_file
-                .clone()
-                .map(std::fs::read_to_string)
-                .unwrap_or(Ok(config.cscs.sbatch_script_template))?;
-            let mut tera = tera::Tera::default();
-            tera.add_raw_template("script.sh", &script_template)?;
-            let mut context = tera::Context::new();
-            context.insert("name", &job_name);
-            context.insert(
-                "command",
-                &options.command.clone().unwrap_or(config.cscs.command).join(" "),
-            );
-            context.insert("environment_file", &environment_path);
-            context.insert("container_workdir", &container_workdir);
-            let script = tera.render("script.sh", &context)?;
-            api_client
-                .upload(current_system, script_path.clone(), script.into_bytes())
-                .await?;
+            let script_path = handle_script(
+                &api_client,
+                &job_name,
+                &base_path,
+                current_system,
+                &environment_path,
+                &container_workdir,
+                &options,
+            )
+            .await?;
 
             // start job
             api_client

coman/src/main.rs

@@ -73,7 +73,6 @@ async fn main() -> Result<()> {
                     }
                     cli::CscsJobCommands::Submit {
                         name,
-                        script_file,
                         image,
                         command,
                         workdir,
@@ -82,11 +81,11 @@ async fn main() -> Result<()> {
                         stdout,
                         stderr,
                         edf_spec,
+                        script_spec,
                     } => {
                         cli_cscs_job_start(
                             name,
                             JobStartOptions {
-                                script_file,
                                 image,
                                 command,
                                 container_workdir: workdir,
@@ -95,6 +94,7 @@ async fn main() -> Result<()> {
                                 stdout,
                                 stderr,
                                 edf_spec: edf_spec.unwrap_or_default().into(),
+                                script_spec: script_spec.unwrap_or_default().into(),
                             },
                             system,
                             platform,