studio: classify embedding models from the HF cache and honor offline mode

studio/backend/core/rag/embeddings.py

@@ -22,6 +22,7 @@
 
 from utils.hardware.hardware import DeviceType, get_device
 from utils.transformers_dtype import dtype_kwargs
+from utils.utils import hf_env_offline
 
 from . import config
 
@@ -119,46 +120,86 @@ def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]:
         return ()
 
 
-def _guard_model_security(name: str) -> None:
+def _guard_model_security(name: str, local_only: bool = False) -> None:
     """Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside
     SentenceTransformer regardless of trust_remote_code. Defense in depth behind the
     /settings gate (a name can also arrive via env/default); local paths and unreachable
     scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error.
+
+    ``local_only`` (offline) inspects the local cache; subdir probes are skipped (they'd hit the
+    network and hang, and the offline gate walks the whole snapshot anyway).
     """
     try:
         from utils.security import evaluate_file_security, security_load_subdirs
 
         token = _ambient_hf_token()
-        # Union the audio-model load roots with the ST module dirs so a flagged pickle
-        # directly under a Transformer module dir (0_Transformer/) blocks instead of
-        # passing as an unreferenced nested shard.
-        load_subdirs = tuple(
-            dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token)))
-        )
-        blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked
+        if local_only:
+            load_subdirs = ()
+        else:
+            # Union audio-model load roots with ST module dirs so a flagged pickle under a
+            # Transformer module dir blocks instead of passing as an unreferenced nested shard.
+            load_subdirs = tuple(
+                dict.fromkeys(
+                    (*security_load_subdirs(name, token), *_st_module_subdirs(name, token))
+                )
+            )
+        blocked = evaluate_file_security(
+            name, hf_token = token, load_subdirs = load_subdirs, local_only_load = local_only
+        ).blocked
     except Exception:
         return
     if blocked:
+        reason = (
+            "has cached pickle weights that cannot be security-scanned offline and no "
+            "safetensors alternative"
+            if local_only
+            else "is flagged as unsafe by Hugging Face's security scan"
+        )
         raise UnsafeEmbeddingModelError(
-            f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security "
-            "scan; refusing to load. Set a different RAG embedding model."
+            f"Embedding model {name!r} {reason}; refusing to load. "
+            "Set a different RAG embedding model."
         )
 
 
+def _st_accepts_local_files_only(st_cls) -> bool:
+    """Whether this SentenceTransformer version accepts local_files_only; passing it to an
+    older constructor raises, so gate on the signature."""
+    try:
+        import inspect
+        return "local_files_only" in inspect.signature(st_cls.__init__).parameters
+    except Exception:
+        return False
+
+
 def _get(model_name: str | None = None):
     """Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16
     for a ~1.5x speedup at negligible accuracy loss."""
     global _model, _name
     name = model_name or config.effective_embedding_model()
+    # Capture offline state once so the gate and the load agree (no window where the gate is
+    # skipped as offline but the constructor then reaches the network).
+    local_only = hf_env_offline()
     with _lock:
         if _model is None or _name != name:
             _install_torchao_stub_once()
             from sentence_transformers import SentenceTransformer
 
             device = _device()
             logger.info("loading embedding model %s on %s", name, device)
-            _guard_model_security(name)
-            _model = SentenceTransformer(name, device = device, model_kwargs = dtype_kwargs("float16"))
+            _guard_model_security(name, local_only)
+            st_kwargs = dict(device = device, model_kwargs = dtype_kwargs("float16"))
+            load_target = name
+            if local_only:
+                from utils.utils import hf_cache_snapshot_dir
+                snapshot = hf_cache_snapshot_dir(name)
+                if snapshot is not None:
+                    # Load from the local snapshot dir: a local path never touches the Hub, so
+                    # this is offline-safe on ANY sentence-transformers version (even ones
+                    # predating local_files_only).
+                    load_target = str(snapshot)
+                elif _st_accepts_local_files_only(SentenceTransformer):
+                    st_kwargs["local_files_only"] = True
+            _model = SentenceTransformer(load_target, **st_kwargs)
             _name = name
         return _model
 

studio/backend/routes/settings.py

@@ -416,6 +416,11 @@ def update_embedding_model(
             log = logger,
         ) from exc
     hf_token = (payload.hf_token or "").strip() or None
+    from utils.utils import hf_env_offline
+
+    # Offline, both the Hub malware scan and the is-embedding check are unreachable and degrade
+    # to the local cache below; capture the state once.
+    local_only_load = hf_env_offline()
     # The env/default model needs no verification; saving it is a no-op override.
     # A local GGUF on the llama-server backend is accepted as-is: it is exactly
     # what the backend loads, and HF metadata cannot verify a local path.
