From 3dd0f12294abfb7c9ac3c4c07404d3bbef751cc4 Mon Sep 17 00:00:00 2001 From: mason h <9421505+drawadiagram@users.noreply.github.com> Date: Fri, 11 Sep 2026 14:20:45 -0500 Subject: [PATCH] protein_binding: switch s4 to Boltz-2, remove hardcoded paths, pin GPUs Protein-binding workflow changes extracted from the impress_fixes branch. Depends on the ImpressManager/find_gpus changes in the preceding commit. Structure prediction (s4): AlphaFold multimer is replaced by Boltz-2. - s4_boltz.sh resolves its interpreter from BOLTZ_VENV/VIRTUAL_ENV rather than sourcing a hardcoded /anvil venv, and makes --use_msa_server opt-in (BOLTZ_USE_MSA_SERVER=1) because compute nodes have no internet. - Fix a CCD cache extraction race that killed 13/16 pipelines in a production run: boltz's download_boltz2() skips extraction when the mols/ directory merely exists, but tarfile.extractall() creates that directory entry immediately, so a second concurrent task read a half-populated cache and failed with "CCD component not found". The lock is now held across a completeness check of mols/ against mols.tar, with a .mols_complete marker to skip the O(45k) recount once warmed. - s3 embeds pre-computed MSA paths in the FASTA header, falling back to single-sequence mode when the cache is cold. - af2_multimer_reduced.sh and the commented-out s4_alphafold task are removed along with scripts/s4_alphafold.sh; the AF2 path had been dead for some time and its test-mode stub emitted Boltz-shaped output. Portability: MPNN_PATH now comes from the environment and fails loudly when unset; base_path is split into input_base_path/output_base_path (IMPRESS_BASE_DIR / IMPRESS_OUTPUT_DIR / IMPRESS_SCRIPTS_DIR) so inputs, outputs and scripts can live apart; every scripts/*.sh re-activates $VIRTUAL_ENV instead of a hardcoded cluster path; plddt_extract_pipeline writes its CSV under --path rather than cwd. GPU scheduling: a gpu_id kwarg is threaded into s4_boltz.sh as CUDA_VISIBLE_DEVICES and propagated to child pipelines, with the runner round-robining find_gpus() across pipelines. A per-GPU semaphore caps concurrent Boltz launches at 2 for each GPU, shared across every pipeline pinned to it. Failure handling: s1, s4 and s5 each check for their real output before re-raising, because the execution backend intermittently reports a failure for a task that completed; run() now aborts the pass if every s4 failed instead of proceeding to pLDDT extraction on nothing. Adds delta_env_setup.sh and delta_gpu_run.sh for Delta HPC, with IMPRESS_BACKEND=dragon|local and IMPRESS_TEST_MODE switches. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01Lo8DwSbyvdWZRkkkka6gA2 --- examples/protein_binding/CLAUDE.md | 2 +- examples/protein_binding/README.md | 3 +- .../protein_binding/af2_multimer_reduced.sh | 44 ---- examples/protein_binding/delta_env_setup.sh | 245 ++++++++++++++++++ examples/protein_binding/delta_gpu_run.sh | 117 +++++++++ examples/protein_binding/mpnn_wrapper.py | 8 +- .../protein_binding/plddt_extract_pipeline.py | 3 +- examples/protein_binding/protein_binding.py | 159 +++++++++--- examples/protein_binding/run_nonadaptive.py | 6 +- .../protein_binding/run_protein_binding.py | 78 ++++-- examples/protein_binding/scripts/s1_mpnn.sh | 4 +- .../protein_binding/scripts/s4_alphafold.sh | 25 -- examples/protein_binding/scripts/s4_boltz.sh | 108 +++++++- .../scripts/s5_plddt_extract.sh | 19 +- 14 files changed, 663 insertions(+), 158 deletions(-) delete mode 100644 examples/protein_binding/af2_multimer_reduced.sh create mode 100644 examples/protein_binding/delta_env_setup.sh create mode 100644 examples/protein_binding/delta_gpu_run.sh delete mode 100755 examples/protein_binding/scripts/s4_alphafold.sh diff --git a/examples/protein_binding/CLAUDE.md b/examples/protein_binding/CLAUDE.md index 5642356..21cd725 100644 --- a/examples/protein_binding/CLAUDE.md +++ b/examples/protein_binding/CLAUDE.md @@ -48,7 +48,7 @@ Before running on HPC, edit the path constants in `run_protein_binding.py` (e.g. | `s1` | HPC | `scripts/s1_mpnn.sh` → `mpnn_wrapper.py` (ProteinMPNN) | GPU | | `s2` | local | parses MPNN FASTA output; ranks by score; populates `iter_seqs` | CPU | | `s3` | local | writes paired FASTA (designed sequence + peptide) per structure | CPU | -| `s4` | HPC | `scripts/s4_boltz.sh` (Boltz) or `scripts/s4_alphafold.sh` (AF2, commented out) | GPU | +| `s4` | HPC | `scripts/s4_boltz.sh` (Boltz-2) | GPU | | `s4_post_exec` | HPC | `cp` commands to stage best-model PDB, PTM JSON, and MPNN PDB from Boltz output | CPU | | `s5` | HPC | `scripts/s5_plddt_extract.sh` → `plddt_extract_pipeline.py` (PyRosetta + BioPandas) | CPU | diff --git a/examples/protein_binding/README.md b/examples/protein_binding/README.md index 59fc59e..a1bf6b1 100644 --- a/examples/protein_binding/README.md +++ b/examples/protein_binding/README.md @@ -59,8 +59,7 @@ Writes one paired FASTA file per structure for the structure predictor: ### `s4` — Structure Prediction Predicts the dimer structure for each (designed sequence, peptide) FASTA. All per-structure tasks are launched in parallel with `asyncio.gather`. -- **Default tool**: Boltz (`scripts/s4_boltz.sh`) using MSA server -- **Alternative**: ColabFold/AF2 (`scripts/s4_alphafold.sh`) — commented out in code +- **Tool**: Boltz-2 (`scripts/s4_boltz.sh`). MSA search via the Boltz MSA server is opt-in (`BOLTZ_USE_MSA_SERVER=1`); by default the step uses the MSA cache pre-computed by `delta_env_setup.sh`, since compute nodes have no internet access. - **Output**: `af/prediction/dimer_models//boltz_results_/predictions//` (PDB + PAE files) - **HPC**: 1 GPU per rank diff --git a/examples/protein_binding/af2_multimer_reduced.sh b/examples/protein_binding/af2_multimer_reduced.sh deleted file mode 100644 index ac686ce..0000000 --- a/examples/protein_binding/af2_multimer_reduced.sh +++ /dev/null @@ -1,44 +0,0 @@ -#!/bin/bash - -# have this *first* since they change the WORK env variable in jobs -# module reset -# module load cuda/12.3.0 - -set -e -set -x - -# work and upperdir need to be on same file system -WORK=/tmp/work -UPPER=/tmp/upper -mkdir -p $WORK $UPPER - -export XLA_PYTHON_CLIENT_PREALLOCATE="false" -export XLA_PYTHON_CLIENT_MEM_FRACTION=".75" -export XLA_PYTHON_CLIENT_ALLOCATOR="platform" - -#-B database.squashfs:/database:image-src=/ -INPUT_FASTA_FILE_DIR=$1 -INPUT_FASTA_FILE_NAME=$2 -OUTPUT_DATA_DIR=$3 - -apptainer run --nv \ - --bind $INPUT_FASTA_FILE_DIR:/fasta \ - --bind $OUTPUT_DATA_DIR:/dimer_models \ - --bind /anvil/datasets/alphafold/db_20230311:/database \ - /apps/biocontainers/images/tacc_alphafold:2.3.1.sif \ - --data_dir=/database \ - --uniref90_database_path=/database/uniref90/uniref90.fasta \ - --mgnify_database_path=/database/mgnify/mgy_clusters_2022_05.fa \ - --template_mmcif_dir=/database/pdb_mmcif/mmcif_files/ \ - --obsolete_pdbs_path=/database/pdb_mmcif/obsolete.dat \ - --fasta_paths=/fasta/$INPUT_FASTA_FILE_NAME \ - --output_dir=/dimer_models \ - --model_preset=multimer \ - --db_preset=reduced_dbs \ - --small_bfd_database_path=/database/small_bfd/bfd-first_non_consensus_sequences.fasta \ - --uniprot_database_path=/database/uniprot/uniprot.fasta \ - --pdb_seqres_database_path=/database/pdb_seqres/pdb_seqres.txt \ - --max_template_date=2020-12-01 \ - --use_gpu_relax=False \ - --num_multimer_predictions_per_model=1 \ - --run_relax=False \ No newline at end of file diff --git a/examples/protein_binding/delta_env_setup.sh b/examples/protein_binding/delta_env_setup.sh new file mode 100644 index 0000000..8a2d47f --- /dev/null +++ b/examples/protein_binding/delta_env_setup.sh @@ -0,0 +1,245 @@ +#!/bin/bash +# ============================================================================= +# IMPRESS Protein Binding environment setup — Delta HPC (NCSA) +# +# Creates a Python 3.11+ venv and installs all dependencies. +# +# Usage: +# export SCRATCH=/scratch/ +# bash delta_env_setup.sh [--env-dir DIR] [--impress-dir DIR] [--python PATH] +# +# Defaults: +# ENV_DIR = /u/$USER/ve/impress +# IMPRESS_DIR = $SCRATCH/$USER/IMPRESS +# python = auto-detected via `module load python` (Delta default: 3.13+) +# ============================================================================= +if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then + set -euo pipefail +fi + +# ── Initialize lmod (needed when run as non-interactive bash script) ────────── +if ! declare -f module &>/dev/null; then + _lmod_init=/usr/share/lmod/lmod/init/bash + [ -f "${_lmod_init}" ] && source "${_lmod_init}" +fi + +# ── Require SCRATCH ─────────────────────────────────────────────────────────── +if [[ -z "${SCRATCH:-}" ]]; then + echo "ERROR: set the SCRATCH env var to your allocation scratch root, e.g.:" + echo " export SCRATCH=/scratch/" + echo " bash delta_env_setup.sh" + exit 1 +fi + +# ── Defaults / arg parsing ──────────────────────────────────────────────────── +ENV_DIR="${ENV_DIR:-/u/${USER}/ve/impress}" +IMPRESS_DIR="${IMPRESS_DIR:-${SCRATCH}/${USER}/IMPRESS}" +BASE_PY_OVERRIDE="" + +while [[ $# -gt 0 ]]; do + case $1 in + --env-dir) ENV_DIR="$2"; shift 2 ;; + --impress-dir) IMPRESS_DIR="$2"; shift 2 ;; + --python) BASE_PY_OVERRIDE="$2"; shift 2 ;; + *) echo "Unknown argument: $1"; exit 1 ;; + esac +done + +PY="${ENV_DIR}/bin/python" +PIP="${ENV_DIR}/bin/pip" + +echo "=================================================================" +echo " ENV_DIR = ${ENV_DIR}" +echo " IMPRESS_DIR = ${IMPRESS_DIR}" +echo "=================================================================" + +# ── 1. Create venv ──────────────────────────────────────────────────────────── +echo "" +echo "── Step 1: Creating venv ──" + +if [ -n "${BASE_PY_OVERRIDE}" ]; then + BASE_PY="${BASE_PY_OVERRIDE}" + echo "Using Python override: ${BASE_PY}" +else + # On Delta, `module load python` gives the default Python 3.13+. + module load python 2>/dev/null || true + BASE_PY=$(command -v python3 2>/dev/null || command -v python 2>/dev/null || true) + if [ -z "${BASE_PY}" ]; then + echo "ERROR: no Python found after 'module load python'." + echo " Pass an explicit interpreter: --python /path/to/python3" + exit 1 + fi + ver=$("${BASE_PY}" -c "import sys; v=sys.version_info; print(v.major*100+v.minor)") + if [ "${ver}" -lt 311 ]; then + echo "ERROR: ${BASE_PY} is Python ${ver} — need 3.11+." + echo " Pass an explicit interpreter: --python /path/to/python3.11" + exit 1 + fi +fi +echo "Using Python: ${BASE_PY} ($(${BASE_PY} --version))" + +if [ ! -x "${PY}" ]; then + "${BASE_PY}" -m venv "${ENV_DIR}" +else + echo "venv already exists at ${ENV_DIR}" +fi + +echo "Python: $("${PY}" --version)" + +# ── 2. Bootstrap pip ────────────────────────────────────────────────────────── +echo "" +echo "── Step 2: Bootstrapping pip ──" +"${PY}" -m pip install -q --upgrade pip wheel +"${PIP}" install -q --force-reinstall "setuptools<71" + +# ── 3. radical.asyncflow (PyPI) ────────────────────────────────────────────── +echo "" +echo "── Step 3: radical-asyncflow (PyPI) ──" +"${PIP}" install -q radical-asyncflow + +# ── 4. rhapsody-py (PyPI) ──────────────────────────────────────────────────── +echo "" +echo "── Step 4: rhapsody-py[dragon] (PyPI) ──" +"${PIP}" install -q "rhapsody-py[dragon,telemetry]" + +# ── 5. IMPRESS (local editable) ─────────────────────────────────────────────── +echo "" +echo "── Step 5: IMPRESS (editable) ──" +"${PIP}" install -q -e "${IMPRESS_DIR}" + +# ── 6. PyTorch (CUDA 12.1) ─────────────────────────────────────────────────── +echo "" +echo "── Step 6: PyTorch (cu121) ──" +"${PIP}" install -q torch --index-url https://download.pytorch.org/whl/cu121 + +# ── 7. Additional dependencies ─────────────────────────────────────────────── +echo "" +echo "── Step 7: pandas + biopandas + matplotlib ──" +"${PIP}" install -q pandas biopandas matplotlib + +# ── 8. PyRosetta ───────────────────────────────────────────────────────────── +echo "" +echo "── Step 8: PyRosetta ──" +# pyrosetta_installer handles credential lookup internally. +# Activate the venv in the environment so its subprocess pip installs there. +export VIRTUAL_ENV="${ENV_DIR}" +export PATH="${ENV_DIR}/bin:${PATH}" +"${PIP}" install -q pyrosetta-installer +"${PY}" -c "import pyrosetta_installer; pyrosetta_installer.install_pyrosetta()" + +# ── 9. Boltz (structure prediction — separate conda env) ───────────────────── +echo "" +echo "── Step 9: Boltz (separate conda env) ──" +# boltz 2.x is kept in a separate conda env so its dependency pins (scipy, etc.) +# don't constrain the main IMPRESS venv. Boltz 2.0+ also installs fine via pip +# in Python 3.13, but we keep the separation to avoid pin conflicts. +# s4_boltz.sh activates BOLTZ_VENV (set in delta_gpu_run.sh) instead of VIRTUAL_ENV. +MINIFORGE="${MINIFORGE:-${SCRATCH}/${USER}/miniforge3}" +BOLTZ_ENV="${BOLTZ_ENV:-${HOME}/ve/boltz}" +# Cache lives in home dir — scratch inode quota can't hold the 45K CCD files. +BOLTZ_CACHE="${BOLTZ_CACHE:-${HOME}/boltz}" +if [ ! -x "${BOLTZ_ENV}/bin/python" ]; then + echo " Creating conda env (Python 3.11) at ${BOLTZ_ENV}" + # Use /tmp for package cache to avoid scratch quota exhaustion. + CONDA_PKGS_DIRS=/tmp/conda_pkgs "${MINIFORGE}/bin/conda" create -p "${BOLTZ_ENV}" python=3.11 -y -q +else + echo " boltz conda env already exists at ${BOLTZ_ENV}" +fi +echo " Installing boltz into ${BOLTZ_ENV}" +"${BOLTZ_ENV}/bin/pip" install -q boltz +# Pre-warm the boltz cache so compute nodes (no internet) find weights ready. +# boltz downloads mols.tar (45K CCD pkl files) + model weights on first run. +# Running predict on the login node (which has internet) pre-populates them. +echo " Pre-warming boltz cache at ${BOLTZ_CACHE}" +mkdir -p "${BOLTZ_CACHE}" +_WARM_FA="$(mktemp /tmp/boltz_warmup_XXXXXX.fasta)" +printf ">warmup|A\nGSSGSSGSS\n>warmup|B\nGSSGSSGSS\n" > "${_WARM_FA}" +BOLTZ_CACHE_DIR="${BOLTZ_CACHE}" "${BOLTZ_ENV}/bin/boltz" predict "${_WARM_FA}" \ + --out_dir "$(mktemp -d /tmp/boltz_warmup_out_XXXXXX)" \ + --cache "${BOLTZ_CACHE}" \ + --override 2>&1 | grep -E "Download|Extracting|Error|error" || true +rm -f "${_WARM_FA}" +echo " Boltz cache pre-warm done (model weights cached at ${BOLTZ_CACHE})" + +# Pre-compute MSAs for all input proteins using the MSA server (login node has internet). +# protein_binding.py s3() reads from BOLTZ_MSA_CACHE and embeds paths in the FASTA so +# compute nodes do not need internet access. Entity 0 = receptor, entity 1 = peptide. +BOLTZ_MSA_CACHE="${BOLTZ_CACHE}/msa_cache" +mkdir -p "${BOLTZ_MSA_CACHE}" +echo " Pre-computing MSAs into ${BOLTZ_MSA_CACHE}" +# IMPRESS_BASE_DIR = parent of prod_in/; IMPRESS_OUTPUT_DIR = parent of af_pipeline_outputs_multi/ +_scratch="${SCRATCH}" +_base_dir="${IMPRESS_BASE_DIR:-${_scratch}/IMPRESS_inputs}" +_out_dir="${IMPRESS_OUTPUT_DIR:-${_scratch}/IMPRESS_outputs}" +_msa_inputs_dir="${_base_dir}/prod_in" +if [ -d "${_msa_inputs_dir}" ]; then + for _pdb_dir in "${_msa_inputs_dir}"/p*_in; do + for _pdb in "${_pdb_dir}"/*.pdb; do + [ -f "${_pdb}" ] || continue + _stem="$(basename "${_pdb}" .pdb)" + _msa_csv="${BOLTZ_MSA_CACHE}/boltz_results_${_stem}/msa/${_stem}_0.csv" + if [ -f "${_msa_csv}" ]; then + echo " ${_stem}: MSA already cached, skipping" + continue + fi + echo " ${_stem}: generating MSA via server..." + _tmp_fa="$(mktemp /tmp/boltz_msa_XXXXXX.fasta)" + # Use FASTA from a prior run (sequences match the actual protein); fall back to + # a placeholder that will trigger MSA generation but gives a generic MSA. + _prior_fa="" + for _cand in "${_out_dir}"/af_pipeline_outputs_multi/*/af/fasta/"${_stem}.fa"; do + [ -f "${_cand}" ] && { _prior_fa="${_cand}"; break; } + done + if [ -n "${_prior_fa}" ]; then + cp "${_prior_fa}" "${_tmp_fa}" + else + printf ">pdz|protein\nGSSGSS\n>pep|protein\nGSSG\n" > "${_tmp_fa}" + fi + BOLTZ_CACHE_DIR="${BOLTZ_CACHE}" "${BOLTZ_ENV}/bin/boltz" predict "${_tmp_fa}" \ + --out_dir "${BOLTZ_MSA_CACHE}" \ + --use_msa_server \ + --cache "${BOLTZ_CACHE}" \ + --output_format pdb \ + --override 2>&1 | grep -E "MSA|Generat|Error|error|skip" || true + rm -f "${_tmp_fa}" + done + done +else + echo " prod_in not found at ${_msa_inputs_dir}, skipping MSA pre-compute" +fi +echo " MSA pre-compute done" + +# ── 10. Verify ─────────────────────────────────────────────────────────────── +echo "" +echo "── Step 10: Verifying installation ──" +_check() { + local label="$1"; shift + if out=$("$@" 2>&1); then + echo " ${label}: OK (${out})" + else + echo " WARNING: ${label} failed" + echo " ${out}" | head -3 + fi +} + +_check "radical.asyncflow" "${PY}" -c "import radical.asyncflow; print(radical.asyncflow.__version__)" +_check "rhapsody-py" "${PY}" -c "import rhapsody; print('ok')" +_check "impress" "${PY}" -c "import impress; print('ok')" +_check "torch" "${PY}" -c "import torch; print(torch.__version__)" +_check "pandas" "${PY}" -c "import pandas; print(pandas.__version__)" +BOLTZ_ENV="${BOLTZ_ENV:-${HOME}/ve/boltz}" +_check "boltz" "${BOLTZ_ENV}/bin/python" -c "import boltz; print('ok')" + +echo "" +echo "=================================================================" +echo "Setup complete." +echo "" +echo "Activate with:" +echo " source ${ENV_DIR}/bin/activate" +echo "" +echo "Run the pipeline:" +echo " export SCRATCH=${SCRATCH}" +echo " export SBATCH_ACCOUNT=bblj-delta-gpu" +echo " cd ${IMPRESS_DIR}/examples/protein_binding" +echo " sbatch delta_gpu_run.sh" +echo "=================================================================" diff --git a/examples/protein_binding/delta_gpu_run.sh b/examples/protein_binding/delta_gpu_run.sh new file mode 100644 index 0000000..14161e9 --- /dev/null +++ b/examples/protein_binding/delta_gpu_run.sh @@ -0,0 +1,117 @@ +#!/bin/bash +# +# Protein Binding Pipeline — SLURM batch script (Delta HPC / GPU) +# +# Set before calling sbatch (only SBATCH_ACCOUNT and SCRATCH are required; +# the rest default to standard Delta locations): +# export SBATCH_ACCOUNT=-delta-gpu +# export SCRATCH=/scratch/ +# +# Example: +# sbatch delta_gpu_run.sh +# +# Account: set SBATCH_ACCOUNT=-delta-gpu before calling sbatch +#SBATCH --partition=gpuA40x4 +#SBATCH --nodes=1 +#SBATCH --tasks-per-node=1 +#SBATCH --cpus-per-task=16 +#SBATCH --gpus-per-node=4 +#SBATCH --mem=220G +#SBATCH --time=02:30:00 +#SBATCH --job-name=impress_protein +#SBATCH --mail-user= +#SBATCH --mail-type=ALL +#SBATCH --output=logs/impress_%j.out +#SBATCH --error=logs/impress_%j.err +# NOTE: logs/ must exist before sbatch is called. Create it once with: +# mkdir -p /logs +# NOTE: IMPRESS log output (including errors) goes to .out, not .err. +# On failure, check logs/impress_.out — the .err file will only +# contain Python interpreter crashes or output from non-IMPRESS processes. + +set -e + +# ── Sanity checks ───────────────────────────────────────────────────────────── +if [ -z "${SBATCH_ACCOUNT:-}${SLURM_JOB_ACCOUNT:-}" ]; then + echo "WARNING: SBATCH_ACCOUNT is not set — job may be charged to default account." +fi +echo "Account: ${SLURM_JOB_ACCOUNT:-unknown}" + +if [ -z "${SCRATCH:-}" ]; then + echo "ERROR: SCRATCH is not set." + echo " export SCRATCH=/scratch/ && sbatch delta_gpu_run.sh" + exit 1 +fi + +# ── System library paths (Delta-specific, required by Dragon) ───────────────── +export CUDA_HOME=/opt/nvidia/hpc_sdk/Linux_x86_64/25.3/cuda/12.8 +export MPI_LIB=/opt/cray/pe/mpich/8.1.32/ofi/gnu/11.2/lib-abi-mpich +export FAB_LIB=/opt/cray/libfabric/1.22.0/lib64 +export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${MPI_LIB}:${FAB_LIB}:${LD_LIBRARY_PATH:-} + +# ── Environment ─────────────────────────────────────────────────────────────── +IMPRESS_VENV="${IMPRESS_VENV:-${HOME}/ve/impress_A}" +unset SLURM_EXPORT_ENV +source "${IMPRESS_VENV}/bin/activate" +dragon-config add --ofi-runtime-lib="${FAB_LIB}" + +# ── Tool paths (adjust for your allocation) ─────────────────────────────────── +export MPNN_PATH="${MPNN_PATH:-${SCRATCH}/${USER}/ProteinMPNN}" +export AF2_DATABASE="${AF2_DATABASE:-${SCRATCH}/${USER}/alphafold_database}" +export AF2_SIF="${AF2_SIF:-${SCRATCH}/${USER}/alphafold.sif}" +# Boltz lives in a separate Python 3.12 conda env (boltz 2.x requires numpy<2.0, +# scipy==1.13.1 etc. which have no Python 3.13 wheels). +export BOLTZ_VENV="${BOLTZ_VENV:-${HOME}/ve/boltz}" +# Boltz model weight cache — kept in home dir; scratch inode quota can't hold +# the 45K CCD molecule files that boltz extracts from mols.tar on first run. +export BOLTZ_CACHE_DIR="${BOLTZ_CACHE_DIR:-${HOME}/boltz}" +mkdir -p "${BOLTZ_CACHE_DIR}" + +# ── IMPRESS paths ───────────────────────────────────────────────────────────── +export IMPRESS_SCRIPTS_DIR="${IMPRESS_SCRIPTS_DIR:-${SCRATCH}/${USER}/IMPRESS/examples/protein_binding}" +# IMPRESS_BASE_DIR: parent of prod_in/ — pipeline builds prod_in/_in from here +export IMPRESS_BASE_DIR="${IMPRESS_BASE_DIR:-${SCRATCH}/${USER}/IMPRESS_inputs}" +export IMPRESS_OUTPUT_DIR="${IMPRESS_OUTPUT_DIR:-${SCRATCH}/${USER}/IMPRESS_outputs}" + +# IMPRESS_BACKEND: "dragon" (default, multi-node HPC) or "local" (single-node, +# ProcessPoolExecutor — useful for development / non-Dragon clusters). +# Set before sbatch: IMPRESS_BACKEND=local sbatch delta_gpu_run.sh +export IMPRESS_BACKEND="${IMPRESS_BACKEND:-dragon}" +echo "IMPRESS_BACKEND: ${IMPRESS_BACKEND}" + +# IMPRESS_TEST_MODE=1: 2 pipelines, max_passes=1, no child pipelines. +# Runs a single MPNN → score → AF2 cycle to verify end-to-end path. +# Set before sbatch: IMPRESS_TEST_MODE=1 sbatch delta_gpu_run.sh +export IMPRESS_TEST_MODE="${IMPRESS_TEST_MODE:-0}" +echo "TEST_MODE: ${IMPRESS_TEST_MODE}" + +# ── Working directory ───────────────────────────────────────────────────────── +WORKDIR="${IMPRESS_SCRIPTS_DIR}" +cd "${WORKDIR}" +mkdir -p logs + +# IMPRESS_SESSION_DIR: asyncflow session dir — runinfo, captured task +# stdout/stderr (.stdout/.stderr per task UID). Must be on Lustre so files +# survive the job and can be reviewed after failures. +export IMPRESS_SESSION_DIR="${IMPRESS_SESSION_DIR:-${WORKDIR}/logs/sessions}" +mkdir -p "${IMPRESS_SESSION_DIR}" + +# ── Run ─────────────────────────────────────────────────────────────────────── + +# -s = single-node Dragon runtime; -m = multi-node (uses MPI/OFI fabric). +if [ "${SLURM_NNODES:-1}" -gt 1 ]; then + DRAGON_MODE="-m" +else + DRAGON_MODE="-s" +fi + +if [ "${IMPRESS_BACKEND}" = "dragon" ]; then + rm -f ddict_orc* + echo "Running: dragon ${DRAGON_MODE} run_protein_binding.py (nodes=${SLURM_NNODES:-1})" + dragon ${DRAGON_MODE} run_protein_binding.py +else + echo "Running: python3 run_protein_binding.py (backend=${IMPRESS_BACKEND})" + python3 run_protein_binding.py +fi + +echo "=== Protein Binding pipeline done: $(date) ===" diff --git a/examples/protein_binding/mpnn_wrapper.py b/examples/protein_binding/mpnn_wrapper.py index 1e8942d..811fd4d 100644 --- a/examples/protein_binding/mpnn_wrapper.py +++ b/examples/protein_binding/mpnn_wrapper.py @@ -1,4 +1,4 @@ -#!/bin/sh +#!/usr/bin/env python3 import argparse parser = argparse.ArgumentParser() import subprocess @@ -15,7 +15,7 @@ parser.add_argument("-homo", "--homo", help="Are your input files homomers? (0 or 1) Note: Overrides tied positions. If homomer is specified, all positions on designed chains will be tied. If positions are restricted, lists must be of same length. Example: -index='1 3 5 7, 1 3 5 7'", type=int) parser.add_argument("-bias_AA", "--bias_AA", help="For which amino acids would you like to install bias? Example: -bias_AA='D E H'", type=str) parser.add_argument("-bias_weight", "--bias_weight", help="What weights would you like to install for the biased amino acids? Lists must match in length. Example: -bias_weight='0.3 -0.3 0.5'", type=str) -parser.add_argument("-temp", "--temp", help="What temperature would you like to sample from? Example: 0.3", type=int, default=0.1) +parser.add_argument("-temp", "--temp", help="What temperature would you like to sample from? Example: 0.3", type=float, default=0.1) parser.add_argument("-inter", "--interface", help="Would you like to design the interface? Do not specify indices if designing the interface. (1 or 0)", type=int, default=0) @@ -27,12 +27,12 @@ mpnn_path=args.mpnn_path is_monomer=args.is_monomer #default is_monomer false chains=args.design_chains -if chains == None: +if chains is None: chains='A' #default design chain A index=args.index fix=args.fix #default false, specify non fixed seqs=args.seqs -if seqs == None: +if seqs is None: seqs=1 #default 1 design per structure tie=args.tie homo=args.homo diff --git a/examples/protein_binding/plddt_extract_pipeline.py b/examples/protein_binding/plddt_extract_pipeline.py index d1c8a7d..6dc57bd 100644 --- a/examples/protein_binding/plddt_extract_pipeline.py +++ b/examples/protein_binding/plddt_extract_pipeline.py @@ -59,4 +59,5 @@ print(f"Processed {len(rows)} structure(s)") df = pd.DataFrame(rows, columns=['ID', 'avg_plddt', 'ptm', 'avg_pae']) -df.to_csv('af_stats_' + args.out + '_pass_' + args.iter + '.csv', index=False) +csv_path = os.path.join(args.path, 'af_stats_' + args.out + '_pass_' + args.iter + '.csv') +df.to_csv(csv_path, index=False) diff --git a/examples/protein_binding/protein_binding.py b/examples/protein_binding/protein_binding.py index 64ca80d..81a47c9 100644 --- a/examples/protein_binding/protein_binding.py +++ b/examples/protein_binding/protein_binding.py @@ -6,7 +6,7 @@ from impress.pipelines.impress_pipeline import ImpressBasePipeline -MPNN_PATH = f"/anvil/projects/x-nairr240405/mason/ProteinMPNN" +MPNN_PATH = os.environ.get("MPNN_PATH", "") _BOLTZ_CHAIN_MAP = {'pdz': 'A', 'pep': 'B'} @@ -20,6 +20,10 @@ def _copy_pdb_rename_chains(src, dst, chain_map=_BOLTZ_CHAIN_MAP): line = line[:21] + chain_map[chain] + line[24:] f_out.write(line) +# One semaphore per GPU shared across all pipeline instances — caps concurrent +# Boltz launches per GPU at 2 regardless of how many pipelines share that GPU. +_boltz_sem_per_gpu: dict = {} + class ProteinBindingPipeline(ImpressBasePipeline): def __init__(self, name, flow, configs=None, **kwargs): # Execution metadata @@ -34,7 +38,11 @@ def __init__(self, name, flow, configs=None, **kwargs): self.num_seqs = kwargs.get("num_seqs", 10) self.sub_order = kwargs.get("sub_order", 0) self.max_passes = kwargs.get("max_passes", 10) - self.mpnn_path = kwargs.get("mpnn_path", MPNN_PATH) + self.mpnn_path = kwargs.get("mpnn_path") or MPNN_PATH + if not self.mpnn_path: + raise ValueError("mpnn_path must be supplied via kwarg or MPNN_PATH env var") + self.peptide_seq: str = kwargs.get("peptide_seq", "EGYQDYEPEA") + self.gpu_id = kwargs.get("gpu_id", None) # Sequence and score state self.current_scores = {} @@ -46,12 +54,18 @@ def __init__(self, name, flow, configs=None, **kwargs): # Input-related self.fasta_list_2 = kwargs.get("fasta_list_2", []) self.base_path = kwargs.get("base_path", os.getcwd()) + # input_base_path is the parent of prod_in/; defaults to base_path so + # existing callers that keep input data next to scripts still work. + self.input_base_path = kwargs.get("input_base_path", self.base_path) + # output_base_path is where af_pipeline_outputs_multi/ is written; + # defaults to base_path so existing callers without IMPRESS_OUTPUT_DIR work. + self.output_base_path = kwargs.get("output_base_path", self.base_path) self.scripts_path = os.path.join(self.base_path, "scripts") - self.input_path = os.path.join(self.base_path, f"prod_in/{self.name}_in") + self.input_path = os.path.join(self.input_base_path, f"prod_in/{self.name}_in") # Output paths self.output_path = os.path.join( - self.base_path, "af_pipeline_outputs_multi", self.name + self.output_base_path, "af_pipeline_outputs_multi", self.name ) self.output_path_mpnn = os.path.join(self.output_path, "mpnn") self.output_path_af = os.path.join( @@ -65,12 +79,9 @@ def __init__(self, name, flow, configs=None, **kwargs): def set_up_new_pipeline_dirs(self, new_pipeline_name): base_output = os.path.join( - self.base_path, "af_pipeline_outputs_multi", new_pipeline_name + self.output_base_path, "af_pipeline_outputs_multi", new_pipeline_name ) - input_dir = os.path.join(self.base_path, f"prod_in/{new_pipeline_name}_in") - - if os.path.isdir(base_output): - return # already exists, nothing to do + input_dir = os.path.join(self.input_base_path, f"prod_in/{new_pipeline_name}_in") # all directories to create subdirs = [ @@ -95,15 +106,16 @@ def register_pipeline_tasks(self): """Register all pipeline tasks""" @self.auto_register_task(capture_stdio=True) # MPNN - async def s1(task_description={"gpus_per_rank": 1}): + async def s1(): # noqa: B006 self.step_id += 1 mpnn_script = os.path.join(self.base_path, "mpnn_wrapper.py") output_dir = os.path.join(self.output_path_mpnn, f"job_{self.passes}") + os.makedirs(output_dir, exist_ok=True) chain = "A" input_path = self.input_path if self.passes == 1 else self.output_path_af - return ( + cmd = ( f"bash {self.scripts_path}/s1_mpnn.sh " f"{mpnn_script} " f"{input_path} " @@ -112,6 +124,7 @@ async def s1(task_description={"gpus_per_rank": 1}): f"{self.num_seqs} " f"{chain}" ) + return cmd @self.auto_register_task(local_task=True) async def s2(): @@ -139,38 +152,50 @@ async def s2(): async def s3(): self.step_id += 1 output_dir = os.path.join(self.output_path, "af", "fasta") + # Pre-computed MSA cache (populated by delta_env_setup.sh on the login node, + # which has internet access). Entity 0 = receptor (PDZ), entity 1 = peptide. + msa_cache = os.path.expanduser("~/boltz/msa_cache") fasta_file_to_return = [] for fasta_file in self.fasta_list_2: base_name = fasta_file.split(".")[0] fasta_file_to_return.append(base_name) design_seq = self.iter_seqs[base_name][self.seq_rank][0] - pep_seq = "EGYQDYEPEA" + pep_seq = self.peptide_seq + + # Boltz embeds pre-computed MSAs in the FASTA header rather than + # running MSA search at prediction time (compute nodes have no internet). + # Naming convention: entity 0 = receptor (PDZ), entity 1 = peptide. + pdz_msa = os.path.join( + msa_cache, f"boltz_results_{base_name}", "msa", f"{base_name}_0.csv" + ) + pep_msa = os.path.join( + msa_cache, f"boltz_results_{base_name}", "msa", f"{base_name}_1.csv" + ) + if os.path.exists(pdz_msa) and os.path.exists(pep_msa): + # |path/to/msa.csv| tells Boltz to load the pre-computed MSA + pdz_tag = f">pdz|protein|{pdz_msa}" + pep_tag = f">pep|protein|{pep_msa}" + else: + # |empty| runs single-sequence mode — no MSA, slightly less accurate + pdz_tag = ">pdz|protein|empty" + pep_tag = ">pep|protein|empty" fasta_path = os.path.join(output_dir, f"{base_name}.fa") with open(fasta_path, "w") as f: - f.write(f">pdz|protein\n{design_seq}\n>pep|protein\n{pep_seq}\n") + f.write(f"{pdz_tag}\n{design_seq}\n{pep_tag}\n{pep_seq}\n") return fasta_file_to_return - # alphafold, must be run separately for each structure one at a time! -# @self.auto_register_task() -# async def s4(target_fasta, task_description={"gpus_per_rank": 1}): # noqa: B006 -# return ( -# f"bash {self.scripts_path}/s4_alphafold.sh " -# f"{self.output_path}/af/fasta/{target_fasta}.fa " -# f"{self.output_path}/af/prediction/dimer_models/{target_fasta}" -# ) - @self.auto_register_task(capture_stdio=True) - async def s4(target_fasta, task_description={"gpus_per_rank": 1}): # noqa: B006 + async def s4(target_fasta): # noqa: B006 self.step_id += 1 cmd = ( f"bash {self.scripts_path}/s4_boltz.sh " f"{self.output_path}/af/fasta/{target_fasta}.fa " f"{self.output_path}/af/prediction/dimer_models/{target_fasta}" + + (f" {self.gpu_id}" if self.gpu_id is not None else "") ) - self.logger.pipeline_log(f"s4 command for {target_fasta}: {cmd}") return cmd @self.auto_register_task(local_task=True) @@ -206,7 +231,7 @@ async def s5(): self.step_id += 1 return ( f"bash {self.scripts_path}/s5_plddt_extract.sh " - f"{self.base_path} " + f"{self.output_base_path} " f"{self.passes} " f"{self.name}" ) @@ -217,15 +242,19 @@ async def get_scores_map(self): def finalize(self, sub_iter_seqs): # finalize the "cleanup" of the current pipeline + from pathlib import Path for a in sub_iter_seqs: self.fasta_list_2.remove(f"{a}.pdb") - os.unlink(f"{self.output_path_af}/{a}.pdb") - os.unlink(f"{self.output_path}/af/fasta/{a}.fa") + Path(f"{self.output_path_af}/{a}.pdb").unlink(missing_ok=True) + Path(f"{self.output_path}/af/fasta/{a}.fa").unlink(missing_ok=True) self.previous_scores = copy.deepcopy(self.current_scores) async def run(self): """Main execution logic""" + if self.gpu_id is not None: + self.logger.pipeline_log(f"gpu={self.gpu_id}") + self.logger.pipeline_log(f"Running for a maximum of {self.max_passes} passes") self.set_up_new_pipeline_dirs(self.name) @@ -241,7 +270,20 @@ async def run(self): else: self.logger.pipeline_log("Submitting MPNN task") - await self.s1() + try: + await self.s1() + except Exception as exc: + # The execution backend may report a spurious failure (TypeError/ + # 'NoneType' subscriptable, or ProcessGroup state error) even when + # MPNN completed successfully. Check for output before propagating. + seqs_dir = os.path.join( + self.output_path_mpnn, f"job_{self.passes}", "seqs" + ) + if not (os.path.isdir(seqs_dir) and os.listdir(seqs_dir)): + raise + self.logger.pipeline_log( + f"s1 raised {exc!r} but seqs output exists — treating as success" + ) self.logger.pipeline_log("MPNN task finished") self.logger.pipeline_log("Submitting sequence ranking task") @@ -255,6 +297,29 @@ async def run(self): alphafold_tasks = [] post_exec_tasks = [] + # Shared per-GPU semaphore caps concurrent Boltz launches at 2 per GPU + # across all pipeline instances pinned to the same GPU. + gpu_key = self.gpu_id if self.gpu_id is not None else "default" + if gpu_key not in _boltz_sem_per_gpu: + _boltz_sem_per_gpu[gpu_key] = asyncio.Semaphore(2) + _boltz_sem = _boltz_sem_per_gpu[gpu_key] + + async def _guarded_s4(target_fasta): + async with _boltz_sem: + try: + return await self.s4(target_fasta=target_fasta) + except Exception as exc: + # The execution backend may raise a spurious failure even when + # Boltz completed successfully. Check for the output PDB before propagating. + pred_dir = os.path.join( + self.output_path, "af", "prediction", "dimer_models", + target_fasta, f"boltz_results_{target_fasta}", + "predictions", target_fasta, + ) + if os.path.isfile(os.path.join(pred_dir, f"{target_fasta}_model_0.pdb")): + return None # output exists; treat as success + raise + for target_fasta in fasta_files: models_path = os.path.join( self.output_path, "af", "prediction", "dimer_models", target_fasta, @@ -282,8 +347,8 @@ async def run(self): f"{target_fasta}.pdb", ) - # launch coroutine without awaiting yet - alphafold_tasks.append(self.s4(target_fasta=target_fasta)) + # launch coroutine without awaiting yet (semaphore-gated) + alphafold_tasks.append(_guarded_s4(target_fasta)) post_exec_tasks.append( self.s4_post_exec( target_fasta=target_fasta, @@ -306,26 +371,40 @@ async def run(self): self.logger.pipeline_log(f"s4 DONE for {fasta_name}") s4_post_results = await asyncio.gather(*post_exec_tasks, return_exceptions=True) + any_s4_ok = False for fasta_name, result in zip(fasta_files, s4_post_results): if isinstance(result, Exception): self.logger.pipeline_log(f"s4_post_exec FAILED for {fasta_name}: {result}") else: self.logger.pipeline_log(f"s4_post_exec DONE for {fasta_name}") + any_s4_ok = True + + if not any_s4_ok: + raise RuntimeError("All s4 tasks failed — skipping pLDDT extraction") self.logger.pipeline_log("Submitting pLDTT extraction task") staged_file = f"af_stats_{self.name}_pass_{self.passes}.csv" - await self.s5( - task_description={ - "output_staging": [ - { - "source": f"task:///{staged_file}", - "target": f"client:///{staged_file}", - } - ], - } - ) + try: + await self.s5( + task_description={ + "output_staging": [ + { + "source": f"task:///{staged_file}", + "target": f"client:///{staged_file}", + } + ], + } + ) + except Exception as exc: + # Spurious backend failure: check if s5 wrote the CSV despite the error. + csv_path = os.path.join(self.output_base_path, staged_file) + if not os.path.isfile(csv_path): + raise + self.logger.pipeline_log( + f"s5 raised {exc!r} but CSV exists — treating as success" + ) self.logger.pipeline_log("pLDTT extract finished") await self.run_adaptive_step(wait=True) diff --git a/examples/protein_binding/run_nonadaptive.py b/examples/protein_binding/run_nonadaptive.py index 6eea85e..6a97fe7 100644 --- a/examples/protein_binding/run_nonadaptive.py +++ b/examples/protein_binding/run_nonadaptive.py @@ -1,7 +1,7 @@ import asyncio from typing import List -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend from rhapsody.telemetry import define_event from impress import PipelineSetup @@ -9,7 +9,7 @@ from protein_binding import ProteinBindingPipeline import rhapsody, logging -rhapsody.enable_logging(level=logging.DEBUG) +rhapsody.enable_logging(level=logging.INFO) def _on_task_event(event) -> None: @@ -19,7 +19,7 @@ def _on_task_event(event) -> None: async def impress_protein_bind_nonadaptive() -> None: - backend = await DragonExecutionBackendV3() + backend = await DragonExecutionBackend() manager: ImpressManager = ImpressManager( execution_backend=backend, diff --git a/examples/protein_binding/run_protein_binding.py b/examples/protein_binding/run_protein_binding.py index cf66937..14a3e8a 100644 --- a/examples/protein_binding/run_protein_binding.py +++ b/examples/protein_binding/run_protein_binding.py @@ -1,20 +1,36 @@ import copy +import os import shutil import asyncio from typing import Dict, Any, Optional, List -from rhapsody.backends import DragonExecutionBackendV3 from rhapsody.telemetry import define_event from rhapsody.telemetry.events import make_event -from impress import PipelineSetup -from impress import ImpressManager +from impress import find_gpus, ImpressManager, PipelineSetup from protein_binding import ProteinBindingPipeline import rhapsody, logging rhapsody.enable_logging(level=logging.DEBUG) +# ── Backend / test mode ─────────────────────────────────────────────────── +# IMPRESS_BACKEND: "dragon" (default, multi-node HPC) or "local" (single-node, +# ProcessPoolExecutor — useful for development / non-Dragon clusters). +BACKEND = os.environ.get("IMPRESS_BACKEND", "dragon").lower() + +if BACKEND == "dragon": + from rhapsody.backends import DragonExecutionBackend +else: + from concurrent.futures import ProcessPoolExecutor + from rhapsody.backends import ConcurrentExecutionBackend +TEST_MODE = os.getenv("IMPRESS_TEST_MODE", "0") == "1" +N_PIPELINES = 4 if TEST_MODE else 16 +MAX_PASSES = 10 if TEST_MODE else 10 +MAX_SUB_PIPELINES_OVERRIDE = 3 if TEST_MODE else None # None = use inline default + +print(f"[INFO] IMPRESS_BACKEND={BACKEND} TEST_MODE={TEST_MODE} N_PIPELINES={N_PIPELINES} MAX_PASSES={MAX_PASSES} MAX_SUB_PIPELINES_OVERRIDE={MAX_SUB_PIPELINES_OVERRIDE}") + # --------------------------------------------------------------------------- # Custom application-level telemetry events # --------------------------------------------------------------------------- @@ -52,21 +68,22 @@ # --------------------------------------------------------------------------- def _on_task_event(event) -> None: - if event.event_type == "TaskFailed": - wid = getattr(event, "workflow_id", None) - print(f"[TELEMETRY] TaskFailed task={event.task_id} workflow={wid}") + pass # --------------------------------------------------------------------------- # Adaptive helpers # --------------------------------------------------------------------------- -async def adaptive_criteria(current_score: float, previous_score: float) -> bool: +def adaptive_criteria(current_score: float, previous_score: float) -> bool: return current_score > previous_score async def impress_protein_bind() -> None: - backend = await DragonExecutionBackendV3() + if BACKEND == "dragon": + backend = await DragonExecutionBackend() + else: + backend = await ConcurrentExecutionBackend.create(ProcessPoolExecutor()) manager: ImpressManager = ImpressManager( execution_backend=backend, @@ -80,12 +97,12 @@ async def impress_protein_bind() -> None: # adaptive_decision closes over `manager` so it can emit events via # manager.telemetry, which is set inside start() before any pipeline runs. async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[str, Any]]: - MAX_SUB_PIPELINES: int = 3 + MAX_SUB_PIPELINES: int = MAX_SUB_PIPELINES_OVERRIDE if MAX_SUB_PIPELINES_OVERRIDE is not None else 3 tel = manager.telemetry sid = tel.session_id if tel else None # Read current scores from CSV - file_name = f'af_stats_{pipeline.name}_pass_{pipeline.passes}.csv' + file_name = os.path.join(pipeline.output_base_path, f'af_stats_{pipeline.name}_pass_{pipeline.passes}.csv') with open(file_name) as fd: for line in fd.readlines()[1:]: line = line.strip() @@ -108,13 +125,14 @@ async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[s continue prev_score = pipeline.previous_scores[protein] - decision = await adaptive_criteria(curr_score, prev_score) + decision = adaptive_criteria(curr_score, prev_score) pipeline.logger.pipeline_log(f'Adaptive decision: {decision}') if tel: tel.emit(make_event( ProteinScore, session_id=sid, + backend="rhapsody", protein=protein, pipeline_name=pipeline.name, pass_num=pipeline.passes, @@ -135,13 +153,14 @@ async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[s for protein in sub_iter_seqs: src = f'{pipeline.output_path_af}/{protein}.pdb' - dst = f'{pipeline.base_path}/prod_in/{new_name}_in/{protein}.pdb' + dst = f'{pipeline.input_base_path}/prod_in/{new_name}_in/{protein}.pdb' shutil.copyfile(src, dst) if tel: tel.emit(make_event( ChildPipelineSpawned, session_id=sid, + backend="rhapsody", parent_name=pipeline.name, child_name=new_name, num_proteins=len(sub_iter_seqs), @@ -160,6 +179,9 @@ async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[s 'seq_rank': pipeline.seq_rank + 1, 'sub_order': pipeline.sub_order + 1, 'previous_scores': copy.deepcopy(pipeline.previous_scores), + 'input_base_path': pipeline.input_base_path, + 'output_base_path': pipeline.output_base_path, + 'gpu_id': pipeline.gpu_id, } } @@ -177,6 +199,7 @@ async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[s tel.emit(make_event( PassSummary, session_id=sid, + backend="rhapsody", pipeline_name=pipeline.name, pass_num=pipeline.passes, num_proteins=len(pipeline.current_scores), @@ -184,13 +207,38 @@ async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[s child_spawned=child_spawned, )) + # IMPRESS_SCRIPTS_DIR = protein_binding examples dir (scripts/, mpnn_wrapper.py) + # IMPRESS_BASE_DIR = parent of prod_in/ (input PDB files) + # IMPRESS_OUTPUT_DIR = where af_pipeline_outputs_multi/ is written + scripts_dir = os.environ.get( + "IMPRESS_SCRIPTS_DIR", os.path.dirname(os.path.abspath(__file__)) + ) + input_base_dir = os.environ.get("IMPRESS_BASE_DIR", scripts_dir) + output_base_dir = os.environ.get("IMPRESS_OUTPUT_DIR", scripts_dir) + os.makedirs(output_base_dir, exist_ok=True) + + all_gpus = find_gpus() + + if all_gpus: + print("[INFO] GPU assignment:") + for i in range(1, N_PIPELINES + 1): + gpu_id = all_gpus[(i - 1) % len(all_gpus)] + print(f"[INFO] p{i:>2} -> gpu={gpu_id}") + pipeline_setups: List[PipelineSetup] = [ PipelineSetup( name=f"p{str(i)}", type=ProteinBindingPipeline, - adaptive_fn=adaptive_decision + config={ + "base_path": scripts_dir, + "input_base_path": input_base_dir, + "output_base_path": output_base_dir, + "max_passes": MAX_PASSES, + **({"gpu_id": all_gpus[(i - 1) % len(all_gpus)]} if all_gpus else {}), + }, + adaptive_fn=adaptive_decision, ) - for i in range(1, 17) + for i in range(1, N_PIPELINES + 1) ] await manager.start(pipeline_setups=pipeline_setups) @@ -201,8 +249,8 @@ async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[s dur = summary.get("duration") if dur: print(f"[TELEMETRY] mean task time: {dur['mean_seconds'] * 1000:.1f} ms") + await manager.telemetry.stop() - await manager.flow.shutdown() if __name__ == "__main__": diff --git a/examples/protein_binding/scripts/s1_mpnn.sh b/examples/protein_binding/scripts/s1_mpnn.sh index 67f4110..377a38d 100755 --- a/examples/protein_binding/scripts/s1_mpnn.sh +++ b/examples/protein_binding/scripts/s1_mpnn.sh @@ -11,8 +11,8 @@ mpnn_path="$4" num_seqs="$5" chain="$6" -source /anvil/projects/x-nairr240405/mason/LigandMPNN/.venv/bin/activate -#source /ocean/projects/dmr170002p/hooten/LigandMPNN/.venv/bin/activate +# Re-activate the IMPRESS venv if running inside a subprocess (VIRTUAL_ENV is exported by sbatch). +[ -n "${VIRTUAL_ENV:-}" ] && source "${VIRTUAL_ENV}/bin/activate" python3 "$mpnn_script" \ -pdb="$input_path" \ diff --git a/examples/protein_binding/scripts/s4_alphafold.sh b/examples/protein_binding/scripts/s4_alphafold.sh deleted file mode 100755 index 29b2126..0000000 --- a/examples/protein_binding/scripts/s4_alphafold.sh +++ /dev/null @@ -1,25 +0,0 @@ -#!/bin/bash -set -euo pipefail - -# Step 4: AlphaFold2 multimer prediction via ColabFold -# Args: $1=fasta_path $2=output_dir - -fasta_path="$1" -output_dir="$2" - -module load modtree/gpu -module load cuda/12.8.0 -module load gcc/11.2.0 -#source /anvil/scratch/x-mason/IMPRESS/.venv/bin/activate -source /ocean/projects/dmr170002p/hooten/IMPRESS/.venv/bin/activate - -pixi run --manifest-path /anvil/scratch/x-mason/localcolabfold \ - colabfold_batch \ - --model-type alphafold2_multimer_v3 \ - --max-template-date 2020-12-01 \ - --rank multimer \ - --random-seed 999 \ - --save-all \ - --debug-logging \ - "$fasta_path" \ - "$output_dir" diff --git a/examples/protein_binding/scripts/s4_boltz.sh b/examples/protein_binding/scripts/s4_boltz.sh index bab328c..bb5b50f 100755 --- a/examples/protein_binding/scripts/s4_boltz.sh +++ b/examples/protein_binding/scripts/s4_boltz.sh @@ -2,30 +2,112 @@ set -e # Step 4: Structure prediction via Boltz -# Args: $1=fasta_path $2=output_dir +# Args: $1=fasta_path $2=output_dir $3=gpu_id (optional) fasta_path="$1" output_dir="$2" +# Optional GPU assignment passed by the caller so tasks spread across GPUs. +if [ -n "${3:-}" ]; then + export CUDA_VISIBLE_DEVICES="$3" +fi -#source /ocean/projects/dmr170002p/hooten/IMPRESS/.venv/bin/activate -source /anvil/projects/x-nairr240405/mason/IMPRESS/.venv/bin/activate -module load modtree/gpu +# Boltz requires Python <=3.12 (numpy<2.0 etc.) so it lives in its own env. +# BOLTZ_VENV may point to a conda env (no bin/activate) or a pip venv; prepend +# its bin/ to PATH so the correct python/boltz are found in either case. +_BOLTZ_ENV="${BOLTZ_VENV:-${VIRTUAL_ENV:-}}" +[ -n "${_BOLTZ_ENV}" ] && export PATH="${_BOLTZ_ENV}/bin:${PATH}" export SSL_CERT_FILE=/etc/pki/tls/certs/ca-bundle.crt -# Boltz caches MSA in boltz_results_/msa/ and reuses it across runs even with -# --override, so pass 2+ would fold new MPNN-designed sequences using the pass-1 MSA. -# Delete the stale MSA before each run to force recomputation for the current sequence. -fasta_stem=$(basename "${fasta_path}" .fa) -stale_msa="${output_dir}/boltz_results_${fasta_stem}/msa" -if [ -d "${stale_msa}" ]; then - rm -rf "${stale_msa}" +# Compute nodes typically have no internet access, so --use_msa_server is off +# by default. Set BOLTZ_USE_MSA_SERVER=1 to enable it on nodes with internet. +_MSA_FLAG="" +[ "${BOLTZ_USE_MSA_SERVER:-0}" = "1" ] && _MSA_FLAG="--use_msa_server" + +# ── Test-mode stub (IMPRESS_TEST_MODE=1) ────────────────────────────────── +# Write minimal Boltz-shaped output so downstream steps (plddt_extract_pipeline) +# can run without invoking the real Boltz model. numpy is available because +# BOLTZ_VENV/bin is already on PATH above. +if [ "${IMPRESS_TEST_MODE:-0}" = "1" ]; then + name="$(basename "${fasta_path}" .fa)" + pred_dir="${output_dir}/boltz_results_${name}/predictions/${name}" + mkdir -p "${pred_dir}" + python3 - "${pred_dir}" "${name}" <<'PYEOF' +import sys, json +import numpy as np +pred_dir, name = sys.argv[1], sys.argv[2] +n = 110 # 100 PDZ residues + 10 peptide (PEP_LEN=10 assumed by extractor) +np.savez(f"{pred_dir}/plddt_{name}_model_0.npz", plddt=np.full(n, 0.85)) +np.savez(f"{pred_dir}/pae_{name}_model_0.npz", pae=np.full((n, n), 2.0)) +with open(f"{pred_dir}/confidence_{name}_model_0.json", "w") as f: + json.dump({"iptm": 0.75, "ptm": 0.80}, f) +PYEOF + echo "[MOCK] s4_boltz stub done for ${name}" + exit 0 fi +# ── End test-mode stub ──────────────────────────────────────────────────── + +mkdir -p "${output_dir}" + +_boltz_cache_dir="${BOLTZ_CACHE_DIR:-${HOME}/.boltz}" + +# $_boltz_cache_dir is shared across concurrently-dispatched pipelines. boltz's own +# download_boltz2() checks `mols.exists()` (directory presence), not +# completeness, before skipping extraction -- tarfile.extractall() creates the +# "mols" directory entry immediately, so a *second* concurrent task calling +# download_boltz2() while a first one is still mid-extract sees mols/ already +# existing and skips extraction outright, then reads a half-populated +# directory and fails with "CCD component not found!" for whatever +# hasn't been extracted yet (see examples/small_molecule_binding/scripts/boltz.sh +# for the same fix, ported here after this exact race killed 13/16 pipelines +# in a production run). +# +# Fix: hold the lock for the entire check-and-repair, verify mols/ actually +# contains every file mols.tar lists (not just that the directory exists), +# and if not, delete and re-extract *inside* the lock via boltz's own +# download_boltz2() so no other concurrent task can observe a +# partially-populated mols/ while this one repairs it. A `.mols_complete` +# marker (written only after a verified-complete extraction) lets later +# invocations skip the O(45k) file-count re-check once warmed. +mkdir -p "$_boltz_cache_dir" +( + flock -x 200 + + tar_ok=false + if tar -tf "$_boltz_cache_dir/mols.tar" >/dev/null 2>&1; then + tar_ok=true + fi + + mols_complete=false + if $tar_ok && [ -f "$_boltz_cache_dir/.mols_complete" ]; then + mols_complete=true + elif $tar_ok && [ -d "$_boltz_cache_dir/mols" ]; then + expected=$(tar -tf "$_boltz_cache_dir/mols.tar" | grep -vc '/$') + actual=$(find "$_boltz_cache_dir/mols" -maxdepth 1 -type f | wc -l) + if [ "$actual" -eq "$expected" ]; then + mols_complete=true + touch "$_boltz_cache_dir/.mols_complete" + fi + fi + + if ! $mols_complete; then + rm -rf "$_boltz_cache_dir/mols.tar" "$_boltz_cache_dir/mols" "$_boltz_cache_dir/.mols_complete" + BOLTZ_CACHE_DIR_FOR_PY="$_boltz_cache_dir" python -c " +import os +from pathlib import Path +from boltz.main import download_boltz2 +download_boltz2(Path(os.environ['BOLTZ_CACHE_DIR_FOR_PY'])) +" + touch "$_boltz_cache_dir/.mols_complete" + fi +) 200>"$_boltz_cache_dir/.download.lock" boltz predict \ "${fasta_path}" \ --out_dir "${output_dir}" \ - --use_msa_server \ - --cache /anvil/projects/x-nairr240405/mason/boltz \ + ${_MSA_FLAG} \ + --cache "$_boltz_cache_dir" \ --output_format pdb \ --write_full_pae \ + --no_kernels \ + --devices 1 \ --override diff --git a/examples/protein_binding/scripts/s5_plddt_extract.sh b/examples/protein_binding/scripts/s5_plddt_extract.sh index 38cb576..e4999ac 100755 --- a/examples/protein_binding/scripts/s5_plddt_extract.sh +++ b/examples/protein_binding/scripts/s5_plddt_extract.sh @@ -2,16 +2,19 @@ set -e # Step 5: pLDDT extraction -# Args: $1=base_path $2=iter $3=out_name +# Args: $1=output_base_path $2=iter $3=out_name -base_path="$1" +output_base_path="$1" iter="$2" out_name="$3" -source /anvil/projects/x-nairr240405/mason/IMPRESS/.venv/bin/activate -#source /ocean/projects/dmr170002p/hooten/IMPRESS/.venv/bin/activate +# plddt_extract_pipeline.py lives one level above this scripts/ directory. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -python3 "$base_path/plddt_extract_pipeline.py" \ - --path="$base_path" \ - --iter="$iter" \ - --out="$out_name" +# Re-activate the IMPRESS venv if running inside a subprocess (VIRTUAL_ENV is exported by sbatch). +[ -n "${VIRTUAL_ENV:-}" ] && source "${VIRTUAL_ENV}/bin/activate" + +python3 "${SCRIPT_DIR}/../plddt_extract_pipeline.py" \ + --path="${output_base_path}" \ + --iter="${iter}" \ + --out="${out_name}"