Skip to content

small_molecule_binding: replace AF2 with Boltz-2 co-folding, port to Delta #62

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
drawadiagram wants to merge 1 commit into
base: update/core-manager-0926
Choose a base branch
from
Open

small_molecule_binding: replace AF2 with Boltz-2 co-folding, port to Delta #62

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 44 additions & 19 deletions examples/small_molecule_binding/CLAUDE.md
View file
Open in desktop

Large diffs are not rendered by default.

160 changes: 84 additions & 76 deletions examples/small_molecule_binding/README.md
View file
Open in desktop

Large diffs are not rendered by default.

279 changes: 279 additions & 0 deletions examples/small_molecule_binding/delta_env_setup.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,279 @@
#!/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/<allocation>
# 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 (python/3.11, cray-python/3.11.7, anaconda3)
#
# Tool directories (cloned by this script if absent):
# MPNN_DIR = $SCRATCH/$USER/LigandMPNN
# BOLTZ_CACHE = $SCRATCH/$USER/.cache/boltz (model weights cache)
#
# 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

# ── Require SCRATCH ───────────────────────────────────────────────────────────
if [[ -z "${SCRATCH:-}" ]]; then
echo "ERROR: set the SCRATCH env var to your allocation scratch root, e.g.:"
echo " export SCRATCH=/scratch/<allocation>"
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"

MPNN_DIR="${MPNN_DIR:-${SCRATCH}/${USER}/LigandMPNN}"
BOLTZ_CACHE="${BOLTZ_CACHE:-${SCRATCH}/${USER}/.cache/boltz}"

echo "================================================================="
echo " ENV_DIR = ${ENV_DIR}"
echo " IMPRESS_DIR = ${IMPRESS_DIR}"
echo " MPNN_DIR = ${MPNN_DIR}"
echo " BOLTZ_CACHE = ${BOLTZ_CACHE}"
echo "================================================================="

# ── 1. Create venv ────────────────────────────────────────────────────────────
echo ""
echo "── Step 1: Creating venv ──"

_find_python() {
for candidate in python3.12 python3.11 python3 python; do
local p
p=$(command -v "${candidate}" 2>/dev/null) || continue
local ver
ver=$("${p}" -c "import sys; v=sys.version_info; print(v.major*100+v.minor)" 2>/dev/null) || continue
[ "${ver}" -ge 311 ] && echo "${p}" && return 0
done
return 1
}

if [ -n "${BASE_PY_OVERRIDE}" ]; then
BASE_PY="${BASE_PY_OVERRIDE}"
echo "Using Python override: ${BASE_PY}"
else
BASE_PY=$(_find_python || true)
if [ -z "${BASE_PY}" ]; then
echo "python3.11+ not in PATH — trying modules..."
for mod in python/3.13.5-gcc13.3.1 cray-python/3.12.12 anaconda3; do
module load "${mod}" 2>/dev/null || true
BASE_PY=$(_find_python || true)
[ -n "${BASE_PY}" ] && echo " loaded module: ${mod}" && break
done
fi
if [ -z "${BASE_PY}" ]; then
echo "ERROR: no Python 3.11+ interpreter found."
echo " Pass an explicit interpreter: --python /path/to/python3.11"
echo " Or load a module manually before running this script."
exit 1
fi
fi
echo "Using Python: ${BASE_PY} ($(${BASE_PY} --version))"

if [ ! -x "${PY}" ]; then
"${BASE_PY}" -m venv "${ENV_DIR}"
else
echo "venv already exists at ${ENV_DIR}"
fi

echo "Python: $("${PY}" --version)"

# ── 2. Bootstrap pip ──────────────────────────────────────────────────────────
echo ""
echo "── Step 2: Bootstrapping pip ──"
"${PY}" -m pip install -q --upgrade pip wheel
"${PIP}" install -q --force-reinstall "setuptools<71"

# ── 3. radical.asyncflow (PyPI) ──────────────────────────────────────────────
echo ""
echo "── Step 3: radical-asyncflow (PyPI) ──"
"${PIP}" install -q radical-asyncflow

# ── 4. rhapsody-py (PyPI) ────────────────────────────────────────────────────
echo ""
echo "── Step 4: rhapsody-py[dragon] (PyPI) ──"
"${PIP}" install -q "rhapsody-py[dragon,telemetry]"

# ── 5. IMPRESS (local editable) ───────────────────────────────────────────────
echo ""
echo "── Step 5: IMPRESS (editable) ──"
"${PIP}" install -q -e "${IMPRESS_DIR}"

# ── 6. PyTorch (CUDA 12.1) — required by LigandMPNN ─────────────────────────
echo ""
echo "── Step 6: PyTorch (cu121) ──"
"${PIP}" install -q torch --index-url https://download.pytorch.org/whl/cu121

