add env and workdir options to job submission, add command for setting current cscs system, add command to init coman config

Author: Panaetius

coman/.config/config.toml

@@ -16,8 +16,17 @@ srun --environment={{environment_file}} {{command}}
 edf_file_template = """
 image = "{{edf_image}}"
 mounts = ["${SCRATCH}:/scratch"]
+workdir = "{{container_workdir}}"
+
+[env]
+{% for key, value in env %}
+{{key}} = "{{value}}"
+{% endfor %}
 """
 
+[cscs.env]
+
+
 [cscs.systems]
 
 [cscs.systems.daint]

coman/src/cli.rs

@@ -1,4 +1,4 @@
-use std::path::PathBuf;
+use std::{error::Error, path::PathBuf};
 
 use clap::{Parser, Subcommand};
 
@@ -14,6 +14,11 @@ pub enum CliCommands {
         #[command(subcommand)]
         command: CscsCommands,
     },
+    #[clap(about = "Create a new project configuration file")]
+    Init {
+        #[clap(help = "Destination folder to create config in (default = current directory)")]
+        destination: Option<PathBuf>,
+    },
 }
 
 #[derive(Subcommand, Debug)]
@@ -38,20 +43,40 @@ pub enum CscsJobCommands {
     Get { job_id: i64 },
     #[clap(alias("s"))]
     Submit {
-        #[clap(short, long)]
+        #[clap(short, long, help = "the path to the srun script file to use")]
         script_file: Option<PathBuf>,
-        #[clap(short, long)]
+        #[clap(
+            short,
+            long,
+            help = "the working directory path inside the container (note this is different from the working directory that the srun command is executed from)"
+        )]
+        workdir: Option<String>,
+        #[clap(short='E', value_name="KEY=VALUE", value_parser=parse_key_val::<String,String>, help="Environment variables to set in the container")]
+        env: Vec<(String, String)>,
+        #[clap(short, long, help = "The docker image to use")]
         image: Option<DockerImageUrl>,
-        #[clap(short, long, trailing_var_arg = true)]
+        #[clap(trailing_var_arg = true, help = "The command to run in the container")]
         command: Option<Vec<String>>,
     },
     #[clap(alias("st"))]
     Stop { job_id: i64 },
 }
 #[derive(Subcommand, Debug)]
 pub enum CscsSystemCommands {
-    #[clap(alias("ls"))]
+    #[clap(alias("ls"), about = "List available compute systems")]
     List,
+    #[clap(alias("s"), about = "Set system to use")]
+    Set {
+        #[clap(
+            short,
+            long,
+            action,
+            help = "set in global config instead of project-local one"
+        )]
+        global: bool,
+        #[clap(help = "System name to use")]
+        system_name: String,
+    },
 }
 
 #[derive(Parser, Debug)]
@@ -95,3 +120,16 @@ Config directory: {config_dir_path}
 Data directory: {data_dir_path}"
     )
 }
+
+fn parse_key_val<T, U>(s: &str) -> Result<(T, U), Box<dyn Error + Send + Sync + 'static>>
+where
+    T: std::str::FromStr,
+    T::Err: Error + Send + Sync + 'static,
+    U: std::str::FromStr,
+    U::Err: Error + Send + Sync + 'static,
+{
+    let pos = s
+        .find('=')
+        .ok_or_else(|| format!("invalid KEY=value: no `=` found in `{s}`"))?;
+    Ok((s[..pos].parse()?, s[pos + 1..].parse()?))
+}

coman/src/config.rs

@@ -4,6 +4,7 @@ use std::{collections::HashMap, env, path::PathBuf};
 
 use color_eyre::Result;
 use directories::ProjectDirs;
+use eyre::eyre;
 use lazy_static::lazy_static;
 use serde::{Deserialize, Serialize};
 
