Support for Singularity containers and slurm

README.md

@@ -51,6 +51,71 @@ cd LigandMPNN && bash get_model_params.sh ./model_params && cd ..
 
 ---
 
+## Singularity / Apptainer container
+
+For GPU clusters, a pre-built environment is provided under `containers/switchcraft.def`. The image installs SwitchCraft, Boltz, cuEquivariance, **Boltz weights** (`boltz/boltz1_conf.ckpt`, `boltz/ccd.pkl`), and **LigandMPNN weights** at build time, so routine runs do not need extra weight bind-mounts.
+
+### Build the image
+
+From the repository root (Linux; build requires network access):
+
+```bash
+apptainer build switchcraft.sif containers/switchcraft.def
+# or: singularity build switchcraft.sif containers/switchcraft.def
+```
+
+Build-time `%test` may print `libcuda.so.1` warnings and `cuda available False` without a GPU; that is expected. Use `--nv` at **runtime** on a GPU node.
+
+### Minimal interactive run
+
+On a GPU node, only the **output directory** on the host needs to be bound. Use config paths relative to `/opt/switchcraft` (same as the Quick Start examples under `tasks/`):
+
+```bash
+mkdir -p /scratch/${USER}/switchcraft_out
+
+apptainer exec --nv --pwd /opt/switchcraft \
+  -B /scratch/${USER}/switchcraft_out:/scratch/${USER}/switchcraft_out \
+  switchcraft.sif \
+  python switchcraft.py \
+    --config tasks/pos_allostery.yaml \
+    --outpath /scratch/${USER}/switchcraft_out \
+    --verbose
+```
+
+Built-in configs (`tasks/pos_allostery.yaml`, etc.) and `motifs/` are already in the image; you do not need to copy them unless you are customizing a design.
+
+**Custom YAML on the host:** bind the file into the image and pass its path under `/opt/switchcraft`, e.g. `-B /scratch/${USER}/my_design.yaml:/opt/switchcraft/my_design.yaml` and `--config my_design.yaml` (still use `--pwd /opt/switchcraft`).
+
+### SLURM
+
+`slurm_scripts/launch_switchcraft.slurm` wraps the same invocation. Replace the `#SBATCH` placeholders (`<partition>`, `<gpus_per_node>`, etc.) for your site, then submit:
+
+```bash
+mkdir -p /scratch/${USER}/switchcraft_out
+
+sbatch slurm_scripts/launch_switchcraft.slurm \
+  tasks/pos_allostery.yaml \
+  /path/to/switchcraft.sif \
+  /scratch/${USER}/switchcraft_out
+```
+
+| Argument | Meaning |
+|----------|---------|
+| `CONFIG_FILE` | YAML path inside the image (e.g. `tasks/pos_allostery.yaml`) |
+| `SWITCHCRAFT_IMAGE` | Path to `switchcraft.sif` |
+| `OUTPUT_DIR` | Host directory for results (bind-mounted) |
+
+**Interactive test** (e.g. inside `salloc` on a GPU node, no `sbatch`):
+
+```bash
+bash slurm_scripts/launch_switchcraft.slurm \
+  tasks/pos_allostery.yaml \
+  /path/to/switchcraft.sif \
+  /scratch/${USER}/switchcraft_out
+```
+
+---
+
 ## Quick Start
 
 Each design run is driven by a YAML config file that fully specifies the design problem. Template configs for example tasks live in `tasks/`.
@@ -74,7 +139,7 @@ Command line arguments that can be provided include
 - `--recycles`: Boltz recycling steps during optimization
 - `--ligandmpnn_seqs`: if >0, run LigandMPNN redesign producing N sequences per design
 - `--verbose`: whether to use progress bars (recommended)
-- `--out`: output directory
+- `-o` / `--outpath`: output directory (README shorthand: output dir)
 - `--num_designs`: total number of designs
 
 ---

containers/switchcraft.def