# ── 7. Boltz-2 ────────────────────────────────────────────────────────────────
#
# EMPIRICALLY CONFIRMED: `pip install "boltz[cuda]"` (with or without `-U`)
# thrashes pip's resolver for a very long time (observed: 28GB+ pip cache,
# 60-75+ pip-metadata/pip-unpack temp dirs, no completion after ~1hr each
# attempt) -- boltz pins several dependencies (`numpy<2.0`, `gemmi==0.6.5`,
# `pytorch-lightning==2.5.0`, etc.) that genuinely conflict with what's
# already installed in this venv (numpy 2.x from other packages, gemmi 0.7.5,
# etc. -- see this repo's other steps). Resolving a real, deep conflict like
# this is inherently slow/combinatorial for pip's resolver, `-U` or not.
#
# WORKING APPROACH (installs cleanly in seconds instead of hanging):
# install boltz with --no-deps, then install its actually-imported runtime
# dependencies individually, also with --no-deps, accepting the versions
# already present rather than forcing boltz's exact pins. Verified working:
# boltz 2.2.1 imports and `boltz predict --help` runs correctly against
# numpy 1.26.4 (downgraded from whatever was there before -- re-verify
# pyrosetta/ProDy/impress/asyncflow/rhapsody still import after this step,
# they were confirmed OK against numpy 1.26.4 during initial validation) and
# gemmi 0.6.5 (downgraded from 0.7.5). `pip check` will still report several
# cosmetic mismatches (pytorch-lightning, cuequivariance-ops-torch-cu12,
# colabfold/ml-dtypes leftovers from before ColabFold was removed, torch's
# own sympy/triton/nvidia-cublas sub-pins) -- none of these broke any actual
# import in testing; only re-investigate if a real runtime failure surfaces.
#
echo ""
echo "── Step 7: Boltz-2 ──"
"${PIP}" install -q --no-deps "boltz[cuda]"
"${PIP}" install -q --no-deps \
pytorch_lightning torchmetrics fairscale einops einx mashumaro modelcif \
wandb dm-tree chembl_structure_pipeline \
cuequivariance_ops_cu12 cuequivariance_ops_torch_cu12
"${PY}" -c "import boltz; import torch; print('boltz', boltz.__version__ if hasattr(boltz, '__version__') else '(no __version__)', '+ torch', torch.__version__, 'import OK')"

# ── 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 ────────────────────────
#
# Pinned to 0.6.5, NOT latest: Step 7 installs boltz, which pins gemmi==0.6.5
# exactly. An unpinned `pip install gemmi` here would silently upgrade to
# latest and re-break that pin (this happened during initial validation).
# Verified empirically that 0.6.5 has everything this pipeline's gemmi usage
# needs: mpnn()'s CIF.GZ->PDB conversion (gemmi.cif.read_string,
# make_structure_from_block, write_pdb) and rfd3()'s ligand-normalization
# helper (read_structure, res.het_flag, mutable res.name, write_pdb) both
# round-trip correctly against 0.6.5.
#
echo ""
echo "── Step 9: gemmi ──"
"${PIP}" install -q "gemmi==0.6.5"

# ── 10. Additional dependencies ───────────────────────────────────────────────
echo ""
echo "── Step 10: pandas + biopandas ──"
"${PIP}" install -q pandas biopandas

# ── 11. PyRosetta ─────────────────────────────────────────────────────────────
echo ""
echo "── Step 11: PyRosetta ──"
"${PIP}" install -q pyrosetta-installer
"${PY}" -c "import pyrosetta_installer; pyrosetta_installer.install_pyrosetta()"

# ── 12. Boltz-2 model weights ─────────────────────────────────────────────────
#
# Boltz has no dedicated "download weights" subcommand — weights auto-download
# on first `boltz predict` call. Warm the cache with a trivial CPU prediction
# on a login node so compute nodes (no internet) find them already present at
# BOLTZ_CACHE.
#
echo ""
echo "── Step 12: Boltz-2 model weights (cache warm-up) ──"
BOLTZ_CACHE="${BOLTZ_CACHE:-${SCRATCH}/${USER}/.cache/boltz}"
mkdir -p "${BOLTZ_CACHE}"
_WARM_DIR=$(mktemp -d)
cat > "${_WARM_DIR}/warm.yaml" <<'YAML'
version: 1
sequences:
- protein:
id: [A]
sequence: MAAAAAAAAAAAAAAAAAAA
msa: empty
YAML
"${ENV_DIR}/bin/boltz" predict "${_WARM_DIR}/warm.yaml" \
--out_dir "${_WARM_DIR}/out" --cache "${BOLTZ_CACHE}" \
--devices 1 --accelerator cpu --output_format pdb \
|| echo "WARNING: boltz cache warm-up failed — check login-node internet access"
rm -rf "${_WARM_DIR}"

# ── 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 "boltz" "${PY}" -c "import boltz; 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 "boltz weights" test -f "${BOLTZ_CACHE}/boltz2_conf.ckpt" && 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 "================================================================="
Loading
Loading