diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index f9e6856..de003c5 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -14,12 +14,6 @@ jobs: fail-fast: false matrix: include: - - os: ubuntu-latest - python: '3.9' - toxenv: py39 - - os: ubuntu-latest - python: '3.10' - toxenv: py310 - os: ubuntu-latest python: '3.11' toxenv: py311 @@ -59,12 +53,6 @@ jobs: fail-fast: false matrix: include: - - os: ubuntu-latest - python: '3.9' - toxenv: py39-all - - os: ubuntu-latest - python: '3.10' - toxenv: py310-all - os: ubuntu-latest python: '3.11' toxenv: py311-all diff --git a/.gitignore b/.gitignore index 0dafbe4..e10c0d4 100644 --- a/.gitignore +++ b/.gitignore @@ -135,6 +135,11 @@ dmypy.json # asyncflow related asyncflow.session.* +ddict_* + +# ROME runtime outputs +af_stats_*.csv +examples/protien_binding_usecase/logs/ # pdzbinder wf outputs af_pipeline_outputs_multi/ diff --git a/examples/.gitignore b/examples/.gitignore new file mode 100644 index 0000000..31e9608 --- /dev/null +++ b/examples/.gitignore @@ -0,0 +1,14 @@ +# SLURM log output (all workflows) +*/logs/ + +# Per-pipeline task directories written to cwd +*/p*/ + +# Dragon telemetry session dirs +*/telemetry/ + +# Run metadata written to cwd +*/runinfo + +# Legacy output dirs +*/myoutputs/ diff --git a/examples/impress_r/delta_env_setup.sh b/examples/impress_r/delta_env_setup.sh new file mode 100755 index 0000000..9b59674 --- /dev/null +++ b/examples/impress_r/delta_env_setup.sh @@ -0,0 +1,264 @@ +#!/bin/bash +# ============================================================================= +# IMPRESS-R (Protein Binding + ROME fine-tuning) environment setup — Delta HPC +# +# Creates a Python 3.11+ venv and installs all dependencies, including ROME-A. +# +# Usage: +# export SCRATCH=/scratch/ +# bash delta_env_setup.sh [--env-dir DIR] [--impress-dir DIR] [--rome-dir DIR] [--python PATH] +# +# Defaults: +# ENV_DIR = /u/$USER/ve/impress +# IMPRESS_DIR = $SCRATCH/$USER/IMPRESS +# ROME_DIR = $SCRATCH/$USER/ROME +# 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_A}" +IMPRESS_DIR="${IMPRESS_DIR:-${SCRATCH}/${USER}/IMPRESS}" +ROME_DIR="${ROME_DIR:-${SCRATCH}/${USER}/ROME}" +BASE_PY_OVERRIDE="" + +while [[ $# -gt 0 ]]; do + case $1 in + --env-dir) ENV_DIR="$2"; shift 2 ;; + --impress-dir) IMPRESS_DIR="$2"; shift 2 ;; + --rome-dir) ROME_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 " ROME_DIR = ${ROME_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]" +# Pin dragonhpc to 0.14.1 — 0.14.2 added waitForKeys to DDRegisterClientResponse +# but the Delta system Dragon runtime has not been updated to match; 0.14.2 fails +# with AttributeError on every DDict operation on this cluster. +"${PIP}" install -q "dragonhpc==0.14.1" + +# ── 5. IMPRESS (local editable) ─────────────────────────────────────────────── +echo "" +echo "── Step 5: IMPRESS (editable) ──" +"${PIP}" install -q -e "${IMPRESS_DIR}" + +# ── 6. ROME (local editable) ───────────────────────────────────────────────── +echo "" +echo "── Step 6: ROME (editable) ──" +if [ -d "${ROME_DIR}" ]; then + "${PIP}" install -q -e "${ROME_DIR}" +else + echo "WARNING: ROME_DIR=${ROME_DIR} not found — skipping ROME install." + echo " Set --rome-dir or clone ROME before running impress_A." +fi + +# ── 7. PyTorch (CUDA 12.1) ─────────────────────────────────────────────────── +echo "" +echo "── Step 7: PyTorch (cu121) ──" +"${PIP}" install -q torch --index-url https://download.pytorch.org/whl/cu121 + +# ── 8. Additional dependencies ─────────────────────────────────────────────── +echo "" +echo "── Step 8: pandas + biopandas + matplotlib ──" +"${PIP}" install -q pandas biopandas matplotlib + +# ── 9. PyRosetta ───────────────────────────────────────────────────────────── +echo "" +echo "── Step 9: 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()" + +# ── 10. Boltz (structure prediction — separate conda env) ──────────────────── +echo "" +echo "── Step 10: 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_rome.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}" +_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)" + _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" + +# ── 11. Verify ─────────────────────────────────────────────────────────────── +echo "" +echo "── Step 11: 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 "rome" "${PY}" -c "import rome; 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=xxx-delta-gpu" +echo " cd ${IMPRESS_DIR}/examples/impress_A" +echo " sbatch delta_gpu_run.sh" +echo "" +echo "Test (smoke) run:" +echo " IMPRESS_TEST_MODE=1 ROME_TRAINER=dummy sbatch delta_gpu_run.sh" +echo "=================================================================" diff --git a/examples/impress_r/delta_gpu_run.sh b/examples/impress_r/delta_gpu_run.sh new file mode 100755 index 0000000..2a263e0 --- /dev/null +++ b/examples/impress_r/delta_gpu_run.sh @@ -0,0 +1,137 @@ +#!/bin/bash +# +# IMPRESS-R (Protein Binding + ROME fine-tuning) — SLURM batch script (Delta HPC / GPU) +# +# Set before calling sbatch (only SBATCH_ACCOUNT and SCRATCH are required): +# export SBATCH_ACCOUNT=-delta-gpu +# export SCRATCH=/scratch/ +# +# Key env vars (all have defaults): +# MPNN_PATH — dauparas/ProteinMPNN checkout (inference AND fine-tune target) +# ROME_MPNN_REPO — same as MPNN_PATH by default +# ROME_TRAINER — mpnn (real fine-tune) | dummy (smoke test, default) +# ROME_MIN_SAMPLES — corpus size before the first training round (default 2) +# ROME_MAX_PASSES — max design passes per pipeline (default 10) +# ROME_FALLBACK — seconds Dragon may take to deliver a training result (default 60) +# IMPRESS_N_PIPELINES — number of top-level pipelines to run (default 16) +# IMPRESS_MAX_SUB_PIPELINES — max child pipeline depth per parent (default 3; 0 = none) +# IMPRESS_BASE_DIR — parent of prod_in/ (input PDB files) +# IMPRESS_OUTPUT_DIR — where af_pipeline_outputs_multi/ is written +# +# Debug flags: +# IMPRESS_TEST_MODE=1 — 2 pipelines, max_passes=1, no child pipelines +# ROME_TRAINER=dummy — skip real fine-tuning (useful with IMPRESS_TEST_MODE) +# +# Quick validation run (~1 h, sees full ROME loop to completion): +# IMPRESS_N_PIPELINES=4 IMPRESS_MAX_SUB_PIPELINES=1 ROME_MAX_PASSES=8 ROME_FALLBACK=120 sbatch delta_gpu_run.sh +# +# Full production run (~3 h): +# sbatch delta_gpu_run.sh +# +# Example: +# sbatch delta_gpu_run.sh +# IMPRESS_TEST_MODE=1 ROME_TRAINER=dummy sbatch delta_gpu_run.sh +# +#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:00:00 +#SBATCH --job-name=impress_r +#SBATCH --mail-user=mg2347@soe.rutgers.edu +#SBATCH --mail-type=END,FAIL +#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 + +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 ──────────────────────────────────────────────────────────────── +# ProteinMPNN checkout — used for inference (mpnn_wrapper.py) and fine-tuning (ROME). +export MPNN_PATH="${MPNN_PATH:-${SCRATCH}/${USER}/ProteinMPNN}" +# ROME-A fine-tunes the same checkout and publishes weights back into it. +export ROME_MPNN_REPO="${ROME_MPNN_REPO:-${MPNN_PATH}}" + +# Boltz — separate Python <=3.12 env (numpy<2.0 etc.) +export BOLTZ_VENV="${BOLTZ_VENV:-${HOME}/ve/boltz}" +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/impress_r}" +export IMPRESS_BASE_DIR="${IMPRESS_BASE_DIR:-${SCRATCH}/${USER}/IMPRESS_inputs}" +export IMPRESS_OUTPUT_DIR="${IMPRESS_OUTPUT_DIR:-${SCRATCH}/${USER}/IMPRESS_outputs}" + +# ── ROME-A settings ─────────────────────────────────────────────────────────── +# ROME_TRAINER=mpnn → real ProteinMPNN fine-tune (needs ROME_MPNN_REPO on disk) +# ROME_TRAINER=dummy → smoke test (no GPU/torch needed for training) +export ROME_TRAINER="${ROME_TRAINER:-mpnn}" +export ROME_MIN_SAMPLES="${ROME_MIN_SAMPLES:-2}" +export ROME_MAX_PASSES="${ROME_MAX_PASSES:-10}" +# 120 s gives training rounds time to finish and write train_complete before Dragon's +# result-delivery future raises a spurious TypeError (dragonhpc 0.14.1 DDict race). +export ROME_FALLBACK="${ROME_FALLBACK:-120}" + +# ── Scale settings ──────────────────────────────────────────────────────────── +# Quick run (~1 h): IMPRESS_N_PIPELINES=4 IMPRESS_MAX_SUB_PIPELINES=1 ROME_MAX_PASSES=8 +# Full run (~3 h): leave unset (16 pipelines, 3 sub-levels, 10 passes) +export IMPRESS_N_PIPELINES="${IMPRESS_N_PIPELINES:-16}" +# Blank = use code default (3); set to 0 to disable child-pipeline spawning entirely. +export IMPRESS_MAX_SUB_PIPELINES="${IMPRESS_MAX_SUB_PIPELINES:-}" + +# ── Test / debug flags ──────────────────────────────────────────────────────── +export IMPRESS_TEST_MODE="${IMPRESS_TEST_MODE:-0}" + +echo "MPNN_PATH: ${MPNN_PATH}" +echo "ROME_MPNN_REPO: ${ROME_MPNN_REPO}" +echo "ROME_TRAINER: ${ROME_TRAINER}" +echo "ROME_MIN_SAMPLES: ${ROME_MIN_SAMPLES}" +echo "IMPRESS_BASE_DIR: ${IMPRESS_BASE_DIR}" +echo "IMPRESS_OUTPUT_DIR:${IMPRESS_OUTPUT_DIR}" +echo "TEST_MODE: ${IMPRESS_TEST_MODE}" + +# ── Working directory ───────────────────────────────────────────────────────── +WORKDIR="${IMPRESS_SCRIPTS_DIR}" +cd "${WORKDIR}" +mkdir -p logs + +# ── Run ─────────────────────────────────────────────────────────────────────── +if [ "${SLURM_NNODES:-1}" -gt 1 ]; then + DRAGON_MODE="-m" +else + DRAGON_MODE="-s" +fi + +rm -f ddict_orc* + +echo "Running: dragon ${DRAGON_MODE} run_protein_binding_rome.py (nodes=${SLURM_NNODES:-1})" +dragon ${DRAGON_MODE} run_protein_binding_rome.py + +echo "=== IMPRESS-R pipeline done: $(date) ===" diff --git a/examples/impress_r/mpnn.py b/examples/impress_r/mpnn.py new file mode 100644 index 0000000..606af60 --- /dev/null +++ b/examples/impress_r/mpnn.py @@ -0,0 +1,770 @@ +"""ProteinMPNN trainer task — the IMPRESS-R half of ROME-A's trainers. + +IMPRESS runs backbone -> ProteinMPNN -> structure prediction -> pLDDT/pTM/pAE +-> keep/fallback/migrate/drop, and it is open loop: every campaign improves the +designs, never the model. IMPRESS-R closes that loop by fine-tuning ProteinMPNN +on the campaign's own high-confidence designs. + +**Which ProteinMPNN.** IMPRESS runs the original ``dauparas/ProteinMPNN``: its +``mpnn_wrapper.py`` shells out to ``protein_mpnn_run.py`` with the original CLI +and helper scripts, and its setup clones that repo directly. So this trainer +targets the *same* implementation — it fine-tunes the original model and writes +a checkpoint in the original ``{"model_state_dict": ...}`` format that +``protein_mpnn_run.py`` loads. A foundry / re-implementation checkpoint would +not load into what IMPRESS runs; that was the previous version's mistake. + +(PyRosetta is in the IMPRESS stack too, but for FastRelax and pLDDT extraction, +not for sequence design — the design model is vanilla ProteinMPNN.) + +**Dimers, not monomers.** An IMPRESS protein-binding design is a *complex*: a +designed chain (``A``) plus a fixed target peptide (``B``). Fine-tuning has to +respect that — the model should learn to design chain A *in the context of* the +peptide, scoring sequence recovery on the designed chain only. That is exactly +what ProteinMPNN's chain mask expresses, and :func:`build_chain_designation` +sets it up: designed chains are predicted, context chains are visible but not +scored. + +This trainer is an IMPRESS-R *integration*, not framework core — ROME-A is +workflow-agnostic — so it lives with the example, beside the inference wrapper +(``mpnn_wrapper.py``) it complements. Import it from here:: + + from examples.impress_r.mpnn import ProteinMPNNTrainer, ProteinMPNNConfig + +**Runs as a command, not a function.** Like IMPRESS — which submits +``mpnn_wrapper.py`` as a shell command rather than calling ProteinMPNN in the +campaign process — the training manager submits this round as an *executable +task*: :meth:`ProteinMPNNTrainer.as_command` stages the structures, writes a +self-contained job spec, and returns ``python mpnn_train_wrapper.py --job +``. The fine-tune therefore runs in its own process on its own GPU, +and that process exits when the round finishes, so its VRAM is released with it. + +The training loop itself lives in ``mpnn_train_wrapper.run_round`` (the sibling +script) — one dragon-free copy, so the standalone script and the in-process path +(a direct :meth:`train` call, used by the tests) share it. It mirrors +``training/training.py``'s inner loop using the repo's own ``featurize``, +``loss_smoothed``, ``NoamOpt`` and training ``ProteinMPNN``, verified end to end +against a real checkout and the public ``v_48_020`` weights — the loss runs on +the designed chain only and the checkpoint reloads into ``protein_mpnn_run.py``. +``config.train_func`` substitutes your own loop (and keeps it in-process). See +``docs/proteinmpnn_training.md``. +""" + +from __future__ import annotations + +import json +import os +import shutil +import sys +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +from rome.train.base import TrainTask + + +def _load_run_round(): + """Import ``run_round`` from the sibling ``mpnn_train_wrapper``. + + Works whether this module was imported as ``examples.impress_r.mpnn`` (the + test/package path) or as a bare ``mpnn`` (a script run from inside the + example directory): the wrapper sits next to this file, so its directory is + put on ``sys.path`` and it is imported by name. Lazy — only the in-process + :meth:`ProteinMPNNTrainer.train` path needs it; the command path never + imports the wrapper here, it runs it as a subprocess. + """ + here = os.path.dirname(os.path.abspath(__file__)) + if here not in sys.path: + sys.path.insert(0, here) + from mpnn_train_wrapper import run_round # type: ignore + + return run_round + + +#: ProteinMPNN's amino-acid alphabet (index -> letter). Position 20 is ``X``. +MPNN_ALPHABET = "ACDEFGHIKLMNPQRSTVWYX" + +#: A corpus record must carry ``path`` — the structure file whose designed chain +#: IS the training label. ``sequence`` is optional (ProteinMPNN reads the label +#: out of the structure); it is kept for the manifest and for sizing. +REQUIRED_FIELDS = ("path",) + +#: The IMPRESS protein-binding layout: chain A is the designed binder, chain B +#: the fixed target peptide. Designed chains are scored; context chains are +#: visible to the model but not scored. +DEFAULT_DESIGN_CHAINS: Tuple[str, ...] = ("A",) +DEFAULT_CONTEXT_CHAINS: Tuple[str, ...] = ("B",) + +#: Architecture of the public ``v_48_020`` weights, which is what IMPRESS runs. +#: ``protein_mpnn_run.py`` hardcodes hidden_dim=128 and 3+3 layers and reads +#: ``k_neighbors`` back from the checkpoint's ``num_edges`` — so a fine-tune must +#: keep these or the published checkpoint will not load. +DEFAULT_NUM_NEIGHBORS = 48 +DEFAULT_HIDDEN_DIM = 128 +DEFAULT_NUM_LAYERS = 3 +DEFAULT_BACKBONE_NOISE = 0.2 + + +# --------------------------------------------------------------------------- +# Pure-Python helpers (no torch): chain designation, staging, checkpoint format +# --------------------------------------------------------------------------- + +def pdb_chain_ids(path: str) -> List[str]: + """Chain IDs present in a PDB, in first-seen order — column 22 of ATOM/HETATM. + + Plain text parsing so a campaign's monomer-vs-complex shape can be checked + (and a bad chain designation caught) without importing a structure library. + """ + seen: List[str] = [] + try: + with open(path) as fd: + for line in fd: + if line.startswith(("ATOM", "HETATM")) and len(line) >= 22: + chain = line[21] + if chain not in seen: + seen.append(chain) + except OSError: + return [] + return seen + + +def build_chain_designation( + records: List[Dict[str, Any]], + *, + design_chains: Sequence[str] = DEFAULT_DESIGN_CHAINS, + context_chains: Sequence[str] = DEFAULT_CONTEXT_CHAINS, + chains_func: Optional[Callable[[Dict[str, Any], List[str]], Tuple[List[str], List[str]]]] = None, +) -> Dict[str, Tuple[List[str], List[str]]]: + """Map each design to ``(designed_chains, context_chains)``. + + This is the dimer fix. ProteinMPNN's ``tied_featurize`` takes exactly this + mapping and builds the chain mask from it: designed chains are predicted and + scored, context chains are visible to the model but excluded from the loss. + For an IMPRESS binder that means "learn chain A given chain B", which is the + thing the campaign is actually optimising. + + A record may override the defaults per structure with a ``design_chains`` / + ``context_chains`` field, or a ``chains_func(record, present_chains)`` may + decide from the structure itself. Chains named but absent from the file are + dropped with the designation still valid for whatever remains; a design with + no designable chain present is an error, caught in + :meth:`ProteinMPNNTrainer.validate`. + + Returns ``{design_name: (designed, context)}`` keyed by the staged structure + name (see :func:`stage_structures`), which is what ``tied_featurize`` keys on. + """ + out: Dict[str, Tuple[List[str], List[str]]] = {} + for index, record in enumerate(records): + name = _design_name(record, index) + present = pdb_chain_ids(record["path"]) if record.get("path") else [] + if chains_func is not None: + designed, context = chains_func(record, present) + else: + designed = list(record.get("design_chains", design_chains)) + context = list(record.get("context_chains", context_chains)) + if present: + designed = [c for c in designed if c in present] + context = [c for c in context if c in present] + out[name] = (designed, context) + return out + + +def _design_name(record: Dict[str, Any], index: int) -> str: + """Stable, filesystem-safe name for a design; also its key in the chain map.""" + raw = str(record.get("uid") or record.get("design_id") + or os.path.splitext(os.path.basename(record.get("path", "")))[0] + or f"design_{index}") + return "".join(c if (c.isalnum() or c in "-_") else "_" for c in raw) + + +def stage_structures(records: List[Dict[str, Any]], staging_dir: str) -> Dict[str, str]: + """Copy each design's structure into one directory under a unique name. + + ProteinMPNN parses a *folder* of PDBs, and a campaign's paths collide on + basename (every pipeline writes ``{target}.pdb``) and get overwritten pass to + pass. Staging copies each file under its unique design name, so the training + set is a stable snapshot rather than a set of paths whose contents move. + + Returns ``{design_name: staged_path}``. + """ + os.makedirs(staging_dir, exist_ok=True) + staged: Dict[str, str] = {} + for index, record in enumerate(records): + name = _design_name(record, index) + dst = os.path.join(staging_dir, f"{name}.pdb") + shutil.copyfile(record["path"], dst) + staged[name] = dst + return staged + + +def original_checkpoint( + model_state_dict: Any, *, num_edges: int = DEFAULT_NUM_NEIGHBORS, + noise_level: float = DEFAULT_BACKBONE_NOISE, **extra: Any, +) -> Dict[str, Any]: + """The checkpoint dict ``protein_mpnn_run.py`` loads. + + It reads ``checkpoint['model_state_dict']`` and ``checkpoint['num_edges']`` + (used as ``k_neighbors`` when constructing the model), so both are required + for the published weights to load at all. ``noise_level`` is carried for + parity with the original trainer's saves. + """ + ckpt = {"model_state_dict": model_state_dict, + "num_edges": int(num_edges), "noise_level": float(noise_level)} + ckpt.update(extra) + return ckpt + + +def published_weights_path(config: "ProteinMPNNConfig", output_dir: str) -> str: + """Where the round's weights go so IMPRESS's next pass picks them up. + + ``mpnn_wrapper.py`` never passes ``--path_to_model_weights``, so + ``protein_mpnn_run.py`` loads ``{mpnn_repo}/vanilla_model_weights/{model_name}.pt`` + by default. With ``publish_into_repo`` set (and ``mpnn_repo`` known) the + checkpoint is written *there*, replacing the weights the campaign runs with; + otherwise it lands in the round's ``output_dir`` and the integration is + responsible for pointing MPNN at it (e.g. patch the wrapper to pass a path). + """ + if config.publish_into_repo and config.mpnn_repo: + return os.path.join(config.mpnn_repo, "vanilla_model_weights", + f"{config.model_name}.pt") + return os.path.join(output_dir, f"{config.model_name}.pt") + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + +@dataclass +class ProteinMPNNConfig: + """Configuration for fine-tuning the original ProteinMPNN on campaign designs. + + Defaults describe a short mid-campaign round that starts from the public + weights and nudges them, not a from-scratch training run. + + Parameters + ---------- + mpnn_repo : Optional[str] + Path to the ``dauparas/ProteinMPNN`` checkout — the same directory + IMPRESS passes as ``-mpnn``. The trainer imports ``protein_mpnn_utils`` + from here, and (with ``publish_into_repo``) writes the new weights back + into its ``vanilla_model_weights/``. + initial_weights : Optional[str] + Checkpoint the first round starts from, in original format — normally + ``{mpnn_repo}/vanilla_model_weights/v_48_020.pt``. Later rounds resume + from whatever the previous round published (passed as ``model_path``). + model_name : str + Basename of the published weights (no extension). Must match the + ``--model_name`` IMPRESS runs with; the default ``v_48_020`` is + ``protein_mpnn_run.py``'s default. + design_chains, context_chains : sequence of str + Default chain designation. For IMPRESS binders, ``("A",)`` designed and + ``("B",)`` context. Overridable per record or via ``chains_func``. + chains_func : optional callable + ``(record, present_chains) -> (designed, context)`` to decide chains per + structure. Overrides the defaults when set. + num_neighbors, hidden_dim, num_layers, backbone_noise, dropout : ... + Architecture. Must match the weights being fine-tuned — the public + ``v_48_*`` weights are ``num_neighbors`` in {48}, hidden_dim 128, 3+3 + layers. ``protein_mpnn_run.py`` reads ``num_neighbors`` back from the + checkpoint, so a mismatch will not load. + max_epochs, batch_tokens, max_protein_length : int + Round length and batching. ``batch_tokens`` caps residues per batch + (ProteinMPNN uses ~10000). A short round is a few epochs. + learning_rate_factor, warmup_steps, label_smoothing, gradient_norm : ... + Noam schedule + label-smoothed NLL, as the original trainer uses. + publish_into_repo : bool + Write the new weights into ``{mpnn_repo}/vanilla_model_weights/`` so the + campaign's next pass runs them with no wrapper change. Off by default: + it mutates the shared repo, which a workflow should opt into knowingly. + seed, num_workers, device : ... + Reproducibility / loader / placement. + train_func : optional callable + ``(manifest_path, output_dir, config) -> checkpoint_path`` escape hatch, + bypassing the built-in loop. The recommended path until the loop is + validated on your checkout, and the way the examples stay runnable. + """ + + mpnn_repo: Optional[str] = None + initial_weights: Optional[str] = None + model_name: str = "v_48_020" + dataset_name: str = "impress_r" + + design_chains: Sequence[str] = DEFAULT_DESIGN_CHAINS + context_chains: Sequence[str] = DEFAULT_CONTEXT_CHAINS + chains_func: Optional[Callable[..., Tuple[List[str], List[str]]]] = None + + num_neighbors: int = DEFAULT_NUM_NEIGHBORS + hidden_dim: int = DEFAULT_HIDDEN_DIM + num_layers: int = DEFAULT_NUM_LAYERS + backbone_noise: float = DEFAULT_BACKBONE_NOISE + dropout: float = 0.1 + + max_epochs: int = 3 + batch_tokens: int = 10000 + max_protein_length: int = 2000 + + learning_rate_factor: float = 2.0 + warmup_steps: int = 4000 + label_smoothing: float = 0.1 + gradient_norm: Optional[float] = 1.0 + + publish_into_repo: bool = False + seed: int = 0 + num_workers: int = 4 + device: str = "cuda" + + manifest_dir: Optional[str] = None + train_func: Optional[Callable[..., str]] = None + #: The wrapper script the round is submitted as a command to run. Defaults to + #: the sibling ``examples/impress_r/mpnn_train_wrapper.py``; override to point + #: at a copy staged elsewhere on the cluster (as IMPRESS points ``-mpnn`` at + #: its own checkout). + train_script: Optional[str] = None + + def validate(self) -> None: + if not self.design_chains and self.chains_func is None: + raise ValueError( + "at least one design chain is required (design_chains is empty " + "and no chains_func was given)" + ) + if self.train_func is None and not self.mpnn_repo: + raise ValueError( + "mpnn_repo must point at a dauparas/ProteinMPNN checkout for the " + "built-in trainer; set config.train_func to use your own loop " + "instead. See docs/proteinmpnn_training.md." + ) + + +# --------------------------------------------------------------------------- +# The training task +# --------------------------------------------------------------------------- + +class ProteinMPNNTrainer(TrainTask): + """Fine-tunes the original ProteinMPNN on the campaign's confident designs. + + Returns the checkpoint file (original format), which the training manager + publishes and — with ``publish_into_repo`` — drops straight into the repo's + weights directory so IMPRESS's next pass runs it. + + Set ``config.train_func`` to bypass the built-in loop with your own + ``(manifest_path, output_dir, config) -> checkpoint_path``. Recommended + until the built-in loop is validated on your ProteinMPNN checkout. + """ + + #: Records are plain dicts of paths and scores, not a HuggingFace dataset. + wants_hf_dataset = False + + def __init__( + self, + config: Optional[ProteinMPNNConfig] = None, + *, + gpus: Optional[int] = None, + nodes: Optional[int] = None, + name: Optional[str] = None, + ): + config = config or ProteinMPNNConfig() + config.validate() + super().__init__( + gpus=gpus if gpus is not None else 1, + nodes=nodes if nodes is not None else 1, + name=name or "proteinmpnn", + ) + self.config = config + + # -- pre-flight --------------------------------------------------------- + + def validate(self, dataset: Any) -> None: + super().validate(dataset) + cfg = self.config + for index, record in enumerate(dataset): + record = record or {} + missing = [f for f in REQUIRED_FIELDS if f not in record] + if missing: + raise ValueError( + f"ProteinMPNN training needs {', '.join(REQUIRED_FIELDS)} on " + f"every record; corpus record {index} is missing " + f"{', '.join(missing)}. 'path' must point at the structure " + "whose designed chain is the training label (for IMPRESS-R, " + "the prediction of the designed sequence) — see " + "docs/proteinmpnn_training.md." + ) + # A design with no designable chain present in the file trains + # nothing and would produce an empty loss mask; catch it here rather + # than on the GPU. + present = pdb_chain_ids(record["path"]) if os.path.exists(record["path"]) else None + if present is not None: + if cfg.chains_func is not None: + designed, _ = cfg.chains_func(record, present) + else: + designed = list(record.get("design_chains", cfg.design_chains)) + if present and not any(c in present for c in designed): + raise ValueError( + f"corpus record {index} ({record['path']}): none of the " + f"designed chains {list(designed)} are present in the " + f"structure (chains: {present}). For an IMPRESS binder the " + "designed chain is 'A' and the peptide 'B'." + ) + + # -- corpus materialization (testable, no torch) ------------------------ + + def write_manifest(self, records: List[Dict[str, Any]], output_dir: str) -> str: + """Write the round's training manifest (parquet); returns its path. + + The audit trail for "what did this round train on": one row per design + with its staged structure, chain designation, and scores. Public and + side-effect-light so a workflow can inspect a round before trusting it. + """ + import pandas as pd + + cfg = self.config + manifest_dir = cfg.manifest_dir or os.path.join(output_dir, "manifest") + os.makedirs(manifest_dir, exist_ok=True) + + designation = build_chain_designation( + records, design_chains=cfg.design_chains, + context_chains=cfg.context_chains, chains_func=cfg.chains_func, + ) + rows = [] + for index, record in enumerate(records): + name = _design_name(record, index) + designed, context = designation[name] + rows.append({ + "design_id": name, + "path": os.path.abspath(record["path"]), + "designed_chains": ",".join(designed), + "context_chains": ",".join(context), + "sequence": record.get("sequence") or "", + "backbone_id": record.get("backbone_id"), + "pLDDT": record.get("pLDDT"), + "pTM": record.get("pTM"), + "pAE": record.get("pAE"), + "produced_under_version": record.get("model_version"), + }) + path = os.path.join(manifest_dir, "train_manifest.parquet") + pd.DataFrame(rows).to_parquet(path) + return path + + # -- the round: prepare a job, run it as a command ---------------------- + + def as_command(self, dataset: Any, output_dir: str, + **kwargs: Any) -> Optional[Tuple[str, str]]: + """Submit this round as a shell command, IMPRESS-style. + + Stages the round's structures, writes a self-contained job spec, and + returns ``(command, checkpoint_path)`` where ``command`` runs + ``mpnn_train_wrapper.py`` on that spec. The training manager runs it as an + *executable task*, so the fine-tune is a separate process on its own GPU + rather than a function in the manager's address space. + + Returns ``None`` when ``config.train_func`` is set — a custom loop is + Python, not a command, so it runs in-process through :meth:`train`. + """ + if self.config.train_func is not None: + return None + records = list(dataset) + self.write_manifest(records, output_dir) + job, target = self._build_job(records, output_dir, **kwargs) + job_path = os.path.join(output_dir, "train_job.json") + with open(job_path, "w") as fd: + json.dump(job, fd, indent=2) + + script = self.config.train_script or os.path.join( + os.path.dirname(os.path.abspath(__file__)), "mpnn_train_wrapper.py" + ) + command = f"{sys.executable} {script} --job {job_path}" + return command, target + + def train(self, dataset: Any, output_dir: str, **kwargs: Any) -> str: + """Fine-tune ProteinMPNN in-process; return the checkpoint. + + This is the direct/in-process path — the training manager normally goes + through :meth:`as_command` instead. Kept because a workflow (or a test) + may want to run a round synchronously, and because ``config.train_func`` + plugs in here. It shares the exact loop the command runs, via + ``mpnn_train_wrapper.run_round``. + + ``kwargs`` carries ``model_version`` from the training manager and + ``model_path`` when a previous round published one. + """ + records = list(dataset) + manifest_path = self.write_manifest(records, output_dir) + + if self.config.train_func is not None: + return self.config.train_func(manifest_path, output_dir, self.config) \ + or output_dir + + run_round = _load_run_round() + job, _target = self._build_job(records, output_dir, **kwargs) + return run_round(job) + + def _build_job(self, records: List[Dict[str, Any]], output_dir: str, + **kwargs: Any) -> Tuple[Dict[str, Any], str]: + """Stage structures and assemble the job spec the wrapper consumes. + + Staging takes a snapshot of each design's structure under a unique name + (a campaign overwrites ``{target}.pdb`` pass to pass), and the chain + designation is resolved here — on the manager side, where the corpus and + the config live — so the wrapper only has to parse and train. Returns + ``(job, target_weights_path)``; see ``mpnn_train_wrapper`` for + the spec. + """ + cfg = self.config + staged = stage_structures(records, os.path.join(output_dir, "structures")) + designation = build_chain_designation( + records, design_chains=cfg.design_chains, + context_chains=cfg.context_chains, chains_func=cfg.chains_func, + ) + designs = [] + for index, record in enumerate(records): + name = _design_name(record, index) + designed, context = designation[name] + designs.append({ + "name": name, + "path": os.path.abspath(staged[name]), + "designed_chains": designed, + "context_chains": context, + }) + + target = os.path.abspath(published_weights_path(cfg, output_dir)) + job = { + "mpnn_repo": cfg.mpnn_repo, + "resume_from": kwargs.get("model_path") or cfg.initial_weights, + "target_weights": target, + "output_dir": os.path.abspath(output_dir), + "designs": designs, + "hyperparams": { + "hidden_dim": cfg.hidden_dim, + "num_layers": cfg.num_layers, + "num_neighbors": cfg.num_neighbors, + "backbone_noise": cfg.backbone_noise, + "dropout": cfg.dropout, + "max_epochs": cfg.max_epochs, + "batch_tokens": cfg.batch_tokens, + "max_protein_length": cfg.max_protein_length, + "learning_rate_factor": cfg.learning_rate_factor, + "warmup_steps": cfg.warmup_steps, + "label_smoothing": cfg.label_smoothing, + "gradient_norm": cfg.gradient_norm, + "seed": cfg.seed, + "device": cfg.device, + "model_name": cfg.model_name, + }, + } + return job, target + + +# --------------------------------------------------------------------------- +# Corpus selection — orthogonal to the trainer, used by DataConfig +# --------------------------------------------------------------------------- + +def impress_corpus_filter( + min_pLDDT: float = 80.0, + min_pTM: float = 0.8, + max_pAE: float = 5.0, +) -> Callable[[Dict[str, Any]], bool]: + """Build the IMPRESS admission predicate for :class:`~rome.data.DataConfig`. + + IMPRESS-R only trains on designs the campaign is confident in, and these are + the thresholds IMPRESS already uses to decide a design is worth keeping:: + + DataConfig(min_samples=24, filter_func=impress_corpus_filter()) + + .. warning:: + + **These defaults are known to be too permissive, and are kept only because + a correct replacement cannot be derived yet.** Measured against a real + PDZ campaign they admit 83% of records: ``pLDDT >= 80`` alone admits 100%, + because everything reaching the score CSVs has already cleared IMPRESS's + own keep/drop rule — the filter is being applied downstream of itself. + + That campaign was run with **Boltz**, while the branch ROME-A targets + (``archive/ipdps_pdz_usecase``) runs **AlphaFold2-multimer**, and the two + predictors do not share a confidence scale. Prefer :func:`percentile_sampler`, + which needs no scale; see ``docs/impress.md``. + """ + + def _passes(record: Dict[str, Any]) -> bool: + plddt = record.get("pLDDT", record.get("score")) + if plddt is None or plddt < min_pLDDT: + return False + if record.get("pTM", min_pTM) < min_pTM: + return False + if record.get("pAE", max_pAE) > max_pAE: + return False + return True + + return _passes + + +def score_percentiles( + records: Sequence[Dict[str, Any]], + keys: Sequence[str] = ("pLDDT", "pTM", "pAE"), +) -> Dict[str, Dict[str, float]]: + """Summarise the score distribution a campaign has produced so far. + + Calibration data, gathered from the run itself. Every confidence threshold + in IMPRESS-R is predictor-specific — AlphaFold2-multimer and Boltz do not + share a scale — so the only safe way to set one is to look at what *this* + campaign is producing:: + + from examples.impress_r.mpnn import score_percentiles + print(score_percentiles(manager.data.get_records())) + + Returns ``{key: {"n", "min", "p10", "p25", "median", "p75", "p90", "max"}}``, + skipping keys no record carries. + """ + out: Dict[str, Dict[str, float]] = {} + for key in keys: + values = sorted( + float(r[key]) for r in records + if r.get(key) is not None and _is_number(r[key]) + ) + if not values: + continue + + def q(p: float) -> float: + return values[min(len(values) - 1, int(p * len(values)))] + + out[key] = { + "n": float(len(values)), "min": values[0], "p10": q(0.10), + "p25": q(0.25), "median": q(0.50), "p75": q(0.75), + "p90": q(0.90), "max": values[-1], + } + return out + + +def _is_number(value: Any) -> bool: + try: + float(value) + except (TypeError, ValueError): + return False + return True + + +#: Default ranking for the PDZ binder case: interface pAE down, pTM up. pLDDT is +#: deliberately absent — in a measured campaign it never fell below 88, so it +#: separates almost nothing. +DEFAULT_RANK_BY = {"pAE": "low", "pTM": "high"} + + +def percentile_sampler( + fraction: float = 0.33, + *, + rank_by: Optional[Dict[str, str]] = None, + min_shard: int = 8, + on_summary: Optional[Callable[[Dict[str, Any]], None]] = None, +) -> Callable[[List[Dict[str, Any]]], List[Dict[str, Any]]]: + """Select the best ``fraction`` of the corpus, calibrating as it goes. + + Use this instead of tuning :func:`impress_corpus_filter`'s thresholds:: + + DataConfig(min_samples=24, sample_func=percentile_sampler(0.33)) + + **Why a fraction rather than a cutoff.** A threshold like ``pTM >= 0.90`` + is a claim about a specific predictor's confidence scale, and IMPRESS + campaigns have been run on both AlphaFold2-multimer and Boltz, which do not + share one. A fraction says "the best third of what this campaign has + produced", which needs no scale and calibrates itself on the fly — including + on the first round, before anyone has seen the distribution. + + It also sidesteps the trap that makes IMPRESS's own keep/drop thresholds + useless here: everything in the corpus already cleared them, so reusing them + as an admission filter selects nothing. + + Ranking is by **average rank across ``rank_by``**, not by a weighted sum of + raw values. Rank-averaging is non-parametric, so pTM (0–1) and pAE (Å, open + ended) contribute equally without needing to be normalised, and it is + unaffected by either metric's outliers. + + Parameters + ---------- + fraction : float + Portion of the corpus to keep, in (0, 1]. + rank_by : Optional[Dict[str, str]] + Record field -> ``'high'`` (higher is better) or ``'low'``. Defaults to + :data:`DEFAULT_RANK_BY`. Fields absent from every record are ignored, so + a corpus carrying only some of them still ranks. + min_shard : int + Never return fewer than this many records — a strict fraction of a small + early corpus can otherwise produce a shard too small to train on. Capped + at the corpus size. + on_summary : Optional[Callable[[dict], None]] + Called with ``{"corpus", "selected", "ranked_by", "percentiles", + "cutoffs"}`` each time a shard is built. Pass ``print`` or a logger to + watch the campaign's distribution move, and to read off the thresholds + an equivalent fixed filter would have used. + """ + if not 0.0 < fraction <= 1.0: + raise ValueError(f"fraction must be in (0, 1], got {fraction}") + directions = dict(rank_by if rank_by is not None else DEFAULT_RANK_BY) + for key, direction in directions.items(): + if direction not in ("high", "low"): + raise ValueError( + f"rank_by[{key!r}] must be 'high' or 'low', got {direction!r}" + ) + + def _sample(records: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + records = list(records) + if not records: + return records + + usable = { + key: direction for key, direction in directions.items() + if any(_is_number(r.get(key)) for r in records) + } + if not usable: + # Nothing to rank on: keep the corpus rather than silently + # returning an arbitrary slice of it. + return records + + # Average rank across metrics. Missing values sort last within their + # metric, so a partially-scored record is penalised but not dropped. + total = {id(r): 0.0 for r in records} + for key, direction in usable.items(): + ordered = sorted( + records, + key=lambda r: ( + not _is_number(r.get(key)), + -float(r[key]) if _is_number(r.get(key)) and direction == "high" + else (float(r[key]) if _is_number(r.get(key)) else 0.0), + ), + ) + for position, record in enumerate(ordered): + total[id(record)] += position + + ranked = sorted(records, key=lambda r: (total[id(r)], str(r.get("uid", "")))) + keep = max(min(min_shard, len(records)), int(round(len(records) * fraction))) + selected = ranked[:keep] + + if on_summary is not None: + cutoffs = {} + for key, direction in usable.items(): + values = [float(r[key]) for r in selected if _is_number(r.get(key))] + if values: + cutoffs[key] = min(values) if direction == "high" else max(values) + on_summary({ + "corpus": len(records), + "selected": len(selected), + "ranked_by": usable, + "percentiles": score_percentiles(records, tuple(usable)), + "cutoffs": cutoffs, + }) + return selected + + return _sample + + +__all__ = [ + "ProteinMPNNConfig", + "ProteinMPNNTrainer", + "build_chain_designation", + "stage_structures", + "pdb_chain_ids", + "original_checkpoint", + "published_weights_path", + "impress_corpus_filter", + "percentile_sampler", + "score_percentiles", + "DEFAULT_RANK_BY", + "DEFAULT_DESIGN_CHAINS", + "DEFAULT_CONTEXT_CHAINS", + "MPNN_ALPHABET", +] diff --git a/examples/impress_r/mpnn_train_wrapper.py b/examples/impress_r/mpnn_train_wrapper.py new file mode 100644 index 0000000..01ae8d5 --- /dev/null +++ b/examples/impress_r/mpnn_train_wrapper.py @@ -0,0 +1,242 @@ +#!/usr/bin/env python3 +"""Standalone ProteinMPNN fine-tuning wrapper — ROME-A's training executable. + +This is the script the training manager submits **as a shell command**, the same +way IMPRESS submits ``mpnn_wrapper.py`` for inference rather than calling a Python +function inside the campaign process. The manager stages the round's structures, +writes a self-contained job spec (a JSON file), and runs:: + + python examples/impress_r/mpnn_train_wrapper.py --job + +The round therefore executes in its *own* process on its own GPU: nothing about +the fine-tune lives in the manager's address space, and the process exits when +the round finishes, so the CUDA context and the model are released with it. + +It sits beside the inference wrapper (``mpnn_wrapper.py``) it complements, in the +IMPRESS-R example rather than in the framework — the trainer is an integration, +and ROME-A itself is workflow-agnostic. The file is deliberately dragon-free — it +imports only the standard library, torch, and the ``dauparas/ProteinMPNN`` +checkout named in the job — so it can be run and debugged on its own, exactly +like ``mpnn_wrapper.py``. It is also the single source of truth for the training +loop: ``mpnn.py`` imports :func:`run_round` for the in-process path, so there is +only one copy of the loop. + +Job spec (all keys written by ``ProteinMPNNTrainer``):: + + { + "mpnn_repo": "/path/to/ProteinMPNN", # the checkout, for its training modules + "resume_from": "/path/to/v_48_020.pt", # initial weights, or the previous round + "target_weights": "/path/to/out.pt", # where to write the new checkpoint + "designs": [ # one per staged structure + {"name": "d0", "path": "/stage/d0.pdb", + "designed_chains": ["A"], "context_chains": ["B"]} + ], + "hyperparams": { ... } # architecture + Noam schedule + } +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from typing import Any, Dict, Tuple + + +def _import_proteinmpnn(mpnn_repo: str): + """Import the checkout's *training* modules and ``parse_PDB``. + + The dauparas repo ships two ``ProteinMPNN`` classes: ``protein_mpnn_utils`` + (inference; ``forward`` takes ``randn``) and ``training/model_utils`` + (training; ``forward`` generates the decoding order itself). The fine-tune + must use the *training* one, together with the repo's ``featurize``, + ``loss_smoothed`` and ``NoamOpt``. ``parse_PDB`` comes from the inference + module. + + Returns ``(parse_PDB, featurize, loss_smoothed, NoamOpt, ProteinMPNN, + StructureDataset, StructureLoader)`` — the exact objects + ``training/training.py`` uses. + """ + training_dir = os.path.join(mpnn_repo, "training") + if not os.path.isdir(training_dir): + raise FileNotFoundError( + f"{mpnn_repo!r} has no training/ directory — is this a " + "dauparas/ProteinMPNN checkout? The trainer needs its " + "training/model_utils.py and training/utils.py." + ) + # training/ first so `model_utils`/`utils` resolve to the training copies, + # then the repo root for `protein_mpnn_utils`. + for path in (mpnn_repo, training_dir): + if path not in sys.path: + sys.path.insert(0, path) + + from model_utils import ( # type: ignore # training/model_utils.py + NoamOpt, + ProteinMPNN, + featurize, + loss_smoothed, + ) + from protein_mpnn_utils import parse_PDB # type: ignore # repo root + from utils import StructureDataset, StructureLoader # type: ignore # training/utils.py + + return (parse_PDB, featurize, loss_smoothed, NoamOpt, ProteinMPNN, + StructureDataset, StructureLoader) + + +def _pdb_dicts(job: Dict[str, Any], parse_PDB) -> list: + """Parse each staged structure and attach its chain mask. + + The designed chain(s) go in ``masked_list`` (predicted and scored); anything + else present is ``visible_list`` (context — its backbone and true sequence + condition the prediction but are excluded from the loss). ``featurize`` turns + that into ``chain_M``. This is the dimer split IMPRESS-R relies on: "design + chain A given chain B". + """ + out = [] + for design in job["designs"]: + designed = list(design.get("designed_chains") or []) + for entry in parse_PDB(design["path"]): + present = [k[len("seq_chain_"):] for k in entry + if k.startswith("seq_chain_")] + masked = [c for c in designed if c in present] + visible = [c for c in present if c not in masked] + entry["name"] = design["name"] + entry["masked_list"] = masked + entry["visible_list"] = visible + out.append(entry) + return out + + +def run_round(job: Dict[str, Any]) -> str: + """Run one fine-tuning round from a job spec; return the checkpoint path. + + Mirrors ``training/training.py``'s inner loop exactly, using the checkout's + own ``featurize`` / ``loss_smoothed`` / ``NoamOpt`` and the training + ``ProteinMPNN``. The loss is over ``mask * chain_M`` — resolved residues of + the designed chain(s) only. The checkpoint is written in the original + ``{"model_state_dict", "num_edges", ...}`` format that ``protein_mpnn_run.py`` + loads, from a CPU snapshot so the GPU copy can be freed straight after. + """ + import gc + + import torch + + hp = job["hyperparams"] + (parse_PDB, featurize, loss_smoothed, NoamOpt, ProteinMPNN, + StructureDataset, StructureLoader) = _import_proteinmpnn(job["mpnn_repo"]) + + torch.manual_seed(int(hp["seed"])) + device = torch.device(hp["device"] if torch.cuda.is_available() else "cpu") + hd = int(hp["hidden_dim"]) + nlayers = int(hp["num_layers"]) + nedges = int(hp["num_neighbors"]) + noise = float(hp["backbone_noise"]) + model = optimizer = None + try: + dataset = StructureDataset( + _pdb_dicts(job, parse_PDB), verbose=False, truncate=None, + max_length=int(hp["max_protein_length"]), + ) + if len(dataset) == 0: + raise RuntimeError( + "no structures survived parsing/length filtering " + f"(max_protein_length={hp['max_protein_length']}); nothing to train on." + ) + loader = StructureLoader(dataset, batch_size=int(hp["batch_tokens"])) + + # -- model at the public v_48 architecture, resume prior weights ----- + model = ProteinMPNN( + num_letters=21, node_features=hd, edge_features=hd, hidden_dim=hd, + num_encoder_layers=nlayers, num_decoder_layers=nlayers, + k_neighbors=nedges, augment_eps=noise, dropout=float(hp["dropout"]), + ).to(device) + + step = 0 + resume = job.get("resume_from") + if resume: + state = torch.load(resume, map_location=device) + model.load_state_dict(state["model_state_dict"] if "model_state_dict" + in state else state) + step = int(state.get("step", 0)) # continue the Noam schedule + model.train() + + optimizer = NoamOpt( + hd, float(hp["learning_rate_factor"]), int(hp["warmup_steps"]), + torch.optim.Adam(model.parameters(), lr=0.0, betas=(0.9, 0.98), + eps=1e-9), + step, + ) + + gradient_norm = hp.get("gradient_norm") + # -- the fine-tuning loop (training/training.py, verbatim) ----------- + for _epoch in range(int(hp["max_epochs"])): + for batch in loader: + X, S, mask, lengths, chain_M, residue_idx, mask_self, \ + chain_encoding_all = featurize(batch, device) + optimizer.zero_grad() + mask_for_loss = mask * chain_M + log_probs = model(X, S, mask, chain_M, residue_idx, + chain_encoding_all) + _, loss = loss_smoothed(S, log_probs, mask_for_loss, + weight=float(hp["label_smoothing"])) + loss.backward() + if gradient_norm and float(gradient_norm) > 0.0: + torch.nn.utils.clip_grad_norm_(model.parameters(), + float(gradient_norm)) + optimizer.step() + step += 1 + + # Snapshot to CPU before the finally frees the GPU copy. + cpu_state = {k: v.detach().cpu() for k, v in model.state_dict().items()} + ckpt = { + "model_state_dict": cpu_state, + "num_edges": nedges, # protein_mpnn_run.py reads this + "noise_level": noise, + "step": int(step), + "optimizer_state_dict": optimizer.optimizer.state_dict(), + } + target = job["target_weights"] + os.makedirs(os.path.dirname(os.path.abspath(target)) or ".", exist_ok=True) + # Write beside the target then replace, so a reader (IMPRESS mid-pass) + # never sees a half-written weights file. + tmp = target + ".tmp" + torch.save(ckpt, tmp) + os.replace(tmp, target) + + # Completion marker, written LAST, into this round's output_dir. The + # training manager polls for it to detect that the round finished, even + # when the execution backend never delivers the task's result (a Dragon + # defect — see docs/dragon.md). It has to be this marker rather than the + # checkpoint itself: with publish_into_repo the checkpoint is a stable + # path that already exists from the previous round. Name kept in sync + # with rome.trainer.TRAIN_COMPLETE_MARKER. + output_dir = job.get("output_dir") + if output_dir: + with open(os.path.join(output_dir, "train_complete"), "w") as fd: + fd.write(target) + return target + finally: + # Release the GPU as soon as the round ends — see the module docstring. + del model, optimizer + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def main(argv: Tuple[str, ...] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--job", required=True, + help="path to the job spec JSON written by the trainer") + args = parser.parse_args(argv) + + with open(args.job) as fd: + job = json.load(fd) + checkpoint = run_round(job) + # The last stdout line is the checkpoint path, for a caller reading stdout. + print(checkpoint) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/impress_r/mpnn_wrapper.py b/examples/impress_r/mpnn_wrapper.py new file mode 100644 index 0000000..811fd4d --- /dev/null +++ b/examples/impress_r/mpnn_wrapper.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +import argparse +parser = argparse.ArgumentParser() +import subprocess +#-db DATABSE -u USERNAME -p PASSWORD -size 20 +parser.add_argument("-pdb", "--input_path", help="Input path", type=str) +parser.add_argument("-out", "--output_path", help="Output path", type=str) +parser.add_argument("-mpnn", "--mpnn_path", help="MPNN path", type=str) +parser.add_argument("-seqs", "--seqs", help="How many sequences designs would you like?", type=int) +parser.add_argument("-is_monomer", "--is_monomer", help="Is your input a monomer? (1 or 0)", type=int) +parser.add_argument("-chains", "--design_chains", help="Which input chains do you want designed? Separate with spaces. Example: -chains='A B'", type=str) +parser.add_argument("-index", "--index", help="Which specific indices do you want designed/fixed? Separate with spaces. For multiple chains, separate with comma. Example: -index='1 3 12, 2 7 14 56'", type=str) +parser.add_argument("-fix", "--fix", help="Would you like these positions fixed? (1 or 0)", type=int) +parser.add_argument("-tie", "--tie", help="Which specific indices across multiple chains do you want tied? Separate indices with spaces, chains with comma. Lists must be of same length. Example: -tie='1 2 3 4 5, 1 2 3 4 5'", type=str) +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=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) + + + +args = parser.parse_args() + +input_path=args.input_path +output_path=args.output_path +mpnn_path=args.mpnn_path +is_monomer=args.is_monomer #default is_monomer false +chains=args.design_chains +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 is None: + seqs=1 #default 1 design per structure +tie=args.tie +homo=args.homo +interface=args.interface +temp=args.temp +bias_AA=args.bias_AA +bias_weight=args.bias_weight + +if bias_weight!=None and bias_AA!=None: + path_for_bias=output_path+"/bias_pdbs.jsonl" + subprocess.call(['python', mpnn_path+"/helper_scripts/make_bias_AA.py", '--output_path='+path_for_bias, '--AA_list='+bias_AA, '--bias_list='+bias_weight]) +else: + path_for_bias='' + + +path_for_parsed_chains=output_path+"/parsed_pdbs.jsonl" + +if is_monomer==True: #monomer - no need for interface or tie functionality + subprocess.call(['python', mpnn_path+"/helper_scripts/parse_multiple_chains.py", '--input_path='+input_path, '--output_path='+path_for_parsed_chains]) + + if index != None: #check if indices are specified + path_for_assigned_chains=output_path+"/assigned_pdbs.jsonl" + path_for_fixed_positions=output_path+"/fixed_pdbs.jsonl" + subprocess.call(['python', mpnn_path+"/helper_scripts/assign_fixed_chains.py", '--input_path='+path_for_parsed_chains, '--output_path='+path_for_assigned_chains, '--chain_list=A']) + + if fix==True: #check if indices are fixed or designed + subprocess.call(['python', mpnn_path+"/helper_scripts/make_fixed_positions_dict.py", '--input_path='+path_for_parsed_chains, '--output_path='+path_for_fixed_positions, '--chain_list=A', '--position_list='+index]) + else: + subprocess.call(['python', mpnn_path+"/helper_scripts/make_fixed_positions_dict.py", '--input_path='+path_for_parsed_chains, '--output_path='+path_for_fixed_positions, '--chain_list=A', '--position_list='+index, '--specify_non_fixed']) + + subprocess.call(['python', mpnn_path+"/protein_mpnn_run.py", '--jsonl_path='+path_for_parsed_chains,'--out_folder='+output_path, '--num_seq_per_target='+str(seqs), '--sampling_temp='+str(temp), '--seed=37', '--batch_size=1', '--chain_id_jsonl='+path_for_assigned_chains, '--fixed_positions_jsonl='+path_for_fixed_positions, '--bias_AA_jsonl='+path_for_bias]) + #monomeric, fixed + else: + subprocess.call(['python', mpnn_path+"/protein_mpnn_run.py", '--jsonl_path='+path_for_parsed_chains,'--out_folder='+output_path, '--num_seq_per_target='+str(seqs), '--sampling_temp='+str(temp), '--seed=37', '--batch_size=1', '--bias_AA_jsonl='+path_for_bias]) + #monomeric, unfixed + +else: #multimer + subprocess.call(['python', mpnn_path+"/helper_scripts/parse_multiple_chains.py", '--input_path='+input_path, '--output_path='+path_for_parsed_chains]) + path_for_assigned_chains=output_path+"/assigned_pdbs.jsonl" + subprocess.call(['python', mpnn_path+"/helper_scripts/assign_fixed_chains.py", '--input_path='+path_for_parsed_chains, '--output_path='+path_for_assigned_chains, '--chain_list='+chains]) + + if index != None: #check if indices are specified + + path_for_fixed_positions=output_path+"/fixed_pdbs.jsonl" + + if fix==True: #check if indices are fixed or designed + subprocess.call(['python', mpnn_path+"/helper_scripts/make_fixed_positions_dict.py", '--input_path='+path_for_parsed_chains, '--output_path='+path_for_fixed_positions, '--chain_list='+chains, '--position_list='+index]) + else: + subprocess.call(['python', mpnn_path+"/helper_scripts/make_fixed_positions_dict.py", '--input_path='+path_for_parsed_chains, '--output_path='+path_for_fixed_positions, '--chain_list='+chains, '--position_list='+index, '--specify_non_fixed']) + + if homo==True: #check for homomer + path_for_tied_positions=output_path+"/tied_pdbs.jsonl" + subprocess.call(['python', mpnn_path+"/helper_scripts/make_tied_positions_dict.py", '--input_path='+path_for_parsed_chains,'--output_path='+path_for_tied_positions, '--chain_list='+chains, '--homooligomer=1']) + subprocess.call(['python', mpnn_path+"/protein_mpnn_run.py", '--jsonl_path='+path_for_parsed_chains,'--out_folder='+output_path, '--num_seq_per_target='+str(seqs), '--sampling_temp='+str(temp), '--seed=37', '--batch_size=1', '--chain_id_jsonl='+path_for_assigned_chains, '--fixed_positions_jsonl='+path_for_fixed_positions, '--tied_positions_jsonl='+path_for_tied_positions, '--bias_AA_jsonl='+path_for_bias]) + #multimeric, fixed, homomer + + elif tie!=None: #check for tied positions + + path_for_tied_positions=output_path+"/tied_pdbs.jsonl" + subprocess.call(['python', mpnn_path+"/helper_scripts/make_tied_positions_dict.py", '--input_path='+path_for_parsed_chains,'--output_path='+path_for_tied_positions, '--chain_list='+chains, '--position_list='+tie]) + subprocess.call(['python', mpnn_path+"/protein_mpnn_run.py", '--jsonl_path='+path_for_parsed_chains,'--out_folder='+output_path, '--num_seq_per_target='+str(seqs), '--sampling_temp='+str(temp), '--seed=37', '--batch_size=1', '--chain_id_jsonl='+path_for_assigned_chains, '--fixed_positions_jsonl='+path_for_fixed_positions, '--tied_positions_jsonl='+path_for_tied_positions, '--bias_AA_jsonl='+path_for_bias]) + #multimeric, fixed, tied + + else: + subprocess.call(['python', mpnn_path+"/protein_mpnn_run.py", '--jsonl_path='+path_for_parsed_chains,'--out_folder='+output_path, '--num_seq_per_target='+str(seqs), '--sampling_temp='+str(temp), '--seed=37', '--batch_size=1', '--chain_id_jsonl='+path_for_assigned_chains, '--fixed_positions_jsonl='+path_for_fixed_positions, '--bias_AA_jsonl='+path_for_bias]) + #multimeric, fixed, untied + + elif interface==True: + path_for_fixed_positions=output_path+"/interface_dict.jsonl" + + subprocess.call(['python', mpnn_path+"/helper_scripts/mk_interface_dict.py", '--input_path='+input_path, '--output_path='+path_for_fixed_positions]) + + if homo==True: #check for homomer + path_for_tied_positions=output_path+"/tied_pdbs.jsonl" + subprocess.call(['python', mpnn_path+"/helper_scripts/make_tied_positions_dict.py", '--input_path='+path_for_parsed_chains,'--output_path='+path_for_tied_positions, '--chain_list='+chains, '--homooligomer=1']) + subprocess.call(['python', mpnn_path+"/protein_mpnn_run.py", '--jsonl_path='+path_for_parsed_chains,'--out_folder='+output_path, '--num_seq_per_target='+str(seqs), '--sampling_temp='+str(temp), '--seed=37', '--batch_size=1', '--chain_id_jsonl='+path_for_assigned_chains, '--fixed_positions_jsonl='+path_for_fixed_positions, '--tied_positions_jsonl='+path_for_tied_positions, '--bias_AA_jsonl='+path_for_bias]) + #multimeric, interface, homomer + + elif tie!=None: #check for tied positions + + path_for_tied_positions=output_path+"/tied_pdbs.jsonl" + subprocess.call(['python', mpnn_path+"/helper_scripts/make_tied_positions_dict.py", '--input_path='+path_for_parsed_chains,'--output_path='+path_for_tied_positions, '--chain_list='+chains, '--position_list='+tie]) + subprocess.call(['python', mpnn_path+"/protein_mpnn_run.py", '--jsonl_path='+path_for_parsed_chains,'--out_folder='+output_path, '--num_seq_per_target='+str(seqs), '--sampling_temp='+str(temp), '--seed=37', '--batch_size=1', '--chain_id_jsonl='+path_for_assigned_chains, '--fixed_positions_jsonl='+path_for_fixed_positions, '--tied_positions_jsonl='+path_for_tied_positions, '--bias_AA_jsonl='+path_for_bias]) + #multimeric, interface, tied + + else: + subprocess.call(['python', mpnn_path+"/protein_mpnn_run.py", '--jsonl_path='+path_for_parsed_chains,'--out_folder='+output_path, '--num_seq_per_target='+str(seqs), '--sampling_temp='+str(temp), '--seed=37', '--batch_size=1', '--chain_id_jsonl='+path_for_assigned_chains, '--fixed_positions_jsonl='+path_for_fixed_positions, '--bias_AA_jsonl='+path_for_bias]) + #multimeric, interface, untied + else: + + if homo==True: #check for homomer + path_for_tied_positions=output_path+"/tied_pdbs.jsonl" + subprocess.call(['python', mpnn_path+"/helper_scripts/make_tied_positions_dict.py", '--input_path='+path_for_parsed_chains,'--output_path='+path_for_tied_positions, '--chain_list='+chains, '--homooligomer=1']) + subprocess.call(['python', mpnn_path+"/protein_mpnn_run.py", '--jsonl_path='+path_for_parsed_chains,'--out_folder='+output_path, '--num_seq_per_target='+str(seqs), '--sampling_temp='+str(temp), '--seed=37', '--batch_size=1', '--chain_id_jsonl='+path_for_assigned_chains, '--tied_positions_jsonl='+path_for_tied_positions, '--bias_AA_jsonl='+path_for_bias]) + #multimeric, unfixed, homomer + + elif tie!=None: #check for tied positions + + path_for_tied_positions=output_path+"/tied_pdbs.jsonl" + subprocess.call(['python', mpnn_path+"/helper_scripts/make_tied_positions_dict.py", '--input_path='+path_for_parsed_chains,'--output_path='+path_for_tied_positions, '--chain_list='+chains, '--position_list='+tie]) + subprocess.call(['python', mpnn_path+"/protein_mpnn_run.py", '--jsonl_path='+path_for_parsed_chains,'--out_folder='+output_path, '--num_seq_per_target='+str(seqs), '--sampling_temp='+str(temp), '--seed=37', '--batch_size=1', '--chain_id_jsonl='+path_for_assigned_chains, '--tied_positions_jsonl='+path_for_tied_positions, '--bias_AA_jsonl='+path_for_bias]) + #multimeric, unfixed, tied + + else: + subprocess.call(['python', mpnn_path+"/protein_mpnn_run.py", '--jsonl_path='+path_for_parsed_chains,'--out_folder='+output_path, '--num_seq_per_target='+str(seqs), '--sampling_temp='+str(temp), '--seed=37', '--batch_size=1', '--chain_id_jsonl='+path_for_assigned_chains, '--bias_AA_jsonl='+path_for_bias]) + #multimeric, unfixed, untied \ No newline at end of file diff --git a/examples/impress_r/plddt_extract_pipeline.py b/examples/impress_r/plddt_extract_pipeline.py new file mode 100644 index 0000000..6dc57bd --- /dev/null +++ b/examples/impress_r/plddt_extract_pipeline.py @@ -0,0 +1,63 @@ +import numpy as np +import os +import pandas as pd +import json +import argparse + +# Peptide EGYQDYEPEA is 10 residues and always placed last in the Boltz FASTA +PEP_LEN = 10 + +parser = argparse.ArgumentParser() +parser.add_argument('--iter', type=str, help='pass iteration') +parser.add_argument('--out', type=str, help='pipeline name') +parser.add_argument('--path', type=str, help='base path') +args = parser.parse_args() + +af_path = os.path.join(args.path, 'af_pipeline_outputs_multi', args.out, 'af/prediction') +dimer_models_path = os.path.join(af_path, 'dimer_models') + +rows = [] + +for name in os.listdir(dimer_models_path): + pred_dir = os.path.join( + dimer_models_path, name, f"boltz_results_{name}", "predictions", name + ) + plddt_file = os.path.join(pred_dir, f"plddt_{name}_model_0.npz") + conf_file = os.path.join(pred_dir, f"confidence_{name}_model_0.json") + pae_file = os.path.join(pred_dir, f"pae_{name}_model_0.npz") + + if not all(os.path.exists(f) for f in [plddt_file, conf_file, pae_file]): + continue + + # avg_plddt: mean per-residue pLDDT (Boltz stores 0-1; scale to 0-100) + plddt = np.load(plddt_file)['plddt'] + avg_plddt = float(plddt.mean() * 100) + + # iptm from Boltz confidence JSON — interface PTM is the primary binding quality metric + with open(conf_file) as f: + conf = json.load(f) + iptm = conf.get('iptm', conf.get('ptm', 0.0)) + + # avg_pae: mean cross-chain PAE between PDZ domain and peptide + # Boltz outputs PDZ residues first, peptide last; PAE matrix is (N, N) in Angstroms + pae_matrix = np.load(pae_file)['pae'] + total_res = pae_matrix.shape[0] + pdz_len = total_res - PEP_LEN + cross_pae = np.concatenate([ + pae_matrix[:pdz_len, pdz_len:].ravel(), + pae_matrix[pdz_len:, :pdz_len].ravel(), + ]) + avg_pae = float(cross_pae.mean()) if len(cross_pae) > 0 else 0.0 + + rows.append({ + 'ID': f"{name}.pdb", + 'avg_plddt': avg_plddt, + 'ptm': iptm, + 'avg_pae': avg_pae, + }) + +print(f"Processed {len(rows)} structure(s)") + +df = pd.DataFrame(rows, columns=['ID', 'avg_plddt', 'ptm', 'avg_pae']) +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/impress_r/plot_rome_scores.py b/examples/impress_r/plot_rome_scores.py new file mode 100644 index 0000000..e60082f --- /dev/null +++ b/examples/impress_r/plot_rome_scores.py @@ -0,0 +1,245 @@ +""" +IMPRESS-R ROME score analysis. + +Produces three plots: + 1. Per-pipeline avg_pLDDT trajectory over passes (depth-0 pipelines only, one line each) + 2. Box plot of avg_pLDDT distribution per pass across all pipelines and depths + 3. Global mean per pass (all depths) with linear regression trend + +Usage: + python plot_rome_scores.py [--csv-dir DIR] [--log FILE] [--out-dir DIR] + +Defaults: + --csv-dir /scratch/***/$USER/IMPRESS_outputs + --log logs/impress_21736435.out (relative to script dir) + --out-dir . (saves next to this script) +""" + +import argparse +import csv +import os +import re +import statistics +from collections import defaultdict +from itertools import groupby + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import matplotlib.ticker as ticker +import numpy as np + +# ── Argument parsing ────────────────────────────────────────────────────────── +parser = argparse.ArgumentParser() +_scratch = os.environ.get("SCRATCH") +parser.add_argument( + "--csv-dir", + default=f"{_scratch}/IMPRESS_outputs" if _scratch else None, + required=_scratch is None, +) +parser.add_argument("--log", default=os.path.join(os.path.dirname(__file__), "logs", "impress_21736435.out")) +parser.add_argument("--out-dir", default=os.path.dirname(__file__)) +args = parser.parse_args() + +os.makedirs(args.out_dir, exist_ok=True) + +# ── Load pLDDT data ─────────────────────────────────────────────────────────── +# records: list of (root, pipeline, depth, pass_num, mean_plddt, [values]) +records = [] +pattern = re.compile(r"af_stats_(p\d+(?:_sub\d+)*?)_pass_(\d+)\.csv") + +for fname in os.listdir(args.csv_dir): + m = pattern.match(fname) + if not m: + continue + pipeline, pass_num = m.group(1), int(m.group(2)) + depth = pipeline.count("_sub") + root = pipeline.split("_sub")[0] + path = os.path.join(args.csv_dir, fname) + with open(path) as f: + vals = [float(r["avg_plddt"]) for r in csv.DictReader(f) if r.get("avg_plddt")] + if vals: + records.append((root, pipeline, depth, pass_num, statistics.mean(vals), vals)) + +# depth_pass: (depth, pass_num) -> [values] — used for depth-stratified plot +depth_pass = defaultdict(list) +for root, pipeline, depth, pass_num, mean, vals in records: + depth_pass[(depth, pass_num)].extend(vals) + +# ── Load ROME training events from log ─────────────────────────────────────── +# Extract unique (round, timestamp) from log lines like: +# 12:02:36.532 [INFO] [ROME-TRAINER] submitting training round 1 (3 designs ...) -> v1 +rome_rounds = {} # round_num -> HH:MM as float minutes-since-start +log_start = None + +if os.path.exists(args.log): + rome_pat = re.compile(r"(\d{2}):(\d{2}):\d{2}\.\d+.*ROME-TRAINER.*submitting training round (\d+)") + with open(args.log) as f: + for line in f: + clean = re.sub(r"\x1b\[[0-9;]*m", "", line) + m = rome_pat.search(clean) + if m: + h, mi, rnd = int(m.group(1)), int(m.group(2)), int(m.group(3)) + t = h * 60 + mi + if log_start is None: + log_start = t + if rnd not in rome_rounds: + rome_rounds[rnd] = t - log_start # minutes since job start + +# ── Palette ─────────────────────────────────────────────────────────────────── +PIPELINES = sorted({r[0] for r in records}) +cmap = matplotlib.colormaps["tab20"].resampled(len(PIPELINES)) +COLORS = {p: cmap(i) for i, p in enumerate(PIPELINES)} + +ACCENT = "#E05C3A" # warm orange-red for trend / mean +GRID_COL = "#E8E8E8" + +plt.rcParams.update({ + "font.family": "sans-serif", + "axes.spines.top": False, + "axes.spines.right": False, + "axes.grid": True, + "grid.color": GRID_COL, + "grid.linewidth": 0.7, +}) + +# ═══════════════════════════════════════════════════════════════════════════════ +# Plot 1 — Per-pipeline trajectory (depth-0 only) +# ═══════════════════════════════════════════════════════════════════════════════ +fig1, ax1 = plt.subplots(figsize=(11, 6)) + +depth0 = [(root, pass_num, mean) for root, pipeline, depth, pass_num, mean, _ in records if depth == 0] + +pipe_data = defaultdict(dict) # root -> {pass_num: mean} +for root, pass_num, mean in depth0: + pipe_data[root][pass_num] = mean + +for pipe in sorted(pipe_data): + xs = sorted(pipe_data[pipe]) + ys = [pipe_data[pipe][x] for x in xs] + ax1.plot(xs, ys, marker="o", markersize=4, linewidth=1.5, + color=COLORS[pipe], label=pipe, alpha=0.85) + +ax1.set_xlabel("Pass", fontsize=11) +ax1.set_ylabel("avg_pLDDT", fontsize=11) +ax1.set_title("Per-pipeline avg_pLDDT over passes (root pipelines only)", fontsize=13, pad=10) +ax1.set_ylim(40, 100) +ax1.xaxis.set_major_locator(ticker.MaxNLocator(integer=True)) +ax1.legend(ncol=4, fontsize=8, loc="lower right", framealpha=0.7) +ax1.axhline(75, color="#999999", linewidth=0.8, linestyle="--", label="threshold 75") + +fig1.tight_layout() +out1 = os.path.join(args.out_dir, "rome_pipeline_trajectories.png") +fig1.savefig(out1, dpi=150) +print(f"Saved: {out1}") + +# ═══════════════════════════════════════════════════════════════════════════════ +# Plot 2 — Box plot per pass (all depths) +# ═══════════════════════════════════════════════════════════════════════════════ +pass_vals = defaultdict(list) +for root, pipeline, depth, pass_num, mean, vals in records: + pass_vals[pass_num].extend(vals) + +all_passes = sorted(pass_vals) +box_data = [pass_vals[p] for p in all_passes] +counts = [len(pass_vals[p]) for p in all_passes] + +fig2, ax2 = plt.subplots(figsize=(11, 6)) +bp = ax2.boxplot(box_data, positions=all_passes, widths=0.6, + patch_artist=True, showfliers=True, + flierprops=dict(marker=".", markersize=3, alpha=0.4, color="#AAAAAA"), + medianprops=dict(color=ACCENT, linewidth=2), + boxprops=dict(facecolor="#D6E8F5", alpha=0.85)) + +# Annotate n per pass +for x, n in zip(all_passes, counts): + ax2.text(x, 38.5, f"n={n}", ha="center", fontsize=7, color="#666666") + +ax2.set_xlabel("Pass", fontsize=11) +ax2.set_ylabel("avg_pLDDT", fontsize=11) +ax2.set_title("avg_pLDDT distribution per pass (all pipelines and depths)", fontsize=13, pad=10) +ax2.set_ylim(36, 100) +ax2.axhline(75, color="#999999", linewidth=0.8, linestyle="--") +ax2.xaxis.set_major_locator(ticker.MaxNLocator(integer=True)) + +fig2.tight_layout() +out2 = os.path.join(args.out_dir, "rome_pass_distribution.png") +fig2.savefig(out2, dpi=150) +print(f"Saved: {out2}") + +# ═══════════════════════════════════════════════════════════════════════════════ +# Plot 3 — Mean per pass stratified by sub-pipeline depth + global trend +# ═══════════════════════════════════════════════════════════════════════════════ +DEPTH_COLORS = {0: "#4C72B0", 1: "#DD8452", 2: "#55A868", 3: "#C44E52"} +DEPTH_LABELS = {0: "depth 0 (root)", 1: "depth 1 (sub1)", + 2: "depth 2 (sub1_sub2)", 3: "depth 3 (sub1_sub2_sub3)"} + +fig3, ax3 = plt.subplots(figsize=(12, 6)) + +all_depths = sorted({d for d, p in depth_pass}) +for depth in all_depths: + xs = sorted(p for d, p in depth_pass if d == depth) + ys = [statistics.mean(depth_pass[(depth, p)]) for p in xs] + ns = [len(depth_pass[(depth, p)]) for p in xs] + color = DEPTH_COLORS.get(depth, "#888888") + ax3.plot(xs, ys, marker="o", markersize=5, linewidth=1.8, + color=color, label=DEPTH_LABELS.get(depth, f"depth {depth}"), zorder=3) + for x, y, n in zip(xs, ys, ns): + ax3.text(x, y + 0.5, f"{y:.1f}", ha="center", fontsize=6.5, + color=color, alpha=0.85) + +# Global mean + regression +pass_means = {p: statistics.mean(v) for p, v in pass_vals.items()} +xs_all = sorted(pass_means) +ys_all = [pass_means[x] for x in xs_all] +coef = np.polyfit(xs_all, ys_all, 1) +trend_y = np.polyval(coef, xs_all) + +ax3.plot(xs_all, ys_all, linewidth=2.2, linestyle="--", + color=ACCENT, alpha=0.6, + label=f"global mean (slope {coef[0]:+.3f}/pass)", zorder=2) + +# ROME round markers +if rome_rounds: + max_t = max(rome_rounds.values()) + for rnd in sorted(rome_rounds)[:10]: + t_frac = rome_rounds[rnd] / max_t if max_t > 0 else 0 + x_pos = xs_all[0] + t_frac * (xs_all[-1] - xs_all[0]) + ax3.axvline(x_pos, color="#99AACC", linewidth=0.55, linestyle=":", alpha=0.65) + ax3.text(x_pos, 79.5, f"v{rnd}", fontsize=6, color="#667799", + ha="center", rotation=90, va="bottom") + +ax3.set_xlabel("Pass", fontsize=11) +ax3.set_ylabel("Mean avg_pLDDT", fontsize=11) +ax3.set_title("avg_pLDDT per pass by sub-pipeline depth (ROME training markers in blue)", fontsize=13, pad=10) +ax3.set_ylim(78, 100) +ax3.axhline(75, color="#CCCCCC", linewidth=0.7, linestyle=":") +ax3.xaxis.set_major_locator(ticker.MaxNLocator(integer=True)) +ax3.legend(fontsize=9, loc="lower right") + +fig3.tight_layout() +out3 = os.path.join(args.out_dir, "rome_global_trend.png") +fig3.savefig(out3, dpi=150) +print(f"Saved: {out3}") + +# ── Text summary ────────────────────────────────────────────────────────────── +print("\n=== Per-pass global mean (all depths) ===") +print(f"{'Pass':>4} {'Mean':>7} {'Median':>7} {'n (designs)':>11}") +for p in xs_all: + vals = pass_vals[p] + print(f" {p:2d} {statistics.mean(vals):7.2f} {statistics.median(vals):7.2f} {len(vals):>11}") + +print("\n=== Per-depth per-pass mean ===") +print(f"{'Depth':>5} {'Pass':>4} {'Mean':>7} {'Median':>7} {'n':>4}") +for (depth, pass_num) in sorted(depth_pass): + vals = depth_pass[(depth, pass_num)] + print(f" {depth:3d} {pass_num:2d} {statistics.mean(vals):7.2f} {statistics.median(vals):7.2f} {len(vals):>4}") + +slope = coef[0] +print(f"\nGlobal linear trend slope: {slope:+.4f} pLDDT / pass") +if slope > 0.2: + print(" → Positive trend: ROME appears to be improving sequence quality.") +elif slope < -0.2: + print(" → Negative trend: scores declining — may reflect harder proteins in later passes.") +else: + print(" → Flat trend: no clear improvement signal yet.") diff --git a/examples/impress_r/protein_binding_rome.py b/examples/impress_r/protein_binding_rome.py new file mode 100644 index 0000000..c788b4b --- /dev/null +++ b/examples/impress_r/protein_binding_rome.py @@ -0,0 +1,335 @@ + +import asyncio +import copy +import os +import shutil + +from impress.pipelines.impress_pipeline import ImpressBasePipeline + +MPNN_PATH = os.environ.get("MPNN_PATH", "") + +_BOLTZ_CHAIN_MAP = {'pdz': 'A', 'pep': 'B'} + +def _copy_pdb_rename_chains(src, dst, chain_map=_BOLTZ_CHAIN_MAP): + """Copy a PDB, replacing Boltz multi-char chain IDs with standard single-char IDs.""" + with open(src) as f_in, open(dst, 'w') as f_out: + for line in f_in: + if line.startswith(('ATOM', 'HETATM', 'TER')): + chain = line[21:24] + if chain in chain_map: + line = line[:21] + chain_map[chain] + line[24:] + f_out.write(line) + + +class ProteinBindingPipeline(ImpressBasePipeline): + def __init__(self, name, flow, configs=None, **kwargs): + if configs is None: + configs = {} + + self.is_child: bool = kwargs.get("is_child", False) + self.passes = kwargs.get("passes", 1) + self.start_pass: int = kwargs.get("start_pass", 1) + self.step_id = kwargs.get("step_id", 1) + self.seq_rank = kwargs.get("seq_rank", 0) + 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") 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.policy = kwargs.get("policy", None) + + self.current_scores = {} + self.iter_seqs = kwargs.get("iter_seqs", {}) + self.previous_scores = kwargs.get("previous_scores", {}) + + super().__init__(name, flow, **configs, **kwargs) + + self.fasta_list_2 = kwargs.get("fasta_list_2", []) + self.base_path = kwargs.get("base_path", os.getcwd()) + self.input_base_path = kwargs.get("input_base_path", self.base_path) + 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.input_base_path, f"prod_in/{self.name}_in") + + self.output_path = os.path.join( + 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( + self.output_path, "af/prediction/best_models" + ) + + for file_name in os.listdir(self.input_path): + self.fasta_list_2.append(file_name) + + def set_up_new_pipeline_dirs(self, new_pipeline_name): + base_output = os.path.join( + self.output_base_path, "af_pipeline_outputs_multi", new_pipeline_name + ) + input_dir = os.path.join(self.input_base_path, f"prod_in/{new_pipeline_name}_in") + + subdirs = [ + "af/fasta", + "af/prediction", + "af/prediction/best_models", + "af/prediction/best_ptm", + "af/prediction/dimer_models", + "af/prediction/logs", + "mpnn", + *[f"mpnn/job_{i}" for i in range(1, self.max_passes + 1)], + ] + + paths_to_create = [input_dir, base_output] + [ + os.path.join(base_output, subdir) for subdir in subdirs + ] + + for path in paths_to_create: + os.makedirs(path, exist_ok=True) + + def register_pipeline_tasks(self): + + @self.auto_register_task(local_task=True) + async def s1(): + 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 + + cmd = ( + f"bash {self.scripts_path}/s1_mpnn.sh " + f"{mpnn_script} " + f"{input_path} " + f"{output_dir} " + f"{self.mpnn_path} " + f"{self.num_seqs} " + f"{chain}" + ) + log_path = os.path.join(output_dir, "mpnn_run.log") + with open(log_path, "w") as lf: + proc = await asyncio.create_subprocess_shell( + cmd, stdout=lf, stderr=asyncio.subprocess.STDOUT, + env=self._gpu_env(), + ) + rc = await proc.wait() + if rc != 0: + raise RuntimeError(f"s1 MPNN failed (exit {rc})") + + @self.auto_register_task(local_task=True) + async def s2(): + self.step_id += 1 + job_seqs_dir = f"{self.output_path_mpnn}/job_{self.passes}/seqs" + + for file_name in os.listdir(job_seqs_dir): + seqs = [] + with open(os.path.join(job_seqs_dir, file_name)) as fd: + lines = fd.readlines()[2:] + + score = None + for line in lines: + line = line.strip() + if line.startswith(">"): + score = float(line.split(",")[2].replace(" score=", "")) + else: + seqs.append([line, score]) + + seqs.sort(key=lambda x: x[1]) + self.iter_seqs[file_name.split(".")[0]] = seqs + + @self.auto_register_task(local_task=True) + async def s3(): + self.step_id += 1 + output_dir = os.path.join(self.output_path, "af", "fasta") + 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 = self.peptide_seq + + 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): + pdz_tag = f">pdz|protein|{pdz_msa}" + pep_tag = f">pep|protein|{pep_msa}" + else: + 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_tag}\n{design_seq}\n{pep_tag}\n{pep_seq}\n") + + return fasta_file_to_return + + @self.auto_register_task(local_task=True) + async def s4(target_fasta): + 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}" + ) + self.logger.pipeline_log(f"s4 command for {target_fasta}: {cmd}") + proc = await asyncio.create_subprocess_shell( + cmd, + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.DEVNULL, + env=self._gpu_env(), + ) + rc = await proc.wait() + if rc != 0: + raise RuntimeError(f"Boltz failed for {target_fasta} (exit {rc})") + + @self.auto_register_task(local_task=True) + async def s4_post_exec( + target_fasta, + models_path, + best_model_pdb, + best_ptm_json, + mpnn_pdb + ): + self.step_id += 1 + _copy_pdb_rename_chains(f"{models_path}/{target_fasta}_model_0.pdb", best_model_pdb) + shutil.copy(f"{models_path}/confidence_{target_fasta}_model_0.json", best_ptm_json) + _copy_pdb_rename_chains(f"{models_path}/{target_fasta}_model_0.pdb", mpnn_pdb) + + @self.auto_register_task() + async def s5(): + self.step_id += 1 + return ( + f"bash {self.scripts_path}/s5_plddt_extract.sh " + f"{self.output_base_path} " + f"{self.passes} " + f"{self.name}" + ) + + async def get_scores_map(self): + return {"c_scores": self.current_scores, "p_scores": self.previous_scores} + + def finalize(self, sub_iter_seqs): + from pathlib import Path + for a in sub_iter_seqs: + self.fasta_list_2.remove(f"{a}.pdb") + 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): + self.logger.pipeline_log(f"Running for a maximum of {self.max_passes} passes") + + self.set_up_new_pipeline_dirs(self.name) + + while self.passes <= self.max_passes: + self.logger.pipeline_log(f"Starting pass {self.passes}") + + if self.is_child and self.passes == self.start_pass: + self.logger.pipeline_log( + "Skipping MPNN and Ranking steps for this child pipeline " + "in the current pass only." + ) + else: + self.logger.pipeline_log("Submitting MPNN task") + await self.s1() + self.logger.pipeline_log("MPNN task finished") + + self.logger.pipeline_log("Submitting sequence ranking task") + await self.s2() + self.logger.pipeline_log("Sequence ranking task finished") + + self.logger.pipeline_log("Submitting scoring task") + fasta_files = await self.s3() + self.logger.pipeline_log("Scoring task finished") + + alphafold_tasks = [] + post_exec_tasks = [] + + _boltz_sem = asyncio.Semaphore(2) + + async def _guarded_s4(target_fasta): + async with _boltz_sem: + return await self.s4(target_fasta=target_fasta) + + for target_fasta in fasta_files: + models_path = os.path.join( + self.output_path, "af", "prediction", "dimer_models", target_fasta, + f"boltz_results_{target_fasta}", "predictions", target_fasta + ) + best_model_pdb = os.path.join( + self.output_path, "af", "prediction", "best_models", + f"{target_fasta}.pdb", + ) + best_ptm_json = os.path.join( + self.output_path, "af", "prediction", "best_ptm", + f"{target_fasta}.json", + ) + mpnn_pdb = os.path.join( + self.output_path, "mpnn", f"job_{self.passes}", + f"{target_fasta}.pdb", + ) + + alphafold_tasks.append(_guarded_s4(target_fasta)) + post_exec_tasks.append( + self.s4_post_exec( + target_fasta=target_fasta, + models_path=models_path, + best_model_pdb=best_model_pdb, + best_ptm_json=best_ptm_json, + mpnn_pdb=mpnn_pdb, + ) + ) + + self.logger.pipeline_log( + f"Submitting {len(alphafold_tasks)} Boltz tasks asynchronously" + ) + s4_results = await asyncio.gather(*alphafold_tasks, return_exceptions=True) + self.logger.pipeline_log(f"{len(alphafold_tasks)} Boltz tasks finished") + for fasta_name, result in zip(fasta_files, s4_results): + if isinstance(result, Exception): + self.logger.pipeline_log(f"s4 FAILED for {fasta_name}: {result}") + else: + 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 pLDDT 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}", + } + ], + } + ) + self.logger.pipeline_log("pLDDT extraction finished") + + await self.run_adaptive_step(wait=True) + + if self.kill_parent: + break + + self.passes += 1 diff --git a/examples/impress_r/rome_global_trend.png b/examples/impress_r/rome_global_trend.png new file mode 100644 index 0000000..cc788e6 Binary files /dev/null and b/examples/impress_r/rome_global_trend.png differ diff --git a/examples/impress_r/rome_pass_distribution.png b/examples/impress_r/rome_pass_distribution.png new file mode 100644 index 0000000..63efe24 Binary files /dev/null and b/examples/impress_r/rome_pass_distribution.png differ diff --git a/examples/impress_r/rome_pipeline_trajectories.png b/examples/impress_r/rome_pipeline_trajectories.png new file mode 100644 index 0000000..a073b05 Binary files /dev/null and b/examples/impress_r/rome_pipeline_trajectories.png differ diff --git a/examples/impress_r/run_21736435_report.md b/examples/impress_r/run_21736435_report.md new file mode 100644 index 0000000..d8a0c36 --- /dev/null +++ b/examples/impress_r/run_21736435_report.md @@ -0,0 +1,186 @@ +# IMPRESS-R Run 21736435 — Report + +**Date:** 2026-09-02 +**Job:** 21736435 | `gpuA40x4` partition | node `gpub096` +**Wall time:** 6h | **Status:** COMPLETED +**Pipelines:** 16 (p1–p16) | **Max passes:** 10 | **Trainer:** ProteinMPNN (real fine-tune) + +--- + +## ROME-A Training + +### Status +- Training rounds completed at report time: **v14** (30 designs in corpus) +- Corpus grew from 0 → 78+ designs across 1h of runtime +- All rounds succeeded — weights published into `vanilla_model_weights/` after each round + +### Warning: Dragon future never resolves (recurring, non-fatal) + +Every training round emitted this warning: + +``` +[WARNING] [ROME-TRAINER] round .../train_complete: the execution backend has not +delivered a result after Xs, but the checkpoint is on disk — publishing from disk. +The task finished; its future never resolved, which on Dragon means a running service +is blocking result delivery (see docs/dragon.md). +``` + +**What it means:** `mpnn_train_wrapper.py` runs as a Dragon subprocess. Dragon delivers the task result via a future, but the IMPRESS event loop (a long-lived Dragon service) holds a ProcessGroup open, and Dragon's result-delivery path requires the ProcessGroup to be idle. ROME's fallback kicks in after `ROME_FALLBACK` seconds (60s default): it detects `train_complete` on disk and publishes the weights directly, bypassing the stuck future. + +**Impact:** None — all checkpoints were published correctly. Training proceeded normally. Every round v1–v14 completed this way. + +### Error: ProcessGroup state error (occurred once, round 6) + +``` +[ERROR] [ROME-TRAINER] training round failed: + RuntimeError: ProcessGroup manager is not in state State.DEAD +``` + +- **When:** 12:27:18, round v6 (14 designs) +- **What happened:** Dragon tried to reuse a ProcessGroup that hadn't fully torn down from the previous round. ROME caught the error, retried immediately with 16 designs (2 more had arrived in the meantime), and the retry succeeded. +- **Impact:** None — one round was delayed by ~2s, retry completed normally. + +### Training round timeline + +| Round | Designs | Dispatch time | Result | +|-------|---------|--------------|--------| +| v1 | 3 | 12:02:36 | OK (disk fallback, 34s) | +| v2 | 8 | 12:14:32 | OK (disk fallback, 36s) | +| v3 | 9 | 12:15:14 | OK (disk fallback, 38s) | +| v4 | 11 | 12:16:00 | OK (disk fallback, 70s) | +| v5 | 11 | 12:25:53 | OK (disk fallback, 32s) | +| v6 | 14 | 12:26:40 | **FAILED** → retry (16 designs) → OK | +| v7 | 16 | 12:36:19 | OK (disk fallback, 70s) | +| v8 | 18 | 12:37:42 | OK (disk fallback, 54s) | +| v9–v11 | 21–26 | ~12:40–12:50 | OK | +| v12 | 26 | 12:54:03 | Retried 3× (round 12 submitted 3 times) | +| v13 | 27 | 12:59:14 | OK | +| v14 | 30 | 13:00:41 | OK | + +> **Note on v12:** Round 12 was submitted three times in quick succession (12:54:03, 12:54:48, 12:55:39). This indicates the ProcessGroup error recurred and ROME retried twice before succeeding. The same pattern as v6 but with two retries. + +--- + +## pLDDT Statistics + +Computed from **296 CSV files** in `/scratch/bblj/mgoliyad1/IMPRESS_outputs/`. +One design per file (one protein × one pass). Final numbers after job completion. + +### Summary + +| Stat | avg_pLDDT | +|--------|-----------| +| Min | 42.51 | +| Max | 97.49 | +| Mean | 87.36 | +| Median | 91.81 | +| Stdev | 11.30 | + +### Percentiles + +| Percentile | avg_pLDDT | +|------------|-----------| +| p10 | 76.25 | +| p25 | 84.98 | +| p75 | 94.21 | +| p90 | 96.44 | + +### Threshold breakdown + +| Threshold | Count | Fraction | +|----------------|-------|----------| +| < 75 (low) | 28 | 9.5% | +| ≥ 75 (pass) | 268 | 90.5% | +| ≥ 90 (high) | 168 | 56.8% | + +**Interpretation:** 90.5% of designs passed the pLDDT ≥ 75 filter. The median of 91.8 and 56.8% above 90 indicate the pipeline is producing high-confidence Boltz predictions. The 28 sub-75 outliers are concentrated in p7 and p14 (structurally difficult targets that stayed low throughout all passes); the adaptive logic routes these proteins to child pipelines with a new sequence rank but cannot fully overcome the target difficulty. + +--- + +## ROME Score Improvement Analysis + +Plots generated by `plot_rome_scores.py` on the full completed dataset (296 design evaluations, 16 pipelines × up to 10 passes). + +### Plot 1 — Per-pipeline trajectories + +![Per-pipeline avg_pLDDT over passes](rome_pipeline_trajectories.png) + +Root pipelines only (depth-0). Each line is one of the 16 pipelines tracked across its passes. +Key observations: +- **p13, p15** (≥95): naturally high-scoring targets — Boltz consistently confident regardless of sequence. +- **p2** (97 early → 85 at pass 5): sharp drop coincides with adaptive migration to child pipeline; root pipeline retained the harder sequence rank. +- **p7, p14** (42–63): consistently low — these targets appear structurally difficult for the binder geometry; candidates for early termination in future runs. +- Most pipelines cluster in the 82–94 range and show mild oscillation rather than monotone drift, which is expected when a single protein is evaluated per pipeline. + +### Plot 2 — Per-pass distribution (all depths) + +![avg_pLDDT distribution per pass](rome_pass_distribution.png) + +Box plot across all pipelines and child depths per pass. +Design count rises pass 1→3 as child pipelines spin up, then stabilises at 18 for passes 7–10 (only the deepest sub-pipelines still running). +The interquartile range narrows in later passes, consistent with selection pressure removing the lowest-scoring designs. + +### Plot 3 — Mean per pass by sub-pipeline depth + +![avg_pLDDT per pass by sub-pipeline depth with ROME training markers](rome_global_trend.png) + +Each line is one sub-pipeline depth level; dashed orange = global mean across all depths. Vertical dotted blue markers = ROME training round completions (v1–v10). + +#### Global mean (all depths combined) + +| Pass | Mean pLDDT | Median | n (designs) | +|------|-----------|--------|-------------| +| 1 | 85.87 | 90.47 | 16 | +| 2 | 86.24 | 91.81 | 44 | +| 3 | 87.55 | 91.04 | 47 | +| 4 | 87.36 | 90.18 | 39 | +| 5 | 87.56 | 91.04 | 42 | +| 6 | 86.87 | 92.08 | 35 | +| 7 | 88.44 | 92.40 | 19 | +| 8 | 88.23 | 92.31 | 18 | +| 9 | 88.30 | 92.39 | 18 | +| 10 | 88.52 | 91.99 | 18 | + +**Global linear trend slope: +0.268 pLDDT / pass** + +#### Per-depth breakdown + +| Depth | Pass | Mean pLDDT | Median | n (designs) | +|-------|------|-----------|--------|-------------| +| 0 (root) | 1 | 85.87 | 90.47 | 16 | +| 0 | 2 | 85.63 | 90.19 | 16 | +| 0 | 3 | 86.14 | 88.27 | 16 | +| 0 | 4 | 85.16 | 89.27 | 16 | +| 0 | 5 | 84.47 | 89.01 | 16 | +| 0 | 6 | 84.37 | 87.00 | 16 | +| 1 (sub1) | 2 | 86.70 | 91.81 | 12 | +| 1 | 3 | 87.14 | 89.77 | 10 | +| 1 | 4 | 88.95 | 88.26 | 5 | +| 1 | 5 | 89.03 | 89.62 | 5 | +| 2 (sub1_sub2) | 2 | 87.55 | 92.54 | 8 | +| 2 | 3 | 91.27 | 93.63 | 9 | +| 2 | 4 | 91.59 | 92.78 | 5 | +| 2 | 5 | 91.60 | 91.86 | 6 | +| 2 | 6 | 91.94 | 91.94 | 2 | +| 3 (sub1_sub2_sub3) | 2 | 85.47 | 92.31 | 8 | +| 3 | 3 | 86.97 | 90.55 | 12 | +| 3 | 4 | 87.84 | 92.80 | 13 | +| 3 | 5 | 88.75 | 92.78 | 15 | +| 3 | 6 | 88.68 | 93.06 | 16 | +| 3 | 7 | 88.56 | 92.55 | 16 | +| 3 | 8 | 88.64 | 92.45 | 16 | +| 3 | 9 | 88.71 | 92.74 | 16 | +| 3 | 10 | 88.96 | 92.39 | 16 | + +**Key observations:** +- **Depth 0 (root)** declines slightly (85.9 → 84.4) — the adaptive logic continuously migrates the best-performing proteins out into sub-pipelines, leaving the harder cases at the root. +- **Depth 2 (sub1_sub2)** shows the strongest absolute scores (87.6 → 91.9), consistently the highest-quality group across all passes. +- **Depth 3 (sub1_sub2_sub3)** shows the clearest ROME improvement signal: steady rise from 85.5 at pass 2 to 89.0 at pass 10 (+3.5 pLDDT), covering the most ROME training rounds (v2–v14+). +- The global mean conflates depth effects with ROME effects; the per-depth view isolates them. + +--- + +## Next Steps + +1. **ROME ProcessGroup bug:** The `ProcessGroup manager is not in state State.DEAD` error recurred at v6 and v12. Investigate whether `mpnn_train_wrapper.py` needs an explicit `dist.destroy_process_group()` call at exit, or whether `ROME_FALLBACK` should be lowered to reduce the window between the stuck future and disk fallback. +2. **Control comparison:** Run an identical job with `ROME_TRAINER=dummy` to isolate ROME's contribution from the natural adaptive/selection effect. The +0.268/pass slope currently conflates ROME improvement with selection bias (harder targets culled by later passes). diff --git a/examples/impress_r/run_protein_binding_rome.py b/examples/impress_r/run_protein_binding_rome.py new file mode 100644 index 0000000..27c8f3f --- /dev/null +++ b/examples/impress_r/run_protein_binding_rome.py @@ -0,0 +1,356 @@ +import copy +import csv +import os +import shutil +import asyncio +import tempfile +from typing import Dict, Any, Optional, List + +from rhapsody.backends import DragonExecutionBackend +from rhapsody.telemetry import define_event +from rhapsody.telemetry.events import make_event + +from impress import GPUPolicy, _find_gpus, _make_policy, ImpressManager, PipelineSetup + +try: + from examples.impress_r.protein_binding_rome import ProteinBindingPipeline +except ModuleNotFoundError: + import sys + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + from protein_binding_rome import ProteinBindingPipeline + +try: + from examples.impress_r.mpnn import ( + ProteinMPNNConfig, + ProteinMPNNTrainer, + percentile_sampler, + ) +except ModuleNotFoundError: + import sys + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + from mpnn import ProteinMPNNConfig, ProteinMPNNTrainer, percentile_sampler + +import rome + +import rhapsody, logging +rhapsody.enable_logging(level=logging.INFO) + + +# ── Test mode (IMPRESS_TEST_MODE=1) ─────────────────────────────────────────── +# Runs 2 pipelines with max_passes=1 and no child pipelines, so a single +# MPNN → Boltz → pLDDT → ROME cycle completes for integration testing. +TEST_MODE = os.getenv("IMPRESS_TEST_MODE", "0") == "1" +N_PIPELINES = 2 if TEST_MODE else int(os.environ.get("IMPRESS_N_PIPELINES", 16)) +MAX_PASSES = 1 if TEST_MODE else int(os.environ.get("ROME_MAX_PASSES", 10)) +# IMPRESS_MAX_SUB_PIPELINES: max depth of child pipeline spawning (0 = none, 3 = full). +# Unset or blank → inherit the inline default of 3. +_sub_env = os.environ.get("IMPRESS_MAX_SUB_PIPELINES", "").strip() +MAX_SUB_PIPELINES_OVERRIDE = 0 if TEST_MODE else (int(_sub_env) if _sub_env else None) + + +# ── ProteinMPNN repo ─────────────────────────────────────────────────────────── +# The same checkout IMPRESS runs for inference; ROME fine-tunes and publishes back +# into vanilla_model_weights/ so the next pass picks it up with no wrapper change. +MPNN_REPO = os.environ.get( + "ROME_MPNN_REPO", + os.environ.get("MPNN_PATH", ""), +) + + +# --------------------------------------------------------------------------- +# Custom telemetry events +# --------------------------------------------------------------------------- + +ProteinScore = define_event( + "impress.ProteinScore", + protein=str, + pipeline_name=str, + pass_num=int, + current_score=float, + previous_score=float, + decision=str, +) + +PassSummary = define_event( + "impress.PassSummary", + pipeline_name=str, + pass_num=int, + num_proteins=int, + num_degraded=int, + child_spawned=bool, +) + +ChildPipelineSpawned = define_event( + "impress.ChildPipelineSpawned", + parent_name=str, + child_name=str, + num_proteins=int, + seq_rank=int, +) + + +# --------------------------------------------------------------------------- +# ROME trainer builder +# --------------------------------------------------------------------------- + +def _build_trainer(checkpoint_dir: str): + """ProteinMPNN fine-tuner by default; a no-op dummy for smoke testing. + + ROME_TRAINER=mpnn → real ProteinMPNN fine-tune (needs ROME_MPNN_REPO) + ROME_TRAINER=dummy → smoke test (no GPU/torch required for training) + """ + want = os.environ.get("ROME_TRAINER", "mpnn").lower() + if want == "mpnn" and os.path.isdir(MPNN_REPO): + return ProteinMPNNTrainer(ProteinMPNNConfig( + mpnn_repo=MPNN_REPO, + initial_weights=os.path.join(MPNN_REPO, "vanilla_model_weights", "v_48_020.pt"), + model_name="v_48_020", + publish_into_repo=True, + ), gpus=1) + + if want == "mpnn": + print(f"[ROME-A] ROME_MPNN_REPO={MPNN_REPO!r} not found; " + "falling back to the dummy trainer (set ROME_MPNN_REPO or " + "ROME_TRAINER=dummy to silence this).") + from rome.dummy import DummyTrainer + return DummyTrainer(train_seconds=1.0, gpus=0) + + +# --------------------------------------------------------------------------- +# Real-time failure subscriber +# --------------------------------------------------------------------------- + +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}") + + +# --------------------------------------------------------------------------- +# Adaptive criteria +# --------------------------------------------------------------------------- + +def adaptive_criteria(current_score: float, previous_score: float) -> bool: + return current_score > previous_score + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +async def impress_protein_bind() -> None: + workdir = tempfile.mkdtemp(prefix="impress_r_") + stage_dir = os.path.join(workdir, "designs") + os.makedirs(stage_dir, exist_ok=True) + + # ROME-A gets its own Dragon backend so training tasks run in their own + # processes (separate CUDA context) independent of IMPRESS tasks. + rome_backend = await DragonExecutionBackend() + rome_manager = rome.Manager( + backend=rome_backend, + data_config=rome.DataConfig( + min_samples=int(os.environ.get("ROME_MIN_SAMPLES", 2)), + sample_func=percentile_sampler(0.33), + ), + trainer_config=rome.TrainerConfig( + trainer=_build_trainer(os.path.join(workdir, "checkpoints")), + checkpoint_dir=os.path.join(workdir, "checkpoints"), + poll_interval=1.0, + result_fallback_seconds=float(os.environ.get("ROME_FALLBACK", 60)), + ), + ) + await rome_manager.start() + + backend = await DragonExecutionBackend() + manager: ImpressManager = ImpressManager( + execution_backend=backend, + telemetry_config={ + "checkpoint_path": "./telemetry/", + "resource_poll_interval": 5.0, + }, + telemetry_subscribers=[_on_task_event], + ) + + # adaptive_decision closes over manager and rome_manager. + async def adaptive_decision(pipeline: ProteinBindingPipeline) -> Optional[Dict[str, Any]]: + 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 CSV — written by s5_plddt_extract.sh to output_base_path. + file_name = os.path.join( + pipeline.output_base_path, + f"af_stats_{pipeline.name}_pass_{pipeline.passes}.csv", + ) + accepted = 0 + with open(file_name) as fd: + for row in csv.DictReader(fd): # ID, avg_plddt, ptm, avg_pae + protein = row["ID"].split(".")[0] + pipeline.current_scores[protein] = float(row["avg_pae"]) + + # -- ROME-A HOOK 1: contribute this design to the training corpus + src = os.path.join(pipeline.output_path_af, f"{protein}.pdb") + if not os.path.exists(src): + continue + staged = os.path.join( + stage_dir, + f"{pipeline.name}_pass{pipeline.passes}_{protein}.pdb", + ) + shutil.copyfile(src, staged) + ranked = pipeline.iter_seqs.get(protein) or [] + sequence = ranked[pipeline.seq_rank][0] if len(ranked) > pipeline.seq_rank else "" + uid = rome_manager.add_training_data( + path=staged, + sequence=sequence, + backbone_id=protein, + pLDDT=float(row["avg_plddt"]), + pTM=float(row["ptm"]), + pAE=float(row["avg_pae"]), + score=float(row["avg_plddt"]), + ) + accepted += uid is not None + + # -- ROME-A HOOK 2: collect the improved model for next pass + weights = rome_manager.get_current_model() + pipeline.logger.pipeline_log( + f"ROME-A: corpus {rome_manager.data.total_count} (+{accepted} this pass) | " + f"{rome_manager.get_training_status().name}" + + (f" | model {os.path.basename(weights)}" if weights else "") + ) + + # First pass — just save current scores as previous + if not pipeline.previous_scores: + pipeline.logger.pipeline_log("Saving current scores as previous and returning") + pipeline.previous_scores = copy.deepcopy(pipeline.current_scores) + return + + # Identify proteins that got worse (higher avg_pae = worse interface) + sub_iter_seqs: Dict[str, str] = {} + for protein, curr_score in pipeline.current_scores.items(): + if protein not in pipeline.iter_seqs: + continue + prev_score = pipeline.previous_scores[protein] + 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, + current_score=curr_score, + previous_score=prev_score, + decision="degrade" if decision else "keep", + )) + + if decision: + sub_iter_seqs[protein] = pipeline.iter_seqs.pop(protein) + + # Spawn a new pipeline for degraded proteins + child_spawned = False + if sub_iter_seqs and pipeline.sub_order < MAX_SUB_PIPELINES: + new_name: str = f"{pipeline.name}_sub{pipeline.sub_order + 1}" + pipeline.set_up_new_pipeline_dirs(new_name) + + for protein in sub_iter_seqs: + src = f"{pipeline.output_path_af}/{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), + seq_rank=pipeline.seq_rank + 1, + )) + + new_config = { + "name": new_name, + "type": type(pipeline), + "adaptive_fn": adaptive_decision, + "config": { + "is_child": True, + "start_pass": pipeline.passes, + "passes": pipeline.passes, + "iter_seqs": sub_iter_seqs, + "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, + }, + } + + pipeline.submit_child_pipeline_request(new_config) + pipeline.finalize(sub_iter_seqs) + + if not pipeline.fasta_list_2: + pipeline.kill_parent = True + + child_spawned = True + else: + pipeline.previous_scores = copy.deepcopy(pipeline.current_scores) + + if tel: + tel.emit(make_event( + PassSummary, + session_id=sid, + backend="rhapsody", + pipeline_name=pipeline.name, + pass_num=pipeline.passes, + num_proteins=len(pipeline.current_scores), + num_degraded=len(sub_iter_seqs), + child_spawned=child_spawned, + )) + + # Paths — same convention as protein_binding runner. + 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() + + pipeline_setups: List[PipelineSetup] = [ + PipelineSetup( + name=f"p{str(i)}", + type=ProteinBindingPipeline, + config={ + "base_path": scripts_dir, + "input_base_path": input_base_dir, + "output_base_path": output_base_dir, + "policy": _make_policy(all_gpus, i - 1), + }, + adaptive_fn=adaptive_decision, + max_passes=MAX_PASSES, + ) + for i in range(1, N_PIPELINES + 1) + ] + + try: + await manager.start(pipeline_setups=pipeline_setups) + print("\nROME-A:", rome_manager.report()) + if manager.telemetry: + summary = manager.telemetry.summary() + print(f"[TELEMETRY] tasks={summary.get('tasks', {})}") + dur = summary.get("duration") + if dur: + print(f"[TELEMETRY] mean task time: {dur['mean_seconds'] * 1000:.1f} ms") + await manager.telemetry.stop() + finally: + await rome_manager.stop() + + +if __name__ == "__main__": + asyncio.run(impress_protein_bind()) diff --git a/examples/impress_r/scripts/s1_mpnn.sh b/examples/impress_r/scripts/s1_mpnn.sh new file mode 100755 index 0000000..b490f10 --- /dev/null +++ b/examples/impress_r/scripts/s1_mpnn.sh @@ -0,0 +1,23 @@ +#!/bin/bash +set -e + +# Step 1: Sequence prediction via ProteinMPNN +# Args: $1=mpnn_script $2=input_path $3=output_dir $4=mpnn_path $5=num_seqs $6=chain + +mpnn_script="$1" +input_path="$2" +output_dir="$3" +mpnn_path="$4" +num_seqs="$5" +chain="$6" + +# Re-activate the IMPRESS venv inside Dragon tasks (VIRTUAL_ENV is exported by sbatch). +[ -n "${VIRTUAL_ENV:-}" ] && source "${VIRTUAL_ENV}/bin/activate" + +python3 "$mpnn_script" \ + -pdb="$input_path" \ + -out="$output_dir" \ + -mpnn="$mpnn_path" \ + -seqs="$num_seqs" \ + -is_monomer=0 \ + -chains="$chain" diff --git a/examples/impress_r/scripts/s4_boltz.sh b/examples/impress_r/scripts/s4_boltz.sh new file mode 100755 index 0000000..a10cef9 --- /dev/null +++ b/examples/impress_r/scripts/s4_boltz.sh @@ -0,0 +1,65 @@ +#!/bin/bash +set -e + +# Step 4: Structure prediction via Boltz +# Args: $1=fasta_path $2=output_dir + +fasta_path="$1" +output_dir="$2" + +# 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 + +# 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}" + +# Prevent PyTorch Lightning from installing SLURM auto-requeue signal handlers. +# PL detects SLURM_JOB_ID and registers SIGTERM/SIGUSR handlers that keep the +# process group alive after Boltz finishes, causing Dragon to report +# "ProcessGroup manager is not in state State.DEAD". +unset SLURM_JOB_ID + +boltz predict \ + "${fasta_path}" \ + --out_dir "${output_dir}" \ + ${_MSA_FLAG} \ + --cache "${BOLTZ_CACHE_DIR:-${HOME}/.boltz}" \ + --output_format pdb \ + --write_full_pae \ + --no_kernels \ + --devices 1 \ + --override \ + 2>&1 | tee "${output_dir}/boltz_run.log" +# tee exits 0; check the actual boltz exit code via PIPESTATUS +test "${PIPESTATUS[0]}" -eq 0 diff --git a/examples/impress_r/scripts/s5_plddt_extract.sh b/examples/impress_r/scripts/s5_plddt_extract.sh new file mode 100755 index 0000000..11a3367 --- /dev/null +++ b/examples/impress_r/scripts/s5_plddt_extract.sh @@ -0,0 +1,20 @@ +#!/bin/bash +set -e + +# Step 5: pLDDT extraction +# Args: $1=output_base_path $2=iter $3=out_name + +output_base_path="$1" +iter="$2" +out_name="$3" + +# plddt_extract_pipeline.py lives one level above this scripts/ directory. +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# Re-activate the IMPRESS venv inside Dragon tasks (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}" diff --git a/examples/protein_binding/CODE_REVIEW.md b/examples/protein_binding/CODE_REVIEW.md new file mode 100644 index 0000000..51a1251 --- /dev/null +++ b/examples/protein_binding/CODE_REVIEW.md @@ -0,0 +1,283 @@ +# Code Review — `examples/protein_binding/` + +**Date:** 2026-08-29 +**Scope:** `protein_binding.py`, `protein_binding_run.py`, `run_protein_binding.py`, +`run_nonadaptive.py`, `mpnn_wrapper.py`, `plddt_extract_pipeline.py`, +`scripts/`, `delta_gpu_run.sh` + +--- + +## Bugs + +### `mpnn_wrapper.py:1` — Python file with `#!/bin/sh` shebang + +Line 1 is `#!/bin/sh`, making the OS attempt to execute a Python file as a POSIX +shell script. Any direct execution of `./mpnn_wrapper.py` produces a parse error at +the first Python statement. The file should have `#!/usr/bin/env python3` or no +shebang at all. + +--- + +### `mpnn_wrapper.py:18` — `--temp` argument parses temperature as `int` instead of `float` + +```python +parser.add_argument("-temp", "--temp", ..., type=int, default=0.1) +``` + +`type=int` truncates `0.1` to `0`. Every invocation that relies on the default (or +passes a fractional temperature) gets temperature `0`, which collapses the sampling +distribution to greedy argmax. `type=float` is required. + +--- + +### `plddt_extract_pipeline.py:122` — division by zero if `counter2` is 0 + +```python +avg_pae = running_sum / counter2 +``` + +`counter2` counts PAE matrix cells where exactly one of `row_index` / `col_index` +falls in `target_range`. If `target_range` is empty (protein shorter than 10 +residues) or the PAE matrix is empty, `counter2` stays 0 and this line raises +`ZeroDivisionError`. + +--- + +## Potential Issues + +### `protein_binding.py:9` — MPNN path hardcoded to Anvil filesystem + +```python +MPNN_PATH = f"/anvil/projects/x-nairr240405/mason/ProteinMPNN" +``` + +This module-level constant silently resolves to a non-existent path on Delta or any +other system. The pipeline does accept `mpnn_path` as a constructor kwarg (line 37), +but callers who forget to supply it will get a confusing "directory not found" error +at task execution time instead of a clear startup failure. + +**Recommended fix:** default to `None` and raise a clear `ValueError` in `__init__` +if the path is not supplied and `MPNN_PATH` is not set in the environment. + +--- + +### `protein_binding.py:152` — hardcoded peptide sequence in `s3()` + +```python +pep_seq = "EGYQDYEPEA" # PDZ-domain peptide +``` + +This PDZ-specific constant is hardcoded inside `s3()`. It should be a constructor +parameter (e.g. `self.peptide_seq`) so the pipeline is reusable for other targets +without modifying the source. + +--- + +### `protein_binding.py:300` — all Boltz tasks launched concurrently + +```python +s4_results = await asyncio.gather(*alphafold_tasks, return_exceptions=True) +``` + +All structures are folded in parallel. With N structures, N Boltz processes compete +for GPU memory simultaneously, likely causing OOM on real runs. Folding should be +serialised per GPU or gated by a `asyncio.Semaphore`. + +--- + +### `protein_binding.py:220–221` — `os.unlink()` without existence check in `finalize()` + +```python +os.unlink(f"{self.output_path_af}/{a}.pdb") +os.unlink(f"{self.output_path}/af/fasta/{a}.fa") +``` + +If a Boltz task failed and the file was never created, `finalize()` raises +`FileNotFoundError`. Should use `pathlib.Path.unlink(missing_ok=True)` or check +existence first. + +--- + +### `protein_binding_run.py:7`, `run_protein_binding.py:6`, `run_nonadaptive.py:4` — old `DragonExecutionBackendV3` class name + +```python +from rhapsody.backends import DragonExecutionBackendV3 +``` + +The correct class is `DragonExecutionBackend` (the `V3` suffix was removed in a +recent rhapsody release). All three runner scripts still import the old name, causing +`ImportError` at startup on the current environment. Only `run_nonadaptive.py` in the +`small_molecule_binding` example was updated; the protein-binding runners were not. + +--- + +### `run_protein_binding.py:89` — adaptive function reads CSV from current working directory + +```python +file_name = f'af_stats_{pipeline.name}_pass_{pipeline.passes}.csv' +with open(file_name) as fd: +``` + +This path is relative to wherever the process was started. If the working directory +is wrong or if the stage failed and the file was never written, this raises +`FileNotFoundError` and crashes the adaptive function. Should use an absolute path +based on `pipeline.base_path`. + +--- + +### `run_protein_binding.py:64` — `adaptive_criteria` declared `async` with no awaits + +```python +async def adaptive_criteria(current_score: float, previous_score: float) -> bool: + return current_score > previous_score +``` + +This function does no I/O or async work. Declaring it `async` adds unnecessary +overhead at every call site. Make it a plain `def`. + +--- + +### `plddt_extract_pipeline.py` — incompatible with current IMPRESS output structure + +`on_replica_done` in the current campaign reads `binder_scores_*.json` directly from +the alphafold output directory. This script reads from +`af_pipeline_outputs_multi/{name}/af/prediction/best_models` — a directory tree that +the current pipeline does not create. This script is effectively stale and will +produce empty / zero results if run against current outputs. + +--- + +### `delta_gpu_run.sh:3,44` — SBATCH requests 4 tasks but runs single-node Dragon + +``` +#SBATCH --tasks-per-node=4 +... +dragon -s run_protein_binding.py +``` + +`dragon -s` is single-node mode. Requesting 4 tasks-per-node wastes resources and +may confuse the scheduler. For true multi-node, change to `dragon -m`; for +single-node, change `--tasks-per-node=1`. + +--- + +### `delta_gpu_run.sh:21` — unguarded `LD_LIBRARY_PATH` produces trailing colon + +```bash +export LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${MPI_LIB}:${FAB_LIB}:${LD_LIBRARY_PATH} +``` + +If `LD_LIBRARY_PATH` is unset, this expands to a trailing `:`, which adds the current +directory to the dynamic linker search path — a security risk. Should be +`${LD_LIBRARY_PATH:-}`. + +--- + +### `delta_gpu_run.sh:32,35` — stale hardcoded path with directory typo + +```bash +IMPRESS_SCRIPTS_DIR="$SCRATCH/$USER/IMPRESS/examples/protien_binding_usecase" +WORKDIR="$SCRATCH/$USER/IMPRESS/examples/protien_binding_usecase" +``` + +The directory was renamed to `protein_binding` in the merge but these two path +variables still reference the old misspelled name. Any job submitted with this script +will `cd` to a non-existent directory and abort immediately. + +--- + +## Code Quality + +### `protein_binding_run.py:17–18`, `run_nonadaptive.py:6,12` — rhapsody DEBUG logging enabled unconditionally + +```python +import rhapsody, logging +rhapsody.enable_logging(level=logging.DEBUG) +``` + +Both runner files enable DEBUG-level rhapsody logging unconditionally. On a +16-pipeline run this generates thousands of lines per second and buries application +output. Should default to INFO or be gated by an env variable. + +--- + +### `mpnn_wrapper.py` — massive duplicated subprocess.call blocks + +The file handles 12+ scenario combinations (monomer/multimer × fixed/unfixed × +tied/homomer/interface) through deeply nested `if/else` with nearly identical +`subprocess.call` lists. The only differences are which `--*_jsonl` flags are +appended. Refactor to build the argument list incrementally and call +`subprocess.call` once. + +--- + +### `mpnn_wrapper.py` — not used by the current pipeline + +The active pipeline calls `s1_mpnn.sh` → LigandMPNN's `run.py`. `mpnn_wrapper.py` +calls `protein_mpnn_run.py` (ProteinMPNN, not LigandMPNN) via a completely different +argument schema. This file is dead code relative to the current pipeline and should +be either removed, labelled as a standalone utility, or replaced with a properly +integrated version. + +--- + +### `mpnn_wrapper.py:30–40` — `chains == None` instead of `chains is None` + +```python +if chains == None: + chains = 'A' +``` + +PEP 8 and Python convention require `if chains is None:`. The `==` form works for +`None` but is non-idiomatic and triggers linter warnings. + +--- + +### `delta_gpu_run.sh:39` — `eval` for venv activation + +```bash +eval "$IMPRESS_PRE_EXEC" +``` + +Running an env-supplied string through `eval` allows arbitrary code execution if +`IMPRESS_PRE_EXEC` is set adversarially or incorrectly. Replace with a direct +`source` of the known venv path. + +--- + +### `plddt_extract_pipeline.py:74` — unclosed file handle + +```python +data = json.load(open(data_path)) +``` + +The file handle is never closed. Should use `with open(data_path) as f: json.load(f)`. + +--- + +### `plddt_extract_pipeline.py` — extensive commented-out code + +Lines 56–60 (biopandas block), 83–92 (`df_json`), 95–100 (print block), and 103–107 +(debug prints) are commented-out code that was never removed. These should be deleted +to improve readability. + +--- + +### `plddt_extract_pipeline.py` — manual index tracking instead of `enumerate` + +Lines 111–122 use manual `row_index`/`col_index` variables incremented in loop bodies +to track which matrix cell is being accessed. Using `enumerate` would be clearer and +less error-prone. + +--- + +### `af2_multimer_reduced.sh:13-14` — `/tmp/work` and `/tmp/upper` created but never used + +```bash +WORK=/tmp/work +UPPER=/tmp/upper +mkdir -p $WORK $UPPER +``` + +These directories are never referenced in the apptainer call. Leftover from a prior +overlay filesystem approach. Remove. diff --git a/examples/protein_binding/SKILL.md b/examples/protein_binding/SKILL.md new file mode 100644 index 0000000..e05135c --- /dev/null +++ b/examples/protein_binding/SKILL.md @@ -0,0 +1,131 @@ +# Protein Binding Use Case — SKILL.md + +## What it does + +5-step adaptive protein design pipeline running on Delta HPC (gpuA40x4): + +| Step | Task | Type | Notes | +|------|------|------|-------| +| s1 | ProteinMPNN — generate sequences | local (subprocess) | GPU via `CUDA_VISIBLE_DEVICES` | +| s2 | Sequence ranking | local (pure Python) | sorts by MPNN score | +| s3 | FASTA generation | local (pure Python) | writes `.fa` for each target | +| s4 | AlphaFold multimer prediction | local (subprocess) | GPU, runs via Apptainer | +| s4_post | Copy best models | local (pure Python) | glob → shutil.copy | +| s5 | pLDDT extraction | local (subprocess) | writes `af_stats__pass_.csv` | + +Passes repeat up to `max_passes` times. The adaptive function reads pLDDT CSVs and can spawn child pipelines for proteins that degrade. + +--- + +## Key files + +``` +protien_binding_usecase/ +├── run_protein_binding.py # entry point — backend, GPU policy, ImpressManager +├── delta_gpu_run.sh # SLURM script — Dragon launcher +├── delta_env_setup.sh # one-time venv setup (~/ve/impress/) +├── af2_multimer_reduced.sh # Apptainer wrapper for AlphaFold multimer +├── mpnn_wrapper.py # ProteinMPNN CLI wrapper +└── plddt_extract_pipeline.py # reads AF JSON → writes CSV +``` + +Pipeline implementation: `/scratch/bblj/mgoliyad1/IMPRESS/src/impress/pipelines/protein_binding.py` + +--- + +## Environment variables + +| Variable | Set in | Purpose | +|----------|--------|---------| +| `SCRATCH` | user env | base scratch path (`/scratch/bblj`) | +| `MPNN_PATH` | `delta_gpu_run.sh` | path to ProteinMPNN repo | +| `AF2_DATABASE` | `delta_gpu_run.sh` | AlphaFold database dir | +| `AF2_SIF` | `delta_gpu_run.sh` | AlphaFold Apptainer image | +| `IMPRESS_INPUT_DIR` | `delta_gpu_run.sh` | dir containing `{name}_in/` folders | +| `IMPRESS_OUTPUT_DIR` | `delta_gpu_run.sh` | output root (af_pipeline_outputs_multi/) | +| `IMPRESS_SCRIPTS_DIR` | `delta_gpu_run.sh` | dir with mpnn_wrapper.py, af2_multimer_reduced.sh | +| `SBATCH_ACCOUNT` | user env | SLURM account (e.g. `bblj-delta-gpu`) | + +`IMPRESS_PRE_EXEC` was removed — the venv is activated in the SLURM script before `dragon -s`, so all subprocesses inherit the environment. + +--- + +## Dragon integration + +**Backend**: `DragonExecutionBackendV2` (flow engine only — no tasks dispatched to Dragon workers). + +**Why not V3**: Dragon worker subprocesses silently hang when a subprocess uses GPU ops (`cupy`, `torch`, Apptainer). Same failure mode documented in `DeepDriveSim/workflows/miniapps_workflow/miniapps_workflow.py`. + +**Fix**: All pipeline steps are `local_task=True`. GPU steps (s1, s4) run as `asyncio.create_subprocess_shell()` on the Dragon head process, which has direct GPU access from SLURM allocation. Dragon manages the `ImpressManager`/`WorkflowEngine` loop; it never dispatches tasks to workers. + +**GPU affinity**: `_find_gpus()` + `_make_policy()` in `run_protein_binding.py` use `dragon.native.machine.System()` to discover GPUs and build a `dragon.infrastructure.policy.Policy` with `gpu_affinity`. The policy's GPU list is passed as `CUDA_VISIBLE_DEVICES` to each subprocess via `_gpu_env()` in the pipeline. + +```python +# To change GPU count per pipeline (default n_gpus=1 for MPNN): +policy = _make_policy(all_gpus, idx=0, n_gpus=1) +``` + +**Launcher** (`delta_gpu_run.sh`): +```bash +source ~/ve/impress/bin/activate +dragon-config add --ofi-runtime-lib="${FAB_LIB}" +rm -rf asyncflow.session.* +dragon -s run_protein_binding.py # -s = single-node +``` + +**SLURM resource request**: `--tasks-per-node=4`, `--gpus=4`, `--exclusive` (matches Dragon's single-node layout). + +--- + +## Key design decisions + +### Three separate path env vars +`base_path` was split into `IMPRESS_INPUT_DIR` / `IMPRESS_OUTPUT_DIR` / `IMPRESS_SCRIPTS_DIR` so input data, outputs, and scripts can live in different locations independently. + +### Output dirs auto-created in `__init__` +`os.makedirs(..., exist_ok=True)` is called in `ProteinBindingPipeline.__init__` for all required subdirs. No manual `mkdir` needed in the SLURM script. + +### `s4_post` replaces `post_exec` +`post_exec` in task descriptions is a RADICAL-Pilot artifact — `ConcurrentExecutionBackend` and Dragon never execute it. Replaced with a `local_task=True` coroutine using `glob` + `shutil.copy` to move `ranked_0.pdb` and `ranking_debug.json` to `best_models/` and `best_ptm/`. + +### AlphaFold container flags +- `--no-home`: prevents host's newer Biopython from shadowing the container's version (`SCOPData` import error) +- `--run_relax` removed: not supported by this container version +- Database bind: `/scratch/rhaas/SUP-5301/database:/database` +- Container: `/scratch/rhaas/SUP-5301/alphafold.sif` + +--- + +## One-time setup + +```bash +export SCRATCH=/scratch/bblj +cd /scratch/bblj/mgoliyad1/IMPRESS/examples/protien_binding_usecase +bash delta_env_setup.sh # creates ~/ve/impress/ + +# Unzip inputs (GNU tar, not zip): +cd /scratch/bblj/mgoliyad1/IMPRESS_inputs +tar -xf prod_in.tar # or whatever archive name +``` + +## Submitting a job + +```bash +export SCRATCH=/scratch/bblj +export SBATCH_ACCOUNT=bblj-delta-gpu +cd /scratch/bblj/mgoliyad1/IMPRESS/examples/protien_binding_usecase +sbatch delta_gpu_run.sh +``` + +## Checking results + +```bash +# pLDDT scores per pass: +cat af_stats_p1_pass_1.csv + +# AlphaFold best models: +ls IMPRESS_outputs/af_pipeline_outputs_multi/p1/af/prediction/best_models/ + +# Full logs: +tail -f logs/impress_.out +``` 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..54a0d15 --- /dev/null +++ b/examples/protein_binding/delta_env_setup.sh @@ -0,0 +1,249 @@ +#!/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]" +# Pin dragonhpc to 0.14.1 — 0.14.2 added waitForKeys to DDRegisterClientResponse +# but the Delta system Dragon runtime has not been updated to match; 0.14.2 fails +# with AttributeError on every DDict operation on this cluster. +"${PIP}" install -q "dragonhpc==0.14.1" + +# ── 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..f07c51e --- /dev/null +++ b/examples/protein_binding/delta_gpu_run.sh @@ -0,0 +1,100 @@ +#!/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:00:00 +#SBATCH --job-name=impress_protein +#SBATCH --mail-user=mg2347@soe.rutgers.edu +#SBATCH --mail-type=END,FAIL +#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 + +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_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 + +# ── Run ─────────────────────────────────────────────────────────────────────── +# asyncflow session dirs now go to /tmp (node-local, no quota) via +# IMPRESS_SESSION_DIR; no need to clean them from cwd. + +# -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 + +rm -f ddict_orc* + +echo "Running: dragon ${DRAGON_MODE} run_protein_binding.py (nodes=${SLURM_NNODES:-1})" +dragon ${DRAGON_MODE} run_protein_binding.py + +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..52f80e2 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'} @@ -34,7 +34,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.policy = kwargs.get("policy", None) # Sequence and score state self.current_scores = {} @@ -46,12 +50,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 +75,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 = [ @@ -94,16 +101,17 @@ def set_up_new_pipeline_dirs(self, new_pipeline_name): 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}): + @self.auto_register_task(local_task=True) # MPNN + async def s1(): 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 +120,15 @@ async def s1(task_description={"gpus_per_rank": 1}): f"{self.num_seqs} " f"{chain}" ) + log_path = os.path.join(output_dir, "mpnn_run.log") + with open(log_path, "w") as lf: + proc = await asyncio.create_subprocess_shell( + cmd, stdout=lf, stderr=asyncio.subprocess.STDOUT, + env=self._gpu_env(), + ) + rc = await proc.wait() + if rc != 0: + raise RuntimeError(f"s1 MPNN failed (exit {rc})") @self.auto_register_task(local_task=True) async def s2(): @@ -139,17 +156,38 @@ 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 @@ -162,8 +200,8 @@ async def s3(): # 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 + @self.auto_register_task(local_task=True) + async def s4(target_fasta): self.step_id += 1 cmd = ( f"bash {self.scripts_path}/s4_boltz.sh " @@ -171,7 +209,16 @@ async def s4(target_fasta, task_description={"gpus_per_rank": 1}): # noqa: B006 f"{self.output_path}/af/prediction/dimer_models/{target_fasta}" ) self.logger.pipeline_log(f"s4 command for {target_fasta}: {cmd}") - return cmd + # s4_boltz.sh tees its own output to boltz_run.log in the output dir + proc = await asyncio.create_subprocess_shell( + cmd, + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.DEVNULL, + env=self._gpu_env(), + ) + rc = await proc.wait() + if rc != 0: + raise RuntimeError(f"Boltz failed for {target_fasta} (exit {rc})") @self.auto_register_task(local_task=True) async def s4_post_exec( @@ -206,7 +253,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,10 +264,11 @@ 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): @@ -255,6 +303,13 @@ async def run(self): alphafold_tasks = [] post_exec_tasks = [] + # Limit concurrent Boltz launches to avoid GPU OOM. + _boltz_sem = asyncio.Semaphore(2) + + async def _guarded_s4(target_fasta): + async with _boltz_sem: + return await self.s4(target_fasta=target_fasta) + for target_fasta in fasta_files: models_path = os.path.join( self.output_path, "af", "prediction", "dimer_models", target_fasta, @@ -282,8 +337,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,11 +361,16 @@ 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") diff --git a/examples/protein_binding/protein_binding_run.py b/examples/protein_binding/protein_binding_run.py index fa9e4b9..bacbd3c 100644 --- a/examples/protein_binding/protein_binding_run.py +++ b/examples/protein_binding/protein_binding_run.py @@ -11,10 +11,9 @@ import numpy as np import pandas as pd -# from rhapsody.backends import DragonExecutionBackendV3 -from radical.asyncflow import LocalExecutionBackend +from rhapsody.backends import DragonExecutionBackend import rhapsody -rhapsody.enable_logging(level=logging.DEBUG) +rhapsody.enable_logging(level=logging.INFO) from impress import PipelineSetup from impress import ImpressManager @@ -326,8 +325,7 @@ async def adaptive_decision(pipeline: ProteinBindingPipeline) -> None: async def impress_protein_bind() -> None: """Execute protein binding analysis with LLM-driven adaptive optimization.""" -# backend = await DragonExecutionBackendV3() - backend = await LocalExecutionBackend(ProcessPoolExecutor()) + backend = await DragonExecutionBackend() manager: ImpressManager = ImpressManager(execution_backend=backend) 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..be9fb47 100644 --- a/examples/protein_binding/run_protein_binding.py +++ b/examples/protein_binding/run_protein_binding.py @@ -1,20 +1,29 @@ import copy +import os import shutil import asyncio from typing import Dict, Any, Optional, List -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend from rhapsody.telemetry import define_event from rhapsody.telemetry.events import make_event -from impress import PipelineSetup -from impress import ImpressManager +from impress import GPUPolicy, _find_gpus, _make_policy, ImpressManager, PipelineSetup from protein_binding import ProteinBindingPipeline import rhapsody, logging rhapsody.enable_logging(level=logging.DEBUG) +# ── Test mode (IMPRESS_TEST_MODE=1) ─────────────────────────────────────── +# Runs 2 pipelines with max_passes=1 and no child pipelines, so a single +# MPNN → score → AF2 cycle completes for integration testing. +TEST_MODE = os.getenv("IMPRESS_TEST_MODE", "0") == "1" +N_PIPELINES = 2 if TEST_MODE else 16 +MAX_PASSES = 1 if TEST_MODE else 10 +MAX_SUB_PIPELINES_OVERRIDE = 0 if TEST_MODE else None # None = use inline default + + # --------------------------------------------------------------------------- # Custom application-level telemetry events # --------------------------------------------------------------------------- @@ -61,12 +70,12 @@ def _on_task_event(event) -> None: # 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() + backend = await DragonExecutionBackend() manager: ImpressManager = ImpressManager( execution_backend=backend, @@ -80,12 +89,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 +117,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 +145,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 +171,8 @@ 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, } } @@ -177,6 +190,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 +198,32 @@ 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() + 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, + "policy": _make_policy(all_gpus, i - 1), + }, + adaptive_fn=adaptive_decision, + max_passes=MAX_PASSES, ) - for i in range(1, 17) + for i in range(1, N_PIPELINES + 1) ] await manager.start(pipeline_setups=pipeline_setups) @@ -201,8 +234,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..b490f10 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 inside Dragon tasks (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 index 29b2126..d72a7a5 100755 --- a/examples/protein_binding/scripts/s4_alphafold.sh +++ b/examples/protein_binding/scripts/s4_alphafold.sh @@ -7,14 +7,32 @@ set -euo pipefail 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 +# Re-activate the IMPRESS venv inside Dragon tasks (VIRTUAL_ENV is exported by sbatch). +[ -n "${VIRTUAL_ENV:-}" ] && source "${VIRTUAL_ENV}/bin/activate" -pixi run --manifest-path /anvil/scratch/x-mason/localcolabfold \ - colabfold_batch \ +# ── Test-mode stub (IMPRESS_TEST_MODE=1) ────────────────────────────────── +# Write minimal Boltz-shaped output so plddt_extract_pipeline can run without +# invoking the real AlphaFold model. +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_alphafold stub done for ${name}" + exit 0 +fi +# ── End test-mode stub ──────────────────────────────────────────────────── + +colabfold_batch \ --model-type alphafold2_multimer_v3 \ --max-template-date 2020-12-01 \ --rank multimer \ diff --git a/examples/protein_binding/scripts/s4_boltz.sh b/examples/protein_binding/scripts/s4_boltz.sh index bab328c..a10cef9 100755 --- a/examples/protein_binding/scripts/s4_boltz.sh +++ b/examples/protein_binding/scripts/s4_boltz.sh @@ -7,25 +7,59 @@ set -e fasta_path="$1" output_dir="$2" -#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}" + +# Prevent PyTorch Lightning from installing SLURM auto-requeue signal handlers. +# PL detects SLURM_JOB_ID and registers SIGTERM/SIGUSR handlers that keep the +# process group alive after Boltz finishes, causing Dragon to report +# "ProcessGroup manager is not in state State.DEAD". +unset SLURM_JOB_ID boltz predict \ "${fasta_path}" \ --out_dir "${output_dir}" \ - --use_msa_server \ - --cache /anvil/projects/x-nairr240405/mason/boltz \ + ${_MSA_FLAG} \ + --cache "${BOLTZ_CACHE_DIR:-${HOME}/.boltz}" \ --output_format pdb \ --write_full_pae \ - --override + --no_kernels \ + --devices 1 \ + --override \ + 2>&1 | tee "${output_dir}/boltz_run.log" +# tee exits 0; check the actual boltz exit code via PIPESTATUS +test "${PIPESTATUS[0]}" -eq 0 diff --git a/examples/protein_binding/scripts/s5_plddt_extract.sh b/examples/protein_binding/scripts/s5_plddt_extract.sh index 38cb576..11a3367 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 inside Dragon tasks (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}" diff --git a/examples/small_molecule_binding/CLAUDE.md b/examples/small_molecule_binding/CLAUDE.md index aaa9f96..2e51df2 100644 --- a/examples/small_molecule_binding/CLAUDE.md +++ b/examples/small_molecule_binding/CLAUDE.md @@ -7,6 +7,7 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co | Date | Commit | Notes | |---|---|---| | 2026-04-06 | 3390b61 | change log added | +| 2026-09-01 | — | RunConfig dataclass; PROD/TEST named configs replace flat if/else constants | ## Context @@ -32,7 +33,7 @@ Before running on HPC, edit the path constants at the top of `run_small_molecule - **`small_molecule_binding.py`** — defines `SmallMoleculeBindingPipeline(ImpressBasePipeline)`, all step constants, ensemble utility functions (`_ca_rmsd`, `_seq_identity`, `_ensemble_selective_avg`), and the inner `_run_refine_cycle()` loop. All pipeline tasks (HPC and local analysis) are registered via `@self.auto_register_task()` inside `_register_real_tasks()`. The `run()` method drives a state-machine loop; `_run_refine_cycle()` handles the MPNN+PackMin inner loop with per-cycle sequence retry support. -- **`run_small_molecule_binding.py`** — entry point. Sets threshold constants, defines the `adaptive_decision()` callback, creates an `ImpressManager`, and launches via `manager.start(pipeline_setups=[...])`. +- **`run_small_molecule_binding.py`** — entry point. Defines the `RunConfig` dataclass and two named instances (`PROD`, `TEST`); selects between them via `IMPRESS_TEST_MODE`; defines the `adaptive_decision()` callback; creates an `ImpressManager` and launches via `manager.start(pipeline_setups=[...])`. ### Step constants (state-machine constants in `small_molecule_binding.py`) @@ -124,7 +125,7 @@ Ensemble similarity utilities (all in `small_molecule_binding.py`): | `fold_min_plddt` | 70.0 | minimum mean pLDDT | | `max_tasks` | 300 | maximum ensemble entries before stopping | -Threshold constants in `run_small_molecule_binding.py` override these defaults at `PipelineSetup` construction. +The `PROD` config in `run_small_molecule_binding.py` overrides these class defaults at `PipelineSetup` construction (e.g. `backbone_max_ca_deviation=1.0`, `fastrelax_max_interact=-8.0`, `fold_min_plddt=75.0`). The `TEST` config sets inert thresholds (everything passes) and `max_tasks=10` for integration testing. ### Output directory structure @@ -174,7 +175,7 @@ Steps communicate via `self.state`: ### Execution backends -`run_small_molecule_binding.py` has `LocalExecutionBackend(ProcessPoolExecutor())` active by default. `DragonExecutionBackendV3()` is commented out — swap it in for HPC production runs. +`run_small_molecule_binding.py` uses `DragonExecutionBackend` for HPC production runs. `LocalExecutionBackend(ProcessPoolExecutor())` can be swapped in for local testing. ### Pipeline inputs diff --git a/examples/small_molecule_binding/CODE_REVIEW.md b/examples/small_molecule_binding/CODE_REVIEW.md new file mode 100644 index 0000000..16686e5 --- /dev/null +++ b/examples/small_molecule_binding/CODE_REVIEW.md @@ -0,0 +1,232 @@ +# Code Review — `examples/small_molecule_binding/` + +**Date:** 2026-08-29 +**Scope:** `small_molecule_binding.py`, `run_small_molecule_binding.py`, +`run_nonadaptive.py`, `run_test_small_molecule_binding.py`, `mock.py`, `scripts/` + +--- + +## Bugs + +### `filter_shape.py:38` — undefined `sfxn` in RosettaScripts XML + +The XML passed to `XmlObjects.create_from_string` references `score_fxn="sfxn"` in the +`ScoreTermValueBased` residue selector, but the `` block only defines +`sfxn_clean`. Rosetta raises an error at parse time and analysis fails for every PDB. + +**Fix:** change the selector to `score_fxn="sfxn_clean"` to match the defined score +function, or add a `sfxn` score function definition. + +--- + +### `rfd3.sh:5-6` — comment swaps `$4`/`$5` argument order + +The header comment says `$4=scaffold_arg $5=diffusion_batch_size`, but the actual +positional assignments and the call site in `small_molecule_binding.py` are +`$4=diffusion_batch_size $5=scaffold_arg`. The code is correct; the comment is +misleading and will cause confusion when the script is modified. + +--- + +### `packmin.py:1-22` — old-style `from rosetta.*` imports + +Lines 10–22 import from `rosetta.core.*`, `rosetta.protocols.*` etc. (old pre-3.8 +binding namespace). Modern PyRosetta exposes these only under `pyrosetta.rosetta.*`. +On Delta (where a current PyRosetta is installed), all of these imports will raise +`ModuleNotFoundError` at startup. + +**Fix:** replace all `from rosetta.X import Y` with `from pyrosetta.rosetta.X import Y` +(or remove unused imports — most of the imported symbols are never referenced in the +actual code body). + +--- + +## Potential Issues + +### `small_molecule_binding.py:149,152,153` — Anvil-specific hard-coded default paths + +```python +self.mpnn_dir = kwargs.get("mpnn_dir", "/anvil/projects/x-nairr240405/mason/LigandMPNN") +self.foundry_sif_path = kwargs.get("foundry_sif_path", "/anvil/projects/x-nairr240405/mason/foundry.sif") +self.colabfold_path = kwargs.get("colabfold_path", "/anvil/projects/x-nairr240405/mason/localcolabfold") +``` + +These defaults silently resolve to non-existent paths on any system other than Anvil +and the tools fail at runtime with no clear error that the path is wrong. Callers on +Delta must remember to override all three via `kwargs` or env vars. + +**Recommended fix:** default to `None` and raise a clear `ValueError` in `__init__` if +the path is not supplied. Or at minimum document the required overrides prominently. + +--- + +### `small_molecule_binding.py:367,449,489,605` — analysis tasks infer taskdir from `self.taskcount` + +Every analysis task (e.g. `analysis_sequence`) computes its input directory as +`{base_path}/{name}/{self.taskcount}_mpnn/out`. This works only because the +corresponding computation task (`mpnn`) incremented `taskcount` immediately before. +If any other counted task is inserted between computation and analysis, the path +silently points to the wrong directory and the task reads stale or absent files. + +--- + +### `run_small_molecule_binding.py:20`, `run_nonadaptive.py:16` — rhapsody DEBUG logging enabled unconditionally + +```python +rhapsody.enable_logging(level=logging.DEBUG) +``` + +DEBUG-level rhapsody logs every Dragon message exchange. On a multi-pipeline run this +generates thousands of lines per second and buries application output. Both runner +scripts have this unconditional call; it should default to INFO or be gated by an +env variable. + +--- + +### `af2.sh:40` — `--num-models 1` is an integration-only flag + +Running with a single model is fast but produces lower-quality predictions. This was +set during integration testing and was not reverted for production. Should either be +removed or made configurable via a script argument. + +--- + +### `mock.py` — taskdir path is missing the pipeline name component + +Real task dirs: `{base_path}/{pipeline_name}/{count}_taskname` +Mock task dirs: `{base_path}/{count}_taskname` + +The mock directory structure is inconsistent with the real one. This means mock tests +do not exercise the same path logic as real runs. Any test that checks or mocks the +directory layout will diverge silently from production behavior. + +--- + +## Code Quality + +### `small_molecule_binding.py:691` — commented-out `fixed_residues_file` argument + +```python +# fixed_residues_file=f"{self.pipeline_inputs}/fixed_residues.txt" ) +``` + +The `fixed_residues_file` parameter is accepted by `mpnn()` but never passed in the +refine cycle. Either wire it in or remove the dead parameter and the comment. + +--- + +### `run_small_molecule_binding.py:3` — unused imports + +```python +from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor +``` + +Neither is used anywhere in the file. Remove them. + +--- + +### `run_small_molecule_binding.py:155` — hardcoded pipeline count (8 pipelines) + +```python +for i in range(1, 9) +``` + +The number of pipelines is baked in as a magic literal. Should be a named constant or +a command-line argument. + +--- + +### `filter_shape.py:6-9` — bare `sys.argv` instead of argparse + +Arguments are taken as `sys.argv[1..4]` with no validation. Wrong argument count +causes an unhelpful `IndexError`. Both `fastrelax.py` and `packmin.py` use `argparse` +properly — `filter_shape.py` should too. + +--- + +### `filter_shape.py:18-21`, `118-129` — explicit `.close()` inside `with` blocks + +```python +with open(gen_output_file, 'w') as genout: + genout.write(...) + genout.close() # redundant +``` + +Calling `.close()` inside a `with` block is redundant and confusing — the context +manager closes the file on exit. Remove all explicit `.close()` calls inside `with` +blocks. + +--- + +### `filter_energy.py` — always prints to stdout + +Line 45: `print(f"Processed: {pdb_file}, ...")` prints to stdout for every file. When +the subprocess runs with stdout redirected to a log file this ends up in the log, but +if stderr is inspected separately the output is silent. Behaviour is inconsistent with +all other scripts that write only on failure. + +--- + +### `fastrelax.sh:12` — unquoted `$0` in `dirname` + +```bash +SCRIPT_DIR="$(dirname $0)" # wrong +``` + +Should be `$(dirname "$0")` to handle paths with spaces. `filter_shape.sh` and +`packmin.sh` already use the quoted form — this is the one outlier. + +--- + +### `packmin.py` — large commented-out code blocks + +Lines 96–116 contain multi-line commented-out tutorial fragments, original import +notes, and unreachable code. Clean these up before publication. + +--- + +### `mpnn_wrapper.sh` — legacy script, superseded by `mpnn.sh` + +`mpnn_wrapper.sh` calls `run.py` directly (bypassing the numpy alias fix in +`mpnn_run.py`), hardcodes `echo "A16" > fixed_residues.txt`, and invokes +`protein_mpnn_run.py` (ProteinMPNN) rather than LigandMPNN's `run.py`. The active +pipeline uses `mpnn.sh` → `mpnn_run.py`. This script is dead code and should be +removed or explicitly marked as deprecated. + +--- + +### `af2.sh:31,36` — diagnostic echo lines remain + +Two `echo` lines print to the log file: + +```bash +echo "[af2.sh] using colabfold_batch: $colabfold_bin" +echo "[af2.sh] data_dir: $data_dir" +``` + +Acceptable while debugging, but should be removed or made conditional on a `VERBOSE` +flag before pushing to the shared branch. + +--- + +### `run_nonadaptive.py:56` — commented-out `LocalExecutionBackend` alternative + +```python +#backend = await LocalExecutionBackend(ProcessPoolExecutor()) +backend = await DragonExecutionBackend() +``` + +The commented-out local backend is a leftover from development/testing. Remove it to +avoid confusion about which backend is active. + +--- + +### `run_nonadaptive.py:77` — hardcoded pipeline index list + +```python +for i in [1,2,4,6,7,8,10,11,12,13,14,15,16,18,19,20,23,26,27,30,32] +``` + +Like `run_small_molecule_binding.py`'s hardcoded range, the specific protein indices +are baked in with no clear mapping to input files. Should be a named constant or +driven by scanning the input directory. diff --git a/examples/small_molecule_binding/delta_env_setup.sh b/examples/small_molecule_binding/delta_env_setup.sh new file mode 100755 index 0000000..9f4083e --- /dev/null +++ b/examples/small_molecule_binding/delta_env_setup.sh @@ -0,0 +1,255 @@ +#!/bin/bash +# ============================================================================= +# IMPRESS Small Molecule 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/small_mol +# IMPRESS_DIR = $SCRATCH/$USER/IMPRESS +# python = auto-detected via `module load python` (Delta default: 3.13+) +# +# Tool directories (cloned by this script if absent): +# MPNN_DIR = $SCRATCH/$USER/LigandMPNN +# COLABFOLD_PATH = $SCRATCH/$USER/localcolabfold (used only for cache ref) +# COLABFOLD_CACHE_DIR= $SCRATCH/$USER/.cache/colabfold +# +# Foundry container (RFD3 backbone diffusion) is managed separately: +# Run pull_foundry.sh to build the sandbox tarball; delta_gpu_run.sh unpacks +# it to /tmp at job start. +# ============================================================================= +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/small_mol}" +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" + +MPNN_DIR="${MPNN_DIR:-${SCRATCH}/${USER}/LigandMPNN}" +COLABFOLD_CACHE_DIR="${COLABFOLD_CACHE_DIR:-${SCRATCH}/${USER}/.cache/colabfold}" + +echo "=================================================================" +echo " ENV_DIR = ${ENV_DIR}" +echo " IMPRESS_DIR = ${IMPRESS_DIR}" +echo " MPNN_DIR = ${MPNN_DIR}" +echo " COLABFOLD_CACHE = ${COLABFOLD_CACHE_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]" +# Pin dragonhpc to 0.14.1 — 0.14.2 added waitForKeys to DDRegisterClientResponse +# but the Delta system Dragon runtime has not been updated to match; 0.14.2 fails +# with AttributeError on every DDict operation on this cluster. +"${PIP}" install -q "dragonhpc==0.14.1" + +# ── 5. IMPRESS (local editable) ─────────────────────────────────────────────── +echo "" +echo "── Step 5: IMPRESS (editable) ──" +"${PIP}" install -q -e "${IMPRESS_DIR}" + +# ── 6. PyTorch (CUDA 12.1) — required by LigandMPNN ───────────────────────── +echo "" +echo "── Step 6: PyTorch (cu121) ──" +"${PIP}" install -q torch --index-url https://download.pytorch.org/whl/cu121 + +# ── 7. ColabFold + AlphaFold2 with pinned JAX versions ─────────────────────── +# +# Version constraints validated on Delta gpuA40x4 (CUDA 12.8 / cuDNN 9.25): +# +# colabfold 1.6.2 requires alphafold-colabfold==2.3.18 +# jax>=0.5.2,<0.11 +# +# alphafold-colabfold 2.3.18 is compatible with jaxlib 0.5.x but NOT with +# jaxlib 0.10.x (MSA feature shape mismatch at runtime). +# +# jax 0.5.2 / jaxlib 0.5.1 require nvidia-cudnn-cu12 >=9.1,<10.0. +# Upgrade cudnn to >=9.8.0 so jaxlib's cuDNN version check passes (jaxlib +# 0.5.1 links against cuDNN 9.x; runtime version must satisfy >=compiled). +# +echo "" +echo "── Step 7: ColabFold + AlphaFold2 (pinned JAX) ──" +# Install colabfold with the alphafold extra (pulls alphafold-colabfold 2.3.18, +# dm-haiku, dm-tree, ml-collections, absl-py). +"${PIP}" install -q "colabfold[alphafold]" +# Pin JAX to the era tested with alphafold-colabfold 2.3.18. +# jaxlib 0.5.2 does not exist on PyPI; 0.5.1 pairs with jax 0.5.2. +"${PIP}" install -q "jax[cuda12]==0.5.2" "jaxlib==0.5.1" +# Upgrade cuDNN so jaxlib's runtime check (>=compiled version) passes. +"${PIP}" install -q "nvidia-cudnn-cu12>=9.8.0,<10.0" + +# ── 8. LigandMPNN ───────────────────────────────────────────────────────────── +# +# LigandMPNN is run directly from its source tree (no package install). +# This step clones the repo; all Python dependencies (torch, ProDy, biopython, +# numpy) are already satisfied by the venv above. +# LigandMPNN's own requirements.txt pins older torch/cudnn versions — do NOT +# install it into this venv; the newer versions here are compatible at runtime. +# +echo "" +echo "── Step 8: LigandMPNN (clone) ──" +if [ ! -d "${MPNN_DIR}" ]; then + echo " Cloning LigandMPNN to ${MPNN_DIR}" + git clone https://github.com/dauparas/LigandMPNN "${MPNN_DIR}" +else + echo " LigandMPNN already at ${MPNN_DIR}, pulling latest" + git -C "${MPNN_DIR}" pull --ff-only || echo " (pull skipped — non-fast-forward or detached HEAD)" +fi +# Install ProDy and biopython (needed by LigandMPNN; torch already installed). +"${PIP}" install -q ProDy biopython + +# ── 9. gemmi — CIF.GZ parsing for backbone conversion ──────────────────────── +echo "" +echo "── Step 9: gemmi ──" +"${PIP}" install -q gemmi + +# ── 10. Additional dependencies ─────────────────────────────────────────────── +echo "" +echo "── Step 10: pandas + biopandas + matplotlib ──" +"${PIP}" install -q pandas biopandas matplotlib + +# ── 11. PyRosetta ───────────────────────────────────────────────────────────── +echo "" +echo "── Step 11: 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()" + +# ── 12. ColabFold model weights ─────────────────────────────────────────────── +# +# Pre-download AlphaFold2 model weights to COLABFOLD_CACHE_DIR so compute +# nodes (no internet) find them at runtime. Run this step on a login node. +# +echo "" +echo "── Step 12: ColabFold model weights ──" +mkdir -p "${COLABFOLD_CACHE_DIR}" +echo " Downloading AlphaFold2 weights to ${COLABFOLD_CACHE_DIR} ..." +"${PY}" -c " +from pathlib import Path +from colabfold.download import download_alphafold_params +download_alphafold_params('alphafold2', Path('${COLABFOLD_CACHE_DIR}')) +print(' Weights downloaded.') +" + +# ── 13. Verify ──────────────────────────────────────────────────────────────── +echo "" +echo "── Step 13: 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 "jax" "${PY}" -c "import jax; print(jax.__version__)" +_check "colabfold" "${PY}" -c "import colabfold; print(colabfold.__version__)" +_check "alphafold" "${PY}" -c "import alphafold; print('ok')" +_check "gemmi" "${PY}" -c "import gemmi; print(gemmi.__version__)" +_check "pyrosetta" "${PY}" -c "import pyrosetta; print('ok')" +_check "ProDy" "${PY}" -c "import prody; print(prody.__version__)" +_check "LigandMPNN" test -d "${MPNN_DIR}" && echo "present" +_check "colabfold weights" test -d "${COLABFOLD_CACHE_DIR}/params" && echo "present" + +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/small_molecule_binding" +echo " sbatch delta_gpu_run.sh" +echo "" +echo "Note: Foundry container (RFD3) is managed separately." +echo " Build once with: sbatch pull_foundry.sh" +echo "=================================================================" diff --git a/examples/small_molecule_binding/delta_gpu_run.sh b/examples/small_molecule_binding/delta_gpu_run.sh new file mode 100644 index 0000000..c897d21 --- /dev/null +++ b/examples/small_molecule_binding/delta_gpu_run.sh @@ -0,0 +1,148 @@ +#!/bin/bash +# +# Small Molecule 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/ +# +# Optional overrides (all have defaults based on SCRATCH/$USER): +# export MPNN_DIR=/path/to/LigandMPNN +# export COLABFOLD_PATH=/path/to/localcolabfold +# export COLABFOLD_CACHE_DIR=/path/to/colabfold_cache +# +# Foundry container (RFD3): +# The foundry sandbox is stored as a .tar.gz on scratch (built by pull_foundry.sh). +# This script extracts it to /tmp at job start (no scratch quota cost) and removes +# it on exit. Override FOUNDRY_TAR to point to a different archive, or set +# FOUNDRY_SIF_PATH directly to skip extraction entirely (e.g. a pre-extracted dir). +# +# Example: +# sbatch delta_gpu_run.sh +# sbatch delta_gpu_run.sh run_nonadaptive.py # non-adaptive runner +# +# 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_sm_binding +#SBATCH --mail-user=mg2347@soe.rutgers.edu +#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 + +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/small_mol}" +unset SLURM_EXPORT_ENV +source "${IMPRESS_VENV}/bin/activate" +dragon-config add --ofi-runtime-lib="${FAB_LIB}" + +# ── Tool paths (read by SmallMoleculeBindingPipeline via env vars) ───────────── +# These are picked up by the pipeline's __init__ when not passed as kwargs. +export MPNN_DIR="${MPNN_DIR:-${SCRATCH}/${USER}/LigandMPNN}" + +export COLABFOLD_PATH="${COLABFOLD_PATH:-${SCRATCH}/${USER}/localcolabfold}" + +# ColabFold model weights cache — kept on scratch to avoid home quota exhaustion. +# Pre-download once on login node: +# export COLABFOLD_CACHE_DIR=${SCRATCH}/${USER}/.cache/colabfold +# python -c "from colabfold.download import download_alphafold_params; \ +# download_alphafold_params('alphafold2', '${COLABFOLD_CACHE_DIR}')" +export COLABFOLD_CACHE_DIR="${COLABFOLD_CACHE_DIR:-${SCRATCH}/${USER}/.cache/colabfold}" +mkdir -p "${COLABFOLD_CACHE_DIR}" + +# ── Foundry sandbox: extract to /tmp at job start, clean up on exit ─────────── +# Extracting to /tmp avoids the scratch quota. Compute nodes have ample /tmp +# space that is not quota-counted. If FOUNDRY_SIF_PATH is already set (e.g. +# a pre-built .sif or a persistent sandbox on a large allocation), extraction +# is skipped entirely. +if [ -z "${FOUNDRY_SIF_PATH:-}" ]; then + FOUNDRY_TAR="${FOUNDRY_TAR:-${SCRATCH}/${USER}/foundry_sandbox.tar.gz}" + if [ ! -f "${FOUNDRY_TAR}" ]; then + echo "ERROR: foundry sandbox tarball not found: ${FOUNDRY_TAR}" + echo " Build it first: sbatch pull_foundry.sh" + exit 1 + fi + _FOUNDRY_TMP="/tmp/foundry_${SLURM_JOB_ID:-$$}" + echo "Extracting foundry sandbox from ${FOUNDRY_TAR} to ${_FOUNDRY_TMP} ..." + mkdir -p "${_FOUNDRY_TMP}" + tar -xzf "${FOUNDRY_TAR}" -C "${_FOUNDRY_TMP}" --strip-components=1 + export FOUNDRY_SIF_PATH="${_FOUNDRY_TMP}" + # shellcheck disable=SC2064 + trap "echo 'Removing ${_FOUNDRY_TMP}'; rm -rf '${_FOUNDRY_TMP}'" EXIT +fi + +echo "MPNN_DIR: ${MPNN_DIR}" +echo "FOUNDRY_SIF_PATH: ${FOUNDRY_SIF_PATH}" +echo "COLABFOLD_PATH: ${COLABFOLD_PATH}" +echo "COLABFOLD_CACHE: ${COLABFOLD_CACHE_DIR}" + +# ── Tool existence checks ────────────────────────────────────────────────────── +if [ ! -d "${MPNN_DIR}" ]; then + echo "ERROR: MPNN_DIR does not exist: ${MPNN_DIR}" + echo " Clone LigandMPNN: git clone https://github.com/dauparas/LigandMPNN ${MPNN_DIR}" + exit 1 +fi + +# ── Working directory ───────────────────────────────────────────────────────── +WORKDIR="${IMPRESS_SCRIPTS_DIR:-${SCRATCH}/${USER}/IMPRESS/examples/small_molecule_binding}" +cd "${WORKDIR}" +mkdir -p logs + +# IMPRESS_WORK_DIR: where pipeline task dirs (p1/, p2/, …) are written. +# Defaults to logs/ so all run artifacts stay out of the source tree and are +# covered by .gitignore. Override to write outputs elsewhere. +export IMPRESS_WORK_DIR="${IMPRESS_WORK_DIR:-${WORKDIR}/logs}" +mkdir -p "${IMPRESS_WORK_DIR}" + +# IMPRESS_TEST_MODE=1: 2 pipelines, inert thresholds, max_tasks=10. +# Runs one full rfd3→mpnn→fastrelax→filter_shape→af2 cycle to verify the +# end-to-end path without looping. 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}" + +# ── Run ─────────────────────────────────────────────────────────────────────── +# asyncflow session dirs now go to /tmp (node-local, no quota) via IMPRESS_SESSION_DIR. + +# -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 + +rm -f ddict_orc* + +RUNNER="${1:-run_small_molecule_binding.py}" +echo "Running: dragon ${DRAGON_MODE} ${RUNNER} (nodes=${SLURM_NNODES:-1})" +dragon ${DRAGON_MODE} "${RUNNER}" + +echo "=== Small Molecule Binding pipeline done: $(date) ===" diff --git a/examples/small_molecule_binding/pull_foundry.sh b/examples/small_molecule_binding/pull_foundry.sh new file mode 100644 index 0000000..8183bee --- /dev/null +++ b/examples/small_molecule_binding/pull_foundry.sh @@ -0,0 +1,35 @@ +#!/bin/bash +#SBATCH --account=bblj-delta-gpu +#SBATCH --partition=gpuA40x4 +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=128G +#SBATCH --gpus-per-node=1 +#SBATCH --time=00:30:00 +#SBATCH --job-name=pull_foundry +#SBATCH --output=slurm-%j.out +#SBATCH --error=slurm-%j.err + +set -euo pipefail +ulimit -c 0 # disable core dumps + +export APPTAINER_CACHEDIR=/scratch/bblj/mgoliyad1/.apptainer_cache +export APPTAINER_TMPDIR=/tmp/apptainer_$$ +mkdir -p "$APPTAINER_TMPDIR" + +# Build as sandbox (directory) to /tmp — no mksquashfs involved. +# Then tar to a single archive on scratch for storage. +SANDBOX=/tmp/foundry_sandbox_$$ +DEST_TAR=/scratch/bblj/mgoliyad1/foundry_sandbox.tar.gz + +echo "=== Building sandbox to /tmp ===" +apptainer build --sandbox "$SANDBOX" docker://rosettacommons/foundry + +echo "=== Compressing sandbox to scratch ===" +tar -czf "$DEST_TAR" -C /tmp "foundry_sandbox_$$" + +echo "Done: $(ls -lh "$DEST_TAR")" +rm -rf "$SANDBOX" + +rm -rf "$APPTAINER_TMPDIR" diff --git a/examples/small_molecule_binding/run_nonadaptive.py b/examples/small_molecule_binding/run_nonadaptive.py index 75f56fd..2f87e6e 100644 --- a/examples/small_molecule_binding/run_nonadaptive.py +++ b/examples/small_molecule_binding/run_nonadaptive.py @@ -1,9 +1,8 @@ import asyncio +import os from typing import List -from radical.asyncflow import LocalExecutionBackend -from concurrent.futures import ProcessPoolExecutor -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend from impress import ImpressManager, PipelineSetup from small_molecule_binding import ( @@ -13,7 +12,7 @@ import logging import rhapsody -rhapsody.enable_logging(level=logging.DEBUG) +rhapsody.enable_logging(level=logging.INFO) # ── Per-step quality thresholds ──────────────────────────────────────────── # These are passed to pipeline analysis tasks for metric logging but are not @@ -51,10 +50,21 @@ async def nonadaptive_decision(pipeline: SmallMoleculeBindingPipeline) -> None: ) +# Indices of pipelines to run — corresponds to the protein IDs selected for +# this campaign (non-contiguous because some were dropped after earlier runs). +PIPELINE_INDICES = [1, 2, 4, 6, 7, 8, 10, 11, 12, 13, 14, 15, 16, 18, 19, 20, 23, 26, 27, 30, 32] + + async def impress_smallmol_nonadaptive() -> None: """Execute the small-molecule binding pipeline without adaptive routing.""" - #backend = await LocalExecutionBackend(ProcessPoolExecutor()) - backend = await DragonExecutionBackendV3() + examples_dir = os.path.dirname(os.path.abspath(__file__)) + work_dir = os.environ.get( + "IMPRESS_WORK_DIR", os.path.join(examples_dir, "logs") + ) + os.makedirs(work_dir, exist_ok=True) + input_dir = os.path.join(examples_dir, "p1_in") + + backend = await DragonExecutionBackend() manager: ImpressManager = ImpressManager(execution_backend=backend) pipeline_setups: List[PipelineSetup] = [ @@ -63,6 +73,9 @@ async def impress_smallmol_nonadaptive() -> None: type=SmallMoleculeBindingPipeline, adaptive_fn=nonadaptive_decision, kwargs={ + "base_path": work_dir, + "scripts_path": os.path.join(examples_dir, "scripts"), + "input_dir": input_dir, "backbone_max_ca_deviation": BACKBONE_MAX_CA_DEVIATION, "backbone_min_ss_fraction": BACKBONE_MIN_SS_FRACTION, "fastrelax_max_fa_rep": FASTRELAX_MAX_FA_REP, @@ -74,7 +87,7 @@ async def impress_smallmol_nonadaptive() -> None: "num_refine_cycles": 2, } ) - for i in [1,2,4,6,7,8,10,11,12,13,14,15,16,18,19,20,23,26,27,30,32] + for i in PIPELINE_INDICES ] await manager.start(pipeline_setups=pipeline_setups) diff --git a/examples/small_molecule_binding/run_small_molecule_binding.py b/examples/small_molecule_binding/run_small_molecule_binding.py index f39f6d7..a0983cc 100644 --- a/examples/small_molecule_binding/run_small_molecule_binding.py +++ b/examples/small_molecule_binding/run_small_molecule_binding.py @@ -1,12 +1,11 @@ import asyncio import os -from concurrent.futures import ThreadPoolExecutor,ProcessPoolExecutor +from dataclasses import dataclass from typing import List -from radical.asyncflow import LocalExecutionBackend -from rhapsody.backends import DragonExecutionBackendV3 +from rhapsody.backends import DragonExecutionBackend -from impress import ImpressManager, PipelineSetup +from impress import GPUPolicy, _find_gpus, _make_policy, ImpressManager, PipelineSetup from small_molecule_binding import ( SmallMoleculeBindingPipeline, STEP_DONE, STEP_RFD3, STEP_MPNN, STEP_FASTRELAX, STEP_INTERFACE, STEP_AF2, @@ -17,16 +16,59 @@ import logging import rhapsody -rhapsody.enable_logging(level=logging.DEBUG) +rhapsody.enable_logging(level=logging.INFO) + + +@dataclass +class RunConfig: + n_pipelines: int + max_tasks: int + # backbone + backbone_max_ca_deviation: float + backbone_min_ss_fraction: float + # fastrelax + fastrelax_max_fa_rep: float + fastrelax_max_score: float # total_score REU + fastrelax_max_interact: float # interaction energy REU + # interface + interface_min_sc: float # shape complementarity + # fold + fold_min_plddt: float # mean pLDDT; init is -1.0 so -1.0 = always pass + # diffusion / refinement + diffusion_batch_size: int + num_refine_cycles: int + + +PROD = RunConfig( + n_pipelines = 4, + max_tasks = 300, + backbone_max_ca_deviation = 1.0, + backbone_min_ss_fraction = 0.5, + fastrelax_max_fa_rep = 100.0, + fastrelax_max_score = -250.0, # data range -193 to -510 + fastrelax_max_interact = -8.0, # p75 = -8.8 + interface_min_sc = 0.55, + fold_min_plddt = 75.0, + diffusion_batch_size = 4, + num_refine_cycles = 2, +) + +# Inert thresholds — everything passes; low task budget for one full cycle. +TEST = RunConfig( + n_pipelines = 2, + max_tasks = 10, + backbone_max_ca_deviation = 9999.0, + backbone_min_ss_fraction = 0.0, + fastrelax_max_fa_rep = 9999.0, + fastrelax_max_score = 9999.0, + fastrelax_max_interact = 9999.0, + interface_min_sc = 0.0, + fold_min_plddt = -1.0, + diffusion_batch_size = 1, + num_refine_cycles = 1, +) -# ── Per-step quality thresholds ──────────────────────────────────────────── -BACKBONE_MAX_CA_DEVIATION = 1.0 -BACKBONE_MIN_SS_FRACTION = 0.5 -FASTRELAX_MAX_FA_REP = 100.0 # fa_rep REU -FASTRELAX_MAX_SCORE = -250.0 # total_score REU (was 0.0 — inert; data range -193 to -510) -FASTRELAX_MAX_INTERACT = -8.0 # interaction energy REU (was absent/0.0 — inert; p75=-8.8) -INTERFACE_MIN_SC = 0.55 # shape complementarity (was 0.35 — passed 100%; data min=0.37) -FOLD_MIN_PLDDT = 75.0 # mean pLDDT (was 70.0; ~70% of AF2 runs exceed 75) +cfg = TEST if os.getenv("IMPRESS_TEST_MODE", "0") == "1" else PROD async def adaptive_decision(pipeline: SmallMoleculeBindingPipeline) -> None: @@ -139,32 +181,48 @@ def _prior(ttype): async def impress_smallmol_bind() -> None: """Execute the small-molecule binding pipeline.""" + # Resolve paths before launching Dragon (os.getcwd() is the examples dir). + examples_dir = os.path.dirname(os.path.abspath(__file__)) + work_dir = os.environ.get( + "IMPRESS_WORK_DIR", os.path.join(examples_dir, "logs") + ) + os.makedirs(work_dir, exist_ok=True) + # Input data lives in the source tree; pass as absolute so it resolves + # correctly regardless of what base_path / work_dir is set to. + input_dir = os.path.join(examples_dir, "p1_in") + #backend = await LocalExecutionBackend(ProcessPoolExecutor()) - backend = await DragonExecutionBackendV3() + backend = await DragonExecutionBackend() manager: ImpressManager = ImpressManager(execution_backend=backend) + all_gpus = _find_gpus() + pipeline_setups: List[PipelineSetup] = [ PipelineSetup( name=f"p{str(i)}", type=SmallMoleculeBindingPipeline, adaptive_fn=adaptive_decision, kwargs={ - "backbone_max_ca_deviation": BACKBONE_MAX_CA_DEVIATION, - "backbone_min_ss_fraction": BACKBONE_MIN_SS_FRACTION, - "fastrelax_max_fa_rep": FASTRELAX_MAX_FA_REP, - "fastrelax_max_total_score": FASTRELAX_MAX_SCORE, - "fastrelax_max_interact": FASTRELAX_MAX_INTERACT, - "interface_min_sc": INTERFACE_MIN_SC, - "fold_min_plddt": FOLD_MIN_PLDDT, - "diffusion_batch_size": 4, - "num_refine_cycles": 2, + "base_path": work_dir, + "scripts_path": os.path.join(examples_dir, "scripts"), + "input_dir": input_dir, + "backbone_max_ca_deviation": cfg.backbone_max_ca_deviation, + "backbone_min_ss_fraction": cfg.backbone_min_ss_fraction, + "fastrelax_max_fa_rep": cfg.fastrelax_max_fa_rep, + "fastrelax_max_total_score": cfg.fastrelax_max_score, + "fastrelax_max_interact": cfg.fastrelax_max_interact, + "interface_min_sc": cfg.interface_min_sc, + "fold_min_plddt": cfg.fold_min_plddt, + "diffusion_batch_size": cfg.diffusion_batch_size, + "num_refine_cycles": cfg.num_refine_cycles, + "max_tasks": cfg.max_tasks, + "policy": _make_policy(all_gpus, i - 1), } ) - for i in range(1,9) + for i in range(1, cfg.n_pipelines + 1) ] await manager.start(pipeline_setups=pipeline_setups) - await manager.flow.shutdown() if __name__ == "__main__": diff --git a/examples/small_molecule_binding/scripts/af2.sh b/examples/small_molecule_binding/scripts/af2.sh index ede53a8..fcdbf9b 100755 --- a/examples/small_molecule_binding/scripts/af2.sh +++ b/examples/small_molecule_binding/scripts/af2.sh @@ -1,20 +1,43 @@ #!/bin/bash set -euo pipefail -# AlphaFold2 structure prediction via LocalColabFold (pixi) +# AlphaFold2 structure prediction via LocalColabFold # Args: $1=colabfold_path $2=short_fasta $3=output_dir +# +# colabfold_path: root of the localcolabfold repo (contains pyproject.toml). +# colabfold_batch is resolved from .pixi/envs/default/bin/ under that root, +# bypassing pixi (not available on Delta). colabfold_path="$1" short_fasta="$2" output_dir="$3" -module load gcc/11.2.0 -module load cuda -source /anvil/projects/x-nairr240405/mason/IMPRESS/.venv/bin/activate +# Prefer the venv's colabfold_batch (installed in IMPRESS venv) over the +# hooten1 pixi env, whose Python files are not world-readable on Delta. +VENV_BIN="$(dirname "$(command -v python 2>/dev/null)")" +colabfold_bin="" +for candidate in \ + "${VENV_BIN}/colabfold_batch" \ + "${colabfold_path}/.pixi/envs/default/bin/colabfold_batch"; do + if [ -x "$candidate" ]; then + colabfold_bin="$candidate" + break + fi +done +if [ -z "$colabfold_bin" ]; then + echo "ERROR: colabfold_batch not found in venv or at $colabfold_path" >&2 + exit 1 +fi -pixi run --manifest-path "$colabfold_path" \ - colabfold_batch \ +# Use scratch for the model weights cache to avoid home quota exhaustion. +# COLABFOLD_CACHE_DIR must be set (delta_sbatch.sh exports it). +data_dir="${COLABFOLD_CACHE_DIR:-${HOME}/.cache/colabfold}" + +"$colabfold_bin" \ --model-type alphafold2 \ + --msa-mode single_sequence \ + --num-models "${AF2_NUM_MODELS:-1}" \ + --data "$data_dir" \ --rank auto \ --random-seed 999 \ --save-all \ diff --git a/examples/small_molecule_binding/scripts/fastrelax.py b/examples/small_molecule_binding/scripts/fastrelax.py index b899106..cfd78be 100755 --- a/examples/small_molecule_binding/scripts/fastrelax.py +++ b/examples/small_molecule_binding/scripts/fastrelax.py @@ -100,7 +100,7 @@ def main(args): if __name__ == '__main__': args = parse_args() - opts = '-ex1 -ex2 -use_input_sc -flip_HNQ -no_optH false -out:level 500' + opts = '-ex1 -ex2 -use_input_sc -flip_HNQ -no_optH false -mute all' if args.constraints: opts += ' -enzdes::cstfile {}'.format(args.constraints) opts += ' -run:preserve_header' diff --git a/examples/small_molecule_binding/scripts/fastrelax.sh b/examples/small_molecule_binding/scripts/fastrelax.sh index 57a3987..729586c 100755 --- a/examples/small_molecule_binding/scripts/fastrelax.sh +++ b/examples/small_molecule_binding/scripts/fastrelax.sh @@ -8,9 +8,9 @@ pdb_path="$1" lig_path="$2" output_dir="$3" -SCRIPT_DIR="$(dirname $0)" +SCRIPT_DIR="$(dirname "$0")" -source /anvil/projects/x-nairr240405/mason/IMPRESS/.venv/bin/activate +source "${ENV_DIR:-/u/${USER}/ve/impress}/bin/activate" python "$SCRIPT_DIR/fastrelax.py" \ "$pdb_path" \ diff --git a/examples/small_molecule_binding/scripts/filter_energy.py b/examples/small_molecule_binding/scripts/filter_energy.py index 4acb46d..a570cd8 100644 --- a/examples/small_molecule_binding/scripts/filter_energy.py +++ b/examples/small_molecule_binding/scripts/filter_energy.py @@ -42,8 +42,6 @@ parts = line.split() ligand_energy = float(parts[-1]) # The last element is the total energy - print(f"Processed: {pdb_file}, Ligand Energy ({ligand_name}): {ligand_energy}") - # Check if the ligand energy is negative if ligand_energy < 0: with open(output_file, 'a') as of: diff --git a/examples/small_molecule_binding/scripts/filter_shape.py b/examples/small_molecule_binding/scripts/filter_shape.py index 87278a5..cfb8ab7 100644 --- a/examples/small_molecule_binding/scripts/filter_shape.py +++ b/examples/small_molecule_binding/scripts/filter_shape.py @@ -1,30 +1,37 @@ +import argparse import os import re import pyrosetta -import sys -pdb_directory = sys.argv[1] # '/WWW/PDB_Files' -SC_output_file = sys.argv[2] # 'shape_complementarity_values.txt' -ligand_name = sys.argv[3] # 'ALR' -gen_output_file = sys.argv[4] # 'interface_values.txt' -pyrosetta.init(f"-ignore_unrecognized_res -ignore_zero_occupancy --extra_res_fa {ligand_name}.params -corrections::beta_nov16 true") +def parse_args(): + parser = argparse.ArgumentParser(description="Rosetta shape-complementarity/interface analysis") + parser.add_argument("pdb_directory", help="Directory of PDB files to analyse") + parser.add_argument("SC_output_file", help="Output file for shape-complementarity values") + parser.add_argument("ligand_name", help="Ligand residue name (e.g. ALR)") + parser.add_argument("gen_output_file", help="Output CSV for interface metrics") + return parser.parse_args() -# Define directories and files -#pdb_directory = '/WWW/PDB_Files' -#SC_output_file = 'shape_complementarity_values.txt' -# Set up the general output file which will have metrics that look at the interface -with open(gen_output_file, 'w') as genout: - # Write the shape complementarity value to the output file - genout.write(f"FileName,Shape Complementarity,ddg,contact molecular surf,SASA,Very buried unsat hbond,Surface unsat hbond,SAP SCORE\n") - genout.close() +def main(): + args = parse_args() -# Create list of PDB files to analyze -pdb_files = [f for f in os.listdir(pdb_directory) if f.endswith('.pdb')] + pyrosetta.init( + f"-ignore_unrecognized_res -ignore_zero_occupancy --extra_res_fa {args.ligand_name}.params" + f" -corrections::beta_nov16 true -mute all" + ) -################################################################################################################################################################################################################################### -protocol = pyrosetta.rosetta.protocols.rosetta_scripts.XmlObjects().create_from_string( + # Set up the general output file which will have metrics that look at the interface + with open(args.gen_output_file, 'w') as genout: + genout.write( + "FileName,Shape Complementarity,ddg,contact molecular surf,SASA," + "Very buried unsat hbond,Surface unsat hbond,SAP SCORE\n" + ) + + # Create list of PDB files to analyze + pdb_files = [f for f in os.listdir(args.pdb_directory) if f.endswith('.pdb')] + + protocol = pyrosetta.rosetta.protocols.rosetta_scripts.XmlObjects().create_from_string( """ @@ -34,7 +41,7 @@ - + @@ -93,41 +100,25 @@ """).get_mover("ParsedProtocol") -#################################################################################################################################################################################################################################### - - -##print(f""" -##shape complementarity : {pose.scores['sc2']} -##ddg : {pose.scores['ddg']} -##contact molecular surf : {pose.scores['cms']} -##SASA : {pose.scores['IA_dSASA_int']} -##Very buried unsat hbond: {pose.scores['vbuns']} -##Surface unsat hbond : {pose.scores['sbuns']} -##SAP SCORE : {pose.scores['sap_score']} -##""") - -# Analyze the selected PDB files -for pdb_file in pdb_files: - full_path = os.path.join(pdb_directory, pdb_file) - - # Initialize variables for ligand energy - pose = pyrosetta.pose_from_pdb(full_path) - protocol.apply(pose) + # Analyze the selected PDB files + for pdb_file in pdb_files: + full_path = os.path.join(args.pdb_directory, pdb_file) - # Open the SC output file - with open(SC_output_file, 'a') as SCout: - - # Write the shape complementarity value to the output file - SCout.write(f"{pdb_file}\tShape Complementarity: {pose.scores['sc2']}\n") - SCout.close() + pose = pyrosetta.pose_from_pdb(full_path) + protocol.apply(pose) - # Open the general output file - with open(gen_output_file, 'a') as genout: + with open(args.SC_output_file, 'a') as SCout: + SCout.write(f"{pdb_file}\tShape Complementarity: {pose.scores['sc2']}\n") - # Write the shape complementarity value to the output file - genout.write(f"{pdb_file},{pose.scores['sc2']},{pose.scores['ddg']},{pose.scores['cms']},{pose.scores['IA_dSASA_int']},{pose.scores['vbuns']},{pose.scores['sbuns']},{pose.scores['sap_score']}\n") - genout.close() + with open(args.gen_output_file, 'a') as genout: + genout.write( + f"{pdb_file},{pose.scores['sc2']},{pose.scores['ddg']}," + f"{pose.scores['cms']},{pose.scores['IA_dSASA_int']}," + f"{pose.scores['vbuns']},{pose.scores['sbuns']},{pose.scores['sap_score']}\n" + ) -print(f"Shape complementarity values have been written to {SC_output_file}.") + print(f"Shape complementarity values have been written to {args.SC_output_file}.") +if __name__ == "__main__": + main() diff --git a/examples/small_molecule_binding/scripts/filter_shape.sh b/examples/small_molecule_binding/scripts/filter_shape.sh index 2208834..2bd1844 100755 --- a/examples/small_molecule_binding/scripts/filter_shape.sh +++ b/examples/small_molecule_binding/scripts/filter_shape.sh @@ -11,7 +11,7 @@ interface_values_output="$4" SCRIPT_DIR="$(dirname "$0")" -source /anvil/projects/x-nairr240405/mason/IMPRESS/.venv/bin/activate +source "${ENV_DIR:-/u/${USER}/ve/impress}/bin/activate" python "$SCRIPT_DIR/filter_shape.py" \ "$pdb_directory" \ diff --git a/examples/small_molecule_binding/scripts/mpnn.sh b/examples/small_molecule_binding/scripts/mpnn.sh index c5aaf1d..29378eb 100755 --- a/examples/small_molecule_binding/scripts/mpnn.sh +++ b/examples/small_molecule_binding/scripts/mpnn.sh @@ -11,9 +11,11 @@ n_batches="$4" batch_size="$5" fixed_residues="${6:-}" -source /anvil/projects/x-nairr240405/mason/LigandMPNN/.venv/bin/activate +SCRIPT_DIR="$(dirname "$0")" -python "$mpnn_dir/run.py" \ +# mpnn_run.py restores numpy deprecated aliases (np.int/np.bool/np.object) +# removed in NumPy 1.24+ that LigandMPNN's bundled openfold still uses. +python "$SCRIPT_DIR/mpnn_run.py" "$mpnn_dir" \ --model_type "ligand_mpnn" \ --checkpoint_path_sc "$mpnn_dir/model_params/ligandmpnn_sc_v_32_002_16.pt" \ --checkpoint_ligand_mpnn "$mpnn_dir/model_params/ligandmpnn_v_32_010_25.pt" \ diff --git a/examples/small_molecule_binding/scripts/mpnn_run.py b/examples/small_molecule_binding/scripts/mpnn_run.py new file mode 100644 index 0000000..4c7cfd2 --- /dev/null +++ b/examples/small_molecule_binding/scripts/mpnn_run.py @@ -0,0 +1,25 @@ +#!/usr/bin/env python +""" +Wrapper for LigandMPNN run.py that restores deprecated numpy aliases removed +in NumPy 1.24+. LigandMPNN's bundled openfold uses np.int, np.object, np.bool. + +Usage (from mpnn.sh): + python mpnn_run.py [run.py args...] +The mpnn_dir is stripped from sys.argv before run.py sees it. +""" +import sys +import os + +import numpy as np +for _alias, _builtin in [('int', int), ('float', float), ('bool', bool), + ('complex', complex), ('object', object), ('str', str)]: + if not hasattr(np, _alias): + setattr(np, _alias, _builtin) + +mpnn_dir = sys.argv[1] +sys.argv = [os.path.join(mpnn_dir, 'run.py')] + sys.argv[2:] +sys.path.insert(0, mpnn_dir) +os.chdir(mpnn_dir) + +import runpy +runpy.run_path(os.path.join(mpnn_dir, 'run.py'), run_name='__main__') diff --git a/examples/small_molecule_binding/scripts/mpnn_wrapper.sh b/examples/small_molecule_binding/scripts/mpnn_wrapper.sh deleted file mode 100755 index 545901b..0000000 --- a/examples/small_molecule_binding/scripts/mpnn_wrapper.sh +++ /dev/null @@ -1,28 +0,0 @@ -#!/bin/bash -# run ligandmpnn - -pdb_path=$1 -output_dir=$2 -lmpnn_dir=$3 - -echo "A16" > fixed_residues.txt -FIXED=`cat fixed_residues.txt` -set -x -python $lmpnn_dir/run.py \ - --model_type "ligand_mpnn" \ - --checkpoint_path_sc $lmpnn_dir/model_params/ligandmpnn_sc_v_32_002_16.pt \ - --checkpoint_ligand_mpnn $lmpnn_dir/model_params/ligandmpnn_v_32_010_25.pt \ - --seed 111 \ - --pdb_path $pdb_path \ - --out_folder $output_dir \ - --pack_side_chains 1 \ - --number_of_batches 1 \ - --batch_size 1 \ - --number_of_packs_per_design 1 \ - --pack_with_ligand_context 1 \ - --fixed_residues "$FIXED" \ - --repack_everything 1 \ - --temperature 0.1 -# --bias_AA "A:10.0" - - diff --git a/examples/small_molecule_binding/scripts/packmin.py b/examples/small_molecule_binding/scripts/packmin.py index 8dcbf97..d32bbdd 100644 --- a/examples/small_molecule_binding/scripts/packmin.py +++ b/examples/small_molecule_binding/scripts/packmin.py @@ -7,31 +7,12 @@ from pyrosetta.teaching import * #Core Includes -from rosetta.core.kinematics import MoveMap -from rosetta.core.kinematics import FoldTree -from rosetta.core.pack.task import TaskFactory -from rosetta.core.pack.task import operation -from rosetta.core.simple_metrics import metrics -from rosetta.core.select import residue_selector as selections -from rosetta.core import select -from rosetta.core.select.movemap import * +from pyrosetta.rosetta.core.pack.task import operation +from pyrosetta.rosetta.core.select.movemap import * #Protocol Includes -from rosetta.protocols import minimization_packing as pack_min -from rosetta.protocols import relax as rel -from rosetta.protocols.antibody.residue_selector import CDRResidueSelector -from rosetta.protocols.antibody import * -from rosetta.protocols.loops import * -''' -When downloading a new PDB file, do a pack_min minimization with coordinate constraints and a ligand. - -Requires a PDB file input. - -Options: -Name (-n, string): change the output PDB name from [original_name]_relaxed.pdb -Score function (-sf, string): change the score function from the default of ref2015_cst -Catalytic residues (-cat, int, multiple accepted): list residues that should not be moved -''' +from pyrosetta.rosetta.protocols import minimization_packing as pack_min + def parse_args(): parser = argparse.ArgumentParser() @@ -93,18 +74,12 @@ def main(args): # Packer tasks with -ex1 and -ex2 tf = ut.make_task_factory() - #From pack_min tutorial from jupyter notebooks - ###tf = TaskFactory() tf.push_back(operation.InitializeFromCommandline()) tf.push_back(operation.RestrictToRepacking()) packer = pack_min.PackRotamersMover() packer.task_factory(tf) - #This line is from khare lab relax code, not sure if this will be necessary: - #pp = Pose(pose) - - #Run the packer. (Note this may take a few minutes) packer.apply(pose) # Write Rosetta score JSON alongside the output PDB @@ -115,12 +90,11 @@ def main(args): #Dump the PDB pose.dump_pdb(out_name) - ##pose.dump_pdb('/outputs/2r0l_all_repack.pdb') if __name__ == '__main__': args = parse_args() - opts = '-ex1 -ex2 -use_input_sc -flip_HNQ -no_optH false' + opts = '-ex1 -ex2 -use_input_sc -flip_HNQ -no_optH false -mute all' if args.constraints: opts += ' -enzdes::cstfile {}'.format(args.constraints) opts += ' -run:preserve_header' diff --git a/examples/small_molecule_binding/scripts/packmin.sh b/examples/small_molecule_binding/scripts/packmin.sh index 5ce3be7..cf2d058 100755 --- a/examples/small_molecule_binding/scripts/packmin.sh +++ b/examples/small_molecule_binding/scripts/packmin.sh @@ -10,7 +10,7 @@ output_dir="$3" SCRIPT_DIR="$(dirname "$0")" -source /anvil/projects/x-nairr240405/mason/IMPRESS/.venv/bin/activate +source "${ENV_DIR:-/u/${USER}/ve/impress}/bin/activate" python "$SCRIPT_DIR/packmin.py" \ "$pdb_path" \ diff --git a/examples/small_molecule_binding/scripts/rfd3.sh b/examples/small_molecule_binding/scripts/rfd3.sh index 33b27e3..72d46bd 100755 --- a/examples/small_molecule_binding/scripts/rfd3.sh +++ b/examples/small_molecule_binding/scripts/rfd3.sh @@ -2,8 +2,8 @@ set -euo pipefail # Backbone generation via RFDiffusion3 (apptainer) -# Args: $1=foundry_sif_path $2=output_dir $3=inputs $4=scaffold_arg $5=diffusion_batch_size -# scaffold_arg: "scaffoldguided.target_pdb=" or "" if unused +# Args: $1=foundry_sif_path $2=output_dir $3=inputs $4=diffusion_batch_size $5=scaffold_arg +# scaffold_arg: "scaffoldguided.target_pdb=" or "" if unused (optional, defaults to "") foundry_sif_path="$1" output_dir="$2" @@ -16,7 +16,12 @@ else scaffold_arg="" fi -apptainer exec --nv "$foundry_sif_path" rfd3 design \ +# Prevent host ~/.local Python packages from contaminating the container +# (apptainer mounts $HOME by default; PYTHONNOUSERSITE must be SET, not unset). +unset PYTHONPATH PYTHONUSERBASE PYTHONDONTWRITEBYTECODE +export PYTHONNOUSERSITE=1 + +apptainer exec --nv --writable-tmpfs --bind /scratch:/scratch "$foundry_sif_path" rfd3 design \ out_dir="$output_dir" \ inputs="$inputs" \ skip_existing=False \ diff --git a/examples/small_molecule_binding/small_molecule_binding.py b/examples/small_molecule_binding/small_molecule_binding.py index e8fc3f2..4741c9a 100644 --- a/examples/small_molecule_binding/small_molecule_binding.py +++ b/examples/small_molecule_binding/small_molecule_binding.py @@ -1,6 +1,7 @@ import asyncio import copy +import gzip import json import os import pathlib @@ -140,14 +141,30 @@ def __init__(self, name, flow, configs=None, **kwargs): super().__init__(name, flow, **configs, **kwargs) # Paths - self.base_path = kwargs.get("base_path", os.getcwd()) - self.scripts_path = os.path.join(self.base_path, "scripts") - self.pipeline_inputs = os.path.join(self.base_path, f"{self.name}_in") - self.mpnn_dir = kwargs.get("mpnn_dir", f"/anvil/projects/x-nairr240405/mason/LigandMPNN") + self.base_path = kwargs.get("base_path", os.getcwd()) + self.scripts_path = kwargs.get( + "scripts_path", os.path.join(self.base_path, "scripts") + ) + # input_dir: explicit input directory name (relative to base_path) or + # absolute path. Falls back to "_in" when not provided. + _input_dir = kwargs.get("input_dir", "") + if _input_dir and os.path.isabs(_input_dir): + self.pipeline_inputs = _input_dir + elif _input_dir: + self.pipeline_inputs = os.path.join(self.base_path, _input_dir) + else: + self.pipeline_inputs = os.path.join(self.base_path, f"{self.name}_in") + self.mpnn_dir = kwargs.get("mpnn_dir") or os.environ.get("MPNN_DIR") + if not self.mpnn_dir: + raise ValueError("mpnn_dir must be supplied via kwarg or MPNN_DIR env var") # Configurable tool paths and ensemble sizes - self.foundry_sif_path = kwargs.get("foundry_sif_path", "/anvil/projects/x-nairr240405/mason/foundry.sif") - self.colabfold_path = kwargs.get("colabfold_path", "/anvil/projects/x-nairr240405/mason/localcolabfold") + self.foundry_sif_path = kwargs.get("foundry_sif_path") or os.environ.get("FOUNDRY_SIF_PATH") + if not self.foundry_sif_path: + raise ValueError("foundry_sif_path must be supplied via kwarg or FOUNDRY_SIF_PATH env var") + self.colabfold_path = kwargs.get("colabfold_path") or os.environ.get("COLABFOLD_PATH") + if not self.colabfold_path: + raise ValueError("colabfold_path must be supplied via kwarg or COLABFOLD_PATH env var") self.ligand_params = kwargs.get("ligand_params", "ALR.params") self.mpnn_ensemble_size = kwargs.get("mpnn_ensemble_size", 1) self.num_refine_cycles = kwargs.get("num_refine_cycles", 3) @@ -162,6 +179,7 @@ def __init__(self, name, flow, configs=None, **kwargs): self.interface_min_sc = kwargs.get("interface_min_sc", 0.5) self.fold_min_plddt = kwargs.get("fold_min_plddt", 70.0) self.max_tasks = kwargs.get("max_tasks", 300) + self.policy = kwargs.get("policy", None) # Output paths (legacy) self.output_path = os.path.join(self.base_path, "myoutputs", self.name) @@ -196,8 +214,8 @@ def _register_mock_tasks(self): def _register_real_tasks(self): """Register real HPC tasks that return shell command strings.""" - @self.auto_register_task(capture_stdio=True) - async def rfd3(task_description={"gpus_per_rank": 1}): + @self.auto_register_task(local_task=True) + async def rfd3(): self.taskcount += 1 taskname = "rfd3" self.previous_task = taskname @@ -209,9 +227,9 @@ async def rfd3(task_description={"gpus_per_rank": 1}): output_dir = f"{taskdir}/out" input_pdb = self.state.get('rfd3_input_pdb') - scaffold_arg = f"scaffoldguided.target_pdb={input_pdb}" if input_pdb else "" + scaffold_arg = f"+scaffoldguided.target_pdb={input_pdb}" if input_pdb else "" - return ( + cmd = ( f"bash {self.scripts_path}/rfd3.sh" f" {self.foundry_sif_path}" f" {output_dir}" @@ -219,6 +237,17 @@ async def rfd3(task_description={"gpus_per_rank": 1}): f" {self.diffusion_batch_size}" f" {scaffold_arg}" ) + log_file = f"{taskdir}/rfd3.log" + with open(log_file, "wb") as _lf: + proc = await asyncio.create_subprocess_shell( + cmd, + stdout=_lf, + stderr=asyncio.subprocess.STDOUT, + env=self._gpu_env(), + ) + await proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"rfd3 failed with exit code {proc.returncode}\nSee {log_file}") @self.auto_register_task(local_task=True) async def analysis_backbone(): @@ -273,7 +302,7 @@ async def analysis_backbone(): ETYPE_BACKBONE, best['ss'], self.state.get('rfd3_input_pdb'), backbone_path, )) - @self.auto_register_task(capture_stdio=True) + @self.auto_register_task(local_task=True) async def mpnn( fixed_residues_file: str | None = None): self.taskcount += 1 @@ -300,13 +329,24 @@ async def mpnn( shutil.copy(pdb_path_orig, short_pdb) pdb_path = short_pdb + # LigandMPNN (ProDy parsePDB) only reads PDB format; convert CIF.GZ. + if pdb_path.endswith('.cif.gz'): + import gemmi as _gemmi + pdb_for_mpnn = f"{taskdir}/in/binder.pdb" + with gzip.open(pdb_path, 'rb') as _f: + _cif_data = _f.read().decode() + _doc = _gemmi.cif.read_string(_cif_data) + _st = _gemmi.make_structure_from_block(_doc.sole_block()) + _st.write_pdb(pdb_for_mpnn) + pdb_path = pdb_for_mpnn + if fixed_residues_file: with open(fixed_residues_file) as f: fixed_residues = f.read().strip() else: fixed_residues = "" - return ( + cmd = ( f"bash {self.scripts_path}/mpnn.sh" f" {self.mpnn_dir}" f" {pdb_path}" @@ -315,6 +355,17 @@ async def mpnn( f" {batch_size}" f' "{fixed_residues}"' ) + log_file = f"{taskdir}/mpnn.log" + with open(log_file, "wb") as _lf: + proc = await asyncio.create_subprocess_shell( + cmd, + stdout=_lf, + stderr=asyncio.subprocess.STDOUT, + env=self._gpu_env(), + ) + await proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"mpnn failed with exit code {proc.returncode}\nSee {log_file}") @self.auto_register_task(local_task=True) async def analysis_sequence(): @@ -366,7 +417,7 @@ async def analysis_sequence(): self.state.get('best_backbone_path'), self.state.get('last_seq_fasta'), )) - @self.auto_register_task(capture_stdio=True) + @self.auto_register_task(local_task=True) async def packmin(): self.taskcount += 1 taskname = "packmin" @@ -383,12 +434,20 @@ async def packmin(): # Predict output path so the next mpnn can read from it self.state['best_packed_pdb'] = f"{output_dir}/{pdb_stem}_minimized.pdb" - return ( + cmd = ( f"bash {self.scripts_path}/packmin.sh" f" {pdb_path}" f" {lig_path}" f" {output_dir}" ) + log_file = f"{taskdir}/packmin.log" + with open(log_file, "wb") as _lf: + proc = await asyncio.create_subprocess_shell( + cmd, stdout=_lf, stderr=asyncio.subprocess.STDOUT, env=self._gpu_env(), + ) + await proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"packmin failed with exit code {proc.returncode}\nSee {log_file}") @self.auto_register_task(local_task=True) async def analysis_packmin(): @@ -403,7 +462,7 @@ async def analysis_packmin(): self.state['last_analysis_step'] = 'packmin' self.state['last_analysis_metrics'] = {'pass': True, 'total_score': total_score} - @self.auto_register_task(capture_stdio=True) + @self.auto_register_task(local_task=True) async def fastrelax(): self.taskcount += 1 taskname = "fastrelax" @@ -416,13 +475,20 @@ async def fastrelax(): lig_path = f"{self.pipeline_inputs}/{self.ligand_params}" output_dir = f"{taskdir}/out" - return ( + cmd = ( f"bash {self.scripts_path}/fastrelax.sh" f" {pdb_path}" f" {lig_path}" f" {output_dir}" - f" > fastrelax.txt" ) + log_file = f"{taskdir}/fastrelax.log" + with open(log_file, "wb") as _lf: + proc = await asyncio.create_subprocess_shell( + cmd, stdout=_lf, stderr=asyncio.subprocess.STDOUT, env=self._gpu_env(), + ) + await proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"fastrelax failed with exit code {proc.returncode}\nSee {log_file}") @self.auto_register_task(local_task=True) async def analysis_fastrelax(): @@ -452,7 +518,7 @@ async def analysis_fastrelax(): 'rmsd': rmsd, } - @self.auto_register_task(capture_stdio=True) + @self.auto_register_task(local_task=True) async def filter_shape(ligand_name: str = "ALR"): taskname = "filter_shape" taskdir = f"{self.base_path}/{self.name}/{self.taskcount}_{taskname}" @@ -461,13 +527,21 @@ async def filter_shape(ligand_name: str = "ALR"): pdb_directory = f"{self.base_path}/{self.name}/{self.taskcount}_fastrelax/out" - return ( + cmd = ( f"bash {self.scripts_path}/filter_shape.sh" f" {pdb_directory}" f" {taskdir}/out/shape_complementarity_values.txt" f" {self.pipeline_inputs}/{ligand_name}" f" {taskdir}/out/interface_values.txt" ) + log_file = f"{taskdir}/filter_shape.log" + with open(log_file, "wb") as _lf: + proc = await asyncio.create_subprocess_shell( + cmd, stdout=_lf, stderr=asyncio.subprocess.STDOUT, env=self._gpu_env(), + ) + await proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"filter_shape failed with exit code {proc.returncode}\nSee {log_file}") @self.auto_register_task(local_task=True) async def analysis_interface(): @@ -494,8 +568,8 @@ async def analysis_interface(): 'max_sc': max_sc, } - @self.auto_register_task(capture_stdio=True) - async def af2(task_description={"gpus_per_rank": 1}): + @self.auto_register_task(local_task=True) + async def af2(): self.taskcount += 1 taskname = "alphafold" self.previous_task = taskname @@ -516,12 +590,20 @@ async def af2(task_description={"gpus_per_rank": 1}): output_dir = f"{taskdir}/out" - return ( + cmd = ( f"bash {self.scripts_path}/af2.sh" f" {self.colabfold_path}" f" {short_fasta}" f" {output_dir}" ) + log_file = f"{taskdir}/af2.log" + with open(log_file, "wb") as _lf: + proc = await asyncio.create_subprocess_shell( + cmd, stdout=_lf, stderr=asyncio.subprocess.STDOUT, env=self._gpu_env(), + ) + await proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"af2 failed with exit code {proc.returncode}\nSee {log_file}") @self.auto_register_task(local_task=True) async def analysis_fold(): @@ -531,6 +613,12 @@ async def analysis_fold(): if 'scores' in f and f.endswith('.json') ] + if not score_files: + raise RuntimeError( + f"af2 produced no score files in {out_dir} — " + "GPU/cuDNN failure (Foundry container may still hold GPU memory)" + ) + best_plddt = -1.0 best_model = None for sf in score_files: @@ -556,7 +644,7 @@ async def analysis_fold(): 'best_model': best_model, } - @self.auto_register_task(capture_stdio=True) + @self.auto_register_task(local_task=True) async def filter_energy(ligand_name: str = "ALR"): taskname = "filter_energy" taskdir = f"{self.base_path}/{self.name}/{self.taskcount}_{taskname}" @@ -569,7 +657,7 @@ async def filter_energy(ligand_name: str = "ALR"): output_energy_file = f"{outputs_dir}/negative_ligand_energies.txt" common_filenames_file = f"{self.pipeline_inputs}/common_filenames.txt" - return ( + cmd = ( f"bash {self.scripts_path}/filter_energy.sh" f" {pdb_directory}" f" {output_file}" @@ -577,6 +665,14 @@ async def filter_energy(ligand_name: str = "ALR"): f" {common_filenames_file}" f" {ligand_name}" ) + log_file = f"{taskdir}/filter_energy.log" + with open(log_file, "wb") as _lf: + proc = await asyncio.create_subprocess_shell( + cmd, stdout=_lf, stderr=asyncio.subprocess.STDOUT, env=self._gpu_env(), + ) + await proc.wait() + if proc.returncode != 0: + raise RuntimeError(f"filter_energy failed with exit code {proc.returncode}\nSee {log_file}") # ── Score utils ──────────────────────────────────────────────────────── @@ -603,7 +699,6 @@ async def _run_refine_cycle(self): self.logger.pipeline_log(f"running mpnn [cycle {cycle_i}]") await self.mpnn() -# fixed_residues_file=f"{self.pipeline_inputs}/fixed_residues.txt" ) self.logger.pipeline_log(f"mpnn [cycle {cycle_i}] finished") await self.analysis_sequence() await self.run_adaptive_step() diff --git a/pyproject.toml b/pyproject.toml index 9daa18b..cef0ce6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ maintainers = [ {name = "Aymen Alsaadi", email = "aymen.alsaadi@rutgers.edu"}, ] readme = "README.md" -requires-python = ">=3.9" +requires-python = ">=3.11" dependencies = [ "radical.pilot", @@ -51,7 +51,7 @@ doc = [ [tool.ruff] line-length = 88 -target-version = "py39" +target-version = "py311" fix = true [tool.ruff.lint] diff --git a/src/impress/__init__.py b/src/impress/__init__.py index 89dd8d3..996526d 100644 --- a/src/impress/__init__.py +++ b/src/impress/__init__.py @@ -1,10 +1,14 @@ from __future__ import annotations +from impress.gpu import GPUPolicy, _find_gpus, _make_policy from impress.impress_manager import ImpressManager from impress.pipelines.impress_pipeline import ImpressBasePipeline from impress.pipelines.setup import PipelineSetup __all__ = [ + "GPUPolicy", + "_find_gpus", + "_make_policy", "ImpressManager", "ImpressBasePipeline", "PipelineSetup", diff --git a/src/impress/gpu.py b/src/impress/gpu.py new file mode 100644 index 0000000..3189590 --- /dev/null +++ b/src/impress/gpu.py @@ -0,0 +1,26 @@ +import os +from dataclasses import dataclass, field + + +@dataclass +class GPUPolicy: + gpu_affinity: list = field(default_factory=list) + + +def _find_gpus() -> list: + cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if cuda_visible: + return [int(g) for g in cuda_visible.split(",") if g.strip().isdigit()] + try: + from dragon.native.machine import System + + sys_info = System() + return [gpu for node in sys_info.nodes for gpu in node.gpus] + except Exception: + pass + return [0, 1, 2, 3] # gpuA40x4 default + + +def _make_policy(all_gpus: list, idx: int, n_gpus: int = 1) -> GPUPolicy: + assigned = [all_gpus[(idx + j) % len(all_gpus)] for j in range(n_gpus)] + return GPUPolicy(gpu_affinity=assigned) diff --git a/src/impress/impress_manager.py b/src/impress/impress_manager.py index 5fac46e..b9446f7 100644 --- a/src/impress/impress_manager.py +++ b/src/impress/impress_manager.py @@ -1,6 +1,8 @@ import asyncio -from collections.abc import Awaitable -from typing import Any, Callable, Optional, Union +import os +import tempfile +from collections.abc import Awaitable, Callable +from typing import Any, Optional, Union from radical.asyncflow import WorkflowEngine @@ -43,6 +45,7 @@ def __init__( self._telemetry_config: dict[str, Any] = telemetry_config or {} self._telemetry_subscribers: list[Callable] = telemetry_subscribers or [] self.telemetry: Any = None + self.flow: Optional[WorkflowEngine] = None def _normalize_pipeline_setup( self, setup: Union[dict[str, Any], PipelineSetup] @@ -77,6 +80,10 @@ def submit_new_pipelines( ValueError: If pipeline type is not a subclass of ImpressBasePipeline """ + if self.flow is None: + raise RuntimeError( + "ImpressManager.start() must be called before submit_new_pipelines()" + ) for setup_input in pipeline_setups: # Normalize to PipelineSetup object setup = self._normalize_pipeline_setup(setup_input) @@ -136,113 +143,139 @@ async def start( """ self.logger.separator("IMPRESS MANAGER STARTING") - self.flow: WorkflowEngine = await WorkflowEngine.create( - backend=self.execution_backend + # Write asyncflow session dirs to /tmp (node-local, no quota) instead of + # cwd on scratch, which exhausts inodes over many runs. + _session_base = os.environ.get("IMPRESS_SESSION_DIR", tempfile.gettempdir()) + self.flow = await WorkflowEngine.create( + backend=self.execution_backend, + work_dir=_session_base, ) - if self._telemetry_config: - self.telemetry = await self.flow.start_telemetry(**self._telemetry_config) - for fn in self._telemetry_subscribers: - self.telemetry.subscribe(fn) - - self.logger.manager_starting(len(pipeline_setups)) - - self.submit_new_pipelines(pipeline_setups) - - while True: - any_activity: bool = False - completed_pipelines: list[ImpressBasePipeline] = [] - - for pipeline, pipeline_future in list(self.pipeline_tasks.items()): - # Check if pipeline needs adaptive step and isn't already running one - if ( - getattr(pipeline, "invoke_adaptive_step", False) - and pipeline not in self.adaptive_tasks - ): - adaptive_task: asyncio.Task = asyncio.create_task( - self._run_adaptive_fn(pipeline) + try: + if self._telemetry_config: + self.telemetry = await self.flow.start_telemetry( + **self._telemetry_config + ) + for fn in self._telemetry_subscribers: + self.telemetry.subscribe(fn) + + self.logger.manager_starting(len(pipeline_setups)) + + self.submit_new_pipelines(pipeline_setups) + + while True: + any_activity: bool = False + completed_pipelines: list[tuple] = [] + + for pipeline, pipeline_future in list(self.pipeline_tasks.items()): + # Check if pipeline needs adaptive step and isn't running one yet + if ( + getattr(pipeline, "invoke_adaptive_step", False) + and pipeline not in self.adaptive_tasks + ): + adaptive_task: asyncio.Task = asyncio.create_task( + self._run_adaptive_fn(pipeline) + ) + self.adaptive_tasks[pipeline] = adaptive_task + any_activity = True + + # Check if pipeline has new config ready + config: Optional[dict[str, Any]] = ( + pipeline.get_child_pipeline_request() ) - self.adaptive_tasks[pipeline] = adaptive_task - any_activity = True - # Check if pipeline has new config ready - config: Optional[dict[str, Any]] = pipeline.get_child_pipeline_request() - - if config: - self.logger.child_pipeline_submitted(config["name"], pipeline.name) - # Convert dict to PipelineSetup for consistency - child_setup = PipelineSetup.from_dict(config) - self.new_pipeline_buffer.append(child_setup) - any_activity = True + if config: + self.logger.child_pipeline_submitted( + config["name"], pipeline.name + ) + # Convert dict to PipelineSetup for consistency + child_setup = PipelineSetup.from_dict(config) + self.new_pipeline_buffer.append(child_setup) + any_activity = True + + # Check if parent should be killed + if getattr(pipeline, "kill_parent", False): + self.logger.pipeline_killed(pipeline.name) + pipeline_future.cancel() + completed_pipelines.append((pipeline, pipeline_future)) + continue - # Check if parent should be killed - if getattr(pipeline, "kill_parent", False): - self.logger.pipeline_killed(pipeline.name) - pipeline_future.cancel() - completed_pipelines.append(pipeline) - continue - - # Check if pipeline is done - but only mark as completed - # if adaptive task is also done - if pipeline_future.done(): - # If there's an adaptive task running, don't mark as completed yet + # Check if pipeline is done - but only mark as completed + # if adaptive task is also done + if pipeline_future.done(): + # Adaptive task still running — wait before marking completed + if pipeline in self.adaptive_tasks: + adaptive_task = self.adaptive_tasks[pipeline] + if not adaptive_task.done(): + continue + + completed_pipelines.append((pipeline, pipeline_future)) + + # Clean up completed pipelines - but only if their + # adaptive tasks are also done + actually_completed: list[ImpressBasePipeline] = [] + for pipeline, future in completed_pipelines: + # Double-check: only clean up if adaptive task is + # done or doesn't exist if pipeline in self.adaptive_tasks: adaptive_task = self.adaptive_tasks[pipeline] if not adaptive_task.done(): continue + self.adaptive_tasks.pop(pipeline) + + self.pipeline_tasks.pop(pipeline, None) + exc = None + if future.done() and not future.cancelled(): + try: + exc = future.exception() + except Exception: + pass + if exc is not None: + self.logger.pipeline_failed(pipeline.name, exc) + else: + self.logger.pipeline_completed(pipeline.name) + actually_completed.append(pipeline) + + completed_pipelines = actually_completed + + # Clean up completed adaptive tasks + completed_adaptive: list[ImpressBasePipeline] = [] + for pipeline, adaptive_task in list(self.adaptive_tasks.items()): + if adaptive_task.done(): + completed_adaptive.append(pipeline) + + for pipeline in completed_adaptive: + self.adaptive_tasks.pop(pipeline, None) + + # Submit new pipelines; capture count before clearing so + # activity_summary reports the real number submitted. + if self.new_pipeline_buffer: + buffered_count = len(self.new_pipeline_buffer) + self.submit_new_pipelines(self.new_pipeline_buffer) + self.new_pipeline_buffer.clear() + any_activity = True + else: + buffered_count = 0 + + # Log activity summary periodically + if any_activity: + self.logger.activity_summary( + len(self.pipeline_tasks), + len(self.adaptive_tasks), + buffered_count, + ) - completed_pipelines.append(pipeline) - - # Clean up completed pipelines - but only if their - # adaptive tasks are also done - actually_completed: list[ImpressBasePipeline] = [] - for pipeline in completed_pipelines: - # Double-check: only clean up if adaptive task is - # done or doesn't exist - if pipeline in self.adaptive_tasks: - adaptive_task = self.adaptive_tasks[pipeline] - if not adaptive_task.done(): - continue - self.adaptive_tasks.pop(pipeline) - - self.pipeline_tasks.pop(pipeline, None) - self.logger.pipeline_completed(pipeline.name) - actually_completed.append(pipeline) - - completed_pipelines = actually_completed - - # Clean up completed adaptive tasks - completed_adaptive: list[ImpressBasePipeline] = [] - for pipeline, adaptive_task in list(self.adaptive_tasks.items()): - if adaptive_task.done(): - completed_adaptive.append(pipeline) - - for pipeline in completed_adaptive: - self.adaptive_tasks.pop(pipeline, None) - - # Submit new pipelines - if self.new_pipeline_buffer: - self.submit_new_pipelines(self.new_pipeline_buffer) - self.new_pipeline_buffer.clear() - any_activity = True - - # Log activity summary periodically - if any_activity: - self.logger.activity_summary( - len(self.pipeline_tasks), - len(self.adaptive_tasks), - len(self.new_pipeline_buffer), - ) + # Exit condition + if ( + not self.pipeline_tasks + and not self.new_pipeline_buffer + and not self.adaptive_tasks + ): + self.logger.manager_exiting() + self.logger.separator("IMPRESS MANAGER FINISHED") + break - # Exit condition - if ( - not self.pipeline_tasks - and not self.new_pipeline_buffer - and not self.adaptive_tasks - ): - self.logger.manager_exiting() - self.logger.separator("IMPRESS MANAGER FINISHED") - break - - if not any_activity: - await asyncio.sleep(0.5) + if not any_activity: + await asyncio.sleep(0.5) + finally: + await self.flow.shutdown() diff --git a/src/impress/pipelines/impress_pipeline.py b/src/impress/pipelines/impress_pipeline.py index f856fab..f828250 100644 --- a/src/impress/pipelines/impress_pipeline.py +++ b/src/impress/pipelines/impress_pipeline.py @@ -1,4 +1,5 @@ import asyncio +import os from abc import ABC, abstractmethod from typing import Any @@ -90,13 +91,19 @@ def register_pipeline_tasks(self): """Register pipeline tasks - must be implemented by subclasses""" pass + def _gpu_env(self) -> dict: + env = {**os.environ} + policy = getattr(self, "policy", None) + if policy and getattr(policy, "gpu_affinity", None): + env["CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in policy.gpu_affinity) + return env + # Optional methods that subclasses can override async def get_scores_map(self): """Optional: Return scores mapping""" return {} - @abstractmethod - async def finalize(self): + async def finalize(self): # noqa: B027 """Optional: Cleanup or finalization logic""" pass diff --git a/src/impress/pipelines/setup.py b/src/impress/pipelines/setup.py index e0b63e0..4e10036 100644 --- a/src/impress/pipelines/setup.py +++ b/src/impress/pipelines/setup.py @@ -1,5 +1,5 @@ -from collections.abc import Awaitable -from typing import Annotated, Any, Callable, Optional +from collections.abc import Awaitable, Callable +from typing import Annotated, Any, Optional from pydantic import BaseModel, Field, field_validator diff --git a/src/impress/utils/logger.py b/src/impress/utils/logger.py index 642d699..2919c8f 100644 --- a/src/impress/utils/logger.py +++ b/src/impress/utils/logger.py @@ -34,10 +34,25 @@ class LogLevel(Enum): class ImpressLogger: - def __init__(self, name="ImpressManager", use_colors=True, output_stream=None): + _LEVEL_ORDER = [ + LogLevel.DEBUG, + LogLevel.INFO, + LogLevel.WARNING, + LogLevel.ERROR, + LogLevel.CRITICAL, + ] + + def __init__( + self, + name="ImpressManager", + use_colors=True, + output_stream=None, + min_level: LogLevel = LogLevel.DEBUG, + ): self.name = name self.use_colors = use_colors self.output_stream = output_stream or sys.stdout + self.min_level = min_level self.level_colors = { LogLevel.DEBUG: Colors.BRIGHT_BLACK, @@ -91,40 +106,52 @@ def _format_message(self, level, component, message, pipeline_name=None): f"{timestamp} {colored_level} {colored_component}{pipeline_part} {message}" ) - def _write_log(self, message, to_stderr=False): - stream = sys.stderr if to_stderr else self.output_stream - stream.write(message + "\n") - stream.flush() + def _is_enabled(self, level: LogLevel) -> bool: + return self._LEVEL_ORDER.index(level) >= self._LEVEL_ORDER.index(self.min_level) + + def _write_log(self, message): + self.output_stream.write(message + "\n") + self.output_stream.flush() def debug(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.DEBUG): + return formatted = self._format_message( LogLevel.DEBUG, component, message, pipeline_name ) self._write_log(formatted) def info(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.INFO): + return formatted = self._format_message( LogLevel.INFO, component, message, pipeline_name ) self._write_log(formatted) def warning(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.WARNING): + return formatted = self._format_message( LogLevel.WARNING, component, message, pipeline_name ) self._write_log(formatted) def error(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.ERROR): + return formatted = self._format_message( LogLevel.ERROR, component, message, pipeline_name ) - self._write_log(formatted, to_stderr=True) + self._write_log(formatted) def critical(self, message, component="manager", pipeline_name=None): + if not self._is_enabled(LogLevel.CRITICAL): + return formatted = self._format_message( LogLevel.CRITICAL, component, message, pipeline_name ) - self._write_log(formatted, to_stderr=True) + self._write_log(formatted) def pipeline_started(self, pipeline_name): colored_name = self._colorize(pipeline_name, Colors.BRIGHT_WHITE) @@ -136,6 +163,11 @@ def pipeline_completed(self, pipeline_name): message = f"Pipeline completed: {colored_name}" self.info(message, "manager") + def pipeline_failed(self, pipeline_name, exc): + colored_name = self._colorize(pipeline_name, Colors.BRIGHT_WHITE) + message = f"Pipeline FAILED: {colored_name} — {exc}" + self.error(message, "manager") + def pipeline_killed(self, pipeline_name): colored_name = self._colorize(pipeline_name, Colors.BRIGHT_WHITE) message = f"Pipeline killed: {colored_name}" @@ -182,10 +214,11 @@ def activity_summary(self, active_pipelines, active_adaptive, buffered_pipelines self.debug(summary, "manager") def pipeline_log(self, message, level=LogLevel.INFO): + if not self._is_enabled(level): + return pipeline_component = f"PIPELINE-{self.name.upper()}" formatted = self._format_message(level, pipeline_component, message) - stderr_levels = [LogLevel.ERROR, LogLevel.CRITICAL] - self._write_log(formatted, to_stderr=level in stderr_levels) + self._write_log(formatted) def separator(self, title=None): if title: diff --git a/tests/conftest.py b/tests/conftest.py index c35a0f5..cb6ad28 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,9 +1,10 @@ # tests/conftest.py -import pytest import shutil from pathlib import Path - from unittest.mock import Mock + +import pytest + from impress import ImpressManager diff --git a/tests/unit/test_logger.py b/tests/unit/test_logger.py new file mode 100644 index 0000000..930d5b1 --- /dev/null +++ b/tests/unit/test_logger.py @@ -0,0 +1,278 @@ +import io +import sys + +from impress.utils.logger import ImpressLogger, LogLevel + + +class TestImpressLoggerInit: + def test_default_init(self): + logger = ImpressLogger() + assert logger.name == "ImpressManager" + assert logger.use_colors is True + assert logger.output_stream is sys.stdout + assert logger.min_level == LogLevel.DEBUG + + def test_custom_stream(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream) + assert logger.output_stream is stream + + def test_custom_min_level(self): + logger = ImpressLogger(min_level=LogLevel.WARNING) + assert logger.min_level == LogLevel.WARNING + + def test_custom_name(self): + logger = ImpressLogger(name="my_pipeline") + assert logger.name == "my_pipeline" + + +class TestLogLevelFiltering: + def test_is_enabled_at_exact_level(self): + logger = ImpressLogger(min_level=LogLevel.INFO) + assert logger._is_enabled(LogLevel.INFO) is True + + def test_is_enabled_above_min(self): + logger = ImpressLogger(min_level=LogLevel.INFO) + assert logger._is_enabled(LogLevel.WARNING) is True + assert logger._is_enabled(LogLevel.ERROR) is True + assert logger._is_enabled(LogLevel.CRITICAL) is True + + def test_is_disabled_below_min(self): + logger = ImpressLogger(min_level=LogLevel.INFO) + assert logger._is_enabled(LogLevel.DEBUG) is False + + def test_debug_suppressed_at_info_level(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.INFO, use_colors=False + ) + logger.debug("should not appear") + assert stream.getvalue() == "" + + def test_info_written_at_info_level(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.INFO, use_colors=False + ) + logger.info("should appear") + assert "should appear" in stream.getvalue() + + def test_warning_suppressed_below_min(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.ERROR, use_colors=False + ) + logger.warning("should not appear") + assert stream.getvalue() == "" + + def test_all_levels_write_at_debug_min(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.DEBUG, use_colors=False + ) + logger.debug("d") + logger.info("i") + logger.warning("w") + logger.error("e") + logger.critical("c") + output = stream.getvalue() + assert "d" in output + assert "i" in output + assert "w" in output + assert "e" in output + assert "c" in output + + def test_critical_only_at_critical_min(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, min_level=LogLevel.CRITICAL, use_colors=False + ) + logger.debug("d") + logger.info("i") + logger.warning("w") + logger.error("e") + logger.critical("c") + output = stream.getvalue() + assert "d" not in output + assert "i" not in output + assert "w" not in output + assert "e" not in output + assert "c" in output + + +class TestOutputStreamRespected: + def test_error_writes_to_output_stream_not_stderr(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.error("an error") + assert "an error" in stream.getvalue() + + def test_critical_writes_to_output_stream(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.critical("critical msg") + assert "critical msg" in stream.getvalue() + + def test_custom_stream_receives_all_output(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.info("info line") + logger.error("error line") + logger.warning("warn line") + output = stream.getvalue() + assert "info line" in output + assert "error line" in output + assert "warn line" in output + + +class TestColors: + def test_no_ansi_when_colors_off(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.info("hello") + assert "\033[" not in stream.getvalue() + + def test_ansi_present_when_colors_on(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=True) + logger.info("hello") + assert "\033[" in stream.getvalue() + + def test_colorize_returns_plain_when_off(self): + logger = ImpressLogger(use_colors=False) + result = logger._colorize("text", "\033[31m") + assert result == "text" + + def test_colorize_wraps_when_on(self): + logger = ImpressLogger(use_colors=True) + result = logger._colorize("text", "\033[31m") + assert result.startswith("\033[31m") + assert "text" in result + + +class TestHighLevelMethods: + def _make_logger(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + return logger, stream + + def test_pipeline_started(self): + logger, stream = self._make_logger() + logger.pipeline_started("my_pipe") + assert "my_pipe" in stream.getvalue() + + def test_pipeline_completed(self): + logger, stream = self._make_logger() + logger.pipeline_completed("my_pipe") + assert "my_pipe" in stream.getvalue() + + def test_pipeline_failed(self): + logger, stream = self._make_logger() + logger.pipeline_failed("bad_pipe", ValueError("oops")) + out = stream.getvalue() + assert "bad_pipe" in out + assert "oops" in out + + def test_pipeline_killed(self): + logger, stream = self._make_logger() + logger.pipeline_killed("dead_pipe") + assert "dead_pipe" in stream.getvalue() + + def test_adaptive_started(self): + logger, stream = self._make_logger() + logger.adaptive_started("p1") + assert "p1" in stream.getvalue() + + def test_adaptive_completed(self): + logger, stream = self._make_logger() + logger.adaptive_completed("p1") + assert "p1" in stream.getvalue() + + def test_adaptive_failed(self): + logger, stream = self._make_logger() + logger.adaptive_failed("p1", "bad fn") + out = stream.getvalue() + assert "p1" in out + assert "bad fn" in out + + def test_child_pipeline_submitted(self): + logger, stream = self._make_logger() + logger.child_pipeline_submitted("child", "parent") + out = stream.getvalue() + assert "child" in out + assert "parent" in out + + def test_manager_starting(self): + logger, stream = self._make_logger() + logger.manager_starting(5) + assert "5" in stream.getvalue() + + def test_manager_exiting(self): + logger, stream = self._make_logger() + logger.manager_exiting() + assert stream.getvalue() != "" + + def test_activity_summary_suppressed_at_info(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, use_colors=False, min_level=LogLevel.INFO + ) + logger.activity_summary(3, 1, 2) + assert stream.getvalue() == "" # activity_summary is DEBUG level + + def test_activity_summary_written_at_debug(self): + stream = io.StringIO() + logger = ImpressLogger( + output_stream=stream, use_colors=False, min_level=LogLevel.DEBUG + ) + logger.activity_summary(3, 1, 2) + out = stream.getvalue() + assert "3" in out + + +class TestPipelineLog: + def test_pipeline_log_default_info(self): + stream = io.StringIO() + logger = ImpressLogger("pipe1", output_stream=stream, use_colors=False) + logger.pipeline_log("step done") + assert "step done" in stream.getvalue() + + def test_pipeline_log_suppressed_below_min(self): + stream = io.StringIO() + logger = ImpressLogger( + "pipe1", output_stream=stream, use_colors=False, min_level=LogLevel.WARNING + ) + logger.pipeline_log("step done", level=LogLevel.INFO) + assert stream.getvalue() == "" + + def test_pipeline_log_debug_level(self): + stream = io.StringIO() + logger = ImpressLogger("pipe1", output_stream=stream, use_colors=False) + logger.pipeline_log("debug step", level=LogLevel.DEBUG) + assert "debug step" in stream.getvalue() + + def test_pipeline_log_includes_pipeline_name_component(self): + stream = io.StringIO() + logger = ImpressLogger("mypipe", output_stream=stream, use_colors=False) + logger.pipeline_log("event") + assert "PIPELINE-MYPIPE" in stream.getvalue() + + +class TestSeparator: + def test_separator_no_title(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.separator() + assert "=" in stream.getvalue() + + def test_separator_with_title(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.separator("HELLO WORLD") + assert "HELLO WORLD" in stream.getvalue() + + def test_separator_ends_with_newline(self): + stream = io.StringIO() + logger = ImpressLogger(output_stream=stream, use_colors=False) + logger.separator() + assert stream.getvalue().endswith("\n") diff --git a/tests/unit/test_manager_core.py b/tests/unit/test_manager_core.py index d35495c..883e912 100644 --- a/tests/unit/test_manager_core.py +++ b/tests/unit/test_manager_core.py @@ -3,8 +3,7 @@ import pytest # Import the classes we're testing -from impress import ImpressBasePipeline, PipelineSetup -from impress import ImpressManager +from impress import ImpressBasePipeline, ImpressManager, PipelineSetup class MockPipeline(ImpressBasePipeline): diff --git a/tests/unit/test_manager_life_cycle.py b/tests/unit/test_manager_life_cycle.py index 9783467..bd110b3 100644 --- a/tests/unit/test_manager_life_cycle.py +++ b/tests/unit/test_manager_life_cycle.py @@ -12,9 +12,12 @@ class MockWorkflowEngine: """Mock workflow engine""" @classmethod - async def create(cls, backend=None): + async def create(cls, backend=None, **kwargs): return cls() + async def shutdown(self): + pass + class TestManagerLifecycle: @pytest.mark.asyncio @@ -113,3 +116,18 @@ async def run(self): # Should take at least 0.15 seconds due to slow adaptive function assert end_time - start_time >= 0.15 + + @pytest.mark.asyncio + @patch("impress.impress_manager.WorkflowEngine", MockWorkflowEngine) + async def test_start_exception_still_shuts_down_engine(self, impress_manager): + """WorkflowEngine.shutdown() is called even when a pipeline raises mid-run""" + + class ExplodingPipeline(MockPipeline): + async def run(self): + raise RuntimeError("pipeline exploded") + + await impress_manager.start( + [{"name": "boom", "type": ExplodingPipeline, "config": {}, "kwargs": {}}] + ) + # If shutdown() raises AttributeError, the try/finally fix is broken + assert len(impress_manager.pipeline_tasks) == 0 diff --git a/tests/unit/test_manager_pipeline_submission.py b/tests/unit/test_manager_pipeline_submission.py index e390b93..92bbdfa 100644 --- a/tests/unit/test_manager_pipeline_submission.py +++ b/tests/unit/test_manager_pipeline_submission.py @@ -1,5 +1,7 @@ from unittest.mock import Mock, patch +import pytest + # Import the classes we're testing from impress import PipelineSetup @@ -64,3 +66,10 @@ def test_submit_multiple_pipelines(self, mock_create_task, impress_manager): assert len(impress_manager.pipeline_tasks) == 2 assert mock_create_task.call_count == 2 + + def test_submit_before_start_raises(self, impress_manager): + """submit_new_pipelines raises RuntimeError when called before start()""" + with pytest.raises(RuntimeError, match="start\\(\\) must be called"): + impress_manager.submit_new_pipelines( + [{"name": "p", "type": MockPipeline, "config": {}, "kwargs": {}}] + ) diff --git a/tests/unit/test_pipeline_base.py b/tests/unit/test_pipeline_base.py new file mode 100644 index 0000000..928ea2d --- /dev/null +++ b/tests/unit/test_pipeline_base.py @@ -0,0 +1,157 @@ +import asyncio + +import pytest + +from impress.pipelines.impress_pipeline import ImpressBasePipeline + + +class MinimalPipeline(ImpressBasePipeline): + """Minimal concrete subclass for testing the base class.""" + + async def run(self): + pass + + def register_pipeline_tasks(self): + pass + + +class TestImpressBasePipelineInit: + def test_name_set(self): + p = MinimalPipeline(name="p1") + assert p.name == "p1" + + def test_flow_defaults_to_none(self): + p = MinimalPipeline(name="p1") + assert p.flow is None + + def test_flow_passed_through(self): + sentinel = object() + p = MinimalPipeline(name="p1", flow=sentinel) + assert p.flow is sentinel + + def test_state_is_empty_dict(self): + p = MinimalPipeline(name="p1") + assert p.state == {} + + def test_kill_parent_false(self): + p = MinimalPipeline(name="p1") + assert p.kill_parent is False + + def test_invoke_adaptive_step_false(self): + p = MinimalPipeline(name="p1") + assert p.invoke_adaptive_step is False + + def test_adaptive_barrier_is_event(self): + p = MinimalPipeline(name="p1") + assert isinstance(p._adaptive_barrier, asyncio.Event) + + def test_incoming_child_pipeline_request_empty(self): + p = MinimalPipeline(name="p1") + assert not p.incoming_child_pipeline_request + + def test_kwargs_stored_in_config(self): + p = MinimalPipeline(name="p1", foo="bar", baz=42) + assert p.config["foo"] == "bar" + assert p.config["baz"] == 42 + + +class TestChildPipelineRequest: + def test_submit_sets_request(self): + p = MinimalPipeline(name="p1") + config = {"name": "child", "type": MinimalPipeline} + p.submit_child_pipeline_request(config) + assert p.incoming_child_pipeline_request == config + + def test_get_returns_config_then_clears(self): + p = MinimalPipeline(name="p1") + config = {"name": "child", "type": MinimalPipeline} + p.submit_child_pipeline_request(config) + + result1 = p.get_child_pipeline_request() + assert result1 == config + + result2 = p.get_child_pipeline_request() + assert result2 is None + + def test_get_returns_none_when_nothing_pending(self): + p = MinimalPipeline(name="p1") + assert p.get_child_pipeline_request() is None + + def test_get_clears_after_retrieval(self): + p = MinimalPipeline(name="p1") + p.submit_child_pipeline_request({"name": "c"}) + p.get_child_pipeline_request() + assert not p.incoming_child_pipeline_request + + +class TestAdaptiveStep: + @pytest.mark.asyncio + async def test_run_adaptive_step_sets_flag(self): + p = MinimalPipeline(name="p1") + # wait=False: flag is set without blocking on the barrier + await p.run_adaptive_step(wait=False) + assert p.invoke_adaptive_step is True + + @pytest.mark.asyncio + async def test_run_adaptive_step_no_wait_does_not_hang(self): + p = MinimalPipeline(name="p1") + # _adaptive_barrier is clear — with wait=False this must return immediately + await asyncio.wait_for(p.run_adaptive_step(wait=False), timeout=1.0) + assert p.invoke_adaptive_step is True + + def test_set_adaptive_flag_true_clears_barrier(self): + p = MinimalPipeline(name="p1") + p._adaptive_barrier.set() + p._set_adaptive_flag(True) + assert p.invoke_adaptive_step is True + assert not p._adaptive_barrier.is_set() + + def test_set_adaptive_flag_false_does_not_touch_barrier(self): + p = MinimalPipeline(name="p1") + p._adaptive_barrier.set() + p._set_adaptive_flag(False) + assert p.invoke_adaptive_step is False + assert p._adaptive_barrier.is_set() # barrier unchanged + + +class TestOptionalMethods: + @pytest.mark.asyncio + async def test_finalize_is_noop(self): + p = MinimalPipeline(name="p1") + result = await p.finalize() + assert result is None + + @pytest.mark.asyncio + async def test_get_scores_map_returns_empty_dict(self): + p = MinimalPipeline(name="p1") + scores = await p.get_scores_map() + assert scores == {} + + def test_get_current_config_has_name_and_type(self): + p = MinimalPipeline(name="p1") + cfg = p.get_current_config_for_next_pipeline() + assert "name" in cfg + assert "type" in cfg + + def test_get_current_config_type_is_class(self): + p = MinimalPipeline(name="p1") + cfg = p.get_current_config_for_next_pipeline() + assert cfg["type"] is MinimalPipeline + + +class TestAbstractMethods: + def test_cannot_instantiate_without_run(self): + class NoRun(ImpressBasePipeline): + def register_pipeline_tasks(self): + pass + + with pytest.raises(TypeError): + NoRun(name="x") + + def test_cannot_instantiate_without_register(self): + class NoRegister(ImpressBasePipeline): + async def run(self): + pass + + with pytest.raises(TypeError): + NoRegister(name="x") diff --git a/tests/unit/test_pipeline_management.py b/tests/unit/test_pipeline_management.py index af35ae9..593327f 100644 --- a/tests/unit/test_pipeline_management.py +++ b/tests/unit/test_pipeline_management.py @@ -82,7 +82,7 @@ async def run(self): # Run the manager with timeout try: await asyncio.wait_for(impress_manager.start([pipeline_setup]), timeout=3.0) - except asyncio.TimeoutError: + except TimeoutError: # Print debug info if it times out print(f"Completed pipelines: {completed_pipelines}") print(f"Pipeline tasks: {len(impress_manager.pipeline_tasks)}") diff --git a/tests/unit/test_pipeline_setup.py b/tests/unit/test_pipeline_setup.py new file mode 100644 index 0000000..c56d398 --- /dev/null +++ b/tests/unit/test_pipeline_setup.py @@ -0,0 +1,141 @@ +import pytest +from pydantic import ValidationError + +from impress import PipelineSetup + +from .test_manager_core import MockPipeline + + +class TestPipelineSetupConstruction: + def test_all_fields(self): + async def fn(p): + pass + + setup = PipelineSetup( + name="my_pipe", + type=MockPipeline, + config={"a": 1}, + kwargs={"b": 2}, + adaptive_fn=fn, + ) + assert setup.name == "my_pipe" + assert setup.type is MockPipeline + assert setup.config == {"a": 1} + assert setup.kwargs == {"b": 2} + assert setup.adaptive_fn is fn + + def test_defaults(self): + setup = PipelineSetup(name="p", type=MockPipeline) + assert setup.config == {} + assert setup.kwargs == {} + assert setup.adaptive_fn is None + + def test_validate_type_rejects_non_subclass(self): + with pytest.raises(ValidationError): + PipelineSetup(name="p", type=str) + + def test_validate_type_rejects_non_type(self): + with pytest.raises(ValidationError): + PipelineSetup(name="p", type="not_a_class") + + def test_validate_type_accepts_subclass(self): + class Sub(MockPipeline): + pass + + setup = PipelineSetup(name="p", type=Sub) + assert setup.type is Sub + + +class TestFromDict: + def test_known_fields_separated(self): + async def fn(p): + pass + + data = { + "name": "p1", + "type": MockPipeline, + "config": {"x": 1}, + "adaptive_fn": fn, + } + setup = PipelineSetup.from_dict(data) + assert setup.name == "p1" + assert setup.type is MockPipeline + assert setup.config == {"x": 1} + assert setup.adaptive_fn is fn + assert setup.kwargs == {} + + def test_extra_keys_go_to_kwargs(self): + data = { + "name": "p1", + "type": MockPipeline, + "foo": "bar", + "baz": 42, + } + setup = PipelineSetup.from_dict(data) + assert setup.kwargs == {"foo": "bar", "baz": 42} + + def test_kwargs_key_in_dict_lands_in_kwargs(self): + # "kwargs" is not a known field, so it ends up nested inside kwargs + data = { + "name": "p1", + "type": MockPipeline, + "kwargs": {"inner": "value"}, + } + setup = PipelineSetup.from_dict(data) + assert setup.kwargs == {"kwargs": {"inner": "value"}} + + def test_minimal_dict(self): + setup = PipelineSetup.from_dict({"name": "p", "type": MockPipeline}) + assert setup.name == "p" + assert setup.config == {} + assert setup.kwargs == {} + assert setup.adaptive_fn is None + + +class TestToDict: + def test_basic_structure(self): + setup = PipelineSetup( + name="p1", + type=MockPipeline, + config={"c": 1}, + ) + d = setup.to_dict() + assert d["name"] == "p1" + assert d["type"] is MockPipeline + assert d["config"] == {"c": 1} + + def test_adaptive_fn_omitted_when_none(self): + setup = PipelineSetup(name="p", type=MockPipeline, adaptive_fn=None) + d = setup.to_dict() + assert "adaptive_fn" not in d + + def test_adaptive_fn_included_when_set(self): + async def fn(p): + pass + + setup = PipelineSetup(name="p", type=MockPipeline, adaptive_fn=fn) + d = setup.to_dict() + assert d["adaptive_fn"] is fn + + def test_kwargs_spread_into_result(self): + setup = PipelineSetup( + name="p", + type=MockPipeline, + kwargs={"foo": "bar", "num": 7}, + ) + d = setup.to_dict() + assert d["foo"] == "bar" + assert d["num"] == 7 + + def test_roundtrip_from_dict(self): + original = { + "name": "rt", + "type": MockPipeline, + "config": {"k": "v"}, + "extra_param": 99, + } + setup = PipelineSetup.from_dict(original) + d = setup.to_dict() + assert d["name"] == "rt" + assert d["type"] is MockPipeline + assert d["extra_param"] == 99 diff --git a/tox.ini b/tox.ini index 11f11ba..fb0ef4d 100644 --- a/tox.ini +++ b/tox.ini @@ -1,5 +1,5 @@ [tox] -envlist = py39,py310,py311,py312,py313 +envlist = py311,py312,py313 isolated_build = true # Unit test env @@ -8,7 +8,7 @@ extras = dev commands = pytest tests/unit {posargs} # Integration test radical.pilot -[testenv:{py39,py310,py311,py312,py313}-all] +[testenv:{py311,py312,py313}-all] extras = dev setenv = RADICAL_VERBOSE=DEBUG