Junior B.
08ed713cc8
docs/gpu_plan.md: cost-efficient plan for running AF3-class co-folding (Boltz-2/DiffDock) on a GPU then paying nothing when idle. - Key insight: structure-track data is tiny (MB of PDBs/SMILES); only the GPU + model weights are heavy -> serverless is ideal. - Recommend Modal (per-second billing, scales to zero = nothing to kill); RunPod as the SSH-box alternative with idle auto-terminate. - Lifecycle: image -> weights Volume (cache, don't re-download) -> run -> git push small results -> teardown automatic. - Phase 1 validate on 3 known binders (~$1) before paying for a screen; Boltz-2 (affinity) on an L4/A10 (24-48GB); est total ~$5-15. gpu/modal_app.py: Modal app skeleton (image, weights volume, GPU cofold() function, local entrypoint); boltz invocation stubbed with TODOs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
68 lines
3.0 KiB
Python
68 lines
3.0 KiB
Python
"""Ephemeral GPU runner for AF3-class co-folding (PLAN §12, docs/gpu_plan.md).
|
|
|
|
Serverless: `modal run gpu/modal_app.py` provisions a GPU, runs, and releases it — zero idle cost,
|
|
nothing to remember to kill. Model weights are cached in a persistent Volume so we never re-pay GPU
|
|
time to re-download them. Prep (Meeko/RDKit) and RMSD scoring (spyrmsd) stay light; only the model
|
|
forward pass needs the GPU.
|
|
|
|
Setup (one-time): `pip install modal && modal token new`.
|
|
Run Phase 1 (validate on 3 known binders): `modal run gpu/modal_app.py`.
|
|
|
|
STATUS: scaffold. The boltz invocation (input spec + output parsing) is stubbed where marked TODO;
|
|
wire it after a first `modal run` confirms the image builds and the GPU is reachable.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import modal
|
|
|
|
app = modal.App("reverso-binding")
|
|
|
|
# CUDA image + AF3-class model (Boltz-2) + light prep/scoring deps.
|
|
image = (
|
|
modal.Image.debian_slim(python_version="3.12")
|
|
.apt_install("git", "wget")
|
|
.pip_install("boltz", "rdkit", "meeko", "spyrmsd", "gemmi", "numpy")
|
|
)
|
|
|
|
# Persist model weights across runs so we download them once, not every GPU-billed run.
|
|
weights = modal.Volume.from_name("reverso-binding-weights", create_if_missing=True)
|
|
|
|
# Known binders -> (PDB id, crystal ligand resname, SMILES placeholder filled by caller).
|
|
# Phase 1 validation: does co-folding reproduce these crystal poses where Vina failed?
|
|
KNOWN = {
|
|
"voxelotor_Hb": ("5E83", "5L7"),
|
|
"mitapivat_PKR": ("8XFD", "WV2"),
|
|
"vorinostat_HDAC2": ("4LXZ", "SHH"),
|
|
}
|
|
|
|
|
|
@app.function(gpu="L4", image=image, volumes={"/weights": weights}, timeout=3600)
|
|
def cofold(protein_seq: str, ligand_smiles: str, weights_dir: str = "/weights") -> dict:
|
|
"""Co-fold one protein+ligand complex and return predicted affinity + pose (PDB string).
|
|
|
|
Runs on the GPU only for this call, then the GPU is released. TODO: replace the stub with the
|
|
actual Boltz-2 invocation (write the YAML/FASTA input spec, call `boltz predict
|
|
--use_msa_server --out_dir ... --cache /weights`, parse the predicted structure + affinity).
|
|
"""
|
|
import subprocess # noqa: F401 (used once boltz is wired)
|
|
|
|
# TODO: build boltz input (protein_seq + ligand_smiles), run, parse pose+affinity.
|
|
raise NotImplementedError("Wire Boltz-2 here; see docs/gpu_plan.md Phase 1.")
|
|
|
|
|
|
@app.local_entrypoint()
|
|
def main() -> None:
|
|
"""Phase 1 driver (runs locally; only cofold() touches the GPU).
|
|
|
|
Pulls target sequences + ligand SMILES from the repo, fans out one GPU call per known binder,
|
|
scores redocking RMSD vs the crystal pose locally (spyrmsd), and prints pass/fail. Results are
|
|
tiny — commit a summary into data/processed/binding/.
|
|
"""
|
|
# TODO: load protein sequences from data/raw/structures/<pdb>.pdb (gemmi) and ligand SMILES
|
|
# (PubChem / drug_set), then:
|
|
# results = list(cofold.map(seqs, smiles))
|
|
# and compute in-place spyrmsd RMSD vs the crystal ligand for each.
|
|
print("Scaffold: fill in sequence/SMILES loading + cofold.map, then score RMSD. "
|
|
"See docs/gpu_plan.md.")
|