@@ -0,0 +1,87 @@
+Bootstrap: docker
+From: python:3.12-slim
+
+
+# Singularity Image for SwitchCraft
+#
+%files 
+# Sections where you can specify which files from host you want to copy in the Image.
+# For bigger files (e.g., models weights) you can: 
+# --bind host_path:image_path 
+    . /opt/switchcraft
+
+%post
+    export DEBIAN_FRONTEND=noninteractive
+    apt-get update && apt-get install -y --no-install-recommends \
+        git build-essential wget ca-certificates \
+        && rm -rf /var/lib/apt/lists/*
+
+    pip install --no-cache-dir --upgrade pip
+    pip install --no-cache-dir torch torchvision \
+        --index-url https://download.pytorch.org/whl/cu128
+    # Boltz imports cuequivariance_torch at module load (not in boltz pyproject deps)
+    pip install --no-cache-dir cuequivariance-torch cuequivariance-ops-torch-cu12
+    pip install -e /opt/switchcraft/boltz
+    pip install prody
+
+    python << 'PY'
+from boltz.main import download
+from pathlib import Path
+cache = Path("/opt/switchcraft/boltz")
+cache.mkdir(parents=True, exist_ok=True)
+download(cache)
+PY
+    cd /opt/switchcraft/LigandMPNN && bash get_model_params.sh ./model_params  
+
+    echo "SwitchCraft container built successfully"
+
+%environment
+    export BOLTZ_USE_CUEQ=1
+
+
+%runscript
+    cd /opt/switchcraft || exit 1
+    exec python switchcraft.py "$@"
+
+%test
+# Runs at end of build. Boltz/LigandMPNN weights are baked in under %post.
+    echo "Testing PyTorch and CUDA"
+    python -c "import torch; print('torch', torch.__version__); print('cuda available', torch.cuda.is_available())"
+
+    echo "Testing cuequivariance_torch"
+    python -c "import cuequivariance_torch; print('cuequivariance_torch ok')"
+
+    echo "Testing boltz import"
+    python -c "import boltz; print('boltz ok')"
+
+    echo "Testing switchcraft imports (weights optional at build)"
+    cd /opt/switchcraft
+    python << 'EOF'
+import os
+import yaml
+from boltz.model.model import Boltz1
+from boltz.main import BoltzDiffusionParams
+from utils import motif_utils, mydesign_utils
+from losses import (
+    MotifLoss, AntiMotifLoss, ContactLoss, HelixBiasLoss,
+    ConfChangeLoss, LigandContactLoss, AntiLigandContactLoss,
+    SheetBiasLoss, SequenceSimilarityLoss, RadiusOfGyrationLoss,
+)
+
+ccd = "/opt/switchcraft/boltz/ccd.pkl"
+if os.path.isfile(ccd):
+    from designer import MultistateDesigner
+    print("designer ok (ccd.pkl present)")
+else:
+    print("SKIP designer: boltz/ccd.pkl missing (download step failed?)")
+print("switchcraft import checks passed")
+EOF
+
+%labels
+
+
+
+
+%help
+    SwitchCraft with Boltz + LigandMPNN weights pre-downloaded.
+    Run: singularity exec --nv switchcraft.sif --config tasks/pos_allostery.yaml

slurm_scripts/launch_switchcraft.slurm

@@ -0,0 +1,46 @@
+#!/usr/bin/env bash
+#SBATCH --nodes=<nodes>
+#SBATCH --ntasks-per-node=<ntasks_per_node>
+#SBATCH --cpus-per-task=<cpus_per_task>
+#SBATCH --gpus-per-node=<gpus_per_node>
+#SBATCH --mem=<mem>
+# Partition/account set by caller via sbatch flags (from config)
+#SBATCH --partition=<partition>
+#SBATCH --account=<account>
+#SBATCH --time=01:00:00
+# Log dir from env (caller passes -o)
+#SBATCH --output=<output_dir>%x.%A_%a.out
+
+# Minimal example (repo + Boltz weights baked in SIF; only output dir on host):
+#   mkdir -p /scratch/${USER}/switchcraft_out
+#   sbatch launch_switchcraft.slurm \
+#     tasks/pos_allostery.yaml \
+#     /path/to/switchcraft.sif \
+#     /scratch/${USER}/switchcraft_out
+#
+# Args:
+#   CONFIG_FILE        — YAML under /opt/switchcraft (e.g. tasks/pos_allostery.yaml)
+#   SWITCHCRAFT_IMAGE  — path to .sif
+#   OUTPUT_DIR         — host scratch dir (bind-mounted for writes)
+
+set -euo pipefail
+
+CONFIG_FILE="${1}"
+SWITCHCRAFT_IMAGE="${2}"
+OUTPUT_DIR="${3}"
+
+mkdir -p "${OUTPUT_DIR}"
+
+echo "Running ${SWITCHCRAFT_IMAGE}"
+echo "  config: ${CONFIG_FILE}"
+echo "  output: ${OUTPUT_DIR}"
+
+singularity exec --nv --pwd /opt/switchcraft \
+  -B "${OUTPUT_DIR}:${OUTPUT_DIR}" \
+  "${SWITCHCRAFT_IMAGE}" \
+  python switchcraft.py \
+    --config "${CONFIG_FILE}" \
+    --outpath "${OUTPUT_DIR}" \
+    --verbose
+
+echo "Run completed."