Refactor app specs using with_prefixed_args

Author: lyskov

src/app.rs

@@ -63,7 +63,7 @@ pub enum MountRole {
     Scratch,
 }
 
-enum IoSpec {
+pub enum IoSpec {
     InputDir(PathBuf),
     InputDirOption(String),
 }
@@ -76,8 +76,8 @@ pub struct ContainerRunSpec {
 
 pub struct NativeRunSpec {
     pub pixi: Cow<'static, str>,
+    pub io_spec: IoSpec,
     pub args: Vec<String>,
-    //pub io_spec: IoSpec,
 }
 
 pub struct RunSpec {
@@ -93,6 +93,30 @@ impl ContainerRunSpec {
             mounts: HashMap::new(),
         }
     }
+    pub fn with_prefixed_args<I1, I2, S1, S2>(
+        image: impl Into<String>,
+        prefixes: I1,
+        args: I2,
+    ) -> Self
+    where
+        I1: IntoIterator<Item = S1>,
+        I2: IntoIterator<Item = S2>,
+        S1: Into<String>,
+        S2: Into<String>,
+    {
+        let full_args: Vec<String> = prefixes
+            .into_iter()
+            .map(Into::into)
+            .chain(args.into_iter().map(Into::into))
+            .collect();
+
+        Self {
+            image: Image(image.into()),
+            args: full_args,
+            mounts: HashMap::new(),
+        }
+    }
+
     pub fn scratch(mut self, p: impl Into<String>) -> Self {
         self.mounts.insert(MountRole::Scratch, p.into());
         self
@@ -104,16 +128,17 @@ impl ContainerRunSpec {
 }
 
 impl NativeRunSpec {
-    pub fn new(pixi: impl Into<Cow<'static, str>>, args: Vec<String>) -> Self {
+    pub fn new(pixi: impl Into<Cow<'static, str>>, io_spec: IoSpec, args: Vec<String>) -> Self {
         Self {
             pixi: pixi.into(),
+            io_spec,
             args,
         }
     }
-    pub fn pixi(mut self, p: impl Into<Cow<'static, str>>) -> Self {
-        self.pixi = p.into();
-        self
-    }
+    // pub fn pixi(mut self, p: impl Into<Cow<'static, str>>) -> Self {
+    //     self.pixi = p.into();
+    //     self
+    // }
 }
 
 impl RunSpec {
@@ -136,3 +161,62 @@ impl App {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_with_prefixed_args() {
+        let prefixes = vec!["--prefix1", "--prefix2"];
+        let args = vec!["arg1", "arg2", "arg3"];
+
+        let spec = ContainerRunSpec::with_prefixed_args(
+            "test/image:latest",
+            prefixes.clone(),
+            args.clone(),
+        );
+
+        assert_eq!(spec.image.0, "test/image:latest");
+        assert_eq!(
+            spec.args,
+            vec!["--prefix1", "--prefix2", "arg1", "arg2", "arg3"]
+        );
+        assert_eq!(spec.mounts.len(), 0);
+    }
+
+    #[test]
+    fn test_with_prefixed_args_empty_prefixes() {
+        let prefixes: Vec<String> = vec![];
+        let args = vec!["arg1".to_string(), "arg2".to_string()];
+
+        let spec = ContainerRunSpec::with_prefixed_args("test/image:v1", prefixes, args.clone());
+
+        assert_eq!(spec.image.0, "test/image:v1");
+        assert_eq!(spec.args, vec!["arg1", "arg2"]);
+    }
+
+    #[test]
+    fn test_with_prefixed_args_empty_args() {
+        let prefixes = vec!["--flag".to_string()];
+        let args: Vec<String> = vec![];
+
+        let spec = ContainerRunSpec::with_prefixed_args("test/image", prefixes.clone(), args);
+
+        assert_eq!(spec.image.0, "test/image");
+        assert_eq!(spec.args.len(), 1);
+        assert_eq!(spec.args[0], "--flag");
+    }
+
+    #[test]
+    fn test_with_prefixed_args_both_empty() {
+        let prefixes: Vec<String> = vec![];
+        let args: Vec<String> = vec![];
+
+        let spec = ContainerRunSpec::with_prefixed_args("empty/image", prefixes, args);
+
+        assert_eq!(spec.image.0, "empty/image");
+        assert_eq!(spec.args.len(), 0);
+        assert!(spec.mounts.is_empty());
+    }
+}

src/app/ligandmpnn.rs

@@ -1,17 +1,14 @@
 use crate::app::{ContainerRunSpec, RunSpec};
 
 pub fn spec(app_args: Vec<String>) -> RunSpec {
-    let container = ContainerRunSpec::new(
+    let container = ContainerRunSpec::with_prefixed_args(
         "rosettacommons/ligandmpnn",
         [
             "--out_folder=/w",
             "--checkpoint_protein_mpnn",
             "/app/ligandmpnn/model_params/proteinmpnn_v_48_020.pt",
-        ]
-        .into_iter()
-        .map(Into::into)
-        .chain(app_args)
-        .collect(),
+        ],
+        app_args,
     )
     .working_dir("/w");
 

src/app/picap.rs

@@ -1,17 +1,14 @@
 use crate::app::{ContainerRunSpec, RunSpec};
 
 pub fn spec(app_args: Vec<String>) -> RunSpec {
-    let container = ContainerRunSpec::new(
+    let container = ContainerRunSpec::with_prefixed_args(
         "rosettacommons/picap",
         [
             "--out_folder=/w",
             "--checkpoint_protein_mpnn",
             "/app/ligandmpnn/model_params/proteinmpnn_v_48_020.pt",
-        ]
-        .into_iter()
-        .map(String::from)
-        .chain(app_args)
-        .collect(),
+        ],
+        app_args,
     )
     .working_dir("/w");
 

src/app/proteinmpnn.rs

@@ -1,13 +1,10 @@
 use crate::app::{ContainerRunSpec, RunSpec};
 
 pub fn spec(app_args: Vec<String>) -> RunSpec {
-    let container = ContainerRunSpec::new(
+    let container = ContainerRunSpec::with_prefixed_args(
         "rosettacommons/proteinmpnn",
-        ["--out_folder=/w"]
-            .into_iter()
-            .map(Into::into)
-            .chain(app_args)
-            .collect(),
+        ["--out_folder=/w"],
+        app_args,
     )
     .working_dir("/w");
 

src/app/pyrosetta.rs

@@ -1,15 +1,9 @@
 use crate::app::{ContainerRunSpec, RunSpec};
 
 pub fn spec(app_args: Vec<String>) -> RunSpec {
-    let container = ContainerRunSpec::new(
-        "rosettacommons/rosetta:serial",
-        ["python"]
-            .into_iter()
-            .map(Into::into)
-            .chain(app_args)
-            .collect(),
-    )
-    .working_dir("/w");
+    let container =
+        ContainerRunSpec::with_prefixed_args("rosettacommons/rosetta:serial", ["python"], app_args)
+            .working_dir("/w");
 
     RunSpec::new(container, None)
 }

src/app/rfdiffusion.rs

@@ -1,24 +1,25 @@
 use crate::{
-    app::{ContainerRunSpec, NativeRunSpec, RunSpec},
+    app::{ContainerRunSpec, IoSpec, NativeRunSpec, RunSpec},
     util::include_asset,
 };
 
 pub fn spec(app_args: Vec<String>) -> RunSpec {
-    let container = ContainerRunSpec::new(
+    let container = ContainerRunSpec::with_prefixed_args(
         "rosettacommons/rfdiffusion",
         [
             "inference.output_prefix=/w/",
             "inference.model_directory_path=/app/RFdiffusion/models",
-        ]
-        .into_iter()
-        .map(Into::into)
-        .chain(app_args.clone())
-        .collect(),
+        ],
+        app_args.clone(),
     )
     .scratch("/app/RFdiffusion/schedules")
     .working_dir("/w");
 
-    let native = NativeRunSpec::new(include_asset!("pixi/rfdiffusion.toml"), app_args);
+    let native = NativeRunSpec::new(
+        include_asset!("pixi/rfdiffusion.toml"),
+        IoSpec::InputDir("aaa".into()),
+        app_args,
+    );
 
     RunSpec::new(container, Some(native))
 }

src/app/score.rs

@@ -1,15 +1,9 @@
 use crate::app::{ContainerRunSpec, RunSpec};
 
 pub fn spec(app_args: Vec<String>) -> RunSpec {
-    let container = ContainerRunSpec::new(
-        "rosettacommons/rosetta:serial",
-        ["score"]
-            .into_iter()
-            .map(Into::into)
-            .chain(app_args)
-            .collect(),
-    )
-    .working_dir("/w");
+    let container =
+        ContainerRunSpec::with_prefixed_args("rosettacommons/rosetta:serial", ["score"], app_args)
+            .working_dir("/w");
 
     RunSpec::new(container, None)
 }