@@ -439,26 +444,41 @@ def update_embedding_model(
         # Fall back to the loader's own token so a gated/private repo is actually scanned
         # (a token-less scan fails open for exactly the repo that would still load).
         scan_token = hf_token or _ambient_hf_token()
-        # Include the ST module dirs (0_Transformer/) so a flagged pickle directly under
-        # one blocks instead of passing as an unreferenced nested shard.
-        load_subdirs = tuple(
-            dict.fromkeys(
-                (
-                    *security_load_subdirs(model, scan_token),
-                    *_st_module_subdirs(model, scan_token),
+        # Offline: subdir probes would hit the network and hang; the offline gate walks the
+        # whole cached snapshot, so no load-subdir hints are needed.
+        if local_only_load:
+            load_subdirs = ()
+        else:
+            # Include ST module dirs (0_Transformer/) so a flagged pickle directly under one
+            # blocks instead of passing as an unreferenced nested shard.
+            load_subdirs = tuple(
+                dict.fromkeys(
+                    (
+                        *security_load_subdirs(model, scan_token),
+                        *_st_module_subdirs(model, scan_token),
+                    )
                 )
             )
-        )
-        if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked:
+        if evaluate_file_security(
+            model,
+            hf_token = scan_token,
+            load_subdirs = load_subdirs,
+            local_only_load = local_only_load,
+        ).blocked:
             # 403, not 409: the client routes every 409 into the forceable "save anyway"
             # flow, but this block is a hard, non-forceable security refusal.
-            raise HTTPException(
-                status_code = 403,
+            if local_only_load:
+                detail = (
+                    f"{model!r} has cached pickle weights that cannot be security-scanned "
+                    "offline and no safetensors alternative, so it cannot be used as the "
+                    "embedding model. Re-download it with safetensors weights while online."
+                )
+            else:
                 detail = (
                     f"{model!r} is flagged as unsafe by Hugging Face's security scan and "
                     "cannot be used as the embedding model."
-                ),
-            )
+                )
+            raise HTTPException(status_code = 403, detail = detail)
     if model != default_embedding_model() and not payload.force and not is_local_gguf:
         from core.rag import config as rag_config
 
@@ -468,15 +488,28 @@ def update_embedding_model(
         # which would wrongly 409 a valid online GGUF embedder.
         gguf_named = _llama_backend_active() and rag_config._names_gguf(model)
         if not gguf_named and not is_embedding_model(model, hf_token = hf_token):
-            raise HTTPException(
-                status_code = 409,
-                detail = (
-                    f"Could not verify {model!r} as an embedding model on "
-                    "Hugging Face (it may be the wrong model type, gated, or "
-                    "you may be offline)."
-                ),
-            )
-        gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error(model, hf_token)
+            # Offline, is_embedding_model can only confirm the ST layout (modules.json); a
+            # transformers-native embedder (e.g. gte-modernbert) is unverifiable without Hub
+            # metadata. If already cached and loadable, accept it rather than raising a 409 that
+            # online would not (ST can load any cached encoder). Uncached -> 409.
+            from utils.utils import hf_cache_snapshot_is_loadable
+
+            # Require a genuinely loadable cache (config + weights), not just a resolved refs/main,
+            # so a metadata-only partial cache still gets the forceable 409.
+            offline_cached = local_only_load and hf_cache_snapshot_is_loadable(model)
+            if not offline_cached:
+                raise HTTPException(
+                    status_code = 409,
+                    detail = (
+                        f"Could not verify {model!r} as an embedding model on "
+                        "Hugging Face (it may be the wrong model type, gated, or "
+                        "you may be offline)."
+                    ),
+                )
+        # The Hub GGUF probe (list_repo_files) can hang offline; skip it. Local check stays.
+        gguf_error = _local_gguf_backend_error(model)
+        if gguf_error is None and not local_only_load:
+            gguf_error = _hf_gguf_backend_error(model, hf_token)
         if gguf_error:
             raise HTTPException(status_code = 409, detail = gguf_error)
     set_rag_embedding_model(model)

studio/backend/tests/test_embedding_model_security_gate.py

@@ -106,6 +106,56 @@ def test_hard_block_uses_non_forceable_status(client, monkeypatch):
     assert unverified.status_code == 409
 
 
+def test_offline_cached_non_st_model_is_accepted(client, monkeypatch):
+    # Offline, a cached transformers-native embedder (no modules.json) is unverifiable via HF
+    # metadata, but ST can load any cached encoder, so accept it (no 409).
+    c, saved = client
+    monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
+    monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+    import utils.models as _models
+    import utils.utils as _uu
+
+    monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False)
+    monkeypatch.setattr(_uu, "hf_cache_snapshot_is_loadable", lambda name: True)
+    r = c.put("/embedding-model", json = {"embedding_model": "acme/gte-modernbert"})
+    assert r.status_code == 200
+    assert saved.get("model") == "acme/gte-modernbert"
+
+
+def test_offline_partial_or_uncached_model_still_409(client, monkeypatch):
+    # Offline but not loadable (uncached or metadata-only partial cache): keep the forceable
+    # 409, since the cache-only load would fail anyway.
+    c, _saved = client
+    monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
+    monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+    import utils.models as _models
+    import utils.utils as _uu
+
+    monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False)
+    monkeypatch.setattr(_uu, "hf_cache_snapshot_is_loadable", lambda name: False)
+    r = c.put("/embedding-model", json = {"embedding_model": "acme/uncached-embedder"})
+    assert r.status_code == 409
+
+
+def test_offline_skips_remote_gguf_probe(client, monkeypatch):
+    # Offline + llama backend: the remote GGUF probe (list_repo_files) must be skipped so a
+    # dead-DNS session cannot hang.
+    c, _saved = client
+    monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+    monkeypatch.setattr(settings, "_llama_backend_active", lambda: True)
+    monkeypatch.setattr(settings, "_local_gguf_backend_error", lambda model: None)
+
+    def _boom(*a, **k):
+        raise AssertionError("hit the network for the GGUF probe")
+
+    monkeypatch.setattr(settings, "_hf_gguf_backend_error", _boom)
+    import utils.models as _models
+
+    monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: True)
+    r = c.put("/embedding-model", json = {"embedding_model": "acme/embedder"})
+    assert r.status_code == 200
+
+
 def test_llama_backend_skips_the_st_pickle_scan(monkeypatch):
     # On the llama-server backend the embedder loads GGUF (inert), not the ST repo's
     # pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here.