diff options
Diffstat (limited to 'app/backends/sglomni/models.py')
| -rw-r--r-- | app/backends/sglomni/models.py | 87 |
1 files changed, 27 insertions, 60 deletions
diff --git a/app/backends/sglomni/models.py b/app/backends/sglomni/models.py index 607e36a..990673a 100644 --- a/app/backends/sglomni/models.py +++ b/app/backends/sglomni/models.py @@ -16,31 +16,25 @@ VoiceDesign exist in the qwen backend too) is downloaded once and its deletion affects both — the same convention every backend here accepts. """ -import os import shutil from pathlib import Path from typing import List, Optional from backends import common, envs -from backends.sglomni.catalog import ModelEntry, entry_by_key, \ - entry_by_repo, extra_import_name +from backends.sglomni.catalog import ENTRIES, Extra, ModelEntry, \ + entry_by_key, entry_by_repo, extra_import_name from backends.sglomni.constants import SERVER_NAME, SGLOMNI_PIP_PKG from backends.sglomni.pythonenv import SGLOMNI_ENV, prepare_env -# The cache directory HF keeps repos in (models--<org>--<name> folders). -# Resolution mirrors huggingface_hub.constants: HF_HUB_CACHE beats -# HUGGINGFACE_HUB_CACHE beats HF_HOME/hub beats ~/.cache/huggingface/hub. +# The HF-cache primitives live in backends.common (shared with the qwen +# backend, which fetches into the same cache); the module-level wrappers +# below keep the names internal callers (and the tests) reference. def _hf_cache_dir() -> Path: - 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 cache directory HF keeps repos in (models--<org>--<name> + folders); common.hf_cache_dir resolves the environment overrides.""" + return common.hf_cache_dir() def repo_dir(repo_id: str) -> Path: @@ -55,17 +49,7 @@ def model_repo_dir(entry: ModelEntry) -> 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(entry: ModelEntry) -> bool: @@ -83,7 +67,7 @@ def model_installed(entry: ModelEntry) -> bool: def installed_entries() -> List[ModelEntry]: """The catalog entries whose weights are already on disk.""" - return [entry for entry in _all_entries() if model_installed(entry)] + return [entry for entry in ENTRIES if model_installed(entry)] def installed_keys() -> List[str]: @@ -115,11 +99,6 @@ def preset_voices(entry: ModelEntry) -> List[str]: return [] -def _all_entries() -> List[ModelEntry]: - from backends.sglomni.catalog import ENTRIES - return list(ENTRIES) - - def system_dep_missing(entry: ModelEntry) -> Optional[str]: """Remediation text when ENTRY's system binary is absent (None = ok).""" if entry.system_dep and not shutil.which(entry.system_dep): @@ -147,15 +126,20 @@ def missing_companions(entry: ModelEntry) -> List[Extra]: SGLOMNI_ENV)] -def install_companions(entry: ModelEntry, *, emit=None, cancel=None) -> int: - """pip-install ENTRY's missing companion packages into the venv. +def install_companions(entry: ModelEntry, *, emit=None, cancel=None, + force: bool = False) -> int: + """pip-install ENTRY's companion packages into the venv. The same recipe ``install_model`` runs (the catalog's ``--no-deps`` flags preserved — the Qwen3-TTS companions must not replace the pinned Transformers 5 stack), limited to what the import probe found absent, so a start-time heal touches as little of the pinned environment as - possible. Returns the first failing exit code, 0 when all present.""" - for spec, no_deps in missing_companions(entry): + possible. With FORCE every extra's pip spec re-runs instead — a + satisfied pin is a pip no-op, so the update flow uses that to heal + version drift the import probe cannot see (the protobuf re-pin + especially). Returns the first failing exit code, 0 when all present.""" + wanted = list(entry.extras) if force else missing_companions(entry) + for spec, no_deps in wanted: args = ["--no-deps"] if no_deps else None rc = common.pip_install([spec], emit=emit, cancel=cancel, env_dir=SGLOMNI_ENV, extra_args=args) @@ -237,36 +221,19 @@ def uninstall_model(key: str, *, emit=None, cancel=None) -> int: def delete_model_weights(entries: Optional[List[ModelEntry]] = None) -> int: """Delete the cached HF weight dirs of ENTRIES (every model by default). - Best-effort rmtree of each ``models--<org>--<name>`` 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. + Delegates to common.hf_delete_model_weights: 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. Returns how many were present and removed. """ if entries is None: - entries = _all_entries() - removed = 0 - for entry in entries: - directory = model_repo_dir(entry) - if not directory.is_dir(): - continue - print(f"[INFO] Removing cached {entry.repo} 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 + entries = list(ENTRIES) + return common.hf_delete_model_weights([entry.repo for entry in entries]) def _hf_download_prefix() -> Optional[List[str]]: """The sglang-omni venv's hf CLI argv prefix (None when absent).""" - for name in ("hf", "huggingface-cli"): - candidate = envs.env_script(name, SGLOMNI_ENV) - if candidate.is_file(): - return [str(candidate)] - return None + return common.hf_download_prefix(SGLOMNI_ENV) def _managed_running_repo() -> Optional[str]: @@ -292,7 +259,7 @@ def resolve_model(key: Optional[str]) -> ModelEntry: if key is not None: entry = entry_by_key(key) if entry is None: - known = ", ".join(e.key for e in _all_entries()) + known = ", ".join(e.key for e in ENTRIES) raise RuntimeError( f"Unknown sglang-omni model {key!r} (installed models are " f"picked by catalog key; known keys: {known})") |