@@ -27,10 +28,12 @@ pub struct CscsConfig {
     #[serde(default)]
     pub current_system: String,
     #[serde(default)]
-    pub name: Option<String>,
-    #[serde(default)]
     pub sbatch_script_template: String,
     #[serde(default)]
+    pub workdir: Option<String>,
+    #[serde(default)]
+    pub env: HashMap<String, String>,
+    #[serde(default)]
     pub image: String,
     #[serde(default)]
     pub edf_file_template: String,
@@ -43,6 +46,8 @@ pub struct CscsConfig {
 
 #[derive(Clone, Debug, Default, Serialize, Deserialize)]
 pub struct Config {
+    #[serde(default)]
+    pub name: Option<String>,
     #[serde(default, flatten)]
     pub config: AppConfig,
     #[serde(default)]
@@ -59,71 +64,123 @@ lazy_static! {
         env::var(format!("{}_CONFIG", PROJECT_NAME.clone()))
             .ok()
             .map(PathBuf::from);
+    pub static ref CONFIG_FILE_NAME: String = format!("{}.toml", PROJECT_NAME.clone());
+    pub static ref CONFIG_FORMAT: config::FileFormat = config::FileFormat::Toml;
 }
 
 impl Config {
-    pub fn new() -> Result<Self, config::ConfigError> {
-        let data_dir = get_data_dir();
-        let config_dir = get_config_dir();
-        let mut builder = config::Config::builder()
-            .add_source(config::File::from_str(
-                DEFAULT_CONFIG_TOML,
-                config::FileFormat::Toml,
-            ))
-            .set_default("data_dir", data_dir.to_str().unwrap())?
-            .set_default("config_dir", config_dir.to_str().unwrap())?;
-
-        let config_files = [
-            (
-                format!("{}.toml", PROJECT_NAME.to_lowercase()),
-                config::FileFormat::Toml,
-            ),
-            (
-                format!("{}.json5", PROJECT_NAME.to_lowercase()),
-                config::FileFormat::Json5,
-            ),
-            (
-                format!("{}.json", PROJECT_NAME.to_lowercase()),
-                config::FileFormat::Json,
-            ),
-            (
-                format!("{}.yaml", PROJECT_NAME.to_lowercase()),
-                config::FileFormat::Yaml,
-            ),
-            (
-                format!("{}.ini", PROJECT_NAME.to_lowercase()),
-                config::FileFormat::Ini,
-            ),
-        ];
-        for (file, format) in &config_files {
-            let source = config::File::from(config_dir.join(file))
-                .format(*format)
-                .required(false);
-            builder = builder.add_source(source);
-        }
+    pub fn new() -> Result<Self> {
+        let builder = default_config_builder()?;
+        let builder = global_config_builder(builder)?;
+        let builder = project_local_config_builder(builder)?;
 
-        // find config override in current directory
-        let mut search_path = std::env::current_dir().expect("current directory does not exist");
-        loop {
-            for (file, format) in &config_files {
-                if search_path.join(file).exists() {
-                    let source = config::File::from(search_path.join(file))
-                        .format(*format)
-                        .required(false);
-                    builder = builder.add_source(source);
-                    break;
-                }
-            }
-            if let Some(p) = search_path.parent() {
-                search_path = p.to_path_buf();
-            } else {
-                break;
-            }
-        }
+        let cfg: Self = builder.build()?.try_deserialize()?;
+        Ok(cfg)
+    }
+    pub fn new_global() -> Result<Self> {
+        let builder = default_config_builder()?;
+        let builder = global_config_builder(builder)?;
 
         let cfg: Self = builder.build()?.try_deserialize()?;
         Ok(cfg)
     }
+
+    pub fn write_local(&self) -> Result<()> {
+        match get_project_local_config_file() {
+            Some(path) => {
+                let content = toml::to_string_pretty(self)?;
+                std::fs::write(path, content)?;
+                Ok(())
+            }
+            None => Err(eyre!(
+                "No config file exists in current project. Consider creating one using '{} init",
+                PROJECT_NAME.to_lowercase().clone()
+            )),
+        }
+    }
+
+    pub fn write_global(&self) -> Result<()> {
+        let config_dir = get_config_dir();
+        let path = config_dir.join(CONFIG_FILE_NAME.clone());
+        let content = toml::to_string_pretty(self)?;
+        std::fs::write(path, content)?;
+        Ok(())
+    }
+
+    pub fn create_config(destination: Option<PathBuf>) -> Result<()> {
+        let mut config = Config::new()?;
+        let project_dir = destination
+            .unwrap_or(std::env::current_dir().expect("current directory does not exist"));
+        if !project_dir.exists() || !project_dir.is_dir() {
+            return Err(eyre!(
+                "destination must exist and be a directory, got {}",
+                project_dir.to_string_lossy()
+            ));
+        }
+
+        let name = project_dir
+            .file_name()
+            .expect("could not get base name from destination");
+        config.name = Some(name.to_string_lossy().to_string());
+
+        let config_path = project_dir.join(CONFIG_FILE_NAME.clone());
+        std::fs::write(config_path.clone(), "")?;
+        let content = toml::to_string_pretty(&config)?;
+        std::fs::write(config_path, content)?;
+        Ok(())
+    }
+}
+
+pub fn default_config_builder() -> Result<config::ConfigBuilder<config::builder::DefaultState>> {
+    let data_dir = get_data_dir();
+    let config_dir = get_config_dir();
+    let builder = config::Config::builder()
+        .add_source(config::File::from_str(
+            DEFAULT_CONFIG_TOML,
+            config::FileFormat::Toml,
+        ))
+        .set_default("data_dir", data_dir.to_str().unwrap())?
+        .set_default("config_dir", config_dir.to_str().unwrap())?;
+    Ok(builder)
+}
+
+pub fn global_config_builder(
+    builder: config::ConfigBuilder<config::builder::DefaultState>,
+) -> Result<config::ConfigBuilder<config::builder::DefaultState>> {
+    let config_dir = get_config_dir();
+    let source = config::File::from(config_dir.join(CONFIG_FILE_NAME.clone()))
+        .format(*CONFIG_FORMAT)
+        .required(false);
+    let builder = builder.add_source(source);
+    Ok(builder)
+}
+
+pub fn project_local_config_builder(
+    builder: config::ConfigBuilder<config::builder::DefaultState>,
+) -> Result<config::ConfigBuilder<config::builder::DefaultState>> {
+    if let Some(config_path) = get_project_local_config_file() {
+        let source = config::File::from(config_path)
+            .format(*CONFIG_FORMAT)
+            .required(false);
+        let builder = builder.add_source(source);
+        return Ok(builder);
+    }
+    Ok(builder)
+}
+
+pub fn get_project_local_config_file() -> Option<PathBuf> {
+    let mut search_path = std::env::current_dir().expect("current directory does not exist");
+    loop {
+        if search_path.join(CONFIG_FILE_NAME.clone()).exists() {
+            return Some(search_path.join(CONFIG_FILE_NAME.clone()));
+        }
+        if let Some(p) = search_path.parent() {
+            search_path = p.to_path_buf();
+        } else {
+            break;
+        }
+    }
+    None
 }
 
 pub fn get_data_dir() -> PathBuf {

coman/src/cscs/api_client.rs

@@ -15,7 +15,7 @@ use firecrest_client::{
         JobMetadataModel, JobModelOutput, SchedulerServiceHealth, UserInfoResponse,
     },
 };
-use std::path::PathBuf;
+use std::{collections::HashMap, path::PathBuf};
 use strum::Display;
 
 use crate::trace_dbg;
@@ -263,6 +263,7 @@ impl CscsApi {
         system_name: &str,
         name: &str,
         script_path: PathBuf,
+        envvars: HashMap<String, String>,
     ) -> Result<()> {
         let workingdir = script_path.clone();
         let workingdir = workingdir.parent();
@@ -273,6 +274,7 @@ impl CscsApi {
             None,
             Some(script_path),
             workingdir.map(|p| p.to_path_buf()),
+            envvars,
         )
         .await?;
         let _ = trace_dbg!(result);

coman/src/cscs/cli.rs

@@ -5,7 +5,8 @@ use std::path::PathBuf;
 use crate::{
     cscs::{
         handlers::{
-            cscs_job_stop, cscs_job_details, cscs_job_list, cscs_start_job, cscs_system_list,
+            cscs_job_details, cscs_job_list, cscs_job_stop, cscs_start_job, cscs_system_list,
+            cscs_system_set,
         },
         oauth2::{CLIENT_ID_SECRET_NAME, CLIENT_SECRET_SECRET_NAME, client_credentials_login},
     },
@@ -93,8 +94,10 @@ pub(crate) async fn cli_cscs_job_start(
     script_file: Option<PathBuf>,
     image: Option<DockerImageUrl>,
     command: Option<Vec<String>>,
+    workdir: Option<String>,
+    env: Vec<(String, String)>,
 ) -> Result<()> {
-    cscs_start_job(script_file, image, command).await
+    cscs_start_job(script_file, image, command, workdir, env).await
 }
 
 pub(crate) async fn cli_cscs_job_stop(job_id: i64) -> Result<()> {
@@ -112,3 +115,6 @@ pub(crate) async fn cli_cscs_system_list() -> Result<()> {
         Err(e) => Err(e),
     }
 }
+pub(crate) async fn cli_cscs_set_system(system_name: String, global: bool) -> Result<()> {
+    cscs_system_set(system_name, global).await
+}