From 212cab7c2fb78b4aea216b705948054fd4f42446 Mon Sep 17 00:00:00 2001 From: Luca Pinello Date: Thu, 30 Apr 2026 17:02:07 -0400 Subject: [PATCH] fix(health): recognise slim HF mirror for chrombpnet + add alphagenome_pt probe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two stale checks in `chorus/core/weights_probe.py` were causing `chorus health` to falsely report healthy 0.4.0 installs as "⚠ Not installed": 1. **chrombpnet probe** only looked at the legacy `downloads/chrombpnet/DNASE_K562/` ENCODE-tarball-extracted directory. In 0.3.0+ that directory is empty by default — chrombpnet weights stream from the slim HF mirror at `models--lucapinello--chorus-chrombpnet-slim` (#59 / #60). Fix: accept *either* the slim-mirror snapshot cache (manifest.json present) *or* the legacy ENCODE directory. The runtime path was already 0.3.0-correct; only the probe was stale. 2. **alphagenome_pt** was missing from `_ARTIFACT_PROBES` entirely (only added to the oracle registry in #62; the health probe was never updated). Fix: new `_probe_alphagenome_pt()` checks the HF cache at `models--gtca--alphagenome_pytorch/snapshots// *.safetensors`. Also factored a small helper `_hf_cache_dir()` so both new probes honour `HF_HOME` / `HF_HUB_CACHE` env vars via `huggingface_hub.constants.HF_HUB_CACHE`, with a defensive fallback to the documented default `~/.cache/huggingface/hub`. Verified end-to-end on this machine: a fresh `chorus health` now reports 7/7 ✓ Healthy, where before it reported 5/7 with two false- alarm warnings on chrombpnet + alphagenome_pt. `pytest -m "not integration and not slow"`: 368 passed, 1 skipped, 5 deselected (no regression). Co-Authored-By: Claude Opus 4.7 (1M context) --- chorus/core/weights_probe.py | 73 +++++++++++++++++++++++++++++++++--- 1 file changed, 67 insertions(+), 6 deletions(-) diff --git a/chorus/core/weights_probe.py b/chorus/core/weights_probe.py index 79e36e8..23f81a9 100644 --- a/chorus/core/weights_probe.py +++ b/chorus/core/weights_probe.py @@ -49,13 +49,53 @@ def _probe_legnet() -> Tuple[bool, List[str]]: return (True, []) +def _hf_cache_dir() -> Path: + """Resolve the HuggingFace hub cache directory. + + Prefer the canonical ``huggingface_hub.constants.HF_HUB_CACHE`` (which + honours ``HF_HOME`` / ``HF_HUB_CACHE`` env vars); fall back to the + documented default ``~/.cache/huggingface/hub`` if the import fails + (chorus base env always has huggingface_hub, so the fallback is + defensive). + """ + try: + from huggingface_hub.constants import HF_HUB_CACHE + return Path(HF_HUB_CACHE) + except Exception: + return Path.home() / ".cache" / "huggingface" / "hub" + + def _probe_chrombpnet() -> Tuple[bool, List[str]]: - # ChromBPNet has no constructor-level default; `chorus setup chrombpnet` - # pre-downloads the canonical DNASE / K562 model. - default = CHORUS_DOWNLOADS_DIR / "chrombpnet" / "DNASE_K562" - if not default.exists() or not any(default.iterdir()): - return (False, [str(default)]) - return (True, []) + """Accept either the 0.3.0+ slim HF mirror or the legacy ENCODE-tarball + layout as proof of install. + + Default in 0.3.0+: weights stream from + ``lucapinello/chorus-chrombpnet-slim`` and live in the HF hub cache. + The local ``downloads/chrombpnet/DNASE_K562`` directory is only + populated for users who explicitly request ``model_type='chrombpnet'`` + (bias-aware) or fold ≠ 0, both of which fall back to ENCODE tarballs. + `chorus setup chrombpnet` succeeds on the slim path, so we must + accept either cache as installed. + """ + # 0.3.0+ default: slim HF mirror. + slim_snapshots = ( + _hf_cache_dir() / "models--lucapinello--chorus-chrombpnet-slim" / "snapshots" + ) + if slim_snapshots.exists(): + for snap in slim_snapshots.iterdir(): + if (snap / "manifest.json").exists(): + return (True, []) + # Legacy ENCODE-tarball cache. + legacy = CHORUS_DOWNLOADS_DIR / "chrombpnet" / "DNASE_K562" + if legacy.exists() and any(legacy.iterdir()): + return (True, []) + return ( + False, + [ + f"neither HF slim mirror cache ({slim_snapshots.parent}) " + f"nor legacy {legacy} is populated" + ], + ) def _probe_alphagenome() -> Tuple[bool, List[str]]: @@ -88,11 +128,32 @@ def _probe_library_cached() -> Tuple[bool, List[str]]: return (True, []) +def _probe_alphagenome_pt() -> Tuple[bool, List[str]]: + """Check the HF cache for the upstream PyTorch port's safetensors. + + Weights live at ``models--gtca--alphagenome_pytorch/snapshots// + model_all_folds.safetensors`` after ``chorus setup --oracle + alphagenome_pt`` (or the first ``load_pretrained_model()`` call). + """ + pt_snapshots = ( + _hf_cache_dir() / "models--gtca--alphagenome_pytorch" / "snapshots" + ) + if pt_snapshots.exists(): + for snap in pt_snapshots.iterdir(): + if any(p.suffix == ".safetensors" for p in snap.iterdir()): + return (True, []) + return ( + False, + [f"HF cache not populated: {pt_snapshots.parent}"], + ) + + _ARTIFACT_PROBES: Dict[str, Callable[[], Tuple[bool, List[str]]]] = { "sei": _probe_sei, "legnet": _probe_legnet, "chrombpnet": _probe_chrombpnet, "alphagenome": _probe_alphagenome, + "alphagenome_pt": _probe_alphagenome_pt, "enformer": _probe_library_cached, "borzoi": _probe_library_cached, }