Spaces:
Running on Zero
Running on Zero
| """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 | |
| 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:]}") | |
| 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 | |
| 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 | |
| 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() | |