"""Modal app: Boltz-2 structure prediction for BioDesignBench Phase B. This is the GPU-side companion to `eval_boltz.py`. The HF Space leaderboard runs on cpu-basic, so it cannot host Boltz directly; instead it POSTs design sequences to this Modal app, which spins up an A10G on demand, runs `boltz predict`, and returns confidence metrics. Setup (one-time, on a machine with `pip install modal`): modal token new # if you don't have a token yet cd biodesignbench-leaderboard modal deploy modal_boltz_app.py After deploy Modal prints a URL like https://--bdb-boltz-predict.modal.run Add that URL plus a shared secret to the HF Space secrets: MODAL_BOLTZ_URL = https://--bdb-boltz-predict.modal.run MODAL_BOLTZ_TOKEN = Cost: A10G is billed per-second, container auto-stops after `container_idle_timeout` seconds. With one submission per month and ~76 tasks * ~30s = ~38min GPU per submission, expected spend is well within Modal's free tier. """ from __future__ import annotations import os import modal APP_NAME = "bdb-boltz" ENDPOINT_LABEL = "bdb-boltz-predict" app = modal.App(APP_NAME) # Persistent volume for Boltz-2 model weights (~6GB, downloaded on first call) weights_volume = modal.Volume.from_name( "bdb-boltz-weights", create_if_missing=True ) # Boltz GPU image. Boltz-2 is published on PyPI as `boltz` and pulls a # CUDA-12 torch wheel automatically. gpu_image = ( modal.Image.from_registry( "nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04", add_python="3.11", ) .apt_install("git", "wget", "build-essential") # Boltz-2 (>=2.2) uses NVIDIA cuequivariance for the triangular-multiply # kernel and requires CUDA 12.5+. We let pip pick a torch that matches # cuequivariance's nvidia-cublas-cu12>=12.5 constraint. .pip_install( # Match dev's known-working stack: torch 2.10 ships nvidia-cublas-cu12 # 12.8 which satisfies cuequivariance>=12.5 requirement. "torch==2.10.0", "boltz==2.2.1", "cuequivariance==0.9.0", "cuequivariance-torch==0.9.0", "cuequivariance-ops-cu12==0.9.0", "cuequivariance-ops-torch-cu12==0.9.0", "fastapi[standard]", "pyyaml", "numpy", ) .env( { "BOLTZ_CACHE": "/weights", "TORCH_HOME": "/weights/torch", "HF_HOME": "/weights/hf", } ) ) # --------------------------------------------------------------------------- # Internal: write YAMLs, run boltz predict, parse outputs # --------------------------------------------------------------------------- def _write_yaml(item: dict) -> str: """Render one prediction item to a Boltz YAML string. item shape: {"name": "task_001", "kind": "monomer" | "complex", "sequences": ["MKKL...", ...]} # 1 for monomer, 2 for complex """ seqs = item.get("sequences") or [] chain_ids = ["A", "B", "C", "D", "E"] lines = ["sequences:"] for i, seq in enumerate(seqs): cid = chain_ids[i] if i < len(chain_ids) else f"X{i}" lines.append(" - protein:") lines.append(f" id: {cid}") lines.append(f" sequence: {seq}") return "\n".join(lines) + "\n" def _parse_confidence(pred_dir) -> dict: """Parse a Boltz prediction directory into a flat metric dict.""" import json from pathlib import Path import numpy as np out = { "pLDDT": 0.0, "pTM": 0.0, "ipTM": 0.0, "i_pAE": 0.0, "success": False, } pred_dir = Path(pred_dir) conf_files = list(pred_dir.rglob("confidence*.json")) if conf_files: try: with open(conf_files[0]) as f: c = json.load(f) out["pLDDT"] = round(float(c.get("complex_plddt", 0.0)) * 100, 2) out["pTM"] = round(float(c.get("ptm", 0.0)), 4) out["ipTM"] = round(float(c.get("iptm", 0.0)), 4) out["i_pAE"] = round(float(c.get("complex_ipae", 0.0)), 2) out["success"] = True except Exception: pass if not out["success"]: # Fall back to per-residue plddt npz if confidence.json is missing plddt_files = list(pred_dir.rglob("plddt*.npz")) if plddt_files: try: arr = np.load(plddt_files[0])["plddt"] out["pLDDT"] = round(float(arr.mean()) * 100, 2) out["success"] = True except Exception: pass return out # --------------------------------------------------------------------------- # GPU entry point — single web endpoint handling both monomer and complex # --------------------------------------------------------------------------- @app.function( image=gpu_image, gpu="A10G", volumes={"/weights": weights_volume}, timeout=1800, scaledown_window=300, secrets=[modal.Secret.from_name("bdb-boltz-shared", required_keys=["TOKEN"])], ) @modal.fastapi_endpoint(method="POST", label=ENDPOINT_LABEL) def predict(payload: dict) -> dict: """Run Boltz-2 on a list of prediction items. Body shape: {"token": "", "items": [{"name": "...", "kind": "monomer"|"complex", "sequences": [...]}, ...]} The list is assembled into a single ``boltz predict`` invocation so the model loads only once per call (amortizes ~30s cold start). Returns a dict mapping each item's `name` to a metric dict: {"pLDDT", "pTM", "ipTM", "i_pAE", "success"} """ import shutil import subprocess import tempfile from pathlib import Path expected_token = os.environ.get("TOKEN", "") if expected_token and (payload.get("token") or "") != expected_token: return {"error": "Unauthorized -- bad MODAL_BOLTZ_TOKEN"} items = payload.get("items") or [] if not items: return {"results": {}} work = Path(tempfile.mkdtemp(prefix="bdb_boltz_")) in_dir = work / "inputs" out_dir = work / "out" in_dir.mkdir() out_dir.mkdir() name_to_yaml: dict[str, str] = {} for i, item in enumerate(items): name = str(item.get("name") or f"item_{i:04d}") safe = "".join(c if c.isalnum() else "_" for c in name)[:60] yaml_name = f"{i:04d}_{safe}" (in_dir / f"{yaml_name}.yaml").write_text(_write_yaml(item)) name_to_yaml[name] = yaml_name cmd = [ "boltz", "predict", str(in_dir), "--out_dir", str(out_dir), "--cache", "/weights/boltz_cache", "--diffusion_samples", "1", "--output_format", "pdb", "--use_msa_server", ] proc = subprocess.run( cmd, capture_output=True, text=True, timeout=1700, cwd=str(work), ) # Persist downloaded model weights to the shared volume try: weights_volume.commit() except Exception: pass if proc.returncode != 0: shutil.rmtree(str(work), ignore_errors=True) return { "error": "boltz predict failed", "stderr": proc.stderr[-2000:], "stdout": proc.stdout[-2000:], } # boltz writes outputs to out/boltz_results_inputs/predictions// predictions_root = None for p in out_dir.rglob("predictions"): if p.is_dir(): predictions_root = p break results: dict[str, dict] = {} if predictions_root is not None: for name, yaml_name in name_to_yaml.items(): pred_dirs = [ d for d in predictions_root.iterdir() if d.is_dir() and (d.name.startswith(yaml_name) or d.name == yaml_name) ] if pred_dirs: results[name] = _parse_confidence(pred_dirs[0]) else: results[name] = { "pLDDT": 0.0, "pTM": 0.0, "ipTM": 0.0, "i_pAE": 0.0, "success": False, "error": "prediction missing", } shutil.rmtree(str(work), ignore_errors=True) return {"results": results} # --------------------------------------------------------------------------- # CLI smoke test: modal run modal_boltz_app.py # --------------------------------------------------------------------------- @app.local_entrypoint() def main(): """Quick sanity check — a short ubiquitin-like sequence.""" import json items = [ { "name": "monomer_demo", "kind": "monomer", "sequences": [ "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG" ], }, ] out = predict.remote(items, authorization="") print(json.dumps(out, indent=2))