Add Foundry integration and tests

Author: lyskov

assets/pixi/rfdiffusion.toml

@@ -5,21 +5,29 @@ description = "Pixi environment for https://github.com/RosettaCommons/rc RFdiffu
 platforms = ["linux-64"]
 #platforms = ["linux-64", "osx-64", "osx-arm64", "win-64"]
 
-channels = ["pytorch", "nvidia", "conda-forge", "dglteam"]
+channels = ["conda-forge"]
 
 [dependencies]
 python = "==3.9"
 pip = "*"
 
 numpy = "1.*"
 
-pytorch = { version = "==1.12.1", channel = "pytorch" }
-pytorch-cuda = { version = "==11.6", channel = "pytorch" }
-"dgl-cuda11.6" = "*"
+cudatoolkit = "11.6.*"
 
 mkl = "<2024.1"
 
+[pypi-options]
+index-url = "https://pypi.org/simple"
+extra-index-urls = ["https://download.pytorch.org/whl/cu116"]
+
+# DGL uses a "find-links" wheel repo (equivalent of pip -f ...)
+find-links = [{ url = "https://data.dgl.ai/wheels/cu116/repo.html" }]
+
 [pypi-dependencies]
+torch = "==1.12.1+cu116"
+dgl = "==1.0.2+cu116"
+
 e3nn = "==0.3.3"
 wandb = "==0.12.0"
 pynvml = "==11.0.0"
@@ -28,6 +36,7 @@ decorator = "==5.1.0"
 hydra-core = "==1.3.2"
 pyrsistent = "==0.19.3"
 
+
 [tasks]
 setup = { depends-on = ["install", "download_weights"] }
 
@@ -43,15 +52,19 @@ cmd = """
 ( rm -rf rfdiffusion-repo-clone/models ; mkdir -p rfdiffusion-repo-clone/models ) \
 && wget -O rfdiffusion-repo-clone/models/Base_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/6f5902ac237024bdd0c176cb93063dc4/Base_ckpt.pt \
 && wget -O rfdiffusion-repo-clone/models/Complex_base_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/e29311f6f1bf1af907f9ef9f44b8328b/Complex_base_ckpt.pt \
+&& wget -O rfdiffusion-repo-clone/models/Complex_beta_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/f572d396fae9206628714fb2ce00f72e/Complex_beta_ckpt.pt \
+\
 && wget -O rfdiffusion-repo-clone/models/Complex_Fold_base_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/60f09a193fb5e5ccdc4980417708dbab/Complex_Fold_base_ckpt.pt \
 && wget -O rfdiffusion-repo-clone/models/InpaintSeq_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/74f51cfb8b440f50d70878e05361d8f0/InpaintSeq_ckpt.pt \
 && wget -O rfdiffusion-repo-clone/models/InpaintSeq_Fold_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/76d00716416567174cdb7ca96e208296/InpaintSeq_Fold_ckpt.pt \
 && wget -O rfdiffusion-repo-clone/models/ActiveSite_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/5532d2e1f3a4738decd58b19d633b3c3/ActiveSite_ckpt.pt \
 && wget -O rfdiffusion-repo-clone/models/Base_epoch8_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/12fc204edeae5b57713c5ad7dcb97d39/Base_epoch8_ckpt.pt \
-&& wget -O rfdiffusion-repo-clone/models/Complex_beta_ckpt.pt http://files.ipd.uw.edu/pub/RFdiffusion/f572d396fae9206628714fb2ce00f72e/Complex_beta_ckpt.pt \
 && wget -O rfdiffusion-repo-clone/models/RF_structure_prediction_weights.pt http://files.ipd.uw.edu/pub/RFdiffusion/1befcb9b28e2f778f53d47f18b7597fa/RF_structure_prediction_weights.pt \
 """
 
+[activation]
+env = { DGLBACKEND = "pytorch", LD_LIBRARY_PATH = "$CONDA_PREFIX/lib:$LD_LIBRARY_PATH" }
+
 [tasks.execute]
 args = ["args"]
 cmd = "python scripts/run_inference.py {{args}}"

src/app.rs

@@ -1,3 +1,4 @@
+mod foundry;
 mod ligandmpnn;
 mod picap;
 mod proteinmpnn;
@@ -41,6 +42,10 @@ pub enum App {
     #[value(aliases = ["LigandMPNN"])]
     Ligandmpnn,
 
+    /// Run the Foundry command https://github.com/RosettaCommons/foundry
+    #[value(aliases = ["Foundry"])]
+    Foundry,
+
     /// Run the PiCAP/CAPSIF2 command https://github.com/Graylab/picap
     #[value(aliases = ["PiCAP", "CAPSIF2"])]
     Picap,
@@ -174,6 +179,7 @@ impl App {
             App::Proteinmpnn => proteinmpnn::container_spec(app_args),
             App::ProteinmpnnScript => proteinmpnn_script::container_spec(app_args),
             App::Ligandmpnn => ligandmpnn::container_spec(app_args),
+            App::Foundry => foundry::container_spec(app_args),
             App::Picap => picap::container_spec(app_args),
         }
     }
@@ -187,7 +193,8 @@ impl App {
             App::Proteinmpnn => todo!("not implemented"), // proteinmpnn::native_spec(app_args),
             App::ProteinmpnnScript => todo!("not implemented"), // proteinmpnn_script::native_spec(app_args),
             App::Ligandmpnn => todo!("not implemented"), // ligandmpnn::native_spec(app_args),
-            App::Picap => todo!("not implemented"),      // picap::native_spec(app_args),
+            App::Foundry => foundry::native_spec(app_args, working_dir),
+            App::Picap => todo!("not implemented"), // picap::native_spec(app_args),
         }
     }
 }

src/app/foundry.rs

@@ -0,0 +1,40 @@
+use std::path::Path;
+
+use crate::{
+    app::{ContainerRunSpec, NativeRunSpec},
+    util::include_asset,
+};
+
+pub fn container_spec(app_args: Vec<String>) -> ContainerRunSpec {
+    dbg!(&app_args);
+    assert!(
+        !app_args.is_empty() || !app_args[0].starts_with("-"),
+        "Foundry arguments must include a protocol name as first argument"
+    );
+
+    // match app_args[0].as_str() {
+    //     "rfd3" => {
+    //         // case 1
+    //         // ckpt_path=
+    //     }
+    //     "mpnn" => {
+    //         // case 2
+    //         // --checkpoint_path
+    //     }
+    //     _ => {
+    //         // default case
+    //     }
+    // }
+
+    // ContainerRunSpec::with_prefixed_args(
+    //     "rosettacommons/foundry:weights",
+    //     std::iter::empty::<&str>(),
+    //     app_args.clone(),
+    // )
+
+    ContainerRunSpec::new("rosettacommons/foundry:weights", app_args).working_dir("/w")
+}
+
+pub fn native_spec(app_args: Vec<String>, _working_dir: &Path) -> NativeRunSpec {
+    NativeRunSpec::new(include_asset!("pixi/foundry.toml"), app_args)
+}

src/app/proteinmpnn_script.rs

@@ -2,7 +2,7 @@ use crate::app::ContainerRunSpec;
 
 pub fn container_spec(mut app_args: Vec<String>) -> ContainerRunSpec {
     assert!(
-        app_args.is_empty() || app_args[0].starts_with("-"),
+        !app_args.is_empty() || !app_args[0].starts_with("-"),
         "ProteinmpnnScript arguments must include a script name as first argument"
     );
 

tests/foundry.rs

@@ -0,0 +1,44 @@
+use assert_cmd::{assert::OutputAssertExt, cargo::cargo_bin_cmd};
+use assert_fs::TempDir;
+
+mod common;
+
+common::engine_tests!(foundry);
+
+fn foundry(engine: &str) {
+    use assert_fs::assert::PathAssert;
+
+    let root = std::path::PathBuf::from("target/foundry").join(engine);
+    std::fs::create_dir_all(&root).expect("create engine testing dir");
+    let work_dir = TempDir::new_in(root).expect("create temp dir");
+
+    let json_rfd3_path = work_dir.join("test_rfd3.json");
+    std::fs::write(json_rfd3_path, r#"{ "foundry": { "length": "10" } }"#)
+        .expect("write test_rfd3.json");
+
+    for (i, c) in [
+        "rfd3 out_dir=/w/rfd3_out/ inputs=/w/test_rfd3.json skip_existing=False prevalidate_inputs=True ckpt_path=/weights/rfd3_latest.ckpt n_batches=1 diffusion_batch_size=1 inference_sampler.num_timesteps=10 low_memory_mode=True global_prefix=test_",
+        "mpnn --structure_path /w/rfd3_out/test_foundry_0_model_0.cif.gz --checkpoint_path /weights/ligandmpnn_v_32_010_25.pt --is_legacy_weights True --model_type ligand_mpnn  --out_directory /w/mpnn_out",
+        "rf3 fold inputs=/w/mpnn_out/test_foundry_0_model_0.cif_b0_d0.cif ckpt_path=/weights/rf3_foundry_01_24_latest_remapped.ckpt diffusion_batch_size=1 num_steps=10 out_dir=/w/rf3_out",
+        ].iter()
+        .enumerate() {
+        let cmd = cargo_bin_cmd!()
+            .args([
+                "run",
+                "--container-engine",
+                engine,
+                "-w",
+                work_dir.path().to_str().unwrap(),
+                "foundry",
+            ])
+            .args(c.split_ascii_whitespace())
+            .unwrap();
+        cmd.assert().success();
+
+        use assert_fs::prelude::PathChild;
+
+        work_dir
+            .child(format!(".000{i}.rc.log"))
+            .assert(predicates::path::exists());
+    }
+}