"""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()