From 5263a30356d7a7b39490e9a3cf5f6c179249500c Mon Sep 17 00:00:00 2001 From: historia Date: Fri, 4 Sep 2026 17:41:14 -0400 Subject: refactor: share the HF cache helpers between qwen and sglomni --- app/backends/common.py | 92 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 92 insertions(+) (limited to 'app/backends/common.py') diff --git a/app/backends/common.py b/app/backends/common.py index a12a366..7974b58 100644 --- a/app/backends/common.py +++ b/app/backends/common.py @@ -11,6 +11,7 @@ run. import os import re +import shutil import sys import time import urllib.parse @@ -743,3 +744,94 @@ def pip_uninstall(packages: List[str], *, emit=None, """ from backends import envs return envs.pip_uninstall(packages, emit=emit, env_dir=env_dir) + + +# ---------------------------------------------------------------------- +# HuggingFace hub cache +# +# Every pip-based backend's model weights are fetched by its server into +# the standard hub cache on first start (and pre-downloaded by the +# Install action via the venv's hf CLI). These helpers read and delete +# the same directories ``from_pretrained`` writes, honoring the same +# environment overrides, so per-model (un)installs land exactly where a +# server start would look. The cache is shared across backends: a repo +# two backends host exists once, and its deletion affects both — the +# same convention every backend here accepts. +# ---------------------------------------------------------------------- + +def hf_cache_dir() -> Path: + """The HF hub cache dir servers fetch model weights into. + + Resolution mirrors huggingface_hub.constants: HF_HUB_CACHE beats + HUGGINGFACE_HUB_CACHE beats HF_HOME/hub beats ~/.cache/huggingface/hub. + """ + override = os.environ.get("HF_HUB_CACHE") or os.environ.get( + "HUGGINGFACE_HUB_CACHE") + if override: + return Path(override) + home = os.environ.get("HF_HOME") + if home: + return Path(home) / "hub" + return Path.home() / ".cache" / "huggingface" / "hub" + + +def hf_repo_dir(repo_id: str) -> Path: + """The cache directory HF keeps REPO_ID's weights in.""" + return hf_cache_dir() / ("models--" + repo_id.replace("/", "--")) + + +def hf_tree_has_file(path: Path) -> bool: + """True when any file or symlink exists under PATH (recursively).""" + try: + for item in path.iterdir(): + # Snapshot files are symlinks into blobs/; count them even + # when temporarily broken (presence is what the loader + # checks). + if item.is_symlink() or item.is_file(): + return True + if item.is_dir() and hf_tree_has_file(item): + return True + except OSError: + return False + return False + + +def hf_delete_model_weights(repo_ids: List[str]) -> int: + """Delete the cached HF weight dirs of REPO_IDS; returns how many + were present and removed. + + Best-effort rmtree of each ``models----`` directory; only + those directories are ever touched — the rest of the HF cache may be + shared with unrelated tools. Prints its progress, which streams into + the task view under curses too. + """ + removed = 0 + for repo_id in repo_ids: + directory = hf_repo_dir(repo_id) + if not directory.is_dir(): + continue + print(f"[INFO] Removing cached {repo_id} weights...") + shutil.rmtree(directory, ignore_errors=True) + if directory.exists(): + print(f"[WARNING] Could not fully remove {directory}") + continue + removed += 1 + if removed: + print(f"[OK] Deleted cached weights for {removed} " + f"{'model' if removed == 1 else 'models'}.") + return removed + + +def hf_download_prefix(env_dir: Path) -> Optional[List[str]]: + """The venv ENV_DIR's hf CLI argv prefix (None when absent). + + The hf CLI (or its older ``huggingface-cli`` name) is what a server + start uses to fetch weights, so the install action pre-downloads + through it too — resumable, streamed, cancelable. + """ + from backends import envs + for name in ("hf", "huggingface-cli"): + candidate = envs.env_script(name, env_dir) + if candidate.is_file(): + return [str(candidate)] + return None -- cgit v1.2.3