Blame for bootstrap.py

"""Start-up of the AI server on the Space: weights, helper environments and settings, before the backend starts.

The Space's disk does not survive a restart, so every start (about 1-2 minutes, in parallel):
  * downloads the weights at fixed revisions: projectxddi/bindsite-weights (AdaBind checkpoints, ESM-IF1, baselines,
    precomputed ESMFold structures) and ESM-2 / ProtT5 from their own repositories (byte-identical to training);
  * builds the two helper environments of the backend with uv, at the exact versions verified on the 4090 machine
    (docker/constraints-*.txt): ESM-IF (Python 3.10, torch 2.4.1 CPU, numpy<2) and the baselines (Python 3.11,
    torch 2.2.2 CPU);
  * fetches the IP-location file from the private bucket (the backup machine refreshes it monthly);
  * restores the visit / prediction log from the bucket (done by the backend's Tracker).
Then it sets the BINDSITE_* variables the backend reads (backend/app/config.py). Everything runs on the CPU.
"""
import logging
import os
import shutil
import subprocess
import sys
import tarfile
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

APP = Path(__file__).resolve().parent
ROOT = Path(os.environ.get("BINDSITE_HOME", "/tmp/bindsite"))
WEIGHTS = ROOT / "weights"

WEIGHTS_REPO = ("projectxddi/bindsite-weights", "8548592d5628e0cd8a970993f9a2005c6ede5f7a")
FOUNDATION = {   # foundation_models/MANIFEST.json
    "esm2_t33_650M_UR50D": ("facebook/esm2_t33_650M_UR50D", "08e4846e537177426273712802403f7ba8261b6c",
                            ["config.json", "special_tokens_map.json", "tokenizer_config.json", "vocab.txt",
                             "model.safetensors"]),
    "prot_t5_xl_half_uniref50-enc": ("Rostlab/prot_t5_xl_half_uniref50-enc", "94a6abc029ae13029317b140b7424e012bf8dfbf",
                                     ["config.json", "special_tokens_map.json", "spiece.model", "tokenizer_config.json",
                                      "pytorch_model.bin"]),
}
BUCKET = "projectxddi/bindsite-logs"
STATIC_SITE = "https://projectproteinbinding-bindsite-explorer.static.hf.space"

log = logging.getLogger("bindsite.bootstrap")
report: dict = {}   # step -> {"seconds": .., "message": ..}; shown on the admin page


def _step(name):
    def wrap(fn):
        def run(*args, **kwargs):
            started = time.perf_counter()
            try:
                message = fn(*args, **kwargs) or "ok"
                report[name] = {"status": "ok", "seconds": round(time.perf_counter() - started, 1), "message": message}
            except Exception as exc:
                report[name] = {"status": "error", "seconds": round(time.perf_counter() - started, 1),
                                "message": f"{type(exc).__name__}: {exc}"}
                log.exception("Start-up step %s failed", name)
            log.info("Start-up step %s: %s", name, report[name])
        return run
    return wrap


@_step("weights")
def download_weights() -> str:
    from huggingface_hub import snapshot_download

    repo, revision = WEIGHTS_REPO
    snapshot_download(repo, revision=revision, local_dir=WEIGHTS, max_workers=16)
    with tarfile.open(WEIGHTS / "esmfold_files.tar.gz") as tar:
        tar.extractall(WEIGHTS, filter="data")
    hub = WEIGHTS / "torch_home" / "hub" / "checkpoints"
    hub.mkdir(parents=True, exist_ok=True)
    link = hub / "esm_if1_gvp4_t16_142M_UR50.pt"
    if not link.exists():
        link.symlink_to(WEIGHTS / "esm_if1" / "esm_if1_gvp4_t16_142M_UR50.pt")
    for name, (repo, revision, files) in FOUNDATION.items():
        snapshot_download(repo, revision=revision, allow_patterns=files, local_dir=WEIGHTS / name, max_workers=8)
    n = len(list((WEIGHTS / "esmfold_files").glob("*.pdb")))
    return f"{WEIGHTS_REPO[0]}@{WEIGHTS_REPO[1][:8]} + ESM-2 + ProtT5, {n} precomputed structures"


def _uv(*args: str) -> None:
    out = subprocess.run(["uv", *args], capture_output=True, text=True, timeout=1200)
    if out.returncode != 0:
        raise RuntimeError(f"uv {' '.join(args[:3])} ...: {(out.stdout + out.stderr).strip()[-600:]}")


@_step("esmif environment")
def build_esmif() -> str:
    venv = ROOT / "venv-esmif"
    py = str(venv / "bin" / "python")
    _uv("venv", str(venv), "--python", "3.10.22", "--quiet")
    _uv("pip", "install", "--quiet", "--python", py, "torch==2.4.1", "--index-url", "https://download.pytorch.org/whl/cpu")
    _uv("pip", "install", "--quiet", "--python", py, "torch-scatter==2.1.2", "-f",
        "https://data.pyg.org/whl/torch-2.4.0+cpu.html")
    _uv("pip", "install", "--quiet", "--python", py, "-r", str(APP / "backend" / "requirements-esmif.txt"),
        "-c", str(APP / "docker" / "constraints-esmif.txt"))
    return py


@_step("baselines environment")
def build_baselines() -> str:
    venv = ROOT / "venv-baselines"
    py = str(venv / "bin" / "python")
    _uv("venv", str(venv), "--python", "3.11", "--quiet")
    _uv("pip", "install", "--quiet", "--python", py, "torch==2.2.2", "--index-url", "https://download.pytorch.org/whl/cpu")
    _uv("pip", "install", "--quiet", "--python", py, "-r", str(APP / "backend" / "requirements-baselines.txt"),
        "-c", str(APP / "docker" / "constraints-baselines.txt"))
    return py


@_step("ip location")
def download_geoip() -> str:
    from huggingface_hub import download_bucket_files

    target = ROOT / "dbip-city-lite.mmdb"
    download_bucket_files(BUCKET, files=[("geoip/dbip-city-lite.mmdb", str(target))], token=os.environ.get("HF_TOKEN"))
    if not target.exists():
        raise FileNotFoundError(f"geoip/dbip-city-lite.mmdb not in bucket {BUCKET}")
    return f"{target.stat().st_size / 1e6:.0f} MB"


def baselines_dir() -> Path:
    """Code from this repository + checkpoints from the weights repo, in the layout baseline_worker.py expects."""
    root = ROOT / "baselines"
    pairs = {"bindEmbed21DL/src": APP / "baselines_2124" / "bindEmbed21DL" / "src",
             "bindEmbed21DL/trained_models": WEIGHTS / "baselines" / "bindEmbed21DL" / "trained_models",
             "bindNode24/src": APP / "baselines_2124" / "bindNode24" / "src",
             "bindNode24/trained_model_dicts": WEIGHTS / "baselines" / "bindNode24" / "trained_model_dicts"}
    for rel, target in pairs.items():
        link = root / rel
        link.parent.mkdir(parents=True, exist_ok=True)
        if not link.is_symlink() and not link.exists():
            link.symlink_to(target)
    return root


def prepare() -> None:
    started = time.perf_counter()
    ROOT.mkdir(parents=True, exist_ok=True)
    if os.environ.get("BINDSITE_SKIP_DOWNLOADS") == "1":   # local check of the pages with the mock predictor
        report["skipped"] = {"status": "warning", "seconds": 0, "message": "BINDSITE_SKIP_DOWNLOADS=1"}
    else:
        if shutil.which("uv") is None:
            raise RuntimeError("uv is missing (requirements.txt)")
        with ThreadPoolExecutor(max_workers=4) as pool:
            for job in (download_weights, build_esmif, build_baselines, download_geoip):
                pool.submit(job)
    report["total"] = {"status": "ok", "seconds": round(time.perf_counter() - started, 1), "message": ""}

    admin_password = os.environ.get("ADMIN_PASSWORD", "")
    env = {
        "BINDSITE_WEIGHTS_DIR": str(WEIGHTS),
        "BINDSITE_FINAL_MODEL_DIR": str(WEIGHTS / "final_model"),
        "BINDSITE_GOLDEN_DIR": str(APP / "adabind" / "golden_set"),
        "BINDSITE_ESMIF_PYTHON": str(ROOT / "venv-esmif" / "bin" / "python"),
        "BINDSITE_BASELINES_PYTHON": str(ROOT / "venv-baselines" / "bin" / "python"),
        "BINDSITE_BASELINES_DIR": str(baselines_dir()),
        "BINDSITE_CACHE_DIR": str(ROOT / "cache"),
        "BINDSITE_LOG_DIR": str(ROOT / "logs"),
        "BINDSITE_DEVICE": "cpu",              # outside @spaces.GPU, CUDA is only emulated on ZeroGPU
        "BINDSITE_ESMFOLD_LOCAL": "off",       # ESMFold comes from the backup machine (BINDSITE_WORKER_KEY)
        "BINDSITE_CORS_ORIGINS": f"{STATIC_SITE},http://localhost:4200",
        "BINDSITE_TRACKING_DB": str(ROOT / "tracking.sqlite"),
        "BINDSITE_TRACKING_BUCKET": BUCKET,
        "BINDSITE_GEOIP_DB": str(ROOT / "dbip-city-lite.mmdb"),
        "BINDSITE_WORKER_KEY": os.environ.get("WORKER_KEY", ""),
        "BINDSITE_ADMIN_KEY": admin_password,
        "BINDSITE_RATE_LIMIT": os.environ.get("BINDSITE_RATE_LIMIT", "120"),
        "HF_HUB_OFFLINE": "0",                 # the tracking log is copied to the bucket through the Hub API
    }
    for key, value in env.items():
        os.environ.setdefault(key, value)
    log.info("Start-up done in %.0f s: %s", time.perf_counter() - started,
             {k: v["status"] for k, v in report.items()})


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, stream=sys.stdout)
    prepare()