Reverso/docs/gpu_plan.md

GPU plan — ephemeral AF3-class co-folding for the binding track

Goal: run AF3-class co-folding (Boltz-2 / DiffDock) on a GPU for the structure-binding track (§12), then pay nothing when idle. The work is bursty (a validation run, then a screen), the inputs are tiny, so the design optimises for zero idle cost, not for a persistent box.

What actually has to move (small!)

Thing	Size	How it gets to the GPU
Code	KB	git clone the gitea repo (or mount)
Target structures (PDBs)	a few MB	in the repo / git
Ligands (SMILES, drug set)	KB	in the repo / git
Model weights (Boltz-2 / DiffDock)	~2–6 GB	downloaded once, cached in a persistent volume
Results (poses, scores, RMSDs)	KB–MB	git push back / download

The 27 GB LINCS data is not part of this track — nothing big to upload. The only thing worth persisting is the model-weights cache (so we don't re-download = re-pay GPU time every run).

How the model weights persist (the cost-saver)

A modal.Volume is a named, cloud-backed filesystem that lives independently of any container or GPU — it survives every teardown. Mounted into the function at /weights:

• Run 1: /weights is empty → the model downloads weights there (the one-time slow cost).
• Run 2+: the same Volume mounts with the files already present → download skipped → no GPU-billed seconds wasted re-fetching 5 GB.

Two things make it actually cache:

1. Point the downloader at the mount (weights only persist if written under /weights): HF_HOME=/weights/hf (HuggingFace), TORCH_HOME=/weights/torch, boltz --cache /weights/boltz.
2. Commit semantics: writes persist on weights.commit() (modern Modal also auto-commits on a clean exit); other containers see them after weights.reload(). Pattern: reload() → run → commit().

The Volume itself costs pennies (~$/GB-month of storage), separate from the GPU — so caching ~5 GB of weights is near-free and saves real GPU time on every subsequent run. (Alternative: bake weights into the image at build time via image.run_function(download) — fastest cold start, but the image rebuilds when weights change. The skeleton uses the Volume approach.)

Provider choice

Option	Billing	Idle cost	"Kill" model	Best for
Modal (recommended)	per-second	$0 (scales to zero)	automatic — nothing to remember	bursty batch runs
RunPod	per-minute, on-demand/spot	only while pod exists	manual terminate	interactive SSH box
Vast.ai	per-minute spot	only while rented	manual destroy	cheapest, more fiddly
Lambda / AWS-GCP spot	per-hour/second	until you stop it	manual stop	if you already have credits

Recommend Modal. You define the run as a function with a gpu= decorator; on call it provisions the GPU, runs, and releases it — there is no GPU to forget to kill, and you pay only for the seconds it ran. That is the "kill the GPU to save money" requirement, automated.

The lifecycle (Modal)

1. Define image once: CUDA base + boltz (or DiffDock) + rdkit, meeko, spyrmsd, gemmi.
2. Weights volume: modal.Volume mounted at /weights; the model downloads into it on first run and is cached forever after (no re-download cost).
3. Run: modal run gpu/modal_app.py → provisions GPU → runs the test → returns results to your laptop. GPU released the moment the function returns.
4. Persist results: write the returned scores/poses into data/processed/binding/ and git commit (small, text).
5. Teardown: nothing to do — Modal scaled to zero. modal app stop only if a run is hung.

RunPod alternative (if you want an interactive box): start pod → git clone → run → git push results → runpodctl remove pod <id> (or Stop in the UI). Set an idle auto-terminate so a forgotten box can't bleed money.

What runs on the GPU (in cost order — cheap validation first)

• Phase 1 — modality validation (~minutes, ~$1): co-fold each known binder + 2 negative controls (caffeine, hydroxyurea) into each target (Hb, PKR, HDAC2) with the binding-mode cofactors co-folded in — HDAC2 + catalytic Zn, PKR + FBP/Mg, Hb + heme (as CCD ligand entries) — and check the known binder has the highest Boltz-2 P(binder) for its own target. This is the discrimination Vina couldn't manage precisely because it can't model Zn-chelation / cofactors. (Ranking uses P(binder), not the raw affinity value, whose sign is version-dependent.) Pose-RMSD vs crystal is a deeper check but needs receptor superposition (align predicted protein to crystal, transform ligand) — a later refinement. If Phase 1 passes, the modality is real; if not, stop before paying for a screen.
• Phase 2 — screen (~tens of minutes, a few $): run the ~300-drug set (or a focused subset) against the sickle targets; rank by Boltz-2 predicted affinity; redo the §12.4 positive-control recovery test. Output a ranked CSV, same shape as the connectivity ranked_candidates.

Model choice

• Boltz-2 (MIT, pip-installable) — predicts the protein–ligand complex and a binding affinity → directly gives a rankable score. Primary choice. Fits a 24–40 GB GPU for these single-domain targets.
• DiffDock-L — lighter, pose-only (needs a separate scorer); fallback if Boltz memory is tight.
• GPU: an L4 (24 GB, ~$0.6–0.8/hr) or A10/L40S (24–48 GB) is plenty; no multi-GPU, no A100 needed for these sizes.

Cost controls (the save-money checklist)

• Serverless (Modal) → zero idle cost by construction; or an idle-timeout auto-kill on a box.
• Cache weights in a persistent volume — re-downloading 5 GB on a $1/hr GPU is wasted money.
• Validate on one target before screening — don't pay for a 300-drug screen until Phase 1 passes.
• Prefer spot/interruptible for the batch screen (Phase 2 is restartable).
• Keep prep (Meeko/RDKit) and result-scoring on the laptop; only the model forward pass needs GPU.
• Estimated total to validate + a first screen: ~$5–15, not a standing bill.

Repo integration

• gpu/modal_app.py — the Modal app (skeleton committed alongside this plan).
• Results land in data/processed/binding/ (gitignored) + a small committed summary.
• Pin model + weights version in the image for reproducibility.

Next step

Scaffold gpu/modal_app.py for Phase-1 validation (3 known binders), do a dry run locally (modal run --detach? no — just modal run), confirm cost, then Phase 2. Requires a Modal account

• pip install modal + modal token new (one-time auth).