diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 15be7f1249..0c743e4ea4 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -22,6 +22,7 @@ from typing import Callable 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,30 +120,55 @@ 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: - 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." + 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} {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): @@ -150,6 +176,9 @@ def _get(model_name: str | None = None): 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() @@ -157,8 +186,20 @@ def _get(model_name: str | None = None): 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 diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 17e64df918..f36c8870e3 100644 --- a/studio/backend/routes/settings.py +++ b/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) diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index b3fa98b604..a6c18bd8de 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/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. diff --git a/studio/backend/tests/test_offline_embedding_minimal.py b/studio/backend/tests/test_offline_embedding_minimal.py new file mode 100644 index 0000000000..8862e231e5 --- /dev/null +++ b/studio/backend/tests/test_offline_embedding_minimal.py @@ -0,0 +1,583 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Offline RAG embedding-model handling (issue #6817). + +Offline the studio must never call the Hub (a DNS-dead session hangs on retries). Using a fake +HF cache under a temp HF_HUB_CACHE, assert that offline: is_embedding_model classifies from the +cached modules.json without the Hub; the file-security gate fails CLOSED on an unscanned pickle +weight with no safetensors alternative and allows an inert cache; the embedder threads +local_files_only into the load. Online behavior is unchanged (bounded timeout + cache fallback). +""" + +import sys +import types +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from utils.security import evaluate_file_security +from utils.utils import ( + hf_cache_snapshot_dir, + hf_cache_snapshot_is_loadable, + hf_env_offline, + st_repo_id_candidates, +) + +# Minimal sentence-transformers modules.json (the marker the gate keys on). +MODULES_JSON = ( + '[{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Transformer"}]' +) + + +def _modules_json(*paths): + """modules.json listing one Transformer module per path (a load root).""" + import json + return json.dumps( + [ + { + "idx": i, + "name": str(i), + "path": p, + "type": "sentence_transformers.models.Transformer", + } + for i, p in enumerate(paths) + ] + ) + + +_COMMIT = "0123456789abcdef0123456789abcdef01234567" + + +def _make_cache( + root, + repo_id, + files, + commit = _COMMIT, +): + """Build a canonical HF-cache snapshot (refs/main + snapshots//) for repo_id under + root from {relpath: contents}; returns the snapshot dir.""" + from huggingface_hub.file_download import repo_folder_name + + repo_dir = Path(root) / repo_folder_name(repo_id = repo_id, repo_type = "model") + (repo_dir / "refs").mkdir(parents = True, exist_ok = True) + (repo_dir / "refs" / "main").write_text(commit) + snapshot = repo_dir / "snapshots" / commit + snapshot.mkdir(parents = True, exist_ok = True) + for rel, contents in files.items(): + path = snapshot / rel + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(contents) + return snapshot + + +def _no_network(): + """Patch model_info to fail loudly if any offline path reaches the network.""" + return patch("huggingface_hub.model_info", side_effect = AssertionError("hit the network")) + + +def _is_embedding_model(*args, **kwargs): + from utils.models.model_config import is_embedding_model + return is_embedding_model(*args, **kwargs) + + +@pytest.fixture +def hf_cache(tmp_path, monkeypatch): + """Point the HF cache at a fresh temp dir.""" + root = tmp_path / "hub" + root.mkdir() + monkeypatch.setenv("HF_HOME", str(tmp_path)) + monkeypatch.setenv("HF_HUB_CACHE", str(root)) + return root + + +@pytest.fixture(autouse = True) +def _clean_env(monkeypatch): + """Start each test online with an empty detection cache; offline tests opt in.""" + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + from utils.models import model_config as mc + + mc._embedding_detection_cache.clear() + yield + mc._embedding_detection_cache.clear() + + +# ── hf_env_offline ─────────────────────────────────────────────── + + +@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on", " On "]) +def test_hf_env_offline_true(monkeypatch, value): + monkeypatch.setenv("HF_HUB_OFFLINE", value) + assert hf_env_offline() is True + + +@pytest.mark.parametrize("value", ["0", "false", "no", "off", ""]) +def test_hf_env_offline_false(monkeypatch, value): + monkeypatch.setenv("HF_HUB_OFFLINE", value) + assert hf_env_offline() is False + + +def test_hf_env_offline_honors_transformers_flag(monkeypatch): + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + assert hf_env_offline() is True + + +def test_hf_env_offline_default_false(): + assert hf_env_offline() is False + + +# ── st_repo_id_candidates ──────────────────────────────────────── + + +def test_candidates_slashless_adds_st_alias(): + assert st_repo_id_candidates("all-MiniLM-L6-v2") == [ + "all-MiniLM-L6-v2", + "sentence-transformers/all-MiniLM-L6-v2", + ] + + +def test_candidates_with_org_is_verbatim(): + assert st_repo_id_candidates("org/model") == ["org/model"] + + +def test_candidates_empty_name(): + assert st_repo_id_candidates(" ") == [] + + +# ── hf_cache_snapshot_dir ──────────────────────────────────────── + + +def test_snapshot_dir_resolves_active_commit(hf_cache): + snapshot = _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON}) + assert hf_cache_snapshot_dir("org/emb") == snapshot + + +def test_snapshot_dir_none_when_uncached(hf_cache): + assert hf_cache_snapshot_dir("org/missing") is None + + +def test_snapshot_dir_uses_st_alias_for_slashless(hf_cache): + snapshot = _make_cache( + hf_cache, "sentence-transformers/all-MiniLM-L6-v2", {"modules.json": MODULES_JSON} + ) + assert hf_cache_snapshot_dir("all-MiniLM-L6-v2") == snapshot + + +def test_snapshot_dir_none_when_snapshot_missing(hf_cache): + from huggingface_hub.file_download import repo_folder_name + + repo_dir = hf_cache / repo_folder_name(repo_id = "org/broken", repo_type = "model") + (repo_dir / "refs").mkdir(parents = True) + (repo_dir / "refs" / "main").write_text("deadbeef") # no snapshots/deadbeef dir + assert hf_cache_snapshot_dir("org/broken") is None + + +def test_snapshot_dir_expands_env_vars_in_cache_path(tmp_path, monkeypatch): + # An unexpanded $VAR in HF_HUB_CACHE must resolve where the loader looks. + real = tmp_path / "hub" + real.mkdir() + monkeypatch.setenv("MY_HF_CACHE", str(real)) + monkeypatch.setenv("HF_HUB_CACHE", "$MY_HF_CACHE") + monkeypatch.delenv("HF_HOME", raising = False) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + snapshot = _make_cache(real, "org/emb", {"modules.json": MODULES_JSON}) + assert hf_cache_snapshot_dir("org/emb") == snapshot + + +def test_snapshot_dir_uses_sentence_transformers_home(tmp_path, monkeypatch): + # ST uses SENTENCE_TRANSFORMERS_HOME as its cache_folder, so the gate must inspect it too. + st_home = tmp_path / "st_home" + st_home.mkdir() + monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home)) + monkeypatch.delenv("HF_HUB_CACHE", raising = False) + monkeypatch.delenv("HF_HOME", raising = False) + snapshot = _make_cache(st_home, "org/emb", {"modules.json": MODULES_JSON}) + assert hf_cache_snapshot_dir("org/emb") == snapshot + + +def test_snapshot_dir_st_home_is_exclusive(tmp_path, monkeypatch): + # With SENTENCE_TRANSFORMERS_HOME set, ST loads only from it, so a model living only under + # HF_HUB_CACHE must not be reported. + st_home = tmp_path / "st_home" + st_home.mkdir() + hub = tmp_path / "hub" + hub.mkdir() + monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home)) + monkeypatch.setenv("HF_HUB_CACHE", str(hub)) + monkeypatch.delenv("HF_HOME", raising = False) + _make_cache(hub, "org/emb", {"modules.json": MODULES_JSON}) # only in the HF hub cache + assert hf_cache_snapshot_dir("org/emb") is None + + +def test_snapshot_is_loadable_with_config_and_weights(hf_cache): + _make_cache(hf_cache, "org/emb", {"config.json": "{}", "model.safetensors": "x"}) + assert hf_cache_snapshot_is_loadable("org/emb") is True + + +def test_snapshot_is_not_loadable_when_metadata_only(hf_cache): + # A partial cache (refs/main resolves but no weights) is not loadable. + _make_cache(hf_cache, "org/partial", {"config.json": "{}", "modules.json": MODULES_JSON}) + assert hf_cache_snapshot_is_loadable("org/partial") is False + + +def test_snapshot_is_not_loadable_when_uncached(hf_cache): + assert hf_cache_snapshot_is_loadable("org/missing") is False + + +def test_gate_blocks_pickle_in_sentence_transformers_home(tmp_path, monkeypatch): + # A pickle under SENTENCE_TRANSFORMERS_HOME must still fail closed offline. + st_home = tmp_path / "st_home" + st_home.mkdir() + monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home)) + monkeypatch.delenv("HF_HUB_CACHE", raising = False) + monkeypatch.delenv("HF_HOME", raising = False) + _make_cache(st_home, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"}) + with _no_network(): + assert evaluate_file_security("org/pk", local_only_load = True).blocked is True + + +# ── is_embedding_model: offline (no network) ───────────────────── + + +def test_offline_true_for_cached_st_model(hf_cache, monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON, "config.json": "{}"}) + with _no_network(): + assert _is_embedding_model("org/emb") is True + + +def test_offline_false_for_cached_non_st_model(hf_cache, monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _make_cache(hf_cache, "org/plain", {"config.json": "{}", "model.safetensors": "x"}) + with _no_network(): + assert _is_embedding_model("org/plain") is False + + +def test_offline_false_when_uncached(hf_cache, monkeypatch): + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + with _no_network(): + assert _is_embedding_model("org/missing") is False + + +def test_offline_slashless_resolves_via_alias(hf_cache, monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _make_cache(hf_cache, "sentence-transformers/all-MiniLM-L6-v2", {"modules.json": MODULES_JSON}) + with _no_network(): + assert _is_embedding_model("all-MiniLM-L6-v2") is True + + +def test_offline_ignores_stale_online_memo(hf_cache, monkeypatch): + # An online lookup memoizes True for an UNCACHED repo (tags say embedding, no weights). Once + # offline, is_embedding_model must reclassify from the empty cache and return False, not the + # stale online True that would make settings accept a repo _get() cannot load. + with patch( + "huggingface_hub.model_info", + side_effect = lambda *a, **k: SimpleNamespace( + tags = ["sentence-transformers"], pipeline_tag = None + ), + ): + assert _is_embedding_model("org/uncached-emb") is True # memoized True online + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + with _no_network(): + assert _is_embedding_model("org/uncached-emb") is False # recomputed from empty cache + + +def test_offline_recomputes_after_cache_materializes(hf_cache, monkeypatch): + # Because the offline branch never records a memo, once an uncached repo's snapshot + # materializes (another process populates the cache) the next call re-reports True. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + with _no_network(): + assert _is_embedding_model("org/later") is False # uncached + _make_cache(hf_cache, "org/later", {"modules.json": MODULES_JSON}) + assert _is_embedding_model("org/later") is True # cache now present, no stale negative + + +# ── is_embedding_model: online (bounded + fallback) ────────────── + + +def test_online_passes_bounded_timeout(hf_cache): + seen = {} + + def _mi( + name, + token = None, + timeout = None, + **kw, + ): + seen["timeout"] = timeout + return SimpleNamespace(tags = ["sentence-transformers"], pipeline_tag = None) + + with patch("huggingface_hub.model_info", side_effect = _mi): + assert _is_embedding_model("org/emb") is True + assert seen["timeout"] == 15.0 + + +def test_online_error_falls_back_to_cache_marker(hf_cache): + _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON}) + with patch("huggingface_hub.model_info", side_effect = RuntimeError("dns dead")): + assert _is_embedding_model("org/emb") is True + + +def test_online_error_without_cache_returns_false(hf_cache): + with patch("huggingface_hub.model_info", side_effect = RuntimeError("dns dead")): + assert _is_embedding_model("org/missing") is False + + +# ── evaluate_file_security: offline fail-closed gate ───────────── + + +def _offline_decision(name): + return evaluate_file_security(name, local_only_load = True) + + +def test_gate_allows_safetensors_only(hf_cache): + _make_cache(hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"}) + with _no_network(): + assert _offline_decision("org/st").blocked is False + + +def test_gate_blocks_pickle_without_safetensors(hf_cache): + _make_cache(hf_cache, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"}) + with _no_network(): + decision = _offline_decision("org/pk") + assert decision.blocked is True + assert any(u["path"] == "pytorch_model.bin" for u in decision.unsafe_files) + + +def test_gate_allows_pickle_with_safetensors_sibling(hf_cache): + _make_cache(hf_cache, "org/both", {"pytorch_model.bin": "x", "model.safetensors": "y"}) + with _no_network(): + assert _offline_decision("org/both").blocked is False + + +def test_gate_blocks_sharded_pickle(hf_cache): + _make_cache( + hf_cache, + "org/shard", + { + "pytorch_model-00001-of-00002.bin": "a", + "pytorch_model-00002-of-00002.bin": "b", + }, + ) + with _no_network(): + assert _offline_decision("org/shard").blocked is True + + +def test_gate_allows_nothing_cached(hf_cache): + with _no_network(): + assert _offline_decision("org/missing").blocked is False + + +def test_gate_allows_gguf_only(hf_cache): + _make_cache(hf_cache, "org/gg", {"model.gguf": "x"}) + with _no_network(): + assert _offline_decision("org/gg").blocked is False + + +def test_gate_blocks_pickle_in_module_subdir(hf_cache): + # 0_Transformer is a module load root (listed in modules.json), so its pickle blocks. + _make_cache( + hf_cache, + "org/mod", + {"modules.json": _modules_json("0_Transformer"), "0_Transformer/pytorch_model.bin": "x"}, + ) + with _no_network(): + assert _offline_decision("org/mod").blocked is True + + +def test_gate_allows_pickle_in_subdir_with_safetensors(hf_cache): + _make_cache( + hf_cache, + "org/mod2", + { + "modules.json": _modules_json("0_Transformer"), + "0_Transformer/pytorch_model.bin": "x", + "0_Transformer/model.safetensors": "y", + }, + ) + with _no_network(): + assert _offline_decision("org/mod2").blocked is False + + +def test_gate_allows_unreferenced_nested_pickle(hf_cache): + # A pickle in a dir NOT referenced by modules.json (e.g. nemo/) is never deserialized, so it + # must not block the offline load (matches the online gate). + _make_cache( + hf_cache, + "org/aux", + { + "modules.json": MODULES_JSON, # Transformer at the root only + "model.safetensors": "w", + "nemo/pytorch_model.bin": "x", + }, + ) + with _no_network(): + assert _offline_decision("org/aux").blocked is False + + +def test_gate_blocks_adapter_pickle_without_safetensors(hf_cache): + _make_cache(hf_cache, "org/ad", {"config.json": "{}", "adapter_model.bin": "x"}) + with _no_network(): + decision = _offline_decision("org/ad") + assert decision.blocked is True + assert any(u["path"] == "adapter_model.bin" for u in decision.unsafe_files) + + +def test_gate_allows_adapter_pickle_with_adapter_safetensors(hf_cache): + _make_cache(hf_cache, "org/ad2", {"adapter_model.bin": "x", "adapter_model.safetensors": "y"}) + with _no_network(): + assert _offline_decision("org/ad2").blocked is False + + +def test_gate_blocks_base_pickle_with_only_adapter_safetensors_decoy(hf_cache): + # A decoy adapter_model.safetensors must NOT suppress a base pytorch_model.bin (the base + # loader would still deserialize the unscanned pickle). + _make_cache(hf_cache, "org/decoy", {"pytorch_model.bin": "x", "adapter_model.safetensors": "y"}) + with _no_network(): + assert _offline_decision("org/decoy").blocked is True + + +def test_gate_blocks_adapter_pickle_with_only_base_safetensors_decoy(hf_cache): + # Symmetric: a base model.safetensors must NOT suppress an adapter_model.bin. + _make_cache(hf_cache, "org/decoy2", {"adapter_model.bin": "x", "model.safetensors": "y"}) + with _no_network(): + assert _offline_decision("org/decoy2").blocked is True + + +def test_gate_reports_snapshot_relative_path(hf_cache): + _make_cache( + hf_cache, + "org/mod3", + {"modules.json": _modules_json("0_Transformer"), "0_Transformer/pytorch_model.bin": "x"}, + ) + with _no_network(): + decision = _offline_decision("org/mod3") + assert decision.blocked is True + assert any(u["path"] == "0_Transformer/pytorch_model.bin" for u in decision.unsafe_files) + + +# ── evaluate_file_security: online path unchanged ──────────────── + + +def test_online_default_blocks_unsafe(): + status = { + "scansDone": True, + "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}], + } + with patch( + "huggingface_hub.model_info", + side_effect = lambda *a, **k: SimpleNamespace(security_repo_status = status), + ): + assert evaluate_file_security("org/x").blocked is True + + +def test_online_default_allows_clean(): + status = {"scansDone": True, "filesWithIssues": []} + with patch( + "huggingface_hub.model_info", + side_effect = lambda *a, **k: SimpleNamespace(security_repo_status = status), + ): + assert evaluate_file_security("org/x").blocked is False + + +# ── embeddings guard + loader ──────────────────────────────────── + + +def test_guard_offline_blocks_pickle_only(hf_cache): + from core.rag.embeddings import UnsafeEmbeddingModelError, _guard_model_security + _make_cache(hf_cache, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"}) + with _no_network(): + with pytest.raises(UnsafeEmbeddingModelError): + _guard_model_security("org/pk", local_only = True) + + +def test_guard_offline_allows_safetensors(hf_cache): + from core.rag.embeddings import _guard_model_security + _make_cache(hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"}) + with _no_network(): + _guard_model_security("org/st", local_only = True) # must not raise + + +def _install_fake_sentence_transformers(monkeypatch, captured): + class FakeSentenceTransformer: + def __init__( + self, + name, + *, + device = None, + model_kwargs = None, + local_files_only = False, + **kw, + ): + captured["name"] = name + captured["device"] = device + captured["local_files_only"] = local_files_only + + module = types.ModuleType("sentence_transformers") + module.SentenceTransformer = FakeSentenceTransformer + monkeypatch.setitem(sys.modules, "sentence_transformers", module) + + +def test_get_offline_loads_from_local_snapshot(hf_cache, monkeypatch): + from core.rag import embeddings + + snapshot = _make_cache( + hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"} + ) + # TRANSFORMERS_OFFLINE only: a cached model loads from its local snapshot dir (a local path, + # never the Hub), offline-safe on ANY sentence-transformers version. + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.setattr(embeddings, "_model", None, raising = False) + monkeypatch.setattr(embeddings, "_name", None, raising = False) + monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) + monkeypatch.setattr(embeddings, "_device", lambda: "cpu") + captured = {} + _install_fake_sentence_transformers(monkeypatch, captured) + with _no_network(): + embeddings._get("org/st") + assert captured["name"] == str(snapshot) + + +def test_get_offline_uncached_uses_local_files_only(tmp_path, monkeypatch): + from core.rag import embeddings + + empty = tmp_path / "hub" + empty.mkdir() + monkeypatch.setenv("HF_HUB_CACHE", str(empty)) + monkeypatch.delenv("HF_HOME", raising = False) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.setattr(embeddings, "_model", None, raising = False) + monkeypatch.setattr(embeddings, "_name", None, raising = False) + monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) + monkeypatch.setattr(embeddings, "_device", lambda: "cpu") + # No cache -> repo-id load forced cache-only (fails fast offline, not a hang). + monkeypatch.setattr(embeddings, "_guard_model_security", lambda name, local_only = False: None) + captured = {} + _install_fake_sentence_transformers(monkeypatch, captured) + embeddings._get("org/uncached-xyz") + assert captured["name"] == "org/uncached-xyz" + assert captured["local_files_only"] is True + + +def test_get_online_omits_local_files_only(monkeypatch): + from core.rag import embeddings + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + monkeypatch.setattr(embeddings, "_model", None, raising = False) + monkeypatch.setattr(embeddings, "_name", None, raising = False) + monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) + monkeypatch.setattr(embeddings, "_device", lambda: "cpu") + # Isolate the loader wiring from the online guard's network calls. + monkeypatch.setattr(embeddings, "_guard_model_security", lambda name, local_only = False: None) + captured = {} + _install_fake_sentence_transformers(monkeypatch, captured) + embeddings._get("org/online") + assert captured["local_files_only"] is False diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 821529083d..50a997218f 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2076,6 +2076,24 @@ def download_gguf_file( _embedding_detection_cache: Dict[tuple, bool] = {} +# Bound the Hub lookup so a DNS-dead session fails fast to the cache instead of hanging on retries. +_HUB_MODEL_INFO_TIMEOUT = 15.0 + + +def _embedding_marker_in_hf_cache(model_name: str) -> bool: + """True when model_name's cached snapshot carries a modules.json (the ST marker). + Cache-only, no network; used offline and as a fallback when the Hub lookup times out.""" + from utils.utils import hf_cache_snapshot_dir + + snapshot = hf_cache_snapshot_dir(model_name) + if snapshot is None: + return False + try: + return (snapshot / "modules.json").is_file() + except OSError: + return False + + def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: """Detect embedding/sentence-transformer models via HF metadata. @@ -2090,6 +2108,15 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: Returns: True if embedding model, else False (default for local paths or errors). """ + from utils.utils import hf_env_offline + + # Offline (remote repo): reclassify from the local cache on every call, before/without the + # memo. An online lookup can memoize True from tags with no weights cached, so trusting it once + # the session goes offline would accept a repo _get() cannot load; a cached negative can also be + # invalidated by later cache materialization. The cache probe is local-only, so it's cheap. + if not is_local_path(model_name) and hf_env_offline(): + return _embedding_marker_in_hf_cache(model_name) + cache_key = (model_name, hf_token) if cache_key in _embedding_detection_cache: return _embedding_detection_cache[cache_key] @@ -2104,7 +2131,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: try: from huggingface_hub import model_info as hf_model_info - info = hf_model_info(model_name, token = hf_token) + info = hf_model_info(model_name, token = hf_token, timeout = _HUB_MODEL_INFO_TIMEOUT) tags = set(info.tags or []) pipeline_tag = info.pipeline_tag or "" @@ -2125,9 +2152,11 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: return is_emb except Exception as e: + # Timeout or transient network error: fall back to the local cache marker, don't hard-fail. logger.warning(f"Could not determine if {model_name} is embedding model: {e}") - _embedding_detection_cache[cache_key] = False - return False + is_emb = _embedding_marker_in_hf_cache(model_name) + _embedding_detection_cache[cache_key] = is_emb + return is_emb def _has_model_weight_files(model_dir: Path) -> bool: diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 466f326f18..0490d38d7c 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -29,13 +29,35 @@ Policy: scanned so a repo cannot dodge the gate by suffixing its name. """ +import re from dataclasses import dataclass, field +from pathlib import Path from typing import Optional from loggers import get_logger logger = get_logger(__name__) +# Pickle-format weight files (plain or sharded) that execute code on load; safetensors/gguf +# are inert. Grouped by weight family so an inert safetensors only suppresses the pickle it +# actually replaces: the loader won't use an adapter's safetensors for pytorch_model.bin. +_PICKLE_WEIGHT_RE = re.compile( + r"^(model|pytorch_model|adapter_model|consolidated)(-\d+-of-\d+)?" + r"\.(bin|pt|pth|ckpt|pkl|pickle)$", + re.IGNORECASE, +) +# Base-model safetensors set: HF names the base pickle pytorch_model.bin but the safetensors +# model.safetensors (stems differ), so a base pickle is replaced only by these, not an adapter's. +_BASE_SAFETENSORS_RE = re.compile( + r"^(model(-\d+-of-\d+)?\.safetensors|model\.safetensors\.index\.json)$", + re.IGNORECASE, +) +# Adapter (PEFT) safetensors set: adapter_model.safetensors, its shards, or index. +_ADAPTER_SAFETENSORS_RE = re.compile( + r"^(adapter_model(-\d+-of-\d+)?\.safetensors|adapter_model\.safetensors\.index\.json)$", + re.IGNORECASE, +) + # Non-blocking levels: clean or not-yet-finished. Anything else (unsafe/suspicious/ # malicious or a future label) blocks, so Hub schema drift fails CLOSED. _NONBLOCKING_LEVELS = frozenset( @@ -265,11 +287,105 @@ def _fetch_security_status(model_name: str, hf_token: Optional[str]): return None +def _st_load_roots(snapshot: Path) -> list: + """Directories a SentenceTransformer load deserializes weights from: the snapshot root plus + each module path in modules.json. Local, no network. Mirrors the online gate (which ignores + unreferenced nested pickles ST never loads) so the offline gate doesn't over-block.""" + roots = [snapshot] + try: + import json + modules = json.loads((snapshot / "modules.json").read_text()) + except (OSError, ValueError): + return roots # no / invalid modules.json -> snapshot root is the only load root + for module in modules or (): + path = str((module or {}).get("path", "")).strip().strip("/") + # Relative module path only; ignore a crafted "../" escape. + if path and ".." not in path.split("/"): + candidate = snapshot / path + if candidate not in roots: + roots.append(candidate) + return roots + + +def _cached_pickle_weight_files(snapshot: Path) -> list: + """Pickle weight files in snapshot's ST load roots, EXCLUDING those whose weight family also + ships an inert safetensors in the same dir (the loader prefers it): a base pickle is suppressed + only by a base model.safetensors, an adapter pickle only by adapter_model.safetensors -- an + unrelated safetensors is no substitute. Load roots only. Raises OSError if the snapshot root is + unreadable (caller blocks).""" + blocked = [] + for root in _st_load_roots(snapshot): + try: + entries = [p for p in root.iterdir() if p.is_file()] + except OSError: + if root == snapshot: + raise # top-level unreadable -> fail closed + continue # unreadable module subdir: nothing loadable to attest here + has_base_safetensors = any(_BASE_SAFETENSORS_RE.match(p.name) for p in entries) + has_adapter_safetensors = any(_ADAPTER_SAFETENSORS_RE.match(p.name) for p in entries) + for path in entries: + if not _PICKLE_WEIGHT_RE.match(path.name): + continue + is_adapter = path.name.lower().startswith("adapter_model") + has_alternative = has_adapter_safetensors if is_adapter else has_base_safetensors + if not has_alternative: + blocked.append(path) + return blocked + + +def _evaluate_local_only(model_name: str) -> FileSecurityDecision: + """Offline security gate. The Hub scan is unreachable, so inspect the local cache and fail + CLOSED on an unscanned pickle weight with no inert safetensors alternative, rather than + failing open or hanging. Safetensors/gguf-only cache loads; nothing cached -> allowed.""" + from utils.utils import hf_cache_snapshot_dir + + try: + snapshot = hf_cache_snapshot_dir(model_name) + except Exception: + logger.warning("Offline gate: could not resolve the cache for '%s'; blocking.", model_name) + return FileSecurityDecision( + model_name, True, reason = "offline; could not inspect the local cache" + ) + + if snapshot is None: + return FileSecurityDecision(model_name, False, reason = "offline; nothing cached to load") + + try: + pickles = _cached_pickle_weight_files(snapshot) + except OSError: + logger.warning("Offline gate: could not read the cache for '%s'; blocking.", model_name) + return FileSecurityDecision( + model_name, True, reason = "offline; could not read the local cache" + ) + + if not pickles: + return FileSecurityDecision( + model_name, False, reason = "offline; cached weights are inert (safetensors/gguf)" + ) + + # Snapshot-relative posix paths (match the online gate; disambiguate same-named pickles). + rel_paths = sorted(p.relative_to(snapshot).as_posix() for p in pickles) + names = ", ".join(rel_paths) + logger.warning( + "Blocking offline load of '%s': cached pickle weight(s) cannot be malware-scanned " + "offline and have no safetensors alternative (%s).", + model_name, + names, + ) + return FileSecurityDecision( + model_name, + True, + unsafe_files = [{"path": rel, "level": "unscanned"} for rel in rel_paths], + reason = f"offline; unscanned pickle weights with no safetensors alternative: {names}", + ) + + def evaluate_file_security( model_name: str, hf_token: Optional[str] = None, *, load_subdirs = (), + local_only_load: bool = False, ) -> FileSecurityDecision: """Block a load when HF's security scan flags unsafe serialized files. @@ -280,6 +396,9 @@ def evaluate_file_security( ``load_subdirs`` names subdirs the load calls ``from_pretrained`` on (e.g. ``("LLM",)`` for Spark-TTS / BiCodec, loading ``/LLM``): a flagged file directly under one is root-level there and blocks, and an index inside it is honored when scoping shards. + + ``local_only_load`` marks an offline load: with the Hub scan unreachable, inspect the local + cache and fail CLOSED on an unscanned pickle weight with no safetensors alternative. """ # Scan the repo the load actually fetches, not the literal alias (which 404s and # fails open): the Spark-TTS "/LLM" alias is really unsloth/ from LLM/. @@ -295,6 +414,10 @@ def evaluate_file_security( # Cannot classify the path -> do not block on that account. return FileSecurityDecision(model_name, False, reason = "path check failed; not blocked") + # Offline: inspect the local cache and fail closed rather than hang on model_info or fail open. + if local_only_load: + return _evaluate_local_only(model_name) + status = _fetch_security_status(model_name, hf_token) if not isinstance(status, dict): return FileSecurityDecision( diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index 31f5f31bee..21e11c6706 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -8,6 +8,7 @@ import structlog from loggers import get_logger from contextlib import contextmanager from pathlib import Path +from typing import Optional import shutil import tempfile @@ -15,6 +16,110 @@ import tempfile logger = get_logger(__name__) +# ── Offline / HF-cache helpers ────────────────────────────────── +# An offline load must never touch the network (a DNS-dead session hangs on hub retries); +# these read the local HF cache the load itself uses. + +_HF_OFFLINE_TRUE_VALUES = frozenset({"1", "true", "yes", "on"}) + + +def hf_env_offline() -> bool: + """True when HF_HUB_OFFLINE or TRANSFORMERS_OFFLINE requests offline mode. + + Also honors TRANSFORMERS_OFFLINE (hub honors only HF_HUB_OFFLINE) since users set it + to keep transformers loads local. + """ + for var in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"): + if os.environ.get(var, "").strip().lower() in _HF_OFFLINE_TRUE_VALUES: + return True + return False + + +def st_repo_id_candidates(model_name: str) -> list: + """Repo ids a Sentence-Transformers load may resolve model_name to; a slashless name + also resolves under the sentence-transformers/ namespace, so both are candidates.""" + name = (model_name or "").strip().strip("/") + if not name: + return [] + candidates = [name] + if "/" not in name: + candidates.append(f"sentence-transformers/{name}") + return candidates + + +def _expand_path(raw: str) -> Path: + """Expand ~ and $VARS as huggingface_hub does, so the gate resolves the loader's dir.""" + return Path(os.path.expandvars(os.path.expanduser(raw))) + + +def _hf_cache_roots() -> list: + """The one cache root the loader resolves to, by its own precedence (it picks ONE + cache_folder, no fall-through): SENTENCE_TRANSFORMERS_HOME, else HF_HUB_CACHE, else + HF_HOME/hub, else ~/.cache/huggingface/hub. Expanded, read from env, one-element list.""" + st_home = os.environ.get("SENTENCE_TRANSFORMERS_HOME") + if st_home: + return [_expand_path(st_home)] + hub = os.environ.get("HF_HUB_CACHE") or os.environ.get("HUGGINGFACE_HUB_CACHE") + if hub: + return [_expand_path(hub)] + hf_home = os.environ.get("HF_HOME") + if hf_home: + return [_expand_path(hf_home) / "hub"] + return [Path.home() / ".cache" / "huggingface" / "hub"] + + +def hf_cache_snapshot_dir(model_name: str) -> Optional[Path]: + """Active local snapshot dir for model_name's main revision, or None if not cached. + Reads refs/main then snapshots/; no network. Tries the ST alias for slashless names.""" + try: + from huggingface_hub.file_download import repo_folder_name + except Exception: + repo_folder_name = None + for cache_root in _hf_cache_roots(): + for repo_id in st_repo_id_candidates(model_name): + try: + if repo_folder_name is not None: + folder = repo_folder_name(repo_id = repo_id, repo_type = "model") + else: + folder = "models--" + repo_id.replace("/", "--") + repo_dir = cache_root / folder + ref = repo_dir / "refs" / "main" + if not ref.is_file(): + continue + commit = ref.read_text().strip() + if not commit: + continue + snapshot = repo_dir / "snapshots" / commit + if snapshot.is_dir(): + return snapshot + except OSError: + continue + return None + + +# A weight file plus a config distinguishes a real cached model from a metadata-only +# partial cache that resolves refs/main but would fail at load time. +_LOADABLE_WEIGHT_SUFFIXES = frozenset({".safetensors", ".bin", ".gguf", ".pt", ".pth", ".ckpt"}) + + +def hf_cache_snapshot_is_loadable(model_name: str) -> bool: + """True when model_name's snapshot is cached and loadable: a config (config.json or + modules.json) plus at least one weight file, not a metadata-only partial cache. No network.""" + snapshot = hf_cache_snapshot_dir(model_name) + if snapshot is None: + return False + try: + has_config = (snapshot / "config.json").is_file() or (snapshot / "modules.json").is_file() + if not has_config: + return False + for path in snapshot.rglob("*"): + if path.suffix.lower() in _LOADABLE_WEIGHT_SUFFIXES and path.is_file(): + return True + except OSError: + return False + return False + + # ── Client-safe error helpers ─────────────────────────────────── # Never return raw exception text to clients; log server-side, return generic.