fix: sglomni retry contract, booting detection, CLI keys, companion refresh, ffmpeg depHEADmain

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