diff options
Diffstat (limited to 'app/backends')
| -rw-r--r-- | app/backends/common.py | 92 | ||||
| -rw-r--r-- | app/backends/qwen.py | 69 |
2 files changed, 107 insertions, 54 deletions
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--<org>--<name>`` 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 diff --git a/app/backends/qwen.py b/app/backends/qwen.py index c099c84..45b691b 100644 --- a/app/backends/qwen.py +++ b/app/backends/qwen.py @@ -25,8 +25,6 @@ Usage: """ import argparse -import os -import shutil import sys from pathlib import Path from typing import List, Optional @@ -81,25 +79,15 @@ QWEN_SPEAKERS = QWEN3_TTS_SPEAKERS # -- HuggingFace weight cache ------------------------------------------------ # # The demo servers pull each model's snapshot into huggingface_hub's default -# hub cache on first start (nothing in this project redirects it). 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. +# hub cache on first start (nothing in this project redirects it). The +# cache primitives are shared with the sglang-omni backend (which fetches +# the same repos into the same cache) in backends.common; the wrappers +# keep the module-level names internal callers (and the tests) reference. 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" + """The HF hub cache dir the servers fetch model weights into; + common.hf_cache_dir resolves the environment overrides.""" + return common.hf_cache_dir() def repo_dir(repo_id: str) -> Path: @@ -114,17 +102,7 @@ def model_repo_dir(model: str) -> Path: def _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 _tree_has_file(item): - return True - except OSError: - return False - return False + return common.hf_tree_has_file(path) def model_installed(model: str) -> bool: @@ -148,27 +126,14 @@ def installed_models() -> List[str]: def delete_model_weights(models: Optional[List[str]] = None) -> int: """Delete the cached HF weight dirs of MODELS (every model by default). - Best-effort rmtree of each ``models--Qwen--…`` directory; returns how - many were present and removed. 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. + Delegates to common.hf_delete_model_weights: best-effort rmtree of + each ``models--Qwen--…`` directory; only those directories are ever + touched — the rest of the HF cache may be shared with unrelated tools. + Returns how many were present and removed. """ names = sorted(MODEL_REPOS) if models is None else list(models) - removed = 0 - for name in names: - directory = model_repo_dir(name) - if not directory.is_dir(): - continue - print(f"[INFO] Removing cached {MODEL_REPOS[name]} weights...") - shutil.rmtree(directory) - 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 + return common.hf_delete_model_weights([MODEL_REPOS[name] + for name in names]) def _is_installed() -> bool: @@ -381,11 +346,7 @@ def _detect_remote(managed: bool = False): def _hf_download_prefix() -> Optional[List[str]]: """The qwen env's hf CLI argv prefix (None when neither script is present).""" - for name in ("hf", "huggingface-cli"): - candidate = envs.env_script(name, QWEN_ENV) - if candidate.is_file(): - return [str(candidate)] - return None + return common.hf_download_prefix(QWEN_ENV) def install_model(model: str, *, emit=None, cancel=None) -> int: |