From 033d60720b5e9d049bbec82f98b53cf6b53256e1 Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Fri, 3 Jul 2026 22:34:18 +0300 Subject: [PATCH 01/77] studio: classify embedding models from the HF cache and honor offline mode is_embedding_model() went straight to huggingface_hub.model_info() for any repo id, so in offline mode (no DNS, or HF_HUB_OFFLINE set) selecting an already-downloaded model hung on network retries that could never succeed and training/export never started (#6817). Check the local HF cache first: a sentence-transformers repo carries modules.json in its snapshot (the same marker used for local paths), so a cached model is classified with no network call. When HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE is set, anything not positively an embedding model returns False without a network call instead of retrying a doomed request. Online, uncached lookups still fall through to model_info(), so tag-only embedding models (feature-extraction) are unaffected. Adds _embedding_marker_in_hf_cache() over the existing _iter_hf_cache_snapshots. --- .../tests/test_is_embedding_model_offline.py | 116 ++++++++++++++++++ studio/backend/utils/models/model_config.py | 34 +++++ 2 files changed, 150 insertions(+) create mode 100644 studio/backend/tests/test_is_embedding_model_offline.py diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py new file mode 100644 index 00000000000..ff51061be63 --- /dev/null +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""is_embedding_model must classify an already-downloaded model from the local +HF cache and honour HF_HUB_OFFLINE, instead of making a model_info() network call +that hangs on DNS retries when offline (#6817).""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +_loggers_stub = types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +_structlog_stub = types.ModuleType("structlog") +_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") +sys.modules.setdefault("structlog", _structlog_stub) + +import utils.models.model_config as mc # noqa: E402 + + +@pytest.fixture(autouse = True) +def _clean_state(monkeypatch): + mc._embedding_detection_cache.clear() + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + yield + mc._embedding_detection_cache.clear() + + +def _snapshot(tmp_path, name, *, sentence_transformer): + d = tmp_path / name + d.mkdir(parents = True, exist_ok = True) + (d / "config.json").write_text("{}") + if sentence_transformer: + (d / "modules.json").write_text("[]") + return d + + +def _fake_hf_model_info(monkeypatch, fn): + fake = types.ModuleType("huggingface_hub") + fake.model_info = fn + monkeypatch.setitem(sys.modules, "huggingface_hub", fake) + + +def _no_network(*a, **k): + raise AssertionError("model_info() must not be called") + + +# ── _embedding_marker_in_hf_cache ── + + +def test_marker_true_when_modules_json_present(tmp_path, monkeypatch): + snap = _snapshot(tmp_path, "snap", sentence_transformer = True) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter([snap])) + assert mc._embedding_marker_in_hf_cache("org/emb") is True + + +def test_marker_false_when_cached_without_modules_json(tmp_path, monkeypatch): + snap = _snapshot(tmp_path, "snap", sentence_transformer = False) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter([snap])) + assert mc._embedding_marker_in_hf_cache("org/llm") is False + + +def test_marker_none_when_not_cached(monkeypatch): + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(())) + assert mc._embedding_marker_in_hf_cache("org/llm") is None + + +# ── is_embedding_model ── + + +def test_cached_sentence_transformer_skips_network(tmp_path, monkeypatch): + snap = _snapshot(tmp_path, "snap", sentence_transformer = True) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter([snap])) + _fake_hf_model_info(monkeypatch, _no_network) + assert mc.is_embedding_model("unsloth/bge-small-en-v1.5") is True + + +def test_offline_cached_non_st_returns_false_without_network(tmp_path, monkeypatch): + snap = _snapshot(tmp_path, "snap", sentence_transformer = False) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter([snap])) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _fake_hf_model_info(monkeypatch, _no_network) + assert mc.is_embedding_model("org/gemma-4-e4b") is False + + +def test_offline_not_cached_returns_false_without_network(monkeypatch): + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(())) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _fake_hf_model_info(monkeypatch, _no_network) + assert mc.is_embedding_model("org/never-downloaded") is False + + +def test_online_uncached_still_uses_network(monkeypatch): + # Not offline, not cached: the network model_info path must still run so an + # embedding model that lacks modules.json (feature-extraction tag) is caught. + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(())) + calls = [] + + def _info(model_name, token = None): + calls.append(model_name) + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + + _fake_hf_model_info(monkeypatch, _info) + assert mc.is_embedding_model("org/gte-modernbert") is True + assert calls == ["org/gte-modernbert"] diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 5d8458e5f08..46e683aa877 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1941,6 +1941,23 @@ def download_gguf_file( _embedding_detection_cache: Dict[tuple, bool] = {} +def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: + """Sentence-transformers detection from the local HF cache, no network call. + + True if a cached snapshot carries the ``modules.json`` marker (the same + signal ``is_embedding_model`` uses for local paths), False if the repo is + cached but no snapshot has it, None when the repo is not in the cache.""" + cached = False + for snap in _iter_hf_cache_snapshots(repo_id): + cached = True + try: + if (snap / "modules.json").is_file(): + return True + except OSError: + continue + return False if cached else None + + def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: """Detect embedding/sentence-transformer models via HF metadata. @@ -1966,6 +1983,23 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: _embedding_detection_cache[cache_key] = is_emb return is_emb + # Prefer the local HF cache: a sentence-transformers repo carries + # modules.json in its snapshot, so an already-downloaded model needs no + # network call. This also lets an offline / HF_HUB_OFFLINE session classify a + # cached model instead of hanging on a model_info() request that fails with a + # DNS error and only ever gets retried (#6817). + cache_hit = _embedding_marker_in_hf_cache(model_name) + if cache_hit is True: + _embedding_detection_cache[cache_key] = True + logger.info(f"Model {model_name} detected as embedding model via HF cache (modules.json)") + return True + if _env_offline(): + # Offline: the cache is the only source; anything not positively a + # sentence-transformers model is treated as non-embedding rather than + # making a network call that cannot succeed. + _embedding_detection_cache[cache_key] = False + return False + try: from huggingface_hub import model_info as hf_model_info From aabf0d06fcdfafd28a5a81d0a02de7432913a6a8 Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Fri, 17 Jul 2026 23:23:58 +0300 Subject: [PATCH 02/77] studio: judge the active cached revision, harden the cache probe, stop stub leaks Three review fixes on the cache-first embedding detection: 1. Prefer the revision refs/main resolves to. The HF cache keeps snapshots of older revisions, so an any-snapshot scan could classify a repo by a stale revision -- e.g. a repo that used to be a sentence-transformers model would short-circuit even the online lookup. When refs/main is recorded, only its snapshot is consulted; the newest-first scan remains the fallback for caches with no ref. 2. Keep the cache probe inside the detection error boundary. The snapshot iterator stat()s entries and could raise if a cached model is deleted concurrently, propagating a 500 out of the config/check-embedding routes. _embedding_marker_in_hf_cache now catches everything and reads as not-cached, so callers keep their normal Hub/offline fallback. 3. Stub loggers/structlog in the test only when the real modules are absent (try-import, mirroring test_windows_gpu_detection_mock), so collecting this file first can no longer shadow the real packages for later tests in the same pytest process. --- .../tests/test_is_embedding_model_offline.py | 122 +++++++++++++++--- studio/backend/utils/models/model_config.py | 42 ++++-- 2 files changed, 133 insertions(+), 31 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index ff51061be63..ff92c6eb7f5 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -7,6 +7,7 @@ from __future__ import annotations +import importlib import sys import types from pathlib import Path @@ -17,13 +18,30 @@ if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -_loggers_stub = types.ModuleType("loggers") -_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) -sys.modules.setdefault("loggers", _loggers_stub) -_structlog_stub = types.ModuleType("structlog") -_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") -sys.modules.setdefault("structlog", _structlog_stub) +def _maybe_stub(name: str, builder): + # Stub only if the real module is unavailable, so this file never shadows + # real packages for later tests in the same pytest process. + try: + importlib.import_module(name) + except ImportError: + sys.modules[name] = builder() + + +def _build_loggers_stub(): + m = types.ModuleType("loggers") + m.get_logger = lambda name: __import__("logging").getLogger(name) + return m + + +def _build_structlog_stub(): + m = types.ModuleType("structlog") + m.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") + return m + + +_maybe_stub("loggers", _build_loggers_stub) +_maybe_stub("structlog", _build_structlog_stub) import utils.models.model_config as mc # noqa: E402 @@ -37,13 +55,27 @@ def _clean_state(monkeypatch): mc._embedding_detection_cache.clear() -def _snapshot(tmp_path, name, *, sentence_transformer): - d = tmp_path / name - d.mkdir(parents = True, exist_ok = True) - (d / "config.json").write_text("{}") - if sentence_transformer: - (d / "modules.json").write_text("[]") - return d +def _repo(tmp_path, *snapshots, main_ref = None): + """Fake HF cache repo dir: snapshots/[/modules.json] (+ refs/main). + + ``snapshots``: (name, sentence_transformer) tuples, oldest last (the + iterator under test yields newest first, so pass them in that order). + Returns the snapshot dirs in the given order. + """ + repo = tmp_path / "models--org--model" + dirs = [] + for name, is_st in snapshots: + d = repo / "snapshots" / name + d.mkdir(parents = True, exist_ok = True) + (d / "config.json").write_text("{}") + if is_st: + (d / "modules.json").write_text("[]") + dirs.append(d) + if main_ref is not None: + refs = repo / "refs" + refs.mkdir(parents = True, exist_ok = True) + (refs / "main").write_text(main_ref) + return dirs def _fake_hf_model_info(monkeypatch, fn): @@ -60,14 +92,14 @@ def _no_network(*a, **k): def test_marker_true_when_modules_json_present(tmp_path, monkeypatch): - snap = _snapshot(tmp_path, "snap", sentence_transformer = True) - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter([snap])) + snaps = _repo(tmp_path, ("aaa", True)) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) assert mc._embedding_marker_in_hf_cache("org/emb") is True def test_marker_false_when_cached_without_modules_json(tmp_path, monkeypatch): - snap = _snapshot(tmp_path, "snap", sentence_transformer = False) - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter([snap])) + snaps = _repo(tmp_path, ("aaa", False)) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) assert mc._embedding_marker_in_hf_cache("org/llm") is False @@ -76,19 +108,67 @@ def test_marker_none_when_not_cached(monkeypatch): assert mc._embedding_marker_in_hf_cache("org/llm") is None +def test_marker_prefers_refs_main_revision(tmp_path, monkeypatch): + # The repo USED to be a sentence-transformers model (old snapshot has + # modules.json) but the revision refs/main points at no longer is. The + # active revision must win: an any-snapshot scan would wrongly say True. + snaps = _repo(tmp_path, ("new", False), ("old", True), main_ref = "new") + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + assert mc._embedding_marker_in_hf_cache("org/was-embedder") is False + + +def test_marker_refs_main_st_revision_is_true(tmp_path, monkeypatch): + snaps = _repo(tmp_path, ("new", True), ("old", False), main_ref = "new") + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + assert mc._embedding_marker_in_hf_cache("org/is-embedder") is True + + +def test_marker_missing_ref_falls_back_to_snapshot_scan(tmp_path, monkeypatch): + # No refs/main recorded: keep the newest-first any-snapshot behavior. + snaps = _repo(tmp_path, ("new", False), ("old", True)) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + assert mc._embedding_marker_in_hf_cache("org/no-ref") is True + + +def test_marker_never_raises_when_cache_mutates(monkeypatch): + # A snapshot vanishing mid-iteration (concurrent cached-model deletion) + # must read as not-cached, not propagate a 500 out of the routes. + def _exploding_iter(repo): + raise FileNotFoundError("snapshot removed underneath") + + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", _exploding_iter) + assert mc._embedding_marker_in_hf_cache("org/racing") is None + + +def test_is_embedding_model_survives_cache_race_online(monkeypatch): + # With the cache probe failing, the online path must still resolve via the + # Hub instead of erroring out. + def _exploding_iter(repo): + raise FileNotFoundError("snapshot removed underneath") + + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", _exploding_iter) + _fake_hf_model_info( + monkeypatch, + lambda name, token = None: types.SimpleNamespace( + tags = ["sentence-transformers"], pipeline_tag = None + ), + ) + assert mc.is_embedding_model("org/racing") is True + + # ── is_embedding_model ── def test_cached_sentence_transformer_skips_network(tmp_path, monkeypatch): - snap = _snapshot(tmp_path, "snap", sentence_transformer = True) - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter([snap])) + snaps = _repo(tmp_path, ("aaa", True)) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) _fake_hf_model_info(monkeypatch, _no_network) assert mc.is_embedding_model("unsloth/bge-small-en-v1.5") is True def test_offline_cached_non_st_returns_false_without_network(tmp_path, monkeypatch): - snap = _snapshot(tmp_path, "snap", sentence_transformer = False) - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter([snap])) + snaps = _repo(tmp_path, ("aaa", False)) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) monkeypatch.setenv("HF_HUB_OFFLINE", "1") _fake_hf_model_info(monkeypatch, _no_network) assert mc.is_embedding_model("org/gemma-4-e4b") is False diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 46e683aa877..b4192ee6861 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1944,18 +1944,40 @@ def download_gguf_file( def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: """Sentence-transformers detection from the local HF cache, no network call. - True if a cached snapshot carries the ``modules.json`` marker (the same - signal ``is_embedding_model`` uses for local paths), False if the repo is - cached but no snapshot has it, None when the repo is not in the cache.""" - cached = False - for snap in _iter_hf_cache_snapshots(repo_id): - cached = True + True/False when the ACTIVE cached revision carries / lacks the + ``modules.json`` marker (the same signal ``is_embedding_model`` uses for + local paths), None when the repo is not in the cache (or the cache is + unreadable). The revision ``refs/main`` resolves to is authoritative when + recorded: the cache keeps snapshots of older revisions, and a repo that + later stopped (or started) being a sentence-transformers model must be + judged by its current revision, not any historical one. Snapshots are only + scanned newest-first when no ref exists. Never raises -- a cache mutating + underneath (concurrent model deletion) reads as not-cached so callers keep + their normal fallback.""" + try: + snapshots = list(_iter_hf_cache_snapshots(repo_id)) + if not snapshots: + return None + # Prefer the snapshot refs/main points at (the active revision). + snapshots_dir = snapshots[0].parent try: - if (snap / "modules.json").is_file(): - return True + commit = (snapshots_dir.parent / "refs" / "main").read_text( + encoding = "utf-8" + ).strip() + preferred = snapshots_dir / commit + if commit and preferred.is_dir(): + return (preferred / "modules.json").is_file() except OSError: - continue - return False if cached else None + pass # no ref recorded: fall back to the newest-first scan + for snap in snapshots: + try: + if (snap / "modules.json").is_file(): + return True + except OSError: + continue + return False + except Exception: + return None def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: From 763ef72ceb352a8ae7d4f3b58125756a467087dd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 17 Jul 2026 20:24:50 +0000 Subject: [PATCH 03/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_is_embedding_model_offline.py | 6 +++++- studio/backend/utils/models/model_config.py | 4 +--- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index ff92c6eb7f5..bc5a9b2ce84 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -55,7 +55,11 @@ def _clean_state(monkeypatch): mc._embedding_detection_cache.clear() -def _repo(tmp_path, *snapshots, main_ref = None): +def _repo( + tmp_path, + *snapshots, + main_ref = None, +): """Fake HF cache repo dir: snapshots/[/modules.json] (+ refs/main). ``snapshots``: (name, sentence_transformer) tuples, oldest last (the diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index b4192ee6861..02a07049690 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1961,9 +1961,7 @@ def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: # Prefer the snapshot refs/main points at (the active revision). snapshots_dir = snapshots[0].parent try: - commit = (snapshots_dir.parent / "refs" / "main").read_text( - encoding = "utf-8" - ).strip() + commit = (snapshots_dir.parent / "refs" / "main").read_text(encoding = "utf-8").strip() preferred = snapshots_dir / commit if commit and preferred.is_dir(): return (preferred / "modules.json").is_file() From fb0cfae97f169d9243226e7ae70bcfb379abad0f Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sat, 18 Jul 2026 05:42:29 +0300 Subject: [PATCH 04/77] studio: treat a missing active-ref snapshot as a cache miss, don't cache offline misses Two review fixes on the cache-first embedding detection: 1. When refs/main is recorded but points at a commit whose snapshot dir is absent (partial download / cache pruning), the recorded ref is still authoritative: return None (cache miss) instead of falling through to scan older snapshots, which could report a stale historical revision's modules.json as the active one -- the same stale-cache class this helper avoids. 2. Do not cache the offline negative. When HF_HUB_OFFLINE/TRANSFORMERS_OFFLINE is set and the repo is not positively an ST model from modules.json, is_embedding_model stored False under the (model_name, hf_token) key shared with online lookups; after the env var cleared in the same process, a tag-only (feature-extraction) embedder returned the cached False and never reached model_info(). The offline negative is now returned without caching. --- .../tests/test_is_embedding_model_offline.py | 33 +++++++++++++++++++ studio/backend/utils/models/model_config.py | 20 +++++++---- 2 files changed, 47 insertions(+), 6 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index bc5a9b2ce84..8caa500e034 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -134,6 +134,16 @@ def test_marker_missing_ref_falls_back_to_snapshot_scan(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache("org/no-ref") is True +def test_marker_ref_points_at_absent_snapshot_is_cache_miss(tmp_path, monkeypatch): + # refs/main names a commit whose snapshot dir is absent (partial download / + # pruning). The recorded ref is authoritative, so this is a cache miss + # (None) -- NOT a fall-through to a stale historical snapshot that has + # modules.json. + snaps = _repo(tmp_path, ("old", True), main_ref = "missing_commit") + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + assert mc._embedding_marker_in_hf_cache("org/partial") is None + + def test_marker_never_raises_when_cache_mutates(monkeypatch): # A snapshot vanishing mid-iteration (concurrent cached-model deletion) # must read as not-cached, not propagate a 500 out of the routes. @@ -198,3 +208,26 @@ def _info(model_name, token = None): _fake_hf_model_info(monkeypatch, _info) assert mc.is_embedding_model("org/gte-modernbert") is True assert calls == ["org/gte-modernbert"] + + +def test_offline_negative_is_not_cached_then_online_detects(monkeypatch): + # A tag-only embedder is not identifiable from modules.json. Offline returns + # False WITHOUT caching, so once the env var clears the online model_info + # lookup still runs and detects it -- the negative must not be sticky. + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(())) + calls = [] + + def _info(model_name, token = None): + calls.append(model_name) + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + + _fake_hf_model_info(monkeypatch, _info) + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + assert mc.is_embedding_model("org/gte-modernbert") is False + assert calls == [] # offline: no network + assert ("org/gte-modernbert", None) not in mc._embedding_detection_cache + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + assert mc.is_embedding_model("org/gte-modernbert") is True # now detected online + assert calls == ["org/gte-modernbert"] diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 02a07049690..3c0df04e663 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1962,11 +1962,17 @@ def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: snapshots_dir = snapshots[0].parent try: commit = (snapshots_dir.parent / "refs" / "main").read_text(encoding = "utf-8").strip() - preferred = snapshots_dir / commit - if commit and preferred.is_dir(): - return (preferred / "modules.json").is_file() except OSError: - pass # no ref recorded: fall back to the newest-first scan + commit = "" # no ref recorded: fall back to the newest-first scan + if commit: + # A ref is recorded, so it is authoritative. If its snapshot is not + # materialized (partial download / pruning) treat the repo as not + # cached (None) rather than scanning stale history -- the exact + # stale-cache class this helper avoids. + preferred = snapshots_dir / commit + if not preferred.is_dir(): + return None + return (preferred / "modules.json").is_file() for snap in snapshots: try: if (snap / "modules.json").is_file(): @@ -2016,8 +2022,10 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: if _env_offline(): # Offline: the cache is the only source; anything not positively a # sentence-transformers model is treated as non-embedding rather than - # making a network call that cannot succeed. - _embedding_detection_cache[cache_key] = False + # making a network call that cannot succeed. Do NOT cache this negative: + # it is offline-conditional. A tag-only (feature-extraction) embedder can + # only be confirmed online, and the same (model_name, hf_token) key is + # reused for online lookups after the env var clears in this process. return False try: From 21bef982a5af6c0c8e4a085a97bc7bd0b9b16d46 Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sat, 18 Jul 2026 07:01:36 +0300 Subject: [PATCH 05/77] studio: defer online embedding detection to the Hub, re-probe offline The local modules.json marker short-circuited is_embedding_model() even online, so a repo that dropped (or added) the marker since it was cached was judged by its stale local revision instead of the current remote one. Online now treats model_info() as authoritative and uses the cache marker only as an uncached fallback when the Hub is unreachable, so a transient failure never poisons the memo. Offline re-probes the marker on every call without consulting or populating the memo, so a model downloaded later in the session (or a cached online negative that predates the download) is detected. _embedding_marker_in_hf_cache() now treats an unreadable refs/main (a non-FileNotFoundError OSError) as a cache miss rather than scanning stale history -- only a genuinely missing ref enables the fallback scan. --- .../tests/test_is_embedding_model_offline.py | 71 ++++++++++++++++++- studio/backend/utils/models/model_config.py | 62 +++++++++------- 2 files changed, 108 insertions(+), 25 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 8caa500e034..ee9c6842c5b 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -144,6 +144,21 @@ def test_marker_ref_points_at_absent_snapshot_is_cache_miss(tmp_path, monkeypatc assert mc._embedding_marker_in_hf_cache("org/partial") is None +def test_marker_unreadable_ref_is_cache_miss(tmp_path, monkeypatch): + # refs/main exists but cannot be read (transient I/O error / restrictive + # permissions). Only a genuinely MISSING ref may enable the historical scan; + # an unreadable ref is a cache miss (None), never a fall-through to a stale + # snapshot that happens to carry modules.json. + snaps = _repo(tmp_path, ("old", True)) + # Make refs/main a directory so read_text() raises IsADirectoryError -- an + # OSError that is NOT FileNotFoundError, i.e. "exists but unreadable". + refs_main = snaps[0].parent.parent / "refs" / "main" + refs_main.parent.mkdir(parents = True, exist_ok = True) + refs_main.mkdir() + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + assert mc._embedding_marker_in_hf_cache("org/unreadable-ref") is None + + def test_marker_never_raises_when_cache_mutates(monkeypatch): # A snapshot vanishing mid-iteration (concurrent cached-model deletion) # must read as not-cached, not propagate a 500 out of the routes. @@ -173,13 +188,67 @@ def _exploding_iter(repo): # ── is_embedding_model ── -def test_cached_sentence_transformer_skips_network(tmp_path, monkeypatch): +def test_offline_cached_st_detected_via_marker_no_network(tmp_path, monkeypatch): + # Offline: a downloaded sentence-transformers repo is classified from its + # modules.json marker with no model_info() network call that would hang on + # DNS retries (#6817). snaps = _repo(tmp_path, ("aaa", True)) monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") _fake_hf_model_info(monkeypatch, _no_network) assert mc.is_embedding_model("unsloth/bge-small-en-v1.5") is True +def test_online_defers_to_hub_over_stale_marker(tmp_path, monkeypatch): + # Online: the Hub is authoritative for the current revision. Even with a + # cached modules.json (the repo WAS an embedder), a Hub lookup that no longer + # reports embedding signals wins -- the stale local marker must not + # short-circuit model_info(). + snaps = _repo(tmp_path, ("aaa", True)) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + calls = [] + + def _info(model_name, token = None): + calls.append(model_name) + return types.SimpleNamespace(tags = ["text-generation"], pipeline_tag = "text-generation") + + _fake_hf_model_info(monkeypatch, _info) + assert mc.is_embedding_model("org/was-embedder") is False + assert calls == ["org/was-embedder"] # Hub consulted, not skipped + + +def test_online_hub_failure_falls_back_to_marker_uncached(tmp_path, monkeypatch): + # A transient model_info() failure falls back to the local marker WITHOUT + # caching: the degraded result must not become sticky, so a later successful + # Hub lookup can still override it. + snaps = _repo(tmp_path, ("aaa", True)) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + _fake_hf_model_info(monkeypatch, _no_network) # raises -> Hub "unreachable" + assert mc.is_embedding_model("org/emb") is True # marker fallback + assert ("org/emb", None) not in mc._embedding_detection_cache # not poisoned + + +def test_online_negative_does_not_block_later_offline_download(tmp_path, monkeypatch): + # An online Hub lookup authoritatively reports non-embedding and is memoized. + # The repo is then downloaded WITH modules.json and the session goes offline; + # the offline path re-probes the marker (never consulting the online memo), + # so the freshly downloaded embedder is detected instead of the stale False. + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(())) + + def _info(model_name, token = None): + return types.SimpleNamespace(tags = ["text-generation"], pipeline_tag = "text-generation") + + _fake_hf_model_info(monkeypatch, _info) + assert mc.is_embedding_model("org/late-embedder") is False # online: Hub says no + assert mc._embedding_detection_cache[("org/late-embedder", None)] is False + + # Now the model is downloaded (marker appears) and the session goes offline. + snaps = _repo(tmp_path, ("aaa", True)) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + assert mc.is_embedding_model("org/late-embedder") is True # marker re-probed + + def test_offline_cached_non_st_returns_false_without_network(tmp_path, monkeypatch): snaps = _repo(tmp_path, ("aaa", False)) monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 2ce93642aec..c13c38d51da 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -992,7 +992,7 @@ def _is_vision_model_uncached( and "<|text_start|>" in tokens and "<|text_end|>" in tokens ), - "snac": lambda tokens: (sum(1 for t in tokens if t.startswith(" 10000), + "snac": lambda tokens: sum(1 for t in tokens if t.startswith(" 10000, } @@ -2026,8 +2026,14 @@ def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: snapshots_dir = snapshots[0].parent try: commit = (snapshots_dir.parent / "refs" / "main").read_text(encoding = "utf-8").strip() - except OSError: + except FileNotFoundError: commit = "" # no ref recorded: fall back to the newest-first scan + except OSError: + # A ref exists but is unreadable (transient I/O error, restrictive + # permissions): the contract is that an unreadable cache reads as + # not-cached, so report a miss rather than scanning stale history -- + # only a genuinely missing ref may enable the fallback scan. + return None if commit: # A ref is recorded, so it is authoritative. If its snapshot is not # materialized (partial download / pruning) treat the repo as not @@ -2063,35 +2069,35 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: True if embedding model, else False (default for local paths or errors). """ cache_key = (model_name, hf_token) - if cache_key in _embedding_detection_cache: - return _embedding_detection_cache[cache_key] - # Local paths: check for sentence-transformer marker (modules.json) + # Local paths: check for sentence-transformer marker (modules.json). This is + # authoritative for an explicit path, so memoize it. if is_local_path(model_name): + if cache_key in _embedding_detection_cache: + return _embedding_detection_cache[cache_key] local_dir = normalize_path(model_name) is_emb = os.path.isfile(os.path.join(local_dir, "modules.json")) _embedding_detection_cache[cache_key] = is_emb return is_emb - # Prefer the local HF cache: a sentence-transformers repo carries - # modules.json in its snapshot, so an already-downloaded model needs no - # network call. This also lets an offline / HF_HUB_OFFLINE session classify a - # cached model instead of hanging on a model_info() request that fails with a - # DNS error and only ever gets retried (#6817). - cache_hit = _embedding_marker_in_hf_cache(model_name) - if cache_hit is True: - _embedding_detection_cache[cache_key] = True - logger.info(f"Model {model_name} detected as embedding model via HF cache (modules.json)") - return True if _env_offline(): - # Offline: the cache is the only source; anything not positively a - # sentence-transformers model is treated as non-embedding rather than - # making a network call that cannot succeed. Do NOT cache this negative: - # it is offline-conditional. A tag-only (feature-extraction) embedder can - # only be confirmed online, and the same (model_name, hf_token) key is - # reused for online lookups after the env var clears in this process. - return False - + # Offline: the local HF cache is the only source -- a network call cannot + # succeed and would only hang on a DNS error and get retried (#6817). + # Re-probe the marker on every call and never consult or populate the + # memo: a model downloaded later in this session must be seen, and a miss + # is not a durable negative (a tag-only feature-extraction embedder can + # only be confirmed online, and the same key is reused once the env var + # clears in this process). + return _embedding_marker_in_hf_cache(model_name) is True + + # Online: the Hub is authoritative for the current remote revision. The local + # cache marker reflects only the last-downloaded revision, which may lag the + # Hub (a repo can add or drop modules.json), so model_info() decides and the + # marker is used solely as a fallback when the Hub is unreachable. Only + # Hub-derived results are memoized, so a transient failure never poisons the + # cache and a later fresh download or successful lookup can override it. + if cache_key in _embedding_detection_cache: + return _embedding_detection_cache[cache_key] try: from huggingface_hub import model_info as hf_model_info @@ -2116,8 +2122,16 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: return is_emb except Exception as e: + # Hub unreachable: fall back to the local marker, uncached -- it is a + # degraded signal a later successful Hub call must be able to override. + marker = _embedding_marker_in_hf_cache(model_name) + if marker is True: + logger.info( + f"Model {model_name} detected as embedding model via HF cache " + f"(modules.json) after Hub lookup failed: {e}" + ) + return True logger.warning(f"Could not determine if {model_name} is embedding model: {e}") - _embedding_detection_cache[cache_key] = False return False From e9c7da63fcc24ebd0af677c350a368787d403a52 Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sat, 18 Jul 2026 07:15:53 +0300 Subject: [PATCH 06/77] studio: harden offline embedding detection against empty refs, offline flips, and cache casing - _embedding_marker_in_hf_cache: an existing-but-empty/whitespace refs/main (a partial write or in-progress truncate-and-rewrite) now reads as a cache miss (None) instead of falling through to scan stale snapshots; only a genuinely missing ref enables the historical scan. - is_embedding_model: while offline, retain a positive already confirmed online this session (model_info only ever memoizes Hub-derived results), so _hf_offline_if_dns_dead() flipping the process to offline mid-load can't downgrade a verified tag-only embedder to False. Cached negatives are still bypassed and re-probed. - resolve_cached_repo_casing + settings route: persist the embedding model in the casing its local HF cache dir uses. Validation accepts a case-insensitive cache hit, but an offline SentenceTransformer load resolves the cache by exact case, so storing the requested spelling (baai/bge-m3 vs models--BAAI--bge-m3) made the model fail to load on a case-sensitive filesystem. --- studio/backend/routes/settings.py | 7 ++- .../tests/test_is_embedding_model_offline.py | 58 ++++++++++++++++++ studio/backend/utils/models/__init__.py | 2 + studio/backend/utils/models/model_config.py | 59 ++++++++++++++++--- 4 files changed, 118 insertions(+), 8 deletions(-) diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 1ddfc0eacb3..bf709fc739b 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -387,7 +387,7 @@ def update_embedding_model( A repo flagged unsafe by HF's security scan returns 403 instead: a hard block that ``force`` cannot bypass, so the UI must not offer "save anyway". Documents indexed under the previous model must be re-uploaded.""" - from utils.models import is_embedding_model + from utils.models import is_embedding_model, resolve_cached_repo_casing try: model = validate_embedding_model(payload.embedding_model) @@ -463,6 +463,11 @@ def update_embedding_model( gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error(model, hf_token) if gguf_error: raise HTTPException(status_code = 409, detail = gguf_error) + # Persist the casing the local HF cache actually uses: validation above + # accepts a case-insensitive cache hit, but the offline SentenceTransformer + # load resolves the cache by exact case, so store the cached spelling (a + # no-op when nothing case-matching is cached) to keep the model loadable. + model = resolve_cached_repo_casing(model) set_rag_embedding_model(model) logger.info( "settings.embedding_model_updated subject=%s model=%s forced=%s", diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index ee9c6842c5b..109e1822f31 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -159,6 +159,15 @@ def test_marker_unreadable_ref_is_cache_miss(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache("org/unreadable-ref") is None +def test_marker_empty_ref_is_cache_miss(tmp_path, monkeypatch): + # refs/main exists but is empty / whitespace (a partial write or in-progress + # truncate-and-rewrite): the active revision is unknown, so this is a cache + # miss (None), NOT a fall-through to a stale snapshot that carries modules.json. + snaps = _repo(tmp_path, ("old", True), main_ref = " \n") + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + assert mc._embedding_marker_in_hf_cache("org/empty-ref") is None + + def test_marker_never_raises_when_cache_mutates(monkeypatch): # A snapshot vanishing mid-iteration (concurrent cached-model deletion) # must read as not-cached, not propagate a 500 out of the routes. @@ -249,6 +258,55 @@ def _info(model_name, token = None): assert mc.is_embedding_model("org/late-embedder") is True # marker re-probed +def test_offline_retains_online_confirmed_positive(monkeypatch): + # A tag-only embedder (feature-extraction, no modules.json) is confirmed + # online and cached True. _hf_offline_if_dns_dead() then flips the process to + # offline mid-load; the offline path must RETAIN that positive, not re-probe + # the absent marker and downgrade a model already verified this session. + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(())) + + def _info(model_name, token = None): + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + + _fake_hf_model_info(monkeypatch, _info) + assert mc.is_embedding_model("org/gte-modernbert") is True # online: cached True + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _fake_hf_model_info(monkeypatch, _no_network) # offline must not hit network + assert mc.is_embedding_model("org/gte-modernbert") is True # positive retained + + +# ── resolve_cached_repo_casing ── + + +def _fake_hf_constants(monkeypatch, cache_root): + fake = types.ModuleType("huggingface_hub") + fake.constants = types.SimpleNamespace(HF_HUB_CACHE = str(cache_root)) + monkeypatch.setitem(sys.modules, "huggingface_hub", fake) + monkeypatch.setitem(sys.modules, "huggingface_hub.constants", fake.constants) + + +def test_resolve_cached_repo_casing_normalizes_to_cache_dir(tmp_path, monkeypatch): + # baai/bge-m3 requested, models--BAAI--bge-m3 cached: the persisted id must be + # normalized to the cached casing so an offline exact-case load still finds it. + (tmp_path / "models--BAAI--bge-m3").mkdir() + _fake_hf_constants(monkeypatch, tmp_path) + assert mc.resolve_cached_repo_casing("baai/bge-m3") == "BAAI/bge-m3" + + +def test_resolve_cached_repo_casing_noop_when_uncached(tmp_path, monkeypatch): + _fake_hf_constants(monkeypatch, tmp_path) + assert mc.resolve_cached_repo_casing("org/not-cached") == "org/not-cached" + + +def test_resolve_cached_repo_casing_ignores_local_and_non_repo(tmp_path, monkeypatch): + # A local dir and a bare (no-slash) name are returned unchanged without a + # cache lookup at all. + _fake_hf_constants(monkeypatch, tmp_path) + assert mc.resolve_cached_repo_casing(str(tmp_path)) == str(tmp_path) + assert mc.resolve_cached_repo_casing("bare-name") == "bare-name" + + def test_offline_cached_non_st_returns_false_without_network(tmp_path, monkeypatch): snaps = _repo(tmp_path, ("aaa", False)) monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 4a5fb7274cf..057a938b372 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -8,6 +8,7 @@ GgufVariantInfo, is_vision_model, is_embedding_model, + resolve_cached_repo_casing, detect_audio_type, is_audio_input_type, VALID_AUDIO_TYPES, @@ -32,6 +33,7 @@ "GgufVariantInfo", "is_vision_model", "is_embedding_model", + "resolve_cached_repo_casing", "detect_audio_type", "is_audio_input_type", "VALID_AUDIO_TYPES", diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index c13c38d51da..3a6e3da23db 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2027,14 +2027,22 @@ def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: try: commit = (snapshots_dir.parent / "refs" / "main").read_text(encoding = "utf-8").strip() except FileNotFoundError: - commit = "" # no ref recorded: fall back to the newest-first scan + # No ref recorded at all: fall back to the newest-first scan. + commit = None except OSError: # A ref exists but is unreadable (transient I/O error, restrictive # permissions): the contract is that an unreadable cache reads as # not-cached, so report a miss rather than scanning stale history -- # only a genuinely missing ref may enable the fallback scan. return None - if commit: + else: + if not commit: + # The ref file exists but is empty / whitespace (a partial write + # or an in-progress truncate-and-rewrite). The active revision is + # unknown, so this is a cache miss -- NOT a fall-through to stale + # history, which is the class of bug this helper avoids. + return None + if commit is not None: # A ref is recorded, so it is authoritative. If its snapshot is not # materialized (partial download / pruning) treat the repo as not # cached (None) rather than scanning stale history -- the exact @@ -2054,6 +2062,38 @@ def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: return None +def resolve_cached_repo_casing(repo_id: str) -> str: + """Return *repo_id* in the exact casing of its local HF cache directory. + + HF cache lookups here are case-insensitive (see _iter_hf_cache_snapshots), so + ``baai/bge-m3`` can validate against a ``models--BAAI--bge-m3`` cache dir. But + an offline SentenceTransformer load resolves the cache by EXACT case, so the + requested spelling must be normalized to the cached one before it is persisted + or the load fails on a case-sensitive filesystem even though validation + passed. Returns *repo_id* unchanged for a local path, a non-repo string, or + when nothing case-matching is cached (online the loader re-resolves casing via + the Hub, so leaving it as-is there is harmless).""" + if is_local_path(repo_id) or "/" not in repo_id: + return repo_id + try: + from huggingface_hub import constants as hf_constants + + cache_dir = Path(hf_constants.HF_HUB_CACHE) + if not cache_dir.is_dir(): + return repo_id + prefix = "models--" + target = f"{prefix}{repo_id.replace('/', '--')}".lower() + for entry in cache_dir.iterdir(): + if entry.is_dir() and entry.name.lower() == target: + # HF encodes org/name as models--org--name; reverse it. The `--` + # separator only appears at path boundaries, so splitting on it + # recovers the Hub-canonical casing the cache dir was named with. + return entry.name[len(prefix) :].replace("--", "/") + except OSError: + return repo_id + return repo_id + + def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: """Detect embedding/sentence-transformer models via HF metadata. @@ -2083,11 +2123,16 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: if _env_offline(): # Offline: the local HF cache is the only source -- a network call cannot # succeed and would only hang on a DNS error and get retried (#6817). - # Re-probe the marker on every call and never consult or populate the - # memo: a model downloaded later in this session must be seen, and a miss - # is not a durable negative (a tag-only feature-extraction embedder can - # only be confirmed online, and the same key is reused once the env var - # clears in this process). + # Retain a positive already confirmed online this session: model_info() + # only ever memoizes Hub-derived results, so a cached True is a real + # detection (e.g. a tag-only feature-extraction embedder with no + # modules.json) that _hf_offline_if_dns_dead() flipping the process to + # offline mid-load must NOT downgrade to False. Otherwise re-probe the + # marker every call without consulting or populating the memo: a cached + # negative must not stick (a model downloaded later, or a tag-only + # embedder, could not be confirmed here), so a miss is not durable. + if _embedding_detection_cache.get(cache_key) is True: + return True return _embedding_marker_in_hf_cache(model_name) is True # Online: the Hub is authoritative for the current remote revision. The local From bfc1cff1de622eff0302f7b30556f54a472d91da Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sat, 18 Jul 2026 07:27:31 +0300 Subject: [PATCH 07/77] studio: reuse the exact-match-first case resolver and preserve the default Replace the ad-hoc resolve_cached_repo_casing with the existing resolve_cached_repo_id_case, which already prefers the exact-case cache dir before any case variant and tie-breaks variants deterministically -- so an exact requested id is never rewritten to a differently cased directory just because iterdir() happened to yield it first. Skip the normalization entirely when the submitted model equals the default: rewriting its casing would make set_rag_embedding_model()'s exact-string default comparison treat it as a custom override, pinning it so later changes to the configured default stop taking effect. --- studio/backend/routes/settings.py | 9 ++++-- .../test_embedding_model_security_gate.py | 31 ++++++++++++++++++ .../tests/test_is_embedding_model_offline.py | 31 ------------------ studio/backend/utils/models/__init__.py | 2 -- studio/backend/utils/models/model_config.py | 32 ------------------- 5 files changed, 38 insertions(+), 67 deletions(-) diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index bf709fc739b..45993bd2793 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -387,7 +387,8 @@ def update_embedding_model( A repo flagged unsafe by HF's security scan returns 403 instead: a hard block that ``force`` cannot bypass, so the UI must not offer "save anyway". Documents indexed under the previous model must be re-uploaded.""" - from utils.models import is_embedding_model, resolve_cached_repo_casing + from utils.models import is_embedding_model + from utils.paths import resolve_cached_repo_id_case try: model = validate_embedding_model(payload.embedding_model) @@ -467,7 +468,11 @@ def update_embedding_model( # accepts a case-insensitive cache hit, but the offline SentenceTransformer # load resolves the cache by exact case, so store the cached spelling (a # no-op when nothing case-matching is cached) to keep the model loadable. - model = resolve_cached_repo_casing(model) + # Skip this for the default: rewriting its casing would make the exact-string + # default comparison in set_rag_embedding_model() treat it as a custom + # override, so later changes to the configured default would stop applying. + if model != default_embedding_model(): + model = resolve_cached_repo_id_case(model) set_rag_embedding_model(model) logger.info( "settings.embedding_model_updated subject=%s model=%s forced=%s", diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index 940b35d7bac..ed41346818b 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -259,6 +259,37 @@ def test_clean_repo_saves_under_force(client, monkeypatch): assert saved.get("model") == "acme/clean-embed" +def test_custom_model_saved_in_cache_casing(client, monkeypatch): + # A custom repo is persisted in the casing its local HF cache dir uses, so an + # offline exact-case SentenceTransformer load still finds it (baai/bge-m3 -> + # BAAI/bge-m3). force=True bypasses the verification gate to isolate this. + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + import utils.paths as _paths + + monkeypatch.setattr(_paths, "resolve_cached_repo_id_case", lambda m: "BAAI/bge-m3") + r = c.put("/embedding-model", json = {"embedding_model": "baai/bge-m3", "force": True}) + assert r.status_code == 200 + assert saved.get("model") == "BAAI/bge-m3" + + +def test_default_model_is_not_casing_normalized(client, monkeypatch): + # Submitting the exact default must NOT be run through cache-casing + # normalization: rewriting it would make set_rag_embedding_model()'s exact + # default comparison treat it as a custom override, so later default changes + # would stop applying. resolve_cached_repo_id_case must not be consulted. + c, saved = client + import utils.paths as _paths + + def _must_not_run(m): + raise AssertionError("the default must not be casing-normalized") + + monkeypatch.setattr(_paths, "resolve_cached_repo_id_case", _must_not_run) + r = c.put("/embedding-model", json = {"embedding_model": "unsloth/default-embed"}) + assert r.status_code == 200 + assert saved.get("model") == "unsloth/default-embed" + + def test_load_sink_refuses_flagged_model(monkeypatch): monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) import core.rag.embeddings as embeddings diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 109e1822f31..1a33dc592fd 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -276,37 +276,6 @@ def _info(model_name, token = None): assert mc.is_embedding_model("org/gte-modernbert") is True # positive retained -# ── resolve_cached_repo_casing ── - - -def _fake_hf_constants(monkeypatch, cache_root): - fake = types.ModuleType("huggingface_hub") - fake.constants = types.SimpleNamespace(HF_HUB_CACHE = str(cache_root)) - monkeypatch.setitem(sys.modules, "huggingface_hub", fake) - monkeypatch.setitem(sys.modules, "huggingface_hub.constants", fake.constants) - - -def test_resolve_cached_repo_casing_normalizes_to_cache_dir(tmp_path, monkeypatch): - # baai/bge-m3 requested, models--BAAI--bge-m3 cached: the persisted id must be - # normalized to the cached casing so an offline exact-case load still finds it. - (tmp_path / "models--BAAI--bge-m3").mkdir() - _fake_hf_constants(monkeypatch, tmp_path) - assert mc.resolve_cached_repo_casing("baai/bge-m3") == "BAAI/bge-m3" - - -def test_resolve_cached_repo_casing_noop_when_uncached(tmp_path, monkeypatch): - _fake_hf_constants(monkeypatch, tmp_path) - assert mc.resolve_cached_repo_casing("org/not-cached") == "org/not-cached" - - -def test_resolve_cached_repo_casing_ignores_local_and_non_repo(tmp_path, monkeypatch): - # A local dir and a bare (no-slash) name are returned unchanged without a - # cache lookup at all. - _fake_hf_constants(monkeypatch, tmp_path) - assert mc.resolve_cached_repo_casing(str(tmp_path)) == str(tmp_path) - assert mc.resolve_cached_repo_casing("bare-name") == "bare-name" - - def test_offline_cached_non_st_returns_false_without_network(tmp_path, monkeypatch): snaps = _repo(tmp_path, ("aaa", False)) monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 057a938b372..4a5fb7274cf 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -8,7 +8,6 @@ GgufVariantInfo, is_vision_model, is_embedding_model, - resolve_cached_repo_casing, detect_audio_type, is_audio_input_type, VALID_AUDIO_TYPES, @@ -33,7 +32,6 @@ "GgufVariantInfo", "is_vision_model", "is_embedding_model", - "resolve_cached_repo_casing", "detect_audio_type", "is_audio_input_type", "VALID_AUDIO_TYPES", diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 3a6e3da23db..9358b177b54 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2062,38 +2062,6 @@ def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: return None -def resolve_cached_repo_casing(repo_id: str) -> str: - """Return *repo_id* in the exact casing of its local HF cache directory. - - HF cache lookups here are case-insensitive (see _iter_hf_cache_snapshots), so - ``baai/bge-m3`` can validate against a ``models--BAAI--bge-m3`` cache dir. But - an offline SentenceTransformer load resolves the cache by EXACT case, so the - requested spelling must be normalized to the cached one before it is persisted - or the load fails on a case-sensitive filesystem even though validation - passed. Returns *repo_id* unchanged for a local path, a non-repo string, or - when nothing case-matching is cached (online the loader re-resolves casing via - the Hub, so leaving it as-is there is harmless).""" - if is_local_path(repo_id) or "/" not in repo_id: - return repo_id - try: - from huggingface_hub import constants as hf_constants - - cache_dir = Path(hf_constants.HF_HUB_CACHE) - if not cache_dir.is_dir(): - return repo_id - prefix = "models--" - target = f"{prefix}{repo_id.replace('/', '--')}".lower() - for entry in cache_dir.iterdir(): - if entry.is_dir() and entry.name.lower() == target: - # HF encodes org/name as models--org--name; reverse it. The `--` - # separator only appears at path boundaries, so splitting on it - # recovers the Hub-canonical casing the cache dir was named with. - return entry.name[len(prefix) :].replace("--", "/") - except OSError: - return repo_id - return repo_id - - def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: """Detect embedding/sentence-transformer models via HF metadata. From 3468fa974fb7cce5457e3566efe48c051b0352aa Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sat, 18 Jul 2026 07:35:15 +0300 Subject: [PATCH 08/77] studio: don't let a stale cache marker mask a permanent Hub error is_embedding_model's Hub-failure fallback consulted the local modules.json marker for ANY model_info() exception, so a permanent error -- a deleted repo, a gated repo without credentials, or a typo that matches stale cache casing -- could pass online validation on a stale marker instead of returning the documented 409, and the persisted model could then fail when the loader refreshes from the Hub. Classify permanent Hub errors (RepositoryNotFound, GatedRepo, RevisionNotFound, EntryNotFound) as False, matching the nearby GGUF/vision detectors, and reserve the cache fallback for transient/5xx failures. --- .../tests/test_is_embedding_model_offline.py | 18 ++++++++++++++++++ studio/backend/utils/models/model_config.py | 18 ++++++++++++++++-- 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 1a33dc592fd..824335733e7 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -226,6 +226,24 @@ def _info(model_name, token = None): assert calls == ["org/was-embedder"] # Hub consulted, not skipped +def test_online_permanent_hub_error_ignores_stale_marker(tmp_path, monkeypatch): + # A permanent Hub error (deleted / gated / typo'd repo) is authoritative: + # even with a cached modules.json, validation must NOT pass on the stale + # marker -- return False so the settings route surfaces its 409, and the + # persisted model can't fail later when the loader refreshes from the Hub. + snaps = _repo(tmp_path, ("aaa", True)) + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + + class RepositoryNotFoundError(Exception): + pass + + def _info(model_name, token = None): + raise RepositoryNotFoundError("404 not found") + + _fake_hf_model_info(monkeypatch, _info) + assert mc.is_embedding_model("org/deleted-or-typo") is False + + def test_online_hub_failure_falls_back_to_marker_uncached(tmp_path, monkeypatch): # A transient model_info() failure falls back to the local marker WITHOUT # caching: the degraded result must not become sticky, so a later successful diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 9358b177b54..fed87475e7a 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2135,8 +2135,22 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: return is_emb except Exception as e: - # Hub unreachable: fall back to the local marker, uncached -- it is a - # degraded signal a later successful Hub call must be able to override. + # A permanent Hub error (deleted / gated / bad revision / typo matching + # stale cache casing) is authoritative: the repo is not a usable + # embedding model, so return False and let the settings route surface its + # documented 409. Falling back to a cached modules.json here would wrongly + # pass validation for a repo the loader can no longer fetch. Matches the + # permanent-vs-transient split the nearby GGUF/vision detectors use. + if type(e).__name__ in ( + "RepositoryNotFoundError", + "GatedRepoError", + "RevisionNotFoundError", + "EntryNotFoundError", + ): + logger.warning(f"Could not determine if {model_name} is embedding model: {e}") + return False + # Transient / connectivity / 5xx failure: fall back to the local marker, + # uncached -- a degraded signal a later successful Hub call can override. marker = _embedding_marker_in_hf_cache(model_name) if marker is True: logger.info( From 202b23f5de6f0423272f7afadcbe757e194ead21 Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sat, 18 Jul 2026 10:49:02 +0300 Subject: [PATCH 09/77] studio: honor TRANSFORMERS_OFFLINE in the embedding preflight, skip casing for local paths - The embedding-model save reached the offline-aware is_embedding_model() only after two preflight helpers made direct huggingface_hub calls that honor just HF_HUB_OFFLINE: _st_module_subdirs() downloads modules.json and the security scan fetches Hub metadata twice. In a TRANSFORMERS_OFFLINE-only session those blocked on network timeouts before the offline return, so saving an already cached model stalled. Both now consult a canonical hf_env_offline() helper -- the download passes local_files_only, and the metadata-only security scan short-circuits to its documented fail-open instead of burning both timeouts. - Skip cache-casing normalization for local paths: a relative directory such as "org/model" is loaded from disk, so rewriting it to a case-insensitive HF cache collision ("Org/model") would stop resolving to that directory and be read as a Hub repo id instead. --- studio/backend/core/rag/embeddings.py | 12 +++++++++++- studio/backend/routes/settings.py | 13 ++++++++----- .../test_embedding_model_security_gate.py | 19 +++++++++++++++++++ .../backend/utils/security/file_security.py | 9 +++++++++ studio/backend/utils/utils.py | 16 ++++++++++++++++ 5 files changed, 63 insertions(+), 6 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index b0ecedd5936..b467a42bb11 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -103,9 +103,19 @@ def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]: else: from huggingface_hub import hf_hub_download from huggingface_hub.utils import EntryNotFoundError + from utils.utils import hf_env_offline try: - local = hf_hub_download(name, "modules.json", token = token or None) + # local_files_only when offline: huggingface_hub honors only + # HF_HUB_OFFLINE natively, so a TRANSFORMERS_OFFLINE-only session + # would otherwise block on network timeouts here even though the + # cached snapshot already has modules.json. + local = hf_hub_download( + name, + "modules.json", + token = token or None, + local_files_only = hf_env_offline(), + ) except EntryNotFoundError: return () data = json.loads(open(local).read()) diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 45993bd2793..7089f37f5eb 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -388,7 +388,7 @@ def update_embedding_model( that ``force`` cannot bypass, so the UI must not offer "save anyway". Documents indexed under the previous model must be re-uploaded.""" from utils.models import is_embedding_model - from utils.paths import resolve_cached_repo_id_case + from utils.paths import is_local_path, resolve_cached_repo_id_case try: model = validate_embedding_model(payload.embedding_model) @@ -468,10 +468,13 @@ def update_embedding_model( # accepts a case-insensitive cache hit, but the offline SentenceTransformer # load resolves the cache by exact case, so store the cached spelling (a # no-op when nothing case-matching is cached) to keep the model loadable. - # Skip this for the default: rewriting its casing would make the exact-string - # default comparison in set_rag_embedding_model() treat it as a custom - # override, so later changes to the configured default would stop applying. - if model != default_embedding_model(): + # Skip the default: rewriting its casing would make the exact-string default + # comparison in set_rag_embedding_model() treat it as a custom override, so + # later changes to the configured default would stop applying. Skip local + # paths too: a relative directory like "org/model" is loaded from disk, and + # rewriting it to a case-insensitive HF cache collision ("Org/model") would + # stop resolving to that directory and be read as a Hub repo id instead. + if model != default_embedding_model() and not is_local_path(model): model = resolve_cached_repo_id_case(model) set_rag_embedding_model(model) logger.info( diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index ed41346818b..738f7cd01eb 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -273,6 +273,25 @@ def test_custom_model_saved_in_cache_casing(client, monkeypatch): assert saved.get("model") == "BAAI/bge-m3" +def test_local_path_model_is_not_casing_normalized(client, monkeypatch, tmp_path): + # A local directory (here an existing path) is loaded from disk. Rewriting it + # to a case-insensitive HF cache collision would stop resolving to that + # directory and be read as a Hub repo id, so normalization must be skipped. + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + import utils.paths as _paths + + def _must_not_run(m): + raise AssertionError("a local path must not be casing-normalized") + + monkeypatch.setattr(_paths, "resolve_cached_repo_id_case", _must_not_run) + local_dir = tmp_path / "org" / "model" + local_dir.mkdir(parents = True) + r = c.put("/embedding-model", json = {"embedding_model": str(local_dir), "force": True}) + assert r.status_code == 200 + assert saved.get("model") == str(local_dir) + + def test_default_model_is_not_casing_normalized(client, monkeypatch): # Submitting the exact default must NOT be run through cache-casing # normalization: rewriting it would make set_rag_embedding_model()'s exact diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 466f326f18f..c77038c972b 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -241,6 +241,15 @@ def _fetch_security_status(model_name: str, hf_token: Optional[str]): retries once on a transient error, then returns None so the caller fails open. """ from huggingface_hub import model_info as hf_model_info + from utils.utils import hf_env_offline + + # Offline (including a TRANSFORMERS_OFFLINE-only session, which huggingface_hub + # does not honor natively): this is a pure Hub-metadata lookup with no local + # fallback, so attempting it would just burn both request timeouts before + # failing open anyway. Skip straight to the fail-open path. + if hf_env_offline(): + logger.debug("HF security scan skipped for '%s': offline; failing open.", model_name) + return None token_arg = hf_token if hf_token else False last_exc = None diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index 3818253ac9e..c951f29f5be 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -161,3 +161,19 @@ def format_error_message(error: Exception, model_name: str) -> str: return f"Not enough {device_label} memory to load '{model_short}'. Try a smaller model or free memory." return str(error) + + +_HF_OFFLINE_TRUE_VALUES = {"1", "true", "yes", "on"} + + +def hf_env_offline() -> bool: + """True when either HF offline env var is truthy (strip+lower, on/true/yes/1). + + ``huggingface_hub`` natively honors only ``HF_HUB_OFFLINE``; callers that make + direct Hub calls must consult this so a ``TRANSFORMERS_OFFLINE``-only session + does not block on network timeouts for data the local cache already has. + """ + return ( + os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _HF_OFFLINE_TRUE_VALUES + or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _HF_OFFLINE_TRUE_VALUES + ) From cb1f5381d72458d107284168ad3346ca78155a60 Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sat, 18 Jul 2026 10:58:31 +0300 Subject: [PATCH 10/77] studio: never skip the security scan on TRANSFORMERS_OFFLINE alone The previous commit skipped the Hub security scan whenever either offline flag was set, but huggingface_hub honors only HF_HUB_OFFLINE: under a TRANSFORMERS_OFFLINE-only session the later SentenceTransformer load still reaches the network, so the scan was being skipped while the repo's pickle could still be downloaded and deserialized -- waving through exactly what _guard_model_security exists to block. Split the flags: hf_hub_offline() (HF_HUB_OFFLINE, the only one that actually prevents a fetch) gates the security short-circuit, while hf_env_offline() (either flag, the user's intent) is used only where local-only behavior is forced explicitly. The SentenceTransformer load now passes local_files_only from that intent, so TRANSFORMERS_OFFLINE genuinely stops the loader fetching instead of merely being assumed to. --- studio/backend/core/rag/embeddings.py | 13 +- studio/backend/tests/test_hf_offline_flags.py | 139 ++++++++++++++++++ .../backend/utils/security/file_security.py | 19 ++- studio/backend/utils/utils.py | 29 +++- 4 files changed, 184 insertions(+), 16 deletions(-) create mode 100644 studio/backend/tests/test_hf_offline_flags.py diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index b467a42bb11..8d291a76edc 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -165,10 +165,21 @@ def _get(model_name: str | None = None): _install_torchao_stub_once() from sentence_transformers import SentenceTransformer + from utils.utils import hf_env_offline + 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")) + # Propagate the user's offline intent into the loader: SentenceTransformer + # performs its own Hub operations, and huggingface_hub honors only + # HF_HUB_OFFLINE, so a TRANSFORMERS_OFFLINE-only session would otherwise + # still fetch missing repo files over the network. + _model = SentenceTransformer( + name, + device = device, + model_kwargs = dtype_kwargs("float16"), + local_files_only = hf_env_offline(), + ) _name = name return _model diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py new file mode 100644 index 00000000000..744fb76993f --- /dev/null +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The two HF offline flags are not interchangeable. + +``huggingface_hub`` honors only ``HF_HUB_OFFLINE``. ``TRANSFORMERS_OFFLINE`` +expresses the same user intent but does NOT stop a Hub fetch, so it may be used +to pass ``local_files_only`` (a real guarantee) but must never be used to skip a +security gate -- doing so would wave through the very download the gate exists +to block. +""" + +from __future__ import annotations + +import importlib +import sys +import types +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + + +def _maybe_stub(name: str, builder): + # Stub only if the real module is unavailable, so this file never shadows + # real packages for later tests in the same pytest process. + try: + importlib.import_module(name) + except ImportError: + sys.modules[name] = builder() + + +def _build_loggers_stub(): + m = types.ModuleType("loggers") + m.get_logger = lambda name: __import__("logging").getLogger(name) + return m + + +def _build_structlog_stub(): + m = types.ModuleType("structlog") + m.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") + return m + + +_maybe_stub("loggers", _build_loggers_stub) +_maybe_stub("structlog", _build_structlog_stub) + +from utils.utils import hf_env_offline, hf_hub_offline # noqa: E402 + + +@pytest.fixture(autouse = True) +def _clean_env(monkeypatch): + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + yield + + +# ── flag semantics ── + + +def test_neither_flag_is_online(): + assert hf_hub_offline() is False + assert hf_env_offline() is False + + +def test_transformers_offline_is_intent_but_not_a_fetch_guarantee(monkeypatch): + # The distinction the security gate depends on: intent yes, guarantee no. + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + assert hf_env_offline() is True + assert hf_hub_offline() is False + + +def test_hub_offline_sets_both(monkeypatch): + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + assert hf_hub_offline() is True + assert hf_env_offline() is True + + +@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on", " 1 "]) +def test_truthy_values_parse(monkeypatch, value): + monkeypatch.setenv("HF_HUB_OFFLINE", value) + assert hf_hub_offline() is True + + +@pytest.mark.parametrize("value", ["", "0", "false", "no", "off", "maybe"]) +def test_non_truthy_values_do_not_parse(monkeypatch, value): + monkeypatch.setenv("HF_HUB_OFFLINE", value) + assert hf_hub_offline() is False + + +# ── the security gate must not be skipped on the weaker flag ── + + +def _fake_hub(monkeypatch, calls): + fake = types.ModuleType("huggingface_hub") + + def _model_info(*a, **k): + calls.append(1) + raise RuntimeError("network unreachable") + + fake.model_info = _model_info + monkeypatch.setitem(sys.modules, "huggingface_hub", fake) + + +def test_security_scan_still_runs_under_transformers_offline(monkeypatch): + # TRANSFORMERS_OFFLINE does NOT stop SentenceTransformer from fetching, so + # the scan must still be attempted; skipping it here would let an unscanned + # repo be downloaded and its pickle deserialized anyway. + import utils.security.file_security as fs + + calls: list = [] + _fake_hub(monkeypatch, calls) + monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") + assert fs._fetch_security_status("org/model", None) is None # fails open + assert calls, "the security scan must not be skipped for TRANSFORMERS_OFFLINE" + + +def test_security_scan_short_circuits_under_hub_offline(monkeypatch): + # With HF_HUB_OFFLINE no fetch is possible, so this metadata-only lookup can + # skip straight to its documented fail-open instead of burning both timeouts. + import utils.security.file_security as fs + + calls: list = [] + _fake_hub(monkeypatch, calls) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + assert fs._fetch_security_status("org/model", None) is None + assert calls == [], "no Hub call is possible when HF_HUB_OFFLINE is set" + + +def test_security_scan_runs_when_online(monkeypatch): + import utils.security.file_security as fs + + calls: list = [] + _fake_hub(monkeypatch, calls) + assert fs._fetch_security_status("org/model", None) is None # fails open on error + assert calls, "online must attempt the Hub" diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index c77038c972b..0d469aedef6 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -241,14 +241,17 @@ def _fetch_security_status(model_name: str, hf_token: Optional[str]): retries once on a transient error, then returns None so the caller fails open. """ from huggingface_hub import model_info as hf_model_info - from utils.utils import hf_env_offline - - # Offline (including a TRANSFORMERS_OFFLINE-only session, which huggingface_hub - # does not honor natively): this is a pure Hub-metadata lookup with no local - # fallback, so attempting it would just burn both request timeouts before - # failing open anyway. Skip straight to the fail-open path. - if hf_env_offline(): - logger.debug("HF security scan skipped for '%s': offline; failing open.", model_name) + from utils.utils import hf_hub_offline + + # Only HF_HUB_OFFLINE -- deliberately NOT the weaker hf_env_offline(). This is + # a security gate: huggingface_hub honors only HF_HUB_OFFLINE, so under a + # TRANSFORMERS_OFFLINE-only session the later SentenceTransformer load can + # still fetch and deserialize the repo's pickle. Skipping the scan there would + # wave through exactly the download _guard_model_security exists to block. + # With HF_HUB_OFFLINE set no fetch is possible, so this metadata-only lookup + # would just burn both request timeouts before failing open anyway. + if hf_hub_offline(): + logger.debug("HF security scan skipped for '%s': hub offline; failing open.", model_name) return None token_arg = hf_token if hf_token else False diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index c951f29f5be..909e11790c5 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -166,14 +166,29 @@ def format_error_message(error: Exception, model_name: str) -> str: _HF_OFFLINE_TRUE_VALUES = {"1", "true", "yes", "on"} +def _offline_flag(name: str) -> bool: + return os.environ.get(name, "").strip().lower() in _HF_OFFLINE_TRUE_VALUES + + +def hf_hub_offline() -> bool: + """True when ``HF_HUB_OFFLINE`` is truthy -- the ONLY offline flag + ``huggingface_hub`` honors natively, so it is the one that actually prevents + a Hub fetch. + + Use this (never :func:`hf_env_offline`) to decide whether a network call is + *impossible*. A security gate must not be skipped on the weaker signal: with + only ``TRANSFORMERS_OFFLINE`` set, hub operations still reach the network, so + skipping a scan there would let an unscanned repo be downloaded anyway. + """ + return _offline_flag("HF_HUB_OFFLINE") + + def hf_env_offline() -> bool: """True when either HF offline env var is truthy (strip+lower, on/true/yes/1). - ``huggingface_hub`` natively honors only ``HF_HUB_OFFLINE``; callers that make - direct Hub calls must consult this so a ``TRANSFORMERS_OFFLINE``-only session - does not block on network timeouts for data the local cache already has. + This is the user's *intent* to work offline. ``huggingface_hub`` honors only + ``HF_HUB_OFFLINE``, so callers that must GUARANTEE no fetch have to pass + ``local_files_only`` explicitly (or consult :func:`hf_hub_offline`) rather + than assume this alone stops the network. """ - return ( - os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _HF_OFFLINE_TRUE_VALUES - or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _HF_OFFLINE_TRUE_VALUES - ) + return hf_hub_offline() or _offline_flag("TRANSFORMERS_OFFLINE") From df32b909a207947e8e48a42678116f571cea56fe Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sat, 18 Jul 2026 11:19:55 +0300 Subject: [PATCH 11/77] studio: short-circuit the security preflight under either offline flag With the loader now pinned to the local cache by local_files_only = hf_env_offline(), a TRANSFORMERS_OFFLINE-only session can no longer fetch anything -- yet the preflight still fell through to two model_info() attempts on 10s and 20s timeouts, stalling every save and load of an already-cached embedder for half a minute before failing open anyway. Skip the metadata-only scan whenever either flag is set. The scan's job is to stop a poisoned pickle being downloaded and deserialized, and nothing can be downloaded under that predicate; the residual case -- a model cached BEFORE it was flagged -- is the same fail-open this function has always documented for an unavailable scan, and is exactly what HF_HUB_OFFLINE already did. That safety argument depends on every loader behind the gate honoring the same predicate, so it is pinned as a test invariant instead of a comment: removing local_files_only from the SentenceTransformer construction now fails the suite. Drops the short-lived hf_hub_offline() helper, which no longer has a caller. --- studio/backend/tests/test_hf_offline_flags.py | 85 ++++++++++--------- .../backend/utils/security/file_security.py | 27 +++--- studio/backend/utils/utils.py | 26 ++---- 3 files changed, 69 insertions(+), 69 deletions(-) diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index 744fb76993f..b4746bf8db5 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -1,13 +1,18 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""The two HF offline flags are not interchangeable. - -``huggingface_hub`` honors only ``HF_HUB_OFFLINE``. ``TRANSFORMERS_OFFLINE`` -expresses the same user intent but does NOT stop a Hub fetch, so it may be used -to pass ``local_files_only`` (a real guarantee) but must never be used to skip a -security gate -- doing so would wave through the very download the gate exists -to block. +"""Effective-offline handling for the embedding preflight. + +``huggingface_hub`` honors only ``HF_HUB_OFFLINE``; ``TRANSFORMERS_OFFLINE`` +expresses the same user intent but does not itself stop a fetch. Studio treats +either as offline (``hf_env_offline``) and makes that real by passing +``local_files_only`` to the loader, which lets the metadata-only Hub security +scan skip straight to its documented fail-open instead of burning both request +timeouts on a session the user declared offline. + +That skip is only sound while the loader really is pinned to the local cache, so +the coupling between the two is pinned here as an explicit invariant rather than +left to a comment. """ from __future__ import annotations @@ -48,7 +53,7 @@ def _build_structlog_stub(): _maybe_stub("loggers", _build_loggers_stub) _maybe_stub("structlog", _build_structlog_stub) -from utils.utils import hf_env_offline, hf_hub_offline # noqa: E402 +from utils.utils import hf_env_offline # noqa: E402 @pytest.fixture(autouse = True) @@ -62,36 +67,28 @@ def _clean_env(monkeypatch): def test_neither_flag_is_online(): - assert hf_hub_offline() is False assert hf_env_offline() is False -def test_transformers_offline_is_intent_but_not_a_fetch_guarantee(monkeypatch): - # The distinction the security gate depends on: intent yes, guarantee no. - monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") - assert hf_env_offline() is True - assert hf_hub_offline() is False - - -def test_hub_offline_sets_both(monkeypatch): - monkeypatch.setenv("HF_HUB_OFFLINE", "1") - assert hf_hub_offline() is True +@pytest.mark.parametrize("var", ["HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"]) +def test_either_flag_means_offline(monkeypatch, var): + monkeypatch.setenv(var, "1") assert hf_env_offline() is True @pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on", " 1 "]) def test_truthy_values_parse(monkeypatch, value): monkeypatch.setenv("HF_HUB_OFFLINE", value) - assert hf_hub_offline() is True + assert hf_env_offline() is True @pytest.mark.parametrize("value", ["", "0", "false", "no", "off", "maybe"]) def test_non_truthy_values_do_not_parse(monkeypatch, value): monkeypatch.setenv("HF_HUB_OFFLINE", value) - assert hf_hub_offline() is False + assert hf_env_offline() is False -# ── the security gate must not be skipped on the weaker flag ── +# ── the offline security short-circuit, and the invariant it rests on ── def _fake_hub(monkeypatch, calls): @@ -105,29 +102,18 @@ def _model_info(*a, **k): monkeypatch.setitem(sys.modules, "huggingface_hub", fake) -def test_security_scan_still_runs_under_transformers_offline(monkeypatch): - # TRANSFORMERS_OFFLINE does NOT stop SentenceTransformer from fetching, so - # the scan must still be attempted; skipping it here would let an unscanned - # repo be downloaded and its pickle deserialized anyway. - import utils.security.file_security as fs - - calls: list = [] - _fake_hub(monkeypatch, calls) - monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1") - assert fs._fetch_security_status("org/model", None) is None # fails open - assert calls, "the security scan must not be skipped for TRANSFORMERS_OFFLINE" - - -def test_security_scan_short_circuits_under_hub_offline(monkeypatch): - # With HF_HUB_OFFLINE no fetch is possible, so this metadata-only lookup can - # skip straight to its documented fail-open instead of burning both timeouts. +@pytest.mark.parametrize("var", ["HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"]) +def test_security_scan_short_circuits_when_offline(monkeypatch, var): + # Metadata-only lookup with no local fallback: on a session the user declared + # offline it must skip straight to its documented fail-open rather than burn + # both request timeouts (10s + 20s) on every save and load. import utils.security.file_security as fs calls: list = [] _fake_hub(monkeypatch, calls) - monkeypatch.setenv("HF_HUB_OFFLINE", "1") + monkeypatch.setenv(var, "1") assert fs._fetch_security_status("org/model", None) is None - assert calls == [], "no Hub call is possible when HF_HUB_OFFLINE is set" + assert calls == [], f"the scan must not hit the Hub under {var}" def test_security_scan_runs_when_online(monkeypatch): @@ -137,3 +123,22 @@ def test_security_scan_runs_when_online(monkeypatch): _fake_hub(monkeypatch, calls) assert fs._fetch_security_status("org/model", None) is None # fails open on error assert calls, "online must attempt the Hub" + + +def test_embedding_loader_forces_local_only_when_offline(): + """The invariant the offline scan skip rests on. + + Skipping the Hub security scan offline is only sound while every loader + behind that gate is pinned to the local cache by the SAME predicate -- if the + loader could still fetch, an unscanned repo's pickle would be downloaded and + deserialized. Pinned at source level because importing the loader would drag + in sentence_transformers/torch. + """ + src = (Path(__file__).resolve().parents[1] / "core" / "rag" / "embeddings.py").read_text( + encoding = "utf-8" + ) + assert "local_files_only = hf_env_offline()" in src, ( + "core/rag/embeddings.py must pin SentenceTransformer to local files when " + "hf_env_offline(); without it the offline security-scan skip in " + "utils/security/file_security.py is unsafe" + ) diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 0d469aedef6..706f1852aa6 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -241,17 +241,22 @@ def _fetch_security_status(model_name: str, hf_token: Optional[str]): retries once on a transient error, then returns None so the caller fails open. """ from huggingface_hub import model_info as hf_model_info - from utils.utils import hf_hub_offline - - # Only HF_HUB_OFFLINE -- deliberately NOT the weaker hf_env_offline(). This is - # a security gate: huggingface_hub honors only HF_HUB_OFFLINE, so under a - # TRANSFORMERS_OFFLINE-only session the later SentenceTransformer load can - # still fetch and deserialize the repo's pickle. Skipping the scan there would - # wave through exactly the download _guard_model_security exists to block. - # With HF_HUB_OFFLINE set no fetch is possible, so this metadata-only lookup - # would just burn both request timeouts before failing open anyway. - if hf_hub_offline(): - logger.debug("HF security scan skipped for '%s': hub offline; failing open.", model_name) + from utils.utils import hf_env_offline + + # Effective offline: skip this metadata-only lookup instead of burning both + # request timeouts (10s + 20s) on a session the user declared offline. + # + # SAFETY INVARIANT: this is only sound because every loader behind this gate + # forces local-only loading from the SAME predicate -- see + # core/rag/embeddings.py, which passes local_files_only = hf_env_offline() to + # SentenceTransformer. Nothing can therefore be fetched here, so the scan's + # job (block a poisoned pickle from being downloaded and deserialized) is + # already served; the residual case, a model cached BEFORE it was flagged, is + # the same fail-open this function has always documented for an unavailable + # scan. If a loader ever stops honoring hf_env_offline(), this skip becomes + # unsafe -- tests/test_hf_offline_flags.py pins that coupling. + if hf_env_offline(): + logger.debug("HF security scan skipped for '%s': offline; failing open.", model_name) return None token_arg = hf_token if hf_token else False diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index 909e11790c5..a029191cd45 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -170,25 +170,15 @@ def _offline_flag(name: str) -> bool: return os.environ.get(name, "").strip().lower() in _HF_OFFLINE_TRUE_VALUES -def hf_hub_offline() -> bool: - """True when ``HF_HUB_OFFLINE`` is truthy -- the ONLY offline flag - ``huggingface_hub`` honors natively, so it is the one that actually prevents - a Hub fetch. - - Use this (never :func:`hf_env_offline`) to decide whether a network call is - *impossible*. A security gate must not be skipped on the weaker signal: with - only ``TRANSFORMERS_OFFLINE`` set, hub operations still reach the network, so - skipping a scan there would let an unscanned repo be downloaded anyway. - """ - return _offline_flag("HF_HUB_OFFLINE") - - def hf_env_offline() -> bool: """True when either HF offline env var is truthy (strip+lower, on/true/yes/1). - This is the user's *intent* to work offline. ``huggingface_hub`` honors only - ``HF_HUB_OFFLINE``, so callers that must GUARANTEE no fetch have to pass - ``local_files_only`` explicitly (or consult :func:`hf_hub_offline`) rather - than assume this alone stops the network. + This is the user's *intent* to work offline, which is broader than what + ``huggingface_hub`` enforces: it honors ``HF_HUB_OFFLINE`` natively but + ignores ``TRANSFORMERS_OFFLINE`` entirely. Setting this alone therefore does + NOT stop a fetch -- callers that need that guarantee must pass + ``local_files_only = hf_env_offline()`` to the loader (as + ``core/rag/embeddings.py`` does), and anything that relaxes a check on the + strength of "we are offline" depends on those callers doing so. """ - return hf_hub_offline() or _offline_flag("TRANSFORMERS_OFFLINE") + return _offline_flag("HF_HUB_OFFLINE") or _offline_flag("TRANSFORMERS_OFFLINE") From 38b78c8d723b40ae1fb17e60458e96b4d3219694 Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sat, 18 Jul 2026 11:54:13 +0300 Subject: [PATCH 12/77] studio: scope the offline scan bypass to callers that load local-only The previous commit put the offline short-circuit inside _fetch_security_status, which is the malware gate shared by every loader -- so TRANSFORMERS_OFFLINE=1 disabled it for all of them, while only the RAG embedder had been changed to pass local_files_only. MLX inference (core/inference/worker.py -> FastMLXModel .from_pretrained), training and export call from_pretrained with no local-only argument, and huggingface_hub ignores that flag, so those paths could still fetch and deserialize an unscanned model with the gate switched off. The bypass is now an explicit local_only_load argument, defaulting to False, and only the two RAG embedding callers -- whose loader is pinned to the local cache by the same predicate -- opt in. Tests pin both halves: the shared gate must still scan under either offline flag by default, and no other caller may pass local_only_load without constraining its loader. --- studio/backend/core/rag/embeddings.py | 12 ++- studio/backend/routes/settings.py | 12 ++- studio/backend/tests/test_hf_offline_flags.py | 74 ++++++++++++++----- .../backend/utils/security/file_security.py | 51 ++++++++----- 4 files changed, 111 insertions(+), 38 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 8d291a76edc..f4c00d65e68 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -145,7 +145,17 @@ def _guard_model_security(name: str) -> None: 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 + from utils.utils import hf_env_offline + + # local_only_load: _get() pins SentenceTransformer to the local cache with + # the same predicate, so offline nothing can be fetched and the Hub scan + # would only stall on timeouts before failing open anyway. + blocked = evaluate_file_security( + name, + hf_token = token, + load_subdirs = load_subdirs, + local_only_load = hf_env_offline(), + ).blocked except Exception: return if blocked: diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 7089f37f5eb..49078d989de 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -434,7 +434,17 @@ def update_embedding_model( ) ) ) - if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked: + # local_only_load: this gate covers the RAG embedder, whose loader pins + # SentenceTransformer to the local cache with the same predicate, so + # offline the scan can only stall on timeouts before failing open. + from utils.utils import hf_env_offline + + if evaluate_file_security( + model, + hf_token = scan_token, + load_subdirs = load_subdirs, + local_only_load = hf_env_offline(), + ).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( diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index b4746bf8db5..1211bdee339 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -103,17 +103,29 @@ def _model_info(*a, **k): @pytest.mark.parametrize("var", ["HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"]) -def test_security_scan_short_circuits_when_offline(monkeypatch, var): - # Metadata-only lookup with no local fallback: on a session the user declared - # offline it must skip straight to its documented fail-open rather than burn - # both request timeouts (10s + 20s) on every save and load. +def test_shared_gate_still_scans_when_offline_by_default(monkeypatch, var): + # THE important one. This gate is shared by every loader (training, MLX, + # export, ...), and most do NOT constrain their loaders to the local cache, + # so an offline-looking env var alone must never disable the malware scan -- + # those paths can still fetch and deserialize an unscanned model. import utils.security.file_security as fs calls: list = [] _fake_hub(monkeypatch, calls) monkeypatch.setenv(var, "1") - assert fs._fetch_security_status("org/model", None) is None - assert calls == [], f"the scan must not hit the Hub under {var}" + assert fs._fetch_security_status("org/model", None) is None # fails open on error + assert calls, f"{var} alone must NOT bypass the shared malware gate" + + +def test_security_scan_short_circuits_for_a_local_only_caller(monkeypatch): + # A caller that guarantees a local-only load gets the Hub round-trip skipped + # instead of burning both request timeouts (10s + 20s) before failing open. + import utils.security.file_security as fs + + calls: list = [] + _fake_hub(monkeypatch, calls) + assert fs._fetch_security_status("org/model", None, True) is None + assert calls == [], "a local-only caller must not hit the Hub" def test_security_scan_runs_when_online(monkeypatch): @@ -125,20 +137,46 @@ def test_security_scan_runs_when_online(monkeypatch): assert calls, "online must attempt the Hub" -def test_embedding_loader_forces_local_only_when_offline(): - """The invariant the offline scan skip rests on. +def _read_backend(rel: str) -> str: + return (Path(__file__).resolve().parents[1] / rel).read_text(encoding = "utf-8") - Skipping the Hub security scan offline is only sound while every loader - behind that gate is pinned to the local cache by the SAME predicate -- if the - loader could still fetch, an unscanned repo's pickle would be downloaded and - deserialized. Pinned at source level because importing the loader would drag - in sentence_transformers/torch. + +def test_embedding_loader_forces_local_only_when_offline(): + """The invariant the RAG opt-in rests on. + + core/rag/embeddings.py is allowed to pass local_only_load because its loader + is pinned to the local cache by the SAME predicate. If that pin is removed + the opt-in silently becomes a hole -- an unscanned repo's pickle could be + fetched and deserialized -- so it is enforced here rather than by comment. + Checked at source level because importing the loader drags in + sentence_transformers/torch. """ - src = (Path(__file__).resolve().parents[1] / "core" / "rag" / "embeddings.py").read_text( - encoding = "utf-8" - ) + src = _read_backend("core/rag/embeddings.py") assert "local_files_only = hf_env_offline()" in src, ( "core/rag/embeddings.py must pin SentenceTransformer to local files when " - "hf_env_offline(); without it the offline security-scan skip in " - "utils/security/file_security.py is unsafe" + "hf_env_offline(); without it passing local_only_load to " + "evaluate_file_security is unsafe" ) + + +def test_only_the_rag_embedding_path_opts_into_the_bypass(): + """No other loader may claim local-only without constraining its loader. + + The MLX/inference, training and export gates call from_pretrained without a + local-only argument, so if one of them started passing local_only_load the + malware gate would be disabled for a path that can still fetch. + """ + allowed = {"core/rag/embeddings.py", "routes/settings.py"} + callers = [ + "core/inference/worker.py", + "core/training/worker.py", + "core/export/worker.py", + "routes/models.py", + "routes/inference.py", + ] + for rel in callers: + assert rel not in allowed + assert "local_only_load" not in _read_backend(rel), ( + f"{rel} passes local_only_load but does not pin its loader to the " + "local cache; that would disable the malware gate for a fetching path" + ) diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 706f1852aa6..9c01c793e4e 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -236,27 +236,31 @@ def _load_scan_target(model_name: str, load_subdirs: tuple) -> tuple: return model_name, load_subdirs -def _fetch_security_status(model_name: str, hf_token: Optional[str]): +def _fetch_security_status( + model_name: str, + hf_token: Optional[str], + local_only_load: bool = False, +): """``security_repo_status`` (a dict) or None if unavailable. Hub metadata only; retries once on a transient error, then returns None so the caller fails open. + + ``local_only_load`` is an explicit promise from the caller that the load it is + gating cannot reach the Hub. See :func:`evaluate_file_security`. """ from huggingface_hub import model_info as hf_model_info - from utils.utils import hf_env_offline - - # Effective offline: skip this metadata-only lookup instead of burning both - # request timeouts (10s + 20s) on a session the user declared offline. - # - # SAFETY INVARIANT: this is only sound because every loader behind this gate - # forces local-only loading from the SAME predicate -- see - # core/rag/embeddings.py, which passes local_files_only = hf_env_offline() to - # SentenceTransformer. Nothing can therefore be fetched here, so the scan's - # job (block a poisoned pickle from being downloaded and deserialized) is - # already served; the residual case, a model cached BEFORE it was flagged, is - # the same fail-open this function has always documented for an unavailable - # scan. If a loader ever stops honoring hf_env_offline(), this skip becomes - # unsafe -- tests/test_hf_offline_flags.py pins that coupling. - if hf_env_offline(): - logger.debug("HF security scan skipped for '%s': offline; failing open.", model_name) + + # Only when the CALLER guarantees a local-only load. This is deliberately not + # keyed off hf_env_offline() here: this gate is shared by every loader + # (training, MLX/inference, export, ...), and most of them do not pass + # local_files_only, so an offline-looking session can still fetch and + # deserialize an unscanned model through those paths. Skipping the scan for + # them would disable the malware gate outright, so the bypass has to be + # opted into by the callers that actually hold the invariant. + if local_only_load: + logger.debug( + "HF security scan skipped for '%s': caller loads local-only; failing open.", + model_name, + ) return None token_arg = hf_token if hf_token else False @@ -287,6 +291,7 @@ def evaluate_file_security( 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. @@ -297,6 +302,16 @@ 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`` lets a caller skip the Hub round-trip when it GUARANTEES the + load it is gating cannot fetch -- e.g. the RAG embedder, which passes + ``local_files_only`` to SentenceTransformer from the same predicate. Pass it only + with that guarantee in hand: the scan exists to stop a poisoned pickle being + downloaded and deserialized, so claiming local-only while the loader can still + fetch disables the gate. Default False, because most callers here (training, + MLX/inference, export) do not constrain their loaders. Even when set, an + already-cached repo still loads unscanned -- the same fail-open this function + documents for an unavailable scan. """ # 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/. @@ -312,7 +327,7 @@ 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") - status = _fetch_security_status(model_name, hf_token) + status = _fetch_security_status(model_name, hf_token, local_only_load) if not isinstance(status, dict): return FileSecurityDecision( model_name, False, reason = "scan unavailable; allowed (fail-open)" From 164b5d81b9cc7aac701b94800841849891994310 Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sat, 18 Jul 2026 13:16:24 +0300 Subject: [PATCH 13/77] studio: capture offline state once, and probe the ST cache root Two holes in the offline embedding path: - _get() read hf_env_offline() twice: once inside _guard_model_security and again for local_files_only. _hf_offline_if_dns_dead() mutates the process-wide offline vars and restores them on exit, so a concurrent load could see True in the guard -- skipping the Hub malware scan -- and False by the time the constructor ran, fetching and deserializing the unscanned repo and breaking the very invariant that licenses the bypass. The value is now read once in _get() and passed to both; _guard_model_security takes it as an argument instead of re-deriving it. - The cache probe searched only HF_HUB_CACHE. SentenceTransformer downloads into SENTENCE_TRANSFORMERS_HOME when that is set, using the same models--org--name/snapshots layout under a different root, so a model fully present there looked uncached and was rejected with a 409 offline even though the local-only loader could load it. Snapshot lookup now covers both roots. --- studio/backend/core/rag/embeddings.py | 28 ++++++---- studio/backend/tests/test_hf_offline_flags.py | 30 +++++++---- .../tests/test_is_embedding_model_offline.py | 33 ++++++++++++ studio/backend/utils/models/model_config.py | 54 +++++++++++++------ 4 files changed, 111 insertions(+), 34 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index f4c00d65e68..ecf76709939 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -129,11 +129,19 @@ 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_load: bool) -> 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_load`` MUST be the same value the caller passes to + SentenceTransformer as ``local_files_only`` -- it is what licenses skipping the + Hub scan, so it is taken as an argument rather than re-read from the + environment here. ``_hf_offline_if_dns_dead()`` mutates the process-wide + offline vars and restores them, so two separate reads can disagree: this + could skip the scan while the constructor then loaded with + ``local_files_only=False`` and fetched the unscanned repo. """ try: from utils.security import evaluate_file_security, security_load_subdirs @@ -145,16 +153,11 @@ def _guard_model_security(name: str) -> None: load_subdirs = tuple( dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token))) ) - from utils.utils import hf_env_offline - - # local_only_load: _get() pins SentenceTransformer to the local cache with - # the same predicate, so offline nothing can be fetched and the Hub scan - # would only stall on timeouts before failing open anyway. blocked = evaluate_file_security( name, hf_token = token, load_subdirs = load_subdirs, - local_only_load = hf_env_offline(), + local_only_load = local_only_load, ).blocked except Exception: return @@ -179,7 +182,14 @@ def _get(model_name: str | None = None): device = _device() logger.info("loading embedding model %s on %s", name, device) - _guard_model_security(name) + # Read the offline state ONCE and use that single value for both the + # security gate and the loader. _hf_offline_if_dns_dead() mutates the + # process-wide offline vars and restores them on exit, so re-reading + # for the constructor could yield False after the guard had already + # skipped the Hub scan on True -- and the load would then fetch and + # deserialize the unscanned repo. + local_only = hf_env_offline() + _guard_model_security(name, local_only) # Propagate the user's offline intent into the loader: SentenceTransformer # performs its own Hub operations, and huggingface_hub honors only # HF_HUB_OFFLINE, so a TRANSFORMERS_OFFLINE-only session would otherwise @@ -188,7 +198,7 @@ def _get(model_name: str | None = None): name, device = device, model_kwargs = dtype_kwargs("float16"), - local_files_only = hf_env_offline(), + local_files_only = local_only, ) _name = name return _model diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index 1211bdee339..db71f325f9b 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -145,17 +145,29 @@ def test_embedding_loader_forces_local_only_when_offline(): """The invariant the RAG opt-in rests on. core/rag/embeddings.py is allowed to pass local_only_load because its loader - is pinned to the local cache by the SAME predicate. If that pin is removed - the opt-in silently becomes a hole -- an unscanned repo's pickle could be - fetched and deserialized -- so it is enforced here rather than by comment. - Checked at source level because importing the loader drags in - sentence_transformers/torch. + is pinned to the local cache by the SAME value. It must be read ONCE and + shared: _hf_offline_if_dns_dead() mutates the process-wide offline vars and + restores them, so two separate hf_env_offline() reads can disagree and the + guard could skip the scan on True while the constructor then fetched the + unscanned repo with local_files_only=False. Checked at source level because + importing the loader drags in sentence_transformers/torch. """ src = _read_backend("core/rag/embeddings.py") - assert "local_files_only = hf_env_offline()" in src, ( - "core/rag/embeddings.py must pin SentenceTransformer to local files when " - "hf_env_offline(); without it passing local_only_load to " - "evaluate_file_security is unsafe" + assert "local_only = hf_env_offline()" in src, ( + "core/rag/embeddings.py must capture the offline state once in _get()" + ) + assert "_guard_model_security(name, local_only)" in src, ( + "the security guard must receive the captured value, not re-read the env" + ) + assert "local_files_only = local_only" in src, ( + "SentenceTransformer must be pinned with the SAME captured value; a second " + "hf_env_offline() read can flip to False and fetch the unscanned repo" + ) + # And the guard must not re-derive it internally. + guard = src.split("def _guard_model_security", 1)[1].split("\ndef ", 1)[0] + assert "hf_env_offline()" not in guard, ( + "_guard_model_security must take local_only_load as an argument so it " + "cannot observe a different offline state than the loader" ) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 824335733e7..16967efec66 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -168,6 +168,39 @@ def test_marker_empty_ref_is_cache_miss(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache("org/empty-ref") is None +def test_snapshots_found_in_sentence_transformers_home(tmp_path, monkeypatch): + # SentenceTransformer downloads into SENTENCE_TRANSFORMERS_HOME when set, + # using the same models--org--name layout under a different root. Searching + # only HF_HUB_CACHE would report a model that is fully present there as + # uncached -- offline that is a 409 for a model the local-only loader can load. + hf_root = tmp_path / "hf" + st_root = tmp_path / "st" + hf_root.mkdir() + snap = st_root / "models--org--model" / "snapshots" / "aaa" + snap.mkdir(parents = True) + (snap / "modules.json").write_text("[]") + + fake = types.ModuleType("huggingface_hub") + fake.constants = types.SimpleNamespace(HF_HUB_CACHE = str(hf_root)) + monkeypatch.setitem(sys.modules, "huggingface_hub", fake) + monkeypatch.setitem(sys.modules, "huggingface_hub.constants", fake.constants) + monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_root)) + + assert [p.name for p in mc._iter_hf_cache_snapshots("org/model")] == ["aaa"] + assert mc._embedding_marker_in_hf_cache("org/model") is True + + +def test_sentence_transformers_home_unset_keeps_hf_cache_only(tmp_path, monkeypatch): + hf_root = tmp_path / "hf" + hf_root.mkdir() + fake = types.ModuleType("huggingface_hub") + fake.constants = types.SimpleNamespace(HF_HUB_CACHE = str(hf_root)) + monkeypatch.setitem(sys.modules, "huggingface_hub", fake) + monkeypatch.setitem(sys.modules, "huggingface_hub.constants", fake.constants) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + assert mc._hf_cache_roots() == [Path(str(hf_root))] + + def test_marker_never_raises_when_cache_mutates(monkeypatch): # A snapshot vanishing mid-iteration (concurrent cached-model deletion) # must read as not-cached, not propagate a 500 out of the routes. diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index fed87475e7a..449592671c1 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1643,29 +1643,51 @@ def _local_gguf_companion_search_root(selected_path: str, gguf_file: str) -> str return str(gguf_dir) -def _iter_hf_cache_snapshots(repo_id: str): - """Yield HF cache snapshot dirs for *repo_id*, newest first. - - Empty if HF_HUB_CACHE is missing, the repo isn't cached, or has no - snapshots. Repo name match is case-insensitive to handle casing drift - between download time and lookup. +def _hf_cache_roots() -> list[Path]: + """Cache roots a model may have been downloaded into, HF_HUB_CACHE first. + + SentenceTransformer downloads into SENTENCE_TRANSFORMERS_HOME when that is + set, using the same ``models--org--name/snapshots/`` layout but under + a different root. Ignoring it makes a model that is fully present there look + uncached -- which offline means a 409 for a model the local-only loader could + actually load. """ + roots: list[Path] = [] try: from huggingface_hub import constants as hf_constants + roots.append(Path(hf_constants.HF_HUB_CACHE)) except Exception: - return + pass + st_home = (os.environ.get("SENTENCE_TRANSFORMERS_HOME") or "").strip() + if st_home: + try: + st_root = Path(st_home).expanduser() + if not any(st_root == root for root in roots): + roots.append(st_root) + except Exception: + pass + return roots + + +def _iter_hf_cache_snapshots(repo_id: str): + """Yield cache snapshot dirs for *repo_id*, newest first. - cache_dir = Path(hf_constants.HF_HUB_CACHE) + Searches every root in :func:`_hf_cache_roots` (HF_HUB_CACHE plus a distinct + SENTENCE_TRANSFORMERS_HOME). Empty if no root exists, the repo isn't cached, + or it has no snapshots. Repo name match is case-insensitive to handle casing + drift between download time and lookup. + """ target = f"models--{repo_id.replace('/', '--')}".lower() repo_dirs: list[Path] = [] - try: - if not cache_dir.is_dir(): - return - for entry in cache_dir.iterdir(): - if entry.is_dir() and entry.name.lower() == target: - repo_dirs.append(entry) - except OSError: - return + for cache_dir in _hf_cache_roots(): + try: + if not cache_dir.is_dir(): + continue + for entry in cache_dir.iterdir(): + if entry.is_dir() and entry.name.lower() == target: + repo_dirs.append(entry) + except OSError: + continue if not repo_dirs: return From f94f3ecb86d007c28a9acb04a16cace945ccd73f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 10:17:00 +0000 Subject: [PATCH 14/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_hf_offline_flags.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index db71f325f9b..c363b2cc3c9 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -153,12 +153,12 @@ def test_embedding_loader_forces_local_only_when_offline(): importing the loader drags in sentence_transformers/torch. """ src = _read_backend("core/rag/embeddings.py") - assert "local_only = hf_env_offline()" in src, ( - "core/rag/embeddings.py must capture the offline state once in _get()" - ) - assert "_guard_model_security(name, local_only)" in src, ( - "the security guard must receive the captured value, not re-read the env" - ) + assert ( + "local_only = hf_env_offline()" in src + ), "core/rag/embeddings.py must capture the offline state once in _get()" + assert ( + "_guard_model_security(name, local_only)" in src + ), "the security guard must receive the captured value, not re-read the env" assert "local_files_only = local_only" in src, ( "SentenceTransformer must be pinned with the SAME captured value; a second " "hf_env_offline() read can flip to False and fetch the unscanned repo" From 72eb4f466f17a4aa682721db4ad17e8ef5f171a8 Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sun, 19 Jul 2026 00:53:31 +0300 Subject: [PATCH 15/77] studio: probe the cache the ST loader actually uses, and require it be loadable Adding SENTENCE_TRANSFORMERS_HOME to the shared snapshot iterator was too broad in one direction and too narrow in another: - _get() builds SentenceTransformer with no cache_folder, so with ST_HOME set it searches THAT root only, never the Hub cache. Probing the union let offline validation pass on a repo cached only in the Hub cache, after which the loader looked in ST_HOME and failed. The Sentence-Transformers probe now resolves to exactly one root: ST_HOME when set, the Hub cache otherwise. - The shared iterator is also used by the GGUF detectors, whose downloads go through hf_hub_download with no cache_dir and therefore really do use the Hub cache. It is back to Hub-cache-only so detection cannot pick a snapshot the GGUF load will not find. - Casing normalization ran through resolve_cached_repo_id_case, which scans the Hub cache, so with ST_HOME set the requested spelling was persisted unchanged and the exact-case offline load missed the differently cased directory that detection had just accepted. It now resolves against the same roots detection uses, exact match first. - A snapshot carrying only modules.json no longer counts as cached: the online security preflight downloads that single file itself, and a partial download leaves it behind, so validation passed for a snapshot with no weights and the first RAG load then failed. A hit now requires the marker plus a config and at least one weight file. --- studio/backend/routes/settings.py | 20 +- .../tests/test_is_embedding_model_offline.py | 171 +++++++++++++----- studio/backend/utils/models/__init__.py | 2 + studio/backend/utils/models/model_config.py | 163 +++++++++++++---- 4 files changed, 265 insertions(+), 91 deletions(-) diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 49078d989de..a5d4fc36ba7 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -387,8 +387,8 @@ def update_embedding_model( A repo flagged unsafe by HF's security scan returns 403 instead: a hard block that ``force`` cannot bypass, so the UI must not offer "save anyway". Documents indexed under the previous model must be re-uploaded.""" - from utils.models import is_embedding_model - from utils.paths import is_local_path, resolve_cached_repo_id_case + from utils.models import is_embedding_model, resolve_st_cached_repo_id_case + from utils.paths import is_local_path try: model = validate_embedding_model(payload.embedding_model) @@ -478,14 +478,16 @@ def update_embedding_model( # accepts a case-insensitive cache hit, but the offline SentenceTransformer # load resolves the cache by exact case, so store the cached spelling (a # no-op when nothing case-matching is cached) to keep the model loadable. - # Skip the default: rewriting its casing would make the exact-string default - # comparison in set_rag_embedding_model() treat it as a custom override, so - # later changes to the configured default would stop applying. Skip local - # paths too: a relative directory like "org/model" is loaded from disk, and - # rewriting it to a case-insensitive HF cache collision ("Org/model") would - # stop resolving to that directory and be read as a Hub repo id instead. + # Resolved against the cache the ST loader itself searches (ST_HOME when set, + # else the Hub cache), so the persisted spelling is one the offline + # exact-case load can actually find. Skip the default: rewriting its casing + # would make the exact-string default comparison in set_rag_embedding_model() + # treat it as a custom override, so later changes to the configured default + # would stop applying. Skip local paths too: a relative directory like + # "org/model" is loaded from disk, and rewriting it to a case-insensitive + # cache collision ("Org/model") would stop resolving to that directory. if model != default_embedding_model() and not is_local_path(model): - model = resolve_cached_repo_id_case(model) + model = resolve_st_cached_repo_id_case(model) set_rag_embedding_model(model) logger.info( "settings.embedding_model_updated subject=%s model=%s forced=%s", diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 16967efec66..86e97e45ab2 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -72,6 +72,7 @@ def _repo( d = repo / "snapshots" / name d.mkdir(parents = True, exist_ok = True) (d / "config.json").write_text("{}") + (d / "model.safetensors").write_bytes(b"\0") if is_st: (d / "modules.json").write_text("[]") dirs.append(d) @@ -97,18 +98,18 @@ def _no_network(*a, **k): def test_marker_true_when_modules_json_present(tmp_path, monkeypatch): snaps = _repo(tmp_path, ("aaa", True)) - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) assert mc._embedding_marker_in_hf_cache("org/emb") is True def test_marker_false_when_cached_without_modules_json(tmp_path, monkeypatch): snaps = _repo(tmp_path, ("aaa", False)) - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) assert mc._embedding_marker_in_hf_cache("org/llm") is False def test_marker_none_when_not_cached(monkeypatch): - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(())) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(())) assert mc._embedding_marker_in_hf_cache("org/llm") is None @@ -117,20 +118,20 @@ def test_marker_prefers_refs_main_revision(tmp_path, monkeypatch): # modules.json) but the revision refs/main points at no longer is. The # active revision must win: an any-snapshot scan would wrongly say True. snaps = _repo(tmp_path, ("new", False), ("old", True), main_ref = "new") - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) assert mc._embedding_marker_in_hf_cache("org/was-embedder") is False def test_marker_refs_main_st_revision_is_true(tmp_path, monkeypatch): snaps = _repo(tmp_path, ("new", True), ("old", False), main_ref = "new") - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) assert mc._embedding_marker_in_hf_cache("org/is-embedder") is True def test_marker_missing_ref_falls_back_to_snapshot_scan(tmp_path, monkeypatch): # No refs/main recorded: keep the newest-first any-snapshot behavior. snaps = _repo(tmp_path, ("new", False), ("old", True)) - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) assert mc._embedding_marker_in_hf_cache("org/no-ref") is True @@ -140,7 +141,7 @@ def test_marker_ref_points_at_absent_snapshot_is_cache_miss(tmp_path, monkeypatc # (None) -- NOT a fall-through to a stale historical snapshot that has # modules.json. snaps = _repo(tmp_path, ("old", True), main_ref = "missing_commit") - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) assert mc._embedding_marker_in_hf_cache("org/partial") is None @@ -155,7 +156,7 @@ def test_marker_unreadable_ref_is_cache_miss(tmp_path, monkeypatch): refs_main = snaps[0].parent.parent / "refs" / "main" refs_main.parent.mkdir(parents = True, exist_ok = True) refs_main.mkdir() - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) assert mc._embedding_marker_in_hf_cache("org/unreadable-ref") is None @@ -164,41 +165,125 @@ def test_marker_empty_ref_is_cache_miss(tmp_path, monkeypatch): # truncate-and-rewrite): the active revision is unknown, so this is a cache # miss (None), NOT a fall-through to a stale snapshot that carries modules.json. snaps = _repo(tmp_path, ("old", True), main_ref = " \n") - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) assert mc._embedding_marker_in_hf_cache("org/empty-ref") is None -def test_snapshots_found_in_sentence_transformers_home(tmp_path, monkeypatch): - # SentenceTransformer downloads into SENTENCE_TRANSFORMERS_HOME when set, - # using the same models--org--name layout under a different root. Searching - # only HF_HUB_CACHE would report a model that is fully present there as - # uncached -- offline that is a 409 for a model the local-only loader can load. - hf_root = tmp_path / "hf" - st_root = tmp_path / "st" - hf_root.mkdir() - snap = st_root / "models--org--model" / "snapshots" / "aaa" - snap.mkdir(parents = True) - (snap / "modules.json").write_text("[]") - +def _fake_hf_cache(monkeypatch, root): fake = types.ModuleType("huggingface_hub") - fake.constants = types.SimpleNamespace(HF_HUB_CACHE = str(hf_root)) + fake.constants = types.SimpleNamespace(HF_HUB_CACHE = str(root)) monkeypatch.setitem(sys.modules, "huggingface_hub", fake) monkeypatch.setitem(sys.modules, "huggingface_hub.constants", fake.constants) - monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_root)) - assert [p.name for p in mc._iter_hf_cache_snapshots("org/model")] == ["aaa"] + +def _st_snapshot( + root, + repo_dir, + commit = "aaa", + loadable = True, +): + snap = root / repo_dir / "snapshots" / commit + snap.mkdir(parents = True) + (snap / "modules.json").write_text("[]") + if loadable: + (snap / "config.json").write_text("{}") + (snap / "model.safetensors").write_bytes(b"\0") + return snap + + +def test_st_probe_uses_sentence_transformers_home(tmp_path, monkeypatch): + # _get() builds SentenceTransformer without cache_folder, so with ST_HOME set + # that is the ONLY cache the load searches. The probe must follow it, or a + # model present there is called uncached and rejected with a 409 offline. + hf_root, st_root = tmp_path / "hf", tmp_path / "st" + hf_root.mkdir() + _st_snapshot(st_root, "models--org--model") + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_root)) + assert [p.name for p in mc._iter_st_cache_snapshots("org/model")] == ["aaa"] assert mc._embedding_marker_in_hf_cache("org/model") is True -def test_sentence_transformers_home_unset_keeps_hf_cache_only(tmp_path, monkeypatch): +def test_st_probe_ignores_hub_cache_when_st_home_is_set(tmp_path, monkeypatch): + # The loader searches ST_HOME only, so a repo cached ONLY in the Hub cache + # must not validate -- otherwise validation passes and the load then fails. + hf_root, st_root = tmp_path / "hf", tmp_path / "st" + st_root.mkdir() + _st_snapshot(hf_root, "models--org--model") + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_root)) + assert mc._embedding_marker_in_hf_cache("org/model") is None + + +def test_gguf_probe_never_follows_st_home(tmp_path, monkeypatch): + # The GGUF path downloads with hf_hub_download and no cache_dir, so it uses + # the Hub cache. Letting its probe see ST_HOME would select a file the GGUF + # loader cannot find. + hf_root, st_root = tmp_path / "hf", tmp_path / "st" + hf_root.mkdir() + _st_snapshot(st_root, "models--org--model") + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_root)) + assert list(mc._iter_hf_cache_snapshots("org/model")) == [] + + +def test_st_probe_falls_back_to_hub_cache(tmp_path, monkeypatch): hf_root = tmp_path / "hf" hf_root.mkdir() - fake = types.ModuleType("huggingface_hub") - fake.constants = types.SimpleNamespace(HF_HUB_CACHE = str(hf_root)) - monkeypatch.setitem(sys.modules, "huggingface_hub", fake) - monkeypatch.setitem(sys.modules, "huggingface_hub.constants", fake.constants) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + assert mc._st_cache_roots() == [Path(str(hf_root))] + + +def test_marker_only_snapshot_is_not_loadable(tmp_path, monkeypatch): + # The online security preflight downloads modules.json on its own via + # hf_hub_download, and a partial download leaves it too. Accepting that + # offline passes validation and then fails on the first RAG load. + hf_root = tmp_path / "hf" + _st_snapshot(hf_root, "models--org--model", loadable = False) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + assert mc._embedding_marker_in_hf_cache("org/model") is False + + +def _case_sensitive_fs(tmp_path) -> bool: + probe = tmp_path / "_CaseProbe" + probe.mkdir() + return not (tmp_path / "_caseprobe").is_dir() + + +def test_st_casing_resolves_against_st_home(tmp_path, monkeypatch): + # resolve_cached_repo_id_case scans only the Hub cache, so with ST_HOME set it + # would persist the requested lower-case id while the exact-case offline load + # looks for models--BAAI--bge-m3 in ST_HOME and misses it. + if not _case_sensitive_fs(tmp_path): + pytest.skip("casing only diverges on a case-sensitive filesystem") + hf_root, st_root = tmp_path / "hf", tmp_path / "st" + hf_root.mkdir() + (st_root / "models--BAAI--bge-m3").mkdir(parents = True) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_root)) + assert mc.resolve_st_cached_repo_id_case("baai/bge-m3") == "BAAI/bge-m3" + + +def test_st_casing_prefers_an_exact_match(tmp_path, monkeypatch): + if not _case_sensitive_fs(tmp_path): + pytest.skip("casing only diverges on a case-sensitive filesystem") + hf_root, st_root = tmp_path / "hf", tmp_path / "st" + hf_root.mkdir() + (st_root / "models--baai--bge-m3").mkdir(parents = True) + (st_root / "models--BAAI--bge-m3").mkdir(parents = True) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_root)) + assert mc.resolve_st_cached_repo_id_case("baai/bge-m3") == "baai/bge-m3" + + +def test_st_casing_noop_when_uncached(tmp_path, monkeypatch): + hf_root = tmp_path / "hf" + hf_root.mkdir() + _fake_hf_cache(monkeypatch, hf_root) monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) - assert mc._hf_cache_roots() == [Path(str(hf_root))] + assert mc.resolve_st_cached_repo_id_case("org/not-cached") == "org/not-cached" def test_marker_never_raises_when_cache_mutates(monkeypatch): @@ -207,7 +292,7 @@ def test_marker_never_raises_when_cache_mutates(monkeypatch): def _exploding_iter(repo): raise FileNotFoundError("snapshot removed underneath") - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", _exploding_iter) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", _exploding_iter) assert mc._embedding_marker_in_hf_cache("org/racing") is None @@ -217,7 +302,7 @@ def test_is_embedding_model_survives_cache_race_online(monkeypatch): def _exploding_iter(repo): raise FileNotFoundError("snapshot removed underneath") - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", _exploding_iter) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", _exploding_iter) _fake_hf_model_info( monkeypatch, lambda name, token = None: types.SimpleNamespace( @@ -235,7 +320,7 @@ def test_offline_cached_st_detected_via_marker_no_network(tmp_path, monkeypatch) # modules.json marker with no model_info() network call that would hang on # DNS retries (#6817). snaps = _repo(tmp_path, ("aaa", True)) - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) monkeypatch.setenv("HF_HUB_OFFLINE", "1") _fake_hf_model_info(monkeypatch, _no_network) assert mc.is_embedding_model("unsloth/bge-small-en-v1.5") is True @@ -247,7 +332,7 @@ def test_online_defers_to_hub_over_stale_marker(tmp_path, monkeypatch): # reports embedding signals wins -- the stale local marker must not # short-circuit model_info(). snaps = _repo(tmp_path, ("aaa", True)) - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) calls = [] def _info(model_name, token = None): @@ -265,7 +350,7 @@ def test_online_permanent_hub_error_ignores_stale_marker(tmp_path, monkeypatch): # marker -- return False so the settings route surfaces its 409, and the # persisted model can't fail later when the loader refreshes from the Hub. snaps = _repo(tmp_path, ("aaa", True)) - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) class RepositoryNotFoundError(Exception): pass @@ -282,7 +367,7 @@ def test_online_hub_failure_falls_back_to_marker_uncached(tmp_path, monkeypatch) # caching: the degraded result must not become sticky, so a later successful # Hub lookup can still override it. snaps = _repo(tmp_path, ("aaa", True)) - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) _fake_hf_model_info(monkeypatch, _no_network) # raises -> Hub "unreachable" assert mc.is_embedding_model("org/emb") is True # marker fallback assert ("org/emb", None) not in mc._embedding_detection_cache # not poisoned @@ -293,7 +378,7 @@ def test_online_negative_does_not_block_later_offline_download(tmp_path, monkeyp # The repo is then downloaded WITH modules.json and the session goes offline; # the offline path re-probes the marker (never consulting the online memo), # so the freshly downloaded embedder is detected instead of the stale False. - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(())) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(())) def _info(model_name, token = None): return types.SimpleNamespace(tags = ["text-generation"], pipeline_tag = "text-generation") @@ -304,7 +389,7 @@ def _info(model_name, token = None): # Now the model is downloaded (marker appears) and the session goes offline. snaps = _repo(tmp_path, ("aaa", True)) - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) monkeypatch.setenv("HF_HUB_OFFLINE", "1") assert mc.is_embedding_model("org/late-embedder") is True # marker re-probed @@ -314,7 +399,7 @@ def test_offline_retains_online_confirmed_positive(monkeypatch): # online and cached True. _hf_offline_if_dns_dead() then flips the process to # offline mid-load; the offline path must RETAIN that positive, not re-probe # the absent marker and downgrade a model already verified this session. - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(())) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(())) def _info(model_name, token = None): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) @@ -329,14 +414,14 @@ def _info(model_name, token = None): def test_offline_cached_non_st_returns_false_without_network(tmp_path, monkeypatch): snaps = _repo(tmp_path, ("aaa", False)) - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(snaps)) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) monkeypatch.setenv("HF_HUB_OFFLINE", "1") _fake_hf_model_info(monkeypatch, _no_network) assert mc.is_embedding_model("org/gemma-4-e4b") is False def test_offline_not_cached_returns_false_without_network(monkeypatch): - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(())) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(())) monkeypatch.setenv("HF_HUB_OFFLINE", "1") _fake_hf_model_info(monkeypatch, _no_network) assert mc.is_embedding_model("org/never-downloaded") is False @@ -345,7 +430,7 @@ def test_offline_not_cached_returns_false_without_network(monkeypatch): def test_online_uncached_still_uses_network(monkeypatch): # Not offline, not cached: the network model_info path must still run so an # embedding model that lacks modules.json (feature-extraction tag) is caught. - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(())) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(())) calls = [] def _info(model_name, token = None): @@ -361,7 +446,7 @@ def test_offline_negative_is_not_cached_then_online_detects(monkeypatch): # A tag-only embedder is not identifiable from modules.json. Offline returns # False WITHOUT caching, so once the env var clears the online model_info # lookup still runs and detects it -- the negative must not be sticky. - monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: iter(())) + monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(())) calls = [] def _info(model_name, token = None): diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 4a5fb7274cf..b1d598e2906 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -8,6 +8,7 @@ GgufVariantInfo, is_vision_model, is_embedding_model, + resolve_st_cached_repo_id_case, detect_audio_type, is_audio_input_type, VALID_AUDIO_TYPES, @@ -32,6 +33,7 @@ "GgufVariantInfo", "is_vision_model", "is_embedding_model", + "resolve_st_cached_repo_id_case", "detect_audio_type", "is_audio_input_type", "VALID_AUDIO_TYPES", diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 449592671c1..9706b7e56ae 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1643,43 +1643,40 @@ def _local_gguf_companion_search_root(selected_path: str, gguf_file: str) -> str return str(gguf_dir) -def _hf_cache_roots() -> list[Path]: - """Cache roots a model may have been downloaded into, HF_HUB_CACHE first. - - SentenceTransformer downloads into SENTENCE_TRANSFORMERS_HOME when that is - set, using the same ``models--org--name/snapshots/`` layout but under - a different root. Ignoring it makes a model that is fully present there look - uncached -- which offline means a 409 for a model the local-only loader could - actually load. - """ - roots: list[Path] = [] +def _hf_hub_cache_root() -> list[Path]: + """The Hub cache root, i.e. what ``hf_hub_download`` uses with no ``cache_dir``.""" try: from huggingface_hub import constants as hf_constants - roots.append(Path(hf_constants.HF_HUB_CACHE)) + return [Path(hf_constants.HF_HUB_CACHE)] except Exception: - pass + return [] + + +def _st_cache_roots() -> list[Path]: + """The cache root a SentenceTransformer load will ACTUALLY search. + + ``_get()`` constructs SentenceTransformer without ``cache_folder``, so it uses + SENTENCE_TRANSFORMERS_HOME when that is set and the Hub cache otherwise -- + one or the other, never both. Probing the union would let offline validation + pass on a repo cached only in the Hub cache while the loader then searched + ST_HOME and failed. Kept separate from the Hub-cache probe because the GGUF + path downloads with ``hf_hub_download`` and no ``cache_dir``, so it really + does use the Hub cache and must not be told about ST_HOME. + """ st_home = (os.environ.get("SENTENCE_TRANSFORMERS_HOME") or "").strip() if st_home: try: - st_root = Path(st_home).expanduser() - if not any(st_root == root for root in roots): - roots.append(st_root) + return [Path(st_home).expanduser()] except Exception: - pass - return roots - + return [] + return _hf_hub_cache_root() -def _iter_hf_cache_snapshots(repo_id: str): - """Yield cache snapshot dirs for *repo_id*, newest first. - Searches every root in :func:`_hf_cache_roots` (HF_HUB_CACHE plus a distinct - SENTENCE_TRANSFORMERS_HOME). Empty if no root exists, the repo isn't cached, - or it has no snapshots. Repo name match is case-insensitive to handle casing - drift between download time and lookup. - """ +def _iter_cache_snapshots_in(repo_id: str, roots: list[Path]): + """Snapshot dirs for *repo_id* under *roots*, newest first (shared machinery).""" target = f"models--{repo_id.replace('/', '--')}".lower() repo_dirs: list[Path] = [] - for cache_dir in _hf_cache_roots(): + for cache_dir in roots: try: if not cache_dir.is_dir(): continue @@ -1688,9 +1685,50 @@ def _iter_hf_cache_snapshots(repo_id: str): repo_dirs.append(entry) except OSError: continue - if not repo_dirs: - return + yield from _iter_snapshots_of(repo_dirs) + + +def resolve_st_cached_repo_id_case(repo_id: str) -> str: + """*repo_id* in the casing of its dir in the cache the ST loader will search. + + ``resolve_cached_repo_id_case`` scans only the Hub cache, so with + SENTENCE_TRANSFORMERS_HOME set it would leave the requested spelling + unchanged -- and the offline ``local_files_only`` load, which is exact-case on + a case-sensitive filesystem, would then miss the differently cased dir that + detection had just accepted. Normalizes against the same roots detection uses. + Returns *repo_id* unchanged for a local path, a non-repo string, or when + nothing case-matching is cached; prefers an exact match before any variant and + tie-breaks variants deterministically. + """ + if is_local_path(repo_id) or "/" not in repo_id: + return repo_id + prefix = "models--" + expected = f"{prefix}{repo_id.replace('/', '--')}" + target = expected.lower() + variants: list[str] = [] + for cache_dir in _st_cache_roots(): + try: + if not cache_dir.is_dir(): + continue + if (cache_dir / expected).is_dir(): + return repo_id # exact case already cached: keep the request + for entry in cache_dir.iterdir(): + if entry.is_dir() and entry.name.lower() == target: + variants.append(entry.name) + except OSError: + continue + if variants: + return sorted(variants)[0][len(prefix) :].replace("--", "/") + return repo_id + + +def _iter_st_cache_snapshots(repo_id: str): + """Snapshot dirs for *repo_id* in the cache the ST loader will search.""" + yield from _iter_cache_snapshots_in(repo_id, _st_cache_roots()) + +def _iter_snapshots_of(repo_dirs: list[Path]): + """Snapshot dirs across *repo_dirs*, newest first.""" snap_dirs: list[Path] = [] for repo_dir in repo_dirs: snapshots = repo_dir / "snapshots" @@ -1704,8 +1742,6 @@ def _iter_hf_cache_snapshots(repo_id: str): continue except OSError: continue - if not snap_dirs: - return snap_dirs_with_mtime = [] for snap_dir in snap_dirs: try: @@ -1716,6 +1752,22 @@ def _iter_hf_cache_snapshots(repo_id: str): yield from (snap_dir for _, snap_dir in snap_dirs_with_mtime) +def _iter_hf_cache_snapshots(repo_id: str): + """Yield HUB cache snapshot dirs for *repo_id*, newest first. + + Deliberately the Hub cache ONLY. Its callers (the GGUF detectors) later + download with ``hf_hub_download`` and no ``cache_dir``, so that is the cache + their load actually uses; adding SENTENCE_TRANSFORMERS_HOME here would let + detection pick a file the GGUF loader cannot find. The Sentence-Transformers + probe is :func:`_iter_st_cache_snapshots`. + + Empty if the root does not exist, the repo isn't cached, or it has no + snapshots. Repo name match is case-insensitive to handle casing drift between + download time and lookup. + """ + yield from _iter_cache_snapshots_in(repo_id, _hf_hub_cache_root()) + + def _list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]: """Variants from the local HF cache snapshot, or None if not cached. @@ -2027,13 +2079,46 @@ def download_gguf_file( _embedding_detection_cache: Dict[tuple, bool] = {} -def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: - """Sentence-transformers detection from the local HF cache, no network call. +_ST_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pt", ".onnx") + + +def _snapshot_is_loadable_st_model(snap: Path) -> bool: + """True when *snap* holds a sentence-transformers model that can actually load. - True/False when the ACTIVE cached revision carries / lacks the - ``modules.json`` marker (the same signal ``is_embedding_model`` uses for - local paths), None when the repo is not in the cache (or the cache is - unreadable). The revision ``refs/main`` resolves to is authoritative when + ``modules.json`` alone is not enough: the online security preflight downloads + exactly that one file via ``hf_hub_download``, and a partial download can + leave it behind too, so a snapshot can carry the marker while the weights and + config SentenceTransformer needs are absent. Accepting that offline would + pass validation and then fail on the first RAG load. Requires the marker plus + a config and at least one weight file somewhere in the snapshot. + """ + try: + if not (snap / "modules.json").is_file(): + return False + if not any( + (snap / name).is_file() for name in ("config.json", "config_sentence_transformers.json") + ): + return False + for path in snap.rglob("*"): + try: + if path.is_file() and path.suffix.lower() in _ST_WEIGHT_SUFFIXES: + return True + except OSError: + continue + return False + except OSError: + return False + + +def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: + """Sentence-transformers detection from the local cache, no network call. + + Probes the cache the ST loader will actually search (see + :func:`_st_cache_roots`), so a hit here means the load can find it too. + True/False when the ACTIVE cached revision is / is not a LOADABLE + sentence-transformers snapshot -- marker plus config plus weights, not the + bare ``modules.json`` the security preflight may have fetched on its own -- + and None when the repo is not in that cache (or the cache is unreadable). The revision ``refs/main`` resolves to is authoritative when recorded: the cache keeps snapshots of older revisions, and a repo that later stopped (or started) being a sentence-transformers model must be judged by its current revision, not any historical one. Snapshots are only @@ -2041,7 +2126,7 @@ def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: underneath (concurrent model deletion) reads as not-cached so callers keep their normal fallback.""" try: - snapshots = list(_iter_hf_cache_snapshots(repo_id)) + snapshots = list(_iter_st_cache_snapshots(repo_id)) if not snapshots: return None # Prefer the snapshot refs/main points at (the active revision). @@ -2072,10 +2157,10 @@ def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: preferred = snapshots_dir / commit if not preferred.is_dir(): return None - return (preferred / "modules.json").is_file() + return _snapshot_is_loadable_st_model(preferred) for snap in snapshots: try: - if (snap / "modules.json").is_file(): + if _snapshot_is_loadable_st_model(snap): return True except OSError: continue From 1a2681296e5df498ec11400674e845d23cfffdb3 Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sun, 19 Jul 2026 01:14:18 +0300 Subject: [PATCH 16/77] studio: thread the captured offline state into the module probe, fix the gate shard - _st_module_subdirs() re-read the process env for its local_files_only. With _hf_offline_if_dns_dead() flipping those vars from another thread, a load that captured local_only=False could still force this probe local-only, get () back because modules.json is not cached, and leave the scan with NO module load roots -- a Hub-flagged pickle under 0_Transformer/ would then pass as an unreferenced nested artifact while the loader fetched and deserialized it. It now takes the captured predicate as an argument, and the settings route reads the state once and uses that single value for both the probe and the scan. - Skip ST-cache casing on the llama-server backend. Nothing there loads through SentenceTransformer: the embedder derives a GGUF companion from the saved spelling and fetches it from the HUB cache, so normalizing to an ST_HOME spelling would point it at a repo _hf_gguf_backend_error() never validated (BAAI/bge-m3-GGUF instead of the checked baai/bge-m3-GGUF). - Fix the security-gate shard, which the signature change had broken: the direct _guard_model_security / _st_module_subdirs callers now pass the new argument (they were raising TypeError before reaching any assertion), and the casing tests patch utils.models.resolve_st_cached_repo_id_case, which the route actually calls, instead of the Hub-only resolver it no longer uses -- those patches were being silently ignored. --- studio/backend/core/rag/embeddings.py | 29 +++++-- studio/backend/routes/settings.py | 40 ++++++--- .../test_embedding_model_security_gate.py | 85 ++++++++++++++----- 3 files changed, 112 insertions(+), 42 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index ecf76709939..d2fa115585e 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -80,12 +80,20 @@ def _ambient_hf_token() -> str | None: return None -def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]: +def _st_module_subdirs(name: str, token: str | None, local_only: bool) -> tuple[str, ...]: """The module directories a SentenceTransformer load reads weights from, taken from the repo's ``modules.json`` (each module's non-empty ``path``, e.g. ``0_Transformer``). ST deserializes ``pytorch_model.bin`` from these dirs, so they are load roots for the security scan: a flagged pickle directly under one must block. Returns () on any failure (no modules.json, offline, malformed) so the guard never bricks the embedder. + + ``local_only`` MUST be the value the caller captured for the load, not a fresh + read of the environment. ``_hf_offline_if_dns_dead()`` flips the process-wide + offline vars and restores them, so re-reading here could force this probe + local-only, return () because modules.json is not cached, and leave the scan + with NO module load roots -- a flagged pickle under ``0_Transformer/`` would + then pass as an unreferenced nested artifact while the loader, using the + captured predicate, still fetched and deserialized it. """ try: import json @@ -103,18 +111,18 @@ def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]: else: from huggingface_hub import hf_hub_download from huggingface_hub.utils import EntryNotFoundError - from utils.utils import hf_env_offline try: - # local_files_only when offline: huggingface_hub honors only - # HF_HUB_OFFLINE natively, so a TRANSFORMERS_OFFLINE-only session - # would otherwise block on network timeouts here even though the - # cached snapshot already has modules.json. + # local_files_only from the CAPTURED predicate: huggingface_hub + # honors only HF_HUB_OFFLINE natively, so an offline session would + # otherwise block on network timeouts here even though the cached + # snapshot already has modules.json -- but it must be the same + # value the load uses, never a fresh env read (see the docstring). local = hf_hub_download( name, "modules.json", token = token or None, - local_files_only = hf_env_offline(), + local_files_only = local_only, ) except EntryNotFoundError: return () @@ -151,7 +159,12 @@ def _guard_model_security(name: str, local_only_load: bool) -> None: # 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))) + dict.fromkeys( + ( + *security_load_subdirs(name, token), + *_st_module_subdirs(name, token, local_only_load), + ) + ) ) blocked = evaluate_file_security( name, diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index a5d4fc36ba7..cc46e147716 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -413,6 +413,14 @@ def update_embedding_model( scan_st_pickle = ( model != default_embedding_model() and not is_local_gguf and not _llama_backend_active() ) + # Read the offline state ONCE for this request and reuse it: the module probe + # and the scan must agree, and _hf_offline_if_dns_dead() can flip the + # process-wide vars between two reads (a probe forced local-only can return no + # module roots, letting a flagged 0_Transformer/ pickle pass as a nested + # artifact while the scan was skipped on a different value). + from utils.utils import hf_env_offline + + local_only_load = hf_env_offline() if scan_st_pickle: # Malware/pickle gate before we persist a repo the embedder later loads with # SentenceTransformer. Runs even under force (force only skips the is-embedding @@ -430,20 +438,18 @@ def update_embedding_model( dict.fromkeys( ( *security_load_subdirs(model, scan_token), - *_st_module_subdirs(model, scan_token), + *_st_module_subdirs(model, scan_token, local_only_load), ) ) ) # local_only_load: this gate covers the RAG embedder, whose loader pins # SentenceTransformer to the local cache with the same predicate, so # offline the scan can only stall on timeouts before failing open. - from utils.utils import hf_env_offline - if evaluate_file_security( model, hf_token = scan_token, load_subdirs = load_subdirs, - local_only_load = hf_env_offline(), + 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. @@ -480,13 +486,25 @@ def update_embedding_model( # no-op when nothing case-matching is cached) to keep the model loadable. # Resolved against the cache the ST loader itself searches (ST_HOME when set, # else the Hub cache), so the persisted spelling is one the offline - # exact-case load can actually find. Skip the default: rewriting its casing - # would make the exact-string default comparison in set_rag_embedding_model() - # treat it as a custom override, so later changes to the configured default - # would stop applying. Skip local paths too: a relative directory like - # "org/model" is loaded from disk, and rewriting it to a case-insensitive - # cache collision ("Org/model") would stop resolving to that directory. - if model != default_embedding_model() and not is_local_path(model): + # exact-case load can actually find. Three cases are deliberately left alone: + # + # * the default -- rewriting its casing would make the exact-string default + # comparison in set_rag_embedding_model() treat it as a custom override, + # so later changes to the configured default would stop applying; + # * a local path -- a relative directory like "org/model" is loaded from + # disk, and rewriting it to a case-insensitive cache collision + # ("Org/model") would stop resolving to that directory; + # * the llama-server backend -- it does not load through SentenceTransformer + # at all. It derives a GGUF companion via effective_gguf_repo() from this + # saved spelling and fetches it with hf_hub_download, i.e. from the HUB + # cache. Normalizing to an ST_HOME spelling would change the derived repo + # to one _hf_gguf_backend_error() never validated (BAAI/bge-m3-GGUF rather + # than the checked baai/bge-m3-GGUF) and which may be absent offline. + if ( + model != default_embedding_model() + and not is_local_path(model) + and not _llama_backend_active() + ): model = resolve_st_cached_repo_id_case(model) set_rag_embedding_model(model) logger.info( diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index 738f7cd01eb..5919f0fe981 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -43,7 +43,9 @@ def client(monkeypatch): # offline and deterministic for the endpoint tests that use this fixture. import core.rag.embeddings as embeddings - monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ()) + monkeypatch.setattr( + embeddings, "_st_module_subdirs", lambda name, token = None, local_only = False: () + ) saved: dict = {} monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) @@ -144,7 +146,9 @@ def test_runtime_llama_fallback_skips_the_st_pickle_scan(monkeypatch): # though the auto resolver would still say sentence-transformers. monkeypatch.setattr(embeddings, "_backend", LlamaServerBackend()) monkeypatch.setattr(embeddings, "_resolve_auto", lambda: "sentence-transformers") - monkeypatch.setattr(embeddings, "_st_module_subdirs", lambda name, token = None: ()) + monkeypatch.setattr( + embeddings, "_st_module_subdirs", lambda name, token = None, local_only = False: () + ) saved: dict = {} monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") @@ -227,7 +231,9 @@ def test_settings_scan_scopes_module_subdirs(monkeypatch): import core.rag.embeddings as embeddings monkeypatch.setattr( - embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",) + embeddings, + "_st_module_subdirs", + lambda name, token = None, local_only = False: ("0_Transformer",), ) seen = {} @@ -265,9 +271,9 @@ def test_custom_model_saved_in_cache_casing(client, monkeypatch): # BAAI/bge-m3). force=True bypasses the verification gate to isolate this. c, saved = client monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) - import utils.paths as _paths + import utils.models as _models - monkeypatch.setattr(_paths, "resolve_cached_repo_id_case", lambda m: "BAAI/bge-m3") + monkeypatch.setattr(_models, "resolve_st_cached_repo_id_case", lambda m: "BAAI/bge-m3") r = c.put("/embedding-model", json = {"embedding_model": "baai/bge-m3", "force": True}) assert r.status_code == 200 assert saved.get("model") == "BAAI/bge-m3" @@ -279,12 +285,12 @@ def test_local_path_model_is_not_casing_normalized(client, monkeypatch, tmp_path # directory and be read as a Hub repo id, so normalization must be skipped. c, saved = client monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) - import utils.paths as _paths + import utils.models as _models def _must_not_run(m): raise AssertionError("a local path must not be casing-normalized") - monkeypatch.setattr(_paths, "resolve_cached_repo_id_case", _must_not_run) + monkeypatch.setattr(_models, "resolve_st_cached_repo_id_case", _must_not_run) local_dir = tmp_path / "org" / "model" local_dir.mkdir(parents = True) r = c.put("/embedding-model", json = {"embedding_model": str(local_dir), "force": True}) @@ -292,18 +298,49 @@ def _must_not_run(m): assert saved.get("model") == str(local_dir) +def test_llama_backend_model_is_not_casing_normalized(monkeypatch): + # On the llama-server backend nothing loads through SentenceTransformer: the + # embedder derives a GGUF companion from this saved spelling via + # effective_gguf_repo() and fetches it with hf_hub_download, i.e. from the HUB + # cache. Rewriting to an ST_HOME spelling would point it at a GGUF repo that + # _hf_gguf_backend_error() never validated and that may be absent offline. + saved: dict = {} + monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") + monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) + monkeypatch.setattr(settings, "set_rag_embedding_model", lambda v: saved.setdefault("model", v)) + monkeypatch.setattr(settings, "_llama_backend_active", lambda: True) + monkeypatch.setattr(settings, "_resolves_as_local_gguf", lambda m: False) + monkeypatch.setattr(settings, "get_rag_embedding_model", lambda: saved.get("model", "")) + monkeypatch.setattr(settings, "get_stored_embedding_model", lambda: saved.get("model")) + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + import utils.models as _models + + def _must_not_run(m): + raise AssertionError("the llama backend must not be ST-casing-normalized") + + monkeypatch.setattr(_models, "resolve_st_cached_repo_id_case", _must_not_run) + + app = FastAPI() + app.include_router(settings.router) + app.dependency_overrides[settings.get_current_subject] = lambda: "admin" + c = TestClient(app, raise_server_exceptions = False) + r = c.put("/embedding-model", json = {"embedding_model": "baai/bge-m3", "force": True}) + assert r.status_code == 200 + assert saved.get("model") == "baai/bge-m3" # the spelling the GGUF check validated + + def test_default_model_is_not_casing_normalized(client, monkeypatch): # Submitting the exact default must NOT be run through cache-casing # normalization: rewriting it would make set_rag_embedding_model()'s exact # default comparison treat it as a custom override, so later default changes - # would stop applying. resolve_cached_repo_id_case must not be consulted. + # would stop applying. resolve_st_cached_repo_id_case must not be consulted. c, saved = client - import utils.paths as _paths + import utils.models as _models def _must_not_run(m): raise AssertionError("the default must not be casing-normalized") - monkeypatch.setattr(_paths, "resolve_cached_repo_id_case", _must_not_run) + monkeypatch.setattr(_models, "resolve_st_cached_repo_id_case", _must_not_run) r = c.put("/embedding-model", json = {"embedding_model": "unsloth/default-embed"}) assert r.status_code == 200 assert saved.get("model") == "unsloth/default-embed" @@ -313,13 +350,13 @@ def test_load_sink_refuses_flagged_model(monkeypatch): monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = True)) import core.rag.embeddings as embeddings with pytest.raises(embeddings.UnsafeEmbeddingModelError): - embeddings._guard_model_security("attacker/malicious-embed") + embeddings._guard_model_security("attacker/malicious-embed", False) def test_load_sink_allows_clean_model(monkeypatch): monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) import core.rag.embeddings as embeddings - embeddings._guard_model_security("acme/clean-embed") # no raise + embeddings._guard_model_security("acme/clean-embed", False) # no raise def test_sink_threads_ambient_token_into_scan(monkeypatch): @@ -327,17 +364,17 @@ def test_sink_threads_ambient_token_into_scan(monkeypatch): # loader's own token to the scan, or it fails open for the repo that still loads. seen = {} mod = _types.ModuleType("utils.security") - mod.security_load_subdirs = ( - lambda name, token = None: seen.setdefault("subdirs_token", token) or () + mod.security_load_subdirs = lambda name, token = None: ( + seen.setdefault("subdirs_token", token) or () + ) + mod.evaluate_file_security = lambda *a, **k: ( + seen.setdefault("scan_token", k.get("hf_token")) or _Decision(False) ) - mod.evaluate_file_security = lambda *a, **k: seen.setdefault( - "scan_token", k.get("hf_token") - ) or _Decision(False) monkeypatch.setitem(sys.modules, "utils.security", mod) import core.rag.embeddings as embeddings monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: "hf_ambient") - embeddings._guard_model_security("acme/gated-embed") + embeddings._guard_model_security("acme/gated-embed", False) assert seen["scan_token"] == "hf_ambient" assert seen["subdirs_token"] == "hf_ambient" @@ -353,16 +390,18 @@ def _capture(*a, **k): return _Decision(False) mod = _types.ModuleType("utils.security") - mod.security_load_subdirs = lambda name, token = None: () + mod.security_load_subdirs = lambda name, token = None, local_only = False: () mod.evaluate_file_security = _capture monkeypatch.setitem(sys.modules, "utils.security", mod) import core.rag.embeddings as embeddings monkeypatch.setattr(embeddings, "_ambient_hf_token", lambda: None) monkeypatch.setattr( - embeddings, "_st_module_subdirs", lambda name, token = None: ("0_Transformer",) + embeddings, + "_st_module_subdirs", + lambda name, token = None, local_only = False: ("0_Transformer",), ) - embeddings._guard_model_security("acme/embed-with-module-dir") + embeddings._guard_model_security("acme/embed-with-module-dir", False) assert "0_Transformer" in seen["subdirs"] @@ -381,7 +420,7 @@ def test_st_module_subdirs_reads_local_modules_json(tmp_path, monkeypatch): ] ) ) - subdirs = embeddings._st_module_subdirs(str(tmp_path), None) + subdirs = embeddings._st_module_subdirs(str(tmp_path), None, False) assert subdirs == ("0_Transformer", "1_Pooling") @@ -395,7 +434,7 @@ def _boom(*a, **k): raise RuntimeError("offline") monkeypatch.setattr(huggingface_hub, "hf_hub_download", _boom) - assert embeddings._st_module_subdirs("acme/no-such-repo-xyz", None) == () + assert embeddings._st_module_subdirs("acme/no-such-repo-xyz", None, False) == () def test_security_block_is_not_swallowed_by_llama_fallback(monkeypatch): From a3fb2e95d03686c0dedb63e56f64e1bb1d550469 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 19 Jul 2026 08:11:42 +0000 Subject: [PATCH 17/77] studio: accept only torch-loadable weights in the offline ST probe; fix re-export lint _snapshot_is_loadable_st_model accepted a cached snapshot whose only weights were .onnx (or .pt), but the RAG loader builds SentenceTransformer with the default torch backend, so such a snapshot passed offline validation and then failed on the first load, the exact validate-then-fail this helper exists to prevent. Restrict _ST_WEIGHT_SUFFIXES to .safetensors and .bin and add a regression test for an ONNX-only snapshot. Also teach scripts/verify_import_hoist.py that names listed in a module-level __all__ are uses, so the legitimately added resolve_st_cached_repo_id_case re-export in utils/models/__init__.py no longer trips HOISTED-IMPORT-UNUSED. Covered by two new self-test cases. --- scripts/verify_import_hoist.py | 52 +++++++++++++++++++ .../tests/test_is_embedding_model_offline.py | 17 ++++++ studio/backend/utils/models/model_config.py | 11 +++- 3 files changed, 78 insertions(+), 2 deletions(-) diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index 22a21a2ebce..d46633c066a 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -447,6 +447,34 @@ def _legb_chain(scope: Scope) -> list[Scope]: # ---------------------------------------------------------------- analysis +def _collect_dunder_all(tree: ast.Module) -> set[str]: + """Module-level ``__all__`` string entries (``= [...]``, ``+= [...]``, or an + annotated assign). + + A name listed in ``__all__`` is a public re-export, which is a real use of the + import that binds it. ``__all__`` entries are string constants, not ``Name`` + loads, so the load-based use scan never sees them -- without this a package + ``__init__`` that adds ``from .x import y`` purely to re-export ``y`` looks + like an unused hoist and trips ``HOISTED-IMPORT-UNUSED``. + """ + names: set[str] = set() + for node in tree.body: + if isinstance(node, ast.Assign): + targets = node.targets + elif isinstance(node, (ast.AnnAssign, ast.AugAssign)): + targets = [node.target] + else: + continue + if not any(isinstance(t, ast.Name) and t.id == "__all__" for t in targets): + continue + value = node.value + if isinstance(value, (ast.List, ast.Tuple, ast.Set)): + for elt in value.elts: + if isinstance(elt, ast.Constant) and isinstance(elt.value, str): + names.add(elt.value) + return names + + def _analyze(src: str): tree = ast.parse(src) b = _Builder() @@ -476,6 +504,16 @@ def _analyze(src: str): for n, bs in module.bindings.items() if any(x.kind in ("import", "importfrom") for x in bs) } + # Re-exports count as uses: a name listed in module-level __all__ that is bound + # by a module import is deliberately exported, not a dangling hoist. Fold its + # targets into the used set so HOISTED-IMPORT-UNUSED does not fire on a + # legitimately-added `from .x import y` in a package __init__. + for _n in _collect_dunder_all(tree): + _bs = module_imports.get(_n) + if _bs: + targets_by_scope.setdefault(module.qualname, set()).update( + x.target for x in _bs if x.target + ) module_dup = { n for n, bs in module.bindings.items() @@ -706,6 +744,20 @@ def used_targets(analysis) -> set[str]: "import os\nimport sys\ndef f(x):\n return x._b + sys.argv[0]\n", None, ), + "reexport_in_all_is_used": ( + # a new re-export added to a package __init__ (name in __all__, no load) is + # a deliberate export, NOT a botched hoist -> must not block + 'from pkg import a\n__all__ = ["a"]\n', + 'from pkg import a\nfrom pkg import b\n__all__ = ["a", "b"]\n', + None, + ), + "unused_import_not_in_all_still_blocks": ( + # the fix is precise: a newly-added module import that is neither loaded nor + # listed in __all__ is still a dangling/unused hoist + 'from pkg import a\n__all__ = ["a"]\n', + 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n', + "BLOCKER", + ), } diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 86e97e45ab2..38a88abe786 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -246,6 +246,23 @@ def test_marker_only_snapshot_is_not_loadable(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache("org/model") is False +def test_marker_onnx_only_snapshot_is_not_loadable(tmp_path, monkeypatch): + # A snapshot with the marker and config but ONLY an ONNX export cached is not + # loadable: _get() builds SentenceTransformer with the default Torch backend + # (no backend="onnx"), which reads model.safetensors / pytorch_model.bin, so + # accepting the ONNX offline would pass validation and then fail on the first + # RAG load -- the same validate-then-fail as the marker-only case. + hf_root = tmp_path / "hf" + snap = hf_root / "models--org--model" / "snapshots" / "aaa" + snap.mkdir(parents = True) + (snap / "modules.json").write_text("[]") + (snap / "config.json").write_text("{}") + (snap / "model.onnx").write_bytes(b"\0") + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + assert mc._embedding_marker_in_hf_cache("org/model") is False + + def _case_sensitive_fs(tmp_path) -> bool: probe = tmp_path / "_CaseProbe" probe.mkdir() diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 9706b7e56ae..0a4b6726d87 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2079,7 +2079,13 @@ def download_gguf_file( _embedding_detection_cache: Dict[tuple, bool] = {} -_ST_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pt", ".onnx") +# Only the weight formats the RAG loader's default backend consumes. _get() +# constructs SentenceTransformer without backend="onnx"/"openvino", so it loads +# through Torch (model.safetensors / pytorch_model.bin, incl. sharded). ".onnx" +# (and ".pt", which is not an HF weight filename) would let a snapshot cached +# with only an ONNX export pass offline validation and then fail on the first +# RAG load -- the exact validate-then-fail this helper exists to prevent. +_ST_WEIGHT_SUFFIXES = (".safetensors", ".bin") def _snapshot_is_loadable_st_model(snap: Path) -> bool: @@ -2090,7 +2096,8 @@ def _snapshot_is_loadable_st_model(snap: Path) -> bool: leave it behind too, so a snapshot can carry the marker while the weights and config SentenceTransformer needs are absent. Accepting that offline would pass validation and then fail on the first RAG load. Requires the marker plus - a config and at least one weight file somewhere in the snapshot. + a config and at least one weight file the default Torch backend can load + (``_ST_WEIGHT_SUFFIXES``) somewhere in the snapshot. """ try: if not (snap / "modules.json").is_file(): From ae2f5c0461cc9dd8a67db73cc75e7a2fb3a98bbe Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sun, 19 Jul 2026 12:38:36 +0300 Subject: [PATCH 18/77] studio: probe the exact repo dir and revision an offline load resolves The cache probe modelled the cache loosely rather than modelling what SentenceTransformer actually does with local_files_only=True: - It merged snapshots across every case-variant repo dir and then read refs/main from whichever held the newest one. With both models--baai--bge-m3 and models--BAAI--bge-m3 present, a complete embedding snapshot in the directory the loader opens could be judged by a newer partial snapshot in the other, failing validation for a usable model. It now selects the ONE directory the loader opens, by the same exact-case-first rule resolve_st_cached_repo_id_case uses to choose the spelling that gets persisted. - It fell back to scanning historical snapshots when refs/main was absent. With local_files_only the default revision is resolved THROUGH that ref, so a snapshot directory alone is not discoverable: the settings request succeeded and the loader then failed at first indexing. A missing, empty or unreadable ref is now a cache miss, and the historical scan is gone. The tests exercise the real lookup against a built cache tree instead of patching the snapshot iterator, so they now cover the directory selection and ref resolution the loader depends on. --- .../tests/test_is_embedding_model_offline.py | 156 +++++++++--------- studio/backend/utils/models/model_config.py | 112 +++++++------ 2 files changed, 142 insertions(+), 126 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 38a88abe786..35728b106d0 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -57,16 +57,21 @@ def _clean_state(monkeypatch): def _repo( tmp_path, + monkeypatch, *snapshots, main_ref = None, + repo_id = "org/model", + root = None, ): - """Fake HF cache repo dir: snapshots/[/modules.json] (+ refs/main). + """Build a real cache repo dir and point the ST cache root at it. - ``snapshots``: (name, sentence_transformer) tuples, oldest last (the - iterator under test yields newest first, so pass them in that order). - Returns the snapshot dirs in the given order. + ``snapshots``: (commit, sentence_transformer) tuples. A snapshot is written + fully loadable (config + weights) so only the marker varies. ``main_ref`` + writes refs/main. Returns the snapshot dirs in the given order. """ - repo = tmp_path / "models--org--model" + cache_root = root if root is not None else tmp_path / "cache" + cache_root.mkdir(parents = True, exist_ok = True) + repo = cache_root / f"models--{repo_id.replace('/', '--')}" dirs = [] for name, is_st in snapshots: d = repo / "snapshots" / name @@ -80,9 +85,14 @@ def _repo( refs = repo / "refs" refs.mkdir(parents = True, exist_ok = True) (refs / "main").write_text(main_ref) + monkeypatch.setattr(mc, "_st_cache_roots", lambda: [cache_root]) return dirs +def _no_cache(monkeypatch): + monkeypatch.setattr(mc, "_st_cache_roots", lambda: []) + + def _fake_hf_model_info(monkeypatch, fn): fake = types.ModuleType("huggingface_hub") fake.model_info = fn @@ -96,77 +106,78 @@ def _no_network(*a, **k): # ── _embedding_marker_in_hf_cache ── -def test_marker_true_when_modules_json_present(tmp_path, monkeypatch): - snaps = _repo(tmp_path, ("aaa", True)) - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) - assert mc._embedding_marker_in_hf_cache("org/emb") is True +def test_marker_true_for_a_loadable_st_revision(tmp_path, monkeypatch): + _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa") + assert mc._embedding_marker_in_hf_cache("org/model") is True -def test_marker_false_when_cached_without_modules_json(tmp_path, monkeypatch): - snaps = _repo(tmp_path, ("aaa", False)) - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) - assert mc._embedding_marker_in_hf_cache("org/llm") is False +def test_marker_false_when_active_revision_is_not_st(tmp_path, monkeypatch): + _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa") + assert mc._embedding_marker_in_hf_cache("org/model") is False def test_marker_none_when_not_cached(monkeypatch): - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(())) - assert mc._embedding_marker_in_hf_cache("org/llm") is None + _no_cache(monkeypatch) + assert mc._embedding_marker_in_hf_cache("org/model") is None -def test_marker_prefers_refs_main_revision(tmp_path, monkeypatch): - # The repo USED to be a sentence-transformers model (old snapshot has - # modules.json) but the revision refs/main points at no longer is. The - # active revision must win: an any-snapshot scan would wrongly say True. - snaps = _repo(tmp_path, ("new", False), ("old", True), main_ref = "new") - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) - assert mc._embedding_marker_in_hf_cache("org/was-embedder") is False +def test_marker_judges_the_revision_refs_main_points_at(tmp_path, monkeypatch): + # The repo USED to be a sentence-transformers model (the old snapshot has + # modules.json) but the revision refs/main points at no longer is. The active + # revision must win; scanning any snapshot would wrongly say True. + _repo(tmp_path, monkeypatch, ("new", False), ("old", True), main_ref = "new") + assert mc._embedding_marker_in_hf_cache("org/model") is False -def test_marker_refs_main_st_revision_is_true(tmp_path, monkeypatch): - snaps = _repo(tmp_path, ("new", True), ("old", False), main_ref = "new") - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) - assert mc._embedding_marker_in_hf_cache("org/is-embedder") is True +def test_marker_true_when_refs_main_revision_is_st(tmp_path, monkeypatch): + _repo(tmp_path, monkeypatch, ("new", True), ("old", False), main_ref = "new") + assert mc._embedding_marker_in_hf_cache("org/model") is True -def test_marker_missing_ref_falls_back_to_snapshot_scan(tmp_path, monkeypatch): - # No refs/main recorded: keep the newest-first any-snapshot behavior. - snaps = _repo(tmp_path, ("new", False), ("old", True)) - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) - assert mc._embedding_marker_in_hf_cache("org/no-ref") is True +def test_marker_missing_ref_is_a_cache_miss(tmp_path, monkeypatch): + # With local_files_only=True huggingface_hub resolves the default revision + # THROUGH refs/main, so a snapshot dir alone is not discoverable. Accepting + # one here would pass validation and then fail at first indexing. + _repo(tmp_path, monkeypatch, ("new", False), ("old", True)) # no refs/main + assert mc._embedding_marker_in_hf_cache("org/model") is None def test_marker_ref_points_at_absent_snapshot_is_cache_miss(tmp_path, monkeypatch): # refs/main names a commit whose snapshot dir is absent (partial download / - # pruning). The recorded ref is authoritative, so this is a cache miss - # (None) -- NOT a fall-through to a stale historical snapshot that has - # modules.json. - snaps = _repo(tmp_path, ("old", True), main_ref = "missing_commit") - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) - assert mc._embedding_marker_in_hf_cache("org/partial") is None + # pruning): a miss, never a fall-through to a stale historical snapshot. + _repo(tmp_path, monkeypatch, ("old", True), main_ref = "missing_commit") + assert mc._embedding_marker_in_hf_cache("org/model") is None def test_marker_unreadable_ref_is_cache_miss(tmp_path, monkeypatch): - # refs/main exists but cannot be read (transient I/O error / restrictive - # permissions). Only a genuinely MISSING ref may enable the historical scan; - # an unreadable ref is a cache miss (None), never a fall-through to a stale - # snapshot that happens to carry modules.json. - snaps = _repo(tmp_path, ("old", True)) - # Make refs/main a directory so read_text() raises IsADirectoryError -- an - # OSError that is NOT FileNotFoundError, i.e. "exists but unreadable". - refs_main = snaps[0].parent.parent / "refs" / "main" + # refs/main exists but cannot be read (transient I/O / restrictive + # permissions): the loader cannot resolve it either. + dirs = _repo(tmp_path, monkeypatch, ("old", True)) + refs_main = dirs[0].parent.parent / "refs" / "main" refs_main.parent.mkdir(parents = True, exist_ok = True) - refs_main.mkdir() - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) - assert mc._embedding_marker_in_hf_cache("org/unreadable-ref") is None + refs_main.mkdir() # a directory -> read_text raises IsADirectoryError + assert mc._embedding_marker_in_hf_cache("org/model") is None def test_marker_empty_ref_is_cache_miss(tmp_path, monkeypatch): - # refs/main exists but is empty / whitespace (a partial write or in-progress - # truncate-and-rewrite): the active revision is unknown, so this is a cache - # miss (None), NOT a fall-through to a stale snapshot that carries modules.json. - snaps = _repo(tmp_path, ("old", True), main_ref = " \n") - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) - assert mc._embedding_marker_in_hf_cache("org/empty-ref") is None + # refs/main exists but is empty / whitespace (a partial write): the active + # revision is unknown. + _repo(tmp_path, monkeypatch, ("old", True), main_ref = " \n") + assert mc._embedding_marker_in_hf_cache("org/model") is None + + +def test_marker_scopes_to_the_repo_dir_the_loader_opens(tmp_path, monkeypatch): + # Two cache dirs differing only by case. The loader opens the one the + # settings route persists (exact case first, as resolve_st_cached_repo_id_case + # picks), so a complete model there must validate even when the OTHER variant + # holds a newer, non-ST snapshot -- judging across both would reject it. + if not _case_sensitive_fs(tmp_path): + pytest.skip("duplicate case variants need a case-sensitive filesystem") + root = tmp_path / "cache" + _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa", repo_id = "baai/bge-m3", root = root) + # A newer variant that is NOT a sentence-transformers model. + _repo(tmp_path, monkeypatch, ("zzz", False), main_ref = "zzz", repo_id = "BAAI/bge-m3", root = root) + assert mc._embedding_marker_in_hf_cache("baai/bge-m3") is True def _fake_hf_cache(monkeypatch, root): @@ -188,6 +199,9 @@ def _st_snapshot( if loadable: (snap / "config.json").write_text("{}") (snap / "model.safetensors").write_bytes(b"\0") + refs = root / repo_dir / "refs" + refs.mkdir(parents = True, exist_ok = True) + (refs / "main").write_text(commit) return snap @@ -200,7 +214,7 @@ def test_st_probe_uses_sentence_transformers_home(tmp_path, monkeypatch): _st_snapshot(st_root, "models--org--model") _fake_hf_cache(monkeypatch, hf_root) monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_root)) - assert [p.name for p in mc._iter_st_cache_snapshots("org/model")] == ["aaa"] + assert mc._st_cache_repo_dir("org/model") == st_root / "models--org--model" assert mc._embedding_marker_in_hf_cache("org/model") is True @@ -309,7 +323,7 @@ def test_marker_never_raises_when_cache_mutates(monkeypatch): def _exploding_iter(repo): raise FileNotFoundError("snapshot removed underneath") - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", _exploding_iter) + monkeypatch.setattr(mc, "_st_cache_repo_dir", _exploding_iter) assert mc._embedding_marker_in_hf_cache("org/racing") is None @@ -319,7 +333,7 @@ def test_is_embedding_model_survives_cache_race_online(monkeypatch): def _exploding_iter(repo): raise FileNotFoundError("snapshot removed underneath") - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", _exploding_iter) + monkeypatch.setattr(mc, "_st_cache_repo_dir", _exploding_iter) _fake_hf_model_info( monkeypatch, lambda name, token = None: types.SimpleNamespace( @@ -336,8 +350,7 @@ def test_offline_cached_st_detected_via_marker_no_network(tmp_path, monkeypatch) # Offline: a downloaded sentence-transformers repo is classified from its # modules.json marker with no model_info() network call that would hang on # DNS retries (#6817). - snaps = _repo(tmp_path, ("aaa", True)) - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) + _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa", repo_id = "unsloth/bge-small-en-v1.5") monkeypatch.setenv("HF_HUB_OFFLINE", "1") _fake_hf_model_info(monkeypatch, _no_network) assert mc.is_embedding_model("unsloth/bge-small-en-v1.5") is True @@ -348,8 +361,7 @@ def test_online_defers_to_hub_over_stale_marker(tmp_path, monkeypatch): # cached modules.json (the repo WAS an embedder), a Hub lookup that no longer # reports embedding signals wins -- the stale local marker must not # short-circuit model_info(). - snaps = _repo(tmp_path, ("aaa", True)) - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) + _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa") calls = [] def _info(model_name, token = None): @@ -366,8 +378,7 @@ def test_online_permanent_hub_error_ignores_stale_marker(tmp_path, monkeypatch): # even with a cached modules.json, validation must NOT pass on the stale # marker -- return False so the settings route surfaces its 409, and the # persisted model can't fail later when the loader refreshes from the Hub. - snaps = _repo(tmp_path, ("aaa", True)) - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) + _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa") class RepositoryNotFoundError(Exception): pass @@ -383,8 +394,7 @@ def test_online_hub_failure_falls_back_to_marker_uncached(tmp_path, monkeypatch) # A transient model_info() failure falls back to the local marker WITHOUT # caching: the degraded result must not become sticky, so a later successful # Hub lookup can still override it. - snaps = _repo(tmp_path, ("aaa", True)) - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) + _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa", repo_id = "org/emb") _fake_hf_model_info(monkeypatch, _no_network) # raises -> Hub "unreachable" assert mc.is_embedding_model("org/emb") is True # marker fallback assert ("org/emb", None) not in mc._embedding_detection_cache # not poisoned @@ -395,7 +405,7 @@ def test_online_negative_does_not_block_later_offline_download(tmp_path, monkeyp # The repo is then downloaded WITH modules.json and the session goes offline; # the offline path re-probes the marker (never consulting the online memo), # so the freshly downloaded embedder is detected instead of the stale False. - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(())) + _no_cache(monkeypatch) def _info(model_name, token = None): return types.SimpleNamespace(tags = ["text-generation"], pipeline_tag = "text-generation") @@ -405,8 +415,7 @@ def _info(model_name, token = None): assert mc._embedding_detection_cache[("org/late-embedder", None)] is False # Now the model is downloaded (marker appears) and the session goes offline. - snaps = _repo(tmp_path, ("aaa", True)) - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) + _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa", repo_id = "org/late-embedder") monkeypatch.setenv("HF_HUB_OFFLINE", "1") assert mc.is_embedding_model("org/late-embedder") is True # marker re-probed @@ -416,7 +425,7 @@ def test_offline_retains_online_confirmed_positive(monkeypatch): # online and cached True. _hf_offline_if_dns_dead() then flips the process to # offline mid-load; the offline path must RETAIN that positive, not re-probe # the absent marker and downgrade a model already verified this session. - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(())) + _no_cache(monkeypatch) def _info(model_name, token = None): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) @@ -430,15 +439,14 @@ def _info(model_name, token = None): def test_offline_cached_non_st_returns_false_without_network(tmp_path, monkeypatch): - snaps = _repo(tmp_path, ("aaa", False)) - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(snaps)) + _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa") monkeypatch.setenv("HF_HUB_OFFLINE", "1") _fake_hf_model_info(monkeypatch, _no_network) assert mc.is_embedding_model("org/gemma-4-e4b") is False def test_offline_not_cached_returns_false_without_network(monkeypatch): - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(())) + _no_cache(monkeypatch) monkeypatch.setenv("HF_HUB_OFFLINE", "1") _fake_hf_model_info(monkeypatch, _no_network) assert mc.is_embedding_model("org/never-downloaded") is False @@ -447,7 +455,7 @@ def test_offline_not_cached_returns_false_without_network(monkeypatch): def test_online_uncached_still_uses_network(monkeypatch): # Not offline, not cached: the network model_info path must still run so an # embedding model that lacks modules.json (feature-extraction tag) is caught. - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(())) + _no_cache(monkeypatch) calls = [] def _info(model_name, token = None): @@ -463,7 +471,7 @@ def test_offline_negative_is_not_cached_then_online_detects(monkeypatch): # A tag-only embedder is not identifiable from modules.json. Offline returns # False WITHOUT caching, so once the env var clears the online model_info # lookup still runs and detects it -- the negative must not be sticky. - monkeypatch.setattr(mc, "_iter_st_cache_snapshots", lambda repo: iter(())) + _no_cache(monkeypatch) calls = [] def _info(model_name, token = None): diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 4f70240ad9b..e39ae449927 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1722,9 +1722,36 @@ def resolve_st_cached_repo_id_case(repo_id: str) -> str: return repo_id -def _iter_st_cache_snapshots(repo_id: str): - """Snapshot dirs for *repo_id* in the cache the ST loader will search.""" - yield from _iter_cache_snapshots_in(repo_id, _st_cache_roots()) +def _st_cache_repo_dir(repo_id: str) -> Optional[Path]: + """The ONE cache repo dir the ST loader will open for *repo_id*, or None. + + Uses the same selection rule as :func:`resolve_st_cached_repo_id_case` -- + exact case first, then a deterministic pick among case variants -- because + that is the spelling the settings route persists and therefore the directory + the loader opens. Scoping to a single directory matters when duplicate case + variants exist: judging the repo by whichever variant happens to hold the + newest snapshot could reject a complete model in the directory that will + actually be loaded (or accept one from a directory that will not be). + """ + prefix = "models--" + expected = f"{prefix}{repo_id.replace('/', '--')}" + target = expected.lower() + variants: list[Path] = [] + for cache_dir in _st_cache_roots(): + try: + if not cache_dir.is_dir(): + continue + exact = cache_dir / expected + if exact.is_dir(): + return exact + for entry in cache_dir.iterdir(): + if entry.is_dir() and entry.name.lower() == target: + variants.append(entry) + except OSError: + continue + if variants: + return sorted(variants, key = lambda path: path.name)[0] + return None def _iter_snapshots_of(repo_dirs: list[Path]): @@ -1759,7 +1786,7 @@ def _iter_hf_cache_snapshots(repo_id: str): download with ``hf_hub_download`` and no ``cache_dir``, so that is the cache their load actually uses; adding SENTENCE_TRANSFORMERS_HOME here would let detection pick a file the GGUF loader cannot find. The Sentence-Transformers - probe is :func:`_iter_st_cache_snapshots`. + probe is :func:`_st_cache_repo_dir`. Empty if the root does not exist, the repo isn't cached, or it has no snapshots. Repo name match is case-insensitive to handle casing drift between @@ -2120,58 +2147,39 @@ def _snapshot_is_loadable_st_model(snap: Path) -> bool: def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: """Sentence-transformers detection from the local cache, no network call. - Probes the cache the ST loader will actually search (see - :func:`_st_cache_roots`), so a hit here means the load can find it too. - True/False when the ACTIVE cached revision is / is not a LOADABLE - sentence-transformers snapshot -- marker plus config plus weights, not the - bare ``modules.json`` the security preflight may have fetched on its own -- - and None when the repo is not in that cache (or the cache is unreadable). The revision ``refs/main`` resolves to is authoritative when - recorded: the cache keeps snapshots of older revisions, and a repo that - later stopped (or started) being a sentence-transformers model must be - judged by its current revision, not any historical one. Snapshots are only - scanned newest-first when no ref exists. Never raises -- a cache mutating - underneath (concurrent model deletion) reads as not-cached so callers keep - their normal fallback.""" + Models exactly what an offline ``local_files_only=True`` load resolves, so a + True here means that load can succeed: + + * one repo dir -- the one :func:`_st_cache_repo_dir` selects, i.e. the casing + the settings route persists and the loader opens; + * its ``refs/main``, because with ``local_files_only`` huggingface_hub + resolves the default revision THROUGH that ref. A snapshot directory alone + is not discoverable, so a missing, empty or unreadable ref is a cache MISS + rather than a reason to go scanning historical snapshots -- accepting one + would pass validation and then fail at first indexing; + * that revision's snapshot, which must be materialized and actually loadable + (marker plus config plus weights, not the bare ``modules.json`` the + security preflight fetches on its own). + + True/False when the active revision is / is not a loadable + sentence-transformers snapshot, None when nothing usable is cached. Never + raises -- a cache mutating underneath (concurrent model deletion) reads as + not-cached so callers keep their normal fallback. + """ try: - snapshots = list(_iter_st_cache_snapshots(repo_id)) - if not snapshots: + repo_dir = _st_cache_repo_dir(repo_id) + if repo_dir is None: return None - # Prefer the snapshot refs/main points at (the active revision). - snapshots_dir = snapshots[0].parent try: - commit = (snapshots_dir.parent / "refs" / "main").read_text(encoding = "utf-8").strip() - except FileNotFoundError: - # No ref recorded at all: fall back to the newest-first scan. - commit = None + commit = (repo_dir / "refs" / "main").read_text(encoding = "utf-8").strip() except OSError: - # A ref exists but is unreadable (transient I/O error, restrictive - # permissions): the contract is that an unreadable cache reads as - # not-cached, so report a miss rather than scanning stale history -- - # only a genuinely missing ref may enable the fallback scan. - return None - else: - if not commit: - # The ref file exists but is empty / whitespace (a partial write - # or an in-progress truncate-and-rewrite). The active revision is - # unknown, so this is a cache miss -- NOT a fall-through to stale - # history, which is the class of bug this helper avoids. - return None - if commit is not None: - # A ref is recorded, so it is authoritative. If its snapshot is not - # materialized (partial download / pruning) treat the repo as not - # cached (None) rather than scanning stale history -- the exact - # stale-cache class this helper avoids. - preferred = snapshots_dir / commit - if not preferred.is_dir(): - return None - return _snapshot_is_loadable_st_model(preferred) - for snap in snapshots: - try: - if _snapshot_is_loadable_st_model(snap): - return True - except OSError: - continue - return False + return None # absent or unreadable: the loader cannot resolve it either + if not commit: + return None # empty / whitespace: a partial write, revision unknown + snapshot = repo_dir / "snapshots" / commit + if not snapshot.is_dir(): + return None # ref recorded but not materialized (partial / pruned) + return _snapshot_is_loadable_st_model(snapshot) except Exception: return None From 9e6a3552ea401c387855486dc91da3fb73e36a8e Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sun, 19 Jul 2026 12:39:57 +0300 Subject: [PATCH 19/77] studio: record refs/main in the ONNX-only probe test The ONNX-only regression test predates the refs/main requirement, so after that change it returned None (a cache miss for want of a ref) before ever reaching the weight-format check it exists to make. Recording the ref restores its intent: the snapshot resolves, and the answer is False because an ONNX export is not loadable by the RAG loader's default Torch backend. --- studio/backend/tests/test_is_embedding_model_offline.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 35728b106d0..c395e44ab36 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -267,11 +267,16 @@ def test_marker_onnx_only_snapshot_is_not_loadable(tmp_path, monkeypatch): # accepting the ONNX offline would pass validation and then fail on the first # RAG load -- the same validate-then-fail as the marker-only case. hf_root = tmp_path / "hf" - snap = hf_root / "models--org--model" / "snapshots" / "aaa" + repo = hf_root / "models--org--model" + snap = repo / "snapshots" / "aaa" snap.mkdir(parents = True) (snap / "modules.json").write_text("[]") (snap / "config.json").write_text("{}") (snap / "model.onnx").write_bytes(b"\0") + # refs/main so the probe resolves this revision and reaches the weight check + # (without it the answer would be None -- a cache miss -- for a different reason). + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text("aaa") _fake_hf_cache(monkeypatch, hf_root) monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) assert mc._embedding_marker_in_hf_cache("org/model") is False From 003c9824ceb4961d6d557bb26970c55ddd7630d5 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 19 Jul 2026 10:15:13 +0000 Subject: [PATCH 20/77] studio: recognize base-model weight files and gate the offline positive on a materialized snapshot _snapshot_is_loadable_st_model matched any .safetensors/.bin by suffix, so a partial cache carrying only a commonly published non-weight bin such as training_args.bin (or an adapter-only artifact) passed offline validation and then failed the local_files_only load at first indexing. Match recognized Torch base-model weight filenames (model / pytorch_model, including sharded) by name. is_embedding_model retained an online-confirmed positive offline even when no files were cached, so a metadata-only /check-embedding result let an uncached repo be saved and then fail at first indexing. Retain the positive only when the active revision is materialized locally, which still covers a downloaded tag-only embedder whose snapshot carries no modules.json. --- .../tests/test_is_embedding_model_offline.py | 77 +++++++++++++++++-- studio/backend/utils/models/model_config.py | 44 ++++++----- 2 files changed, 97 insertions(+), 24 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index c395e44ab36..69a2913f41d 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -282,6 +282,49 @@ def test_marker_onnx_only_snapshot_is_not_loadable(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache("org/model") is False +def _cache_repo_with_files(tmp_path, monkeypatch, *files, commit = "aaa"): + """A cache repo whose active snapshot holds modules.json + config + *files*.""" + hf_root = tmp_path / "hf" + repo = hf_root / "models--org--model" + snap = repo / "snapshots" / commit + snap.mkdir(parents = True) + (snap / "modules.json").write_text("[]") + (snap / "config.json").write_text("{}") + for name in files: + target = snap / name + target.parent.mkdir(parents = True, exist_ok = True) + target.write_bytes(b"\0") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + + +def test_marker_rejects_non_base_weight_bins(tmp_path, monkeypatch): + # A partial cache carrying modules.json + config plus only a commonly published + # NON-weight .bin (training_args.bin) is not a loadable base model: the Torch + # backend needs model.safetensors / pytorch_model.bin, so accepting it would + # pass validation and then fail at first indexing. Adapter-only artifacts are + # rejected for the same reason. + _cache_repo_with_files(tmp_path, monkeypatch, "training_args.bin") + assert mc._embedding_marker_in_hf_cache("org/model") is False + _cache_repo_with_files(tmp_path / "b", monkeypatch, "adapter_model.safetensors") + assert mc._embedding_marker_in_hf_cache("org/model") is False + + +@pytest.mark.parametrize("weights", [ + ("pytorch_model.bin",), # torch .bin + ("model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"), # sharded + ("0_Transformer/model.safetensors",), # weight in a module dir +]) +def test_marker_accepts_recognized_torch_weights(tmp_path, monkeypatch, weights): + # The filename recognizer must not over-reject real base-model weights: single + # pytorch_model.bin, sharded model-000NN-of-000NN.safetensors, and weights that + # live inside a module directory all count as loadable. + _cache_repo_with_files(tmp_path, monkeypatch, *weights) + assert mc._embedding_marker_in_hf_cache("org/model") is True + + def _case_sensitive_fs(tmp_path) -> bool: probe = tmp_path / "_CaseProbe" probe.mkdir() @@ -425,12 +468,15 @@ def _info(model_name, token = None): assert mc.is_embedding_model("org/late-embedder") is True # marker re-probed -def test_offline_retains_online_confirmed_positive(monkeypatch): - # A tag-only embedder (feature-extraction, no modules.json) is confirmed - # online and cached True. _hf_offline_if_dns_dead() then flips the process to - # offline mid-load; the offline path must RETAIN that positive, not re-probe - # the absent marker and downgrade a model already verified this session. - _no_cache(monkeypatch) +def test_offline_retains_online_confirmed_positive(tmp_path, monkeypatch): + # A tag-only embedder (feature-extraction, no modules.json) is confirmed online + # and cached True, and its snapshot IS materialized locally. + # _hf_offline_if_dns_dead() then flips the process to offline mid-load; the + # offline path must RETAIN that positive because the files are present, not + # downgrade a model already verified this session. The marker alone reads False + # here (present but no modules.json), so retention rests on the memo plus a + # present active snapshot, not on the marker. + _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") def _info(model_name, token = None): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) @@ -443,6 +489,25 @@ def _info(model_name, token = None): assert mc.is_embedding_model("org/gte-modernbert") is True # positive retained +def test_offline_metadata_only_positive_not_trusted_without_cache(monkeypatch): + # A positive from an online model_info() call (e.g. /check-embedding) proves + # only that the repo is tagged an embedder, not that any files were downloaded. + # With nothing materialized in the cache, the offline path must NOT trust that + # memo: saving the model would then fail the local_files_only RAG load at first + # indexing. So it returns False despite the cached True. + _no_cache(monkeypatch) + + def _info(model_name, token = None): + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + + _fake_hf_model_info(monkeypatch, _info) + assert mc.is_embedding_model("org/uncached-embedder") is True # online: cached True (metadata only) + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _fake_hf_model_info(monkeypatch, _no_network) # offline must not hit network + assert mc.is_embedding_model("org/uncached-embedder") is False # not cached -> not trusted + + def test_offline_cached_non_st_returns_false_without_network(tmp_path, monkeypatch): _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa") monkeypatch.setenv("HF_HUB_OFFLINE", "1") diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index e39ae449927..320eb3f44dd 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2106,13 +2106,16 @@ def download_gguf_file( _embedding_detection_cache: Dict[tuple, bool] = {} -# Only the weight formats the RAG loader's default backend consumes. _get() +# The base-model weight files the RAG loader's default backend consumes. _get() # constructs SentenceTransformer without backend="onnx"/"openvino", so it loads -# through Torch (model.safetensors / pytorch_model.bin, incl. sharded). ".onnx" -# (and ".pt", which is not an HF weight filename) would let a snapshot cached -# with only an ONNX export pass offline validation and then fail on the first -# RAG load -- the exact validate-then-fail this helper exists to prevent. -_ST_WEIGHT_SUFFIXES = (".safetensors", ".bin") +# through Torch: a single model.safetensors / pytorch_model.bin or their sharded +# model-00001-of-000NN forms, at the snapshot root or inside a module dir +# (0_Transformer/, 2_Dense/). Matched by NAME, not suffix: a bare ".safetensors"/ +# ".bin" match would accept the commonly published training_args.bin / optimizer.bin +# or an adapter-only artifact -- files that carry a weight suffix but are not a +# loadable base model -- and let a partial cache pass offline validation and then +# fail on the first RAG load, the exact validate-then-fail this helper prevents. +_ST_WEIGHT_FILE_RE = re.compile(r"^(model|pytorch_model)(-\d+-of-\d+)?\.(safetensors|bin)$") def _snapshot_is_loadable_st_model(snap: Path) -> bool: @@ -2123,8 +2126,8 @@ def _snapshot_is_loadable_st_model(snap: Path) -> bool: leave it behind too, so a snapshot can carry the marker while the weights and config SentenceTransformer needs are absent. Accepting that offline would pass validation and then fail on the first RAG load. Requires the marker plus - a config and at least one weight file the default Torch backend can load - (``_ST_WEIGHT_SUFFIXES``) somewhere in the snapshot. + a config and at least one recognized Torch base-model weight file + (``_ST_WEIGHT_FILE_RE``) somewhere in the snapshot. """ try: if not (snap / "modules.json").is_file(): @@ -2135,7 +2138,7 @@ def _snapshot_is_loadable_st_model(snap: Path) -> bool: return False for path in snap.rglob("*"): try: - if path.is_file() and path.suffix.lower() in _ST_WEIGHT_SUFFIXES: + if path.is_file() and _ST_WEIGHT_FILE_RE.match(path.name): return True except OSError: continue @@ -2213,17 +2216,22 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: if _env_offline(): # Offline: the local HF cache is the only source -- a network call cannot # succeed and would only hang on a DNS error and get retried (#6817). - # Retain a positive already confirmed online this session: model_info() - # only ever memoizes Hub-derived results, so a cached True is a real - # detection (e.g. a tag-only feature-extraction embedder with no - # modules.json) that _hf_offline_if_dns_dead() flipping the process to - # offline mid-load must NOT downgrade to False. Otherwise re-probe the - # marker every call without consulting or populating the memo: a cached - # negative must not stick (a model downloaded later, or a tag-only + # Re-probe the marker every call without consulting or populating the memo: + # a cached negative must not stick (a model downloaded later, or a tag-only # embedder, could not be confirmed here), so a miss is not durable. - if _embedding_detection_cache.get(cache_key) is True: + marker = _embedding_marker_in_hf_cache(model_name) + # Retain a positive confirmed online this session ONLY when the active + # revision is actually materialized locally (marker is not None). A cached + # True proves model_info() tagged the repo an embedder, not that its files + # are on disk -- an online /check-embedding call populates this process-wide + # memo without downloading -- so trusting it for an uncached repo would save + # a model the local_files_only load then fails on. "not None" (not True) + # still covers a downloaded tag-only feature-extraction embedder whose + # snapshot is present but has no modules.json: the _hf_offline_if_dns_dead() + # mid-load flip case that must NOT be downgraded. + if _embedding_detection_cache.get(cache_key) is True and marker is not None: return True - return _embedding_marker_in_hf_cache(model_name) is True + return marker is True # Online: the Hub is authoritative for the current remote revision. The local # cache marker reflects only the last-downloaded revision, which may lag the From 563dd315936d1cb8b8307e040bb90929f34697d5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 19 Jul 2026 10:15:58 +0000 Subject: [PATCH 21/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../tests/test_is_embedding_model_offline.py | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 69a2913f41d..bdec9f8c6e8 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -282,7 +282,12 @@ def test_marker_onnx_only_snapshot_is_not_loadable(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache("org/model") is False -def _cache_repo_with_files(tmp_path, monkeypatch, *files, commit = "aaa"): +def _cache_repo_with_files( + tmp_path, + monkeypatch, + *files, + commit = "aaa", +): """A cache repo whose active snapshot holds modules.json + config + *files*.""" hf_root = tmp_path / "hf" repo = hf_root / "models--org--model" @@ -312,11 +317,14 @@ def test_marker_rejects_non_base_weight_bins(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache("org/model") is False -@pytest.mark.parametrize("weights", [ - ("pytorch_model.bin",), # torch .bin - ("model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"), # sharded - ("0_Transformer/model.safetensors",), # weight in a module dir -]) +@pytest.mark.parametrize( + "weights", + [ + ("pytorch_model.bin",), # torch .bin + ("model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"), # sharded + ("0_Transformer/model.safetensors",), # weight in a module dir + ], +) def test_marker_accepts_recognized_torch_weights(tmp_path, monkeypatch, weights): # The filename recognizer must not over-reject real base-model weights: single # pytorch_model.bin, sharded model-000NN-of-000NN.safetensors, and weights that @@ -501,7 +509,9 @@ def _info(model_name, token = None): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) _fake_hf_model_info(monkeypatch, _info) - assert mc.is_embedding_model("org/uncached-embedder") is True # online: cached True (metadata only) + assert ( + mc.is_embedding_model("org/uncached-embedder") is True + ) # online: cached True (metadata only) monkeypatch.setenv("HF_HUB_OFFLINE", "1") _fake_hf_model_info(monkeypatch, _no_network) # offline must not hit network From 62e57d18c18ff5a3062874703def7aea77f4f56f Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 19 Jul 2026 10:44:03 +0000 Subject: [PATCH 22/77] studio: require a complete weight set offline and persist embedder verdicts across restarts Two follow-ups to the offline embedding-model classifier: - _snapshot_is_loadable_st_model now requires a COMPLETE Torch base-model weight set in one snapshot directory, not just any single recognized weight file. A partially downloaded sharded model (model-00001-of-00002 without its sibling) no longer passes offline validation and then fails at first indexing under local_files_only. Weight files are grouped by directory and a directory counts only when it holds a single model.safetensors / pytorch_model.bin or a full shard set whose indices cover 1..total. - Online-confirmed embedder verdicts are now recorded under the resolved Studio home (embedding_verdicts.json). The session memo is lost on exit, so a downloaded tag-only feature-extraction embedder (snapshot present but no modules.json) was misclassified as non-embedding the first offline call after a restart. The offline branch consults this durable allowlist in addition to the memo, still gated on the active revision being materialized on disk, so an uncached repo is never trusted. Writes are best-effort and only positive verdicts are stored. --- .../tests/test_is_embedding_model_offline.py | 87 +++++++++++- studio/backend/utils/models/model_config.py | 126 ++++++++++++++++-- 2 files changed, 198 insertions(+), 15 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index bdec9f8c6e8..30dfcecfa1c 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -47,10 +47,14 @@ def _build_structlog_stub(): @pytest.fixture(autouse = True) -def _clean_state(monkeypatch): +def _clean_state(tmp_path, monkeypatch): mc._embedding_detection_cache.clear() monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + # Point the persisted embedder-verdict store at a per-test Studio home so the + # cross-restart allowlist never leaks into another test or the real ~/.unsloth. + monkeypatch.delenv("STUDIO_HOME", raising = False) + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "_studio_home")) yield mc._embedding_detection_cache.clear() @@ -317,6 +321,26 @@ def test_marker_rejects_non_base_weight_bins(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache("org/model") is False +@pytest.mark.parametrize( + "weights", + [ + ("model-00001-of-00002.safetensors",), # only shard 1 of 2 downloaded + ("model-00002-of-00003.bin",), # a lone middle shard + ( # two shards present but the set claims three + "model-00001-of-00003.safetensors", + "model-00002-of-00003.safetensors", + ), + ], +) +def test_marker_rejects_incomplete_shard_set(tmp_path, monkeypatch, weights): + # A partially downloaded sharded model must NOT validate: SentenceTransformer's + # Torch backend needs every shard the index names at load time, so accepting an + # incomplete set would pass offline validation and then fail at first indexing -- + # the same validate-then-fail this helper exists to prevent. + _cache_repo_with_files(tmp_path, monkeypatch, *weights) + assert mc._embedding_marker_in_hf_cache("org/model") is False + + @pytest.mark.parametrize( "weights", [ @@ -518,6 +542,67 @@ def _info(model_name, token = None): assert mc.is_embedding_model("org/uncached-embedder") is False # not cached -> not trusted +def test_offline_detects_persisted_tag_only_embedder_after_restart(tmp_path, monkeypatch): + # A tag-only embedder (feature-extraction, no modules.json) is confirmed online + # in one session, which durably records the verdict, and its snapshot is + # materialized on disk. After a RESTART (the session memo is gone) the process + # comes up offline: the marker alone reads False (present but no modules.json), + # so recognition rests on the persisted allowlist plus the present active + # snapshot -- the exact case the memo cannot cover across a restart. + _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") + + def _info(model_name, token = None): + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + + _fake_hf_model_info(monkeypatch, _info) + assert mc.is_embedding_model("org/gte-modernbert") is True # online: confirmed + persisted + assert "org/gte-modernbert" in mc._load_persisted_embedders() # written to disk + + mc._embedding_detection_cache.clear() # simulate a process restart: memo lost + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _fake_hf_model_info(monkeypatch, _no_network) # offline must not hit network + assert mc.is_embedding_model("org/gte-modernbert") is True # recovered from disk + + +def test_offline_persisted_verdict_not_trusted_when_uncached(tmp_path, monkeypatch): + # The persisted allowlist records only that the repo was tagged an embedder, + # not that its files are on disk. After a restart with NOTHING materialized in + # the cache, the offline path must NOT trust the persisted verdict: saving the + # model would then fail the local_files_only RAG load at first indexing. It is + # gated on the active snapshot being present, so an uncached repo returns False. + _no_cache(monkeypatch) + + def _info(model_name, token = None): + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + + _fake_hf_model_info(monkeypatch, _info) + assert mc.is_embedding_model("org/uncached-embedder") is True # online: confirmed + persisted + assert "org/uncached-embedder" in mc._load_persisted_embedders() + + mc._embedding_detection_cache.clear() # restart: memo lost, disk verdict remains + _no_cache(monkeypatch) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _fake_hf_model_info(monkeypatch, _no_network) + assert mc.is_embedding_model("org/uncached-embedder") is False # nothing on disk -> not trusted + + +def test_persist_embedder_is_best_effort_when_home_unwritable(tmp_path, monkeypatch): + # Persistence is an optimization, never a correctness requirement: if the + # Studio home cannot be created/written, the online path must still return its + # verdict rather than raise. Point the home at a path blocked by a file. + blocker = tmp_path / "blocker" + blocker.write_text("not a dir") + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(blocker / "studio")) # parent is a file + _no_cache(monkeypatch) + + def _info(model_name, token = None): + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + + _fake_hf_model_info(monkeypatch, _info) + assert mc.is_embedding_model("org/emb") is True # no exception despite unwritable home + assert mc._load_persisted_embedders() == set() # nothing recorded, silently + + def test_offline_cached_non_st_returns_false_without_network(tmp_path, monkeypatch): _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa") monkeypatch.setenv("HF_HUB_OFFLINE", "1") diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 320eb3f44dd..605ebc845e9 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2106,6 +2106,69 @@ def download_gguf_file( _embedding_detection_cache: Dict[tuple, bool] = {} +def _persisted_embedders_path() -> Path: + """On-disk store of repo ids confirmed embedders via an online Hub lookup. + + Lives under the resolved Studio home so it survives a restart: the session + memo (:data:`_embedding_detection_cache`) is lost on exit, so without this a + tag-only feature-extraction embedder -- one whose cached snapshot carries no + ``modules.json`` and so cannot be recognized from the cache alone -- would be + misclassified as non-embedding the first offline call after a restart. + """ + return _studio_root() / "embedding_verdicts.json" + + +def _load_persisted_embedders() -> set: + """Repo ids recorded as embedders in a prior session; empty set on any error. + + Best-effort and never raises: a missing, empty, corrupt, or non-list file + (and any read/decode error) reads as "nothing recorded" so callers fall back + to their normal cache-marker logic. + """ + try: + with open(_persisted_embedders_path(), encoding = "utf-8") as fh: + data = json.load(fh) + if isinstance(data, list): + return {name for name in data if isinstance(name, str)} + except Exception: + pass + return set() + + +def _persist_embedder(model_name: str) -> None: + """Record *model_name* as an online-confirmed embedder, best-effort. + + Only positive Hub verdicts are written (never a negative or a token): the + file is a durable allowlist the offline branch consults, gated there on the + active revision actually being materialized on disk, so a stale entry for a + since-deleted cache cannot resurrect a False. Any failure (unwritable home, + race) is swallowed -- persistence is an optimization, not a correctness + requirement for the online path that calls it. + """ + try: + current = _load_persisted_embedders() + if model_name in current: + return + current.add(model_name) + path = _persisted_embedders_path() + path.parent.mkdir(parents = True, exist_ok = True) + tmp = path.with_name(path.name + f".{os.getpid()}.tmp") + with open(tmp, "w", encoding = "utf-8") as fh: + json.dump(sorted(current), fh) + os.replace(tmp, path) + except Exception as e: + logger.debug(f"Could not persist embedder verdict for {model_name}: {e}") + + +def _known_embedder(model_name: str, cache_key: tuple) -> bool: + """True when *model_name* was confirmed an embedder online, this session or a + prior one -- the session memo (authoritative, token-scoped) OR the persisted + cross-restart allowlist. Callers still gate this on a materialized snapshot.""" + if _embedding_detection_cache.get(cache_key) is True: + return True + return model_name in _load_persisted_embedders() + + # The base-model weight files the RAG loader's default backend consumes. _get() # constructs SentenceTransformer without backend="onnx"/"openvino", so it loads # through Torch: a single model.safetensors / pytorch_model.bin or their sharded @@ -2116,6 +2179,34 @@ def download_gguf_file( # loadable base model -- and let a partial cache pass offline validation and then # fail on the first RAG load, the exact validate-then-fail this helper prevents. _ST_WEIGHT_FILE_RE = re.compile(r"^(model|pytorch_model)(-\d+-of-\d+)?\.(safetensors|bin)$") +_ST_SHARD_RE = re.compile(r"^(model|pytorch_model)-(\d+)-of-(\d+)\.(safetensors|bin)$") + + +def _dir_has_complete_torch_weights(names: set) -> bool: + """True when the filenames in ONE directory form a COMPLETE Torch base-model + weight set the default SentenceTransformer backend can load: + + * a single ``model.safetensors`` / ``pytorch_model.bin``; or + * a full shard set -- a shard names its own total (``-00001-of-00002``), so + every index ``1..total`` for that (stem, ext, total) must be present. + + A lone shard such as ``model-00001-of-00002.safetensors`` is NOT complete: the + loader needs the missing shard(s) at indexing time, so accepting a partially + downloaded sharded model would validate and then fail under local_files_only. + """ + for stem in ("model", "pytorch_model"): + for ext in ("safetensors", "bin"): + if f"{stem}.{ext}" in names: + return True + shards: dict = {} + for name in names: + m = _ST_SHARD_RE.match(name) + if m: + shards.setdefault((m.group(1), m.group(4), int(m.group(3))), set()).add(int(m.group(2))) + for (_stem, _ext, total), indices in shards.items(): + if total > 0 and indices == set(range(1, total + 1)): + return True + return False def _snapshot_is_loadable_st_model(snap: Path) -> bool: @@ -2126,8 +2217,8 @@ def _snapshot_is_loadable_st_model(snap: Path) -> bool: leave it behind too, so a snapshot can carry the marker while the weights and config SentenceTransformer needs are absent. Accepting that offline would pass validation and then fail on the first RAG load. Requires the marker plus - a config and at least one recognized Torch base-model weight file - (``_ST_WEIGHT_FILE_RE``) somewhere in the snapshot. + a config and a COMPLETE recognized Torch base-model weight set + (``_dir_has_complete_torch_weights``) in one directory of the snapshot. """ try: if not (snap / "modules.json").is_file(): @@ -2136,13 +2227,14 @@ def _snapshot_is_loadable_st_model(snap: Path) -> bool: (snap / name).is_file() for name in ("config.json", "config_sentence_transformers.json") ): return False + by_dir: dict = {} for path in snap.rglob("*"): try: if path.is_file() and _ST_WEIGHT_FILE_RE.match(path.name): - return True + by_dir.setdefault(path.parent, set()).add(path.name) except OSError: continue - return False + return any(_dir_has_complete_torch_weights(names) for names in by_dir.values()) except OSError: return False @@ -2220,16 +2312,18 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: # a cached negative must not stick (a model downloaded later, or a tag-only # embedder, could not be confirmed here), so a miss is not durable. marker = _embedding_marker_in_hf_cache(model_name) - # Retain a positive confirmed online this session ONLY when the active - # revision is actually materialized locally (marker is not None). A cached - # True proves model_info() tagged the repo an embedder, not that its files - # are on disk -- an online /check-embedding call populates this process-wide - # memo without downloading -- so trusting it for an uncached repo would save - # a model the local_files_only load then fails on. "not None" (not True) - # still covers a downloaded tag-only feature-extraction embedder whose - # snapshot is present but has no modules.json: the _hf_offline_if_dns_dead() - # mid-load flip case that must NOT be downgraded. - if _embedding_detection_cache.get(cache_key) is True and marker is not None: + # Retain a positive confirmed online -- this session (memo) or a prior one + # (persisted allowlist) -- ONLY when the active revision is actually + # materialized locally (marker is not None). A recorded True proves + # model_info() tagged the repo an embedder, not that its files are on disk + # (an online /check-embedding call records the verdict without downloading), + # so trusting it for an uncached repo would save a model the local_files_only + # load then fails on. "not None" (not True) still covers a downloaded + # tag-only feature-extraction embedder whose snapshot is present but has no + # modules.json -- recognizable only from that recorded verdict, and the case + # that must survive a restart, not just the _hf_offline_if_dns_dead() mid-load + # flip. + if _known_embedder(model_name, cache_key) and marker is not None: return True return marker is True @@ -2256,6 +2350,10 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: _embedding_detection_cache[cache_key] = is_emb if is_emb: + # Durably record the positive so a later offline session (after a + # restart, when the memo is gone) can still recognize a downloaded + # tag-only embedder, gated there on its snapshot being materialized. + _persist_embedder(model_name) logger.info( f"Model {model_name} detected as embedding model: " f"pipeline_tag={pipeline_tag}, " From 98abd64032d7e14f8ada9bb4ac57e5920b0b8099 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 19 Jul 2026 10:59:00 +0000 Subject: [PATCH 23/77] studio: require complete weights (with shard index) and resolve default casing offline Follow-ups to the offline embedding-model classifier from the latest review: - Trust a recorded embedder verdict (session memo or persisted allowlist) offline only when the active snapshot carries a COMPLETE, loadable weight set, not merely that it is materialized. A partial download (config present, weights missing or an incomplete shard set) makes _embedding_marker_in_hf_cache read False rather than None, so the previous marker-is-not-None gate wrongly returned True and the local_files_only load then failed. Split out _snapshot_has_complete_weights (config plus complete weights, modules.json aside) and _active_snapshot_dir, and gate the known-embedder positive on the weight set. - Require a sharded checkpoint's index map (model.safetensors.index.json / pytorch_model.bin.index.json) in addition to every shard before accepting it: transformers discovers and wires shards through that index, so a complete shard set without it fails the local-only load. - Resolve the embedding model name to its exact cache casing in the RAG loader before constructing SentenceTransformer. The settings route persists that spelling for a custom override but deliberately leaves the configured default verbatim, so a default whose casing differs from the cache dir would miss it and fail offline. Resolving at load time covers the default too; a no-op for a local path or when nothing case-matching is cached, and idempotent for an already-normalized override. Adds regression tests for the partial-snapshot verdict, the missing shard index, and the loader casing resolution; updates the offline-invariant source assertion to the resolved-name variable. --- studio/backend/core/rag/embeddings.py | 21 ++- studio/backend/tests/test_hf_offline_flags.py | 2 +- .../tests/test_is_embedding_model_offline.py | 56 ++++++- studio/backend/tests/test_rag_embeddings.py | 41 +++++ studio/backend/utils/models/model_config.py | 151 ++++++++++++------ 5 files changed, 208 insertions(+), 63 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 62e64bedfde..d65c3910dff 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -194,7 +194,22 @@ def _get(model_name: str | None = None): from utils.utils import hf_env_offline device = _device() - logger.info("loading embedding model %s on %s", name, device) + # Resolve to the exact cache casing the offline local_files_only load + # needs. The /settings route persists that spelling for a custom override + # but deliberately leaves the configured default verbatim (rewriting it + # would turn the default into an override and break later default changes), + # so a default whose casing differs from the cache dir would otherwise + # miss it and fail offline. Resolve here at load time to cover the default + # too; a no-op for a local path or when nothing case-matching is cached, + # and idempotent for an override already normalized at persist time. + load_name = name + from utils.paths import is_local_path + + if not is_local_path(load_name): + from utils.models import resolve_st_cached_repo_id_case + + load_name = resolve_st_cached_repo_id_case(load_name) + logger.info("loading embedding model %s on %s", load_name, device) # Read the offline state ONCE and use that single value for both the # security gate and the loader. _hf_offline_if_dns_dead() mutates the # process-wide offline vars and restores them on exit, so re-reading @@ -202,13 +217,13 @@ def _get(model_name: str | None = None): # skipped the Hub scan on True -- and the load would then fetch and # deserialize the unscanned repo. local_only = hf_env_offline() - _guard_model_security(name, local_only) + _guard_model_security(load_name, local_only) # Propagate the user's offline intent into the loader: SentenceTransformer # performs its own Hub operations, and huggingface_hub honors only # HF_HUB_OFFLINE, so a TRANSFORMERS_OFFLINE-only session would otherwise # still fetch missing repo files over the network. _model = SentenceTransformer( - name, + load_name, device = device, model_kwargs = dtype_kwargs("float16"), local_files_only = local_only, diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index c363b2cc3c9..d96f79162e5 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -157,7 +157,7 @@ def test_embedding_loader_forces_local_only_when_offline(): "local_only = hf_env_offline()" in src ), "core/rag/embeddings.py must capture the offline state once in _get()" assert ( - "_guard_model_security(name, local_only)" in src + "_guard_model_security(load_name, local_only)" in src ), "the security guard must receive the captured value, not re-read the env" assert "local_files_only = local_only" in src, ( "SentenceTransformer must be pinned with the SAME captured value; a second " diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 30dfcecfa1c..dd00a408df9 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -324,19 +324,25 @@ def test_marker_rejects_non_base_weight_bins(tmp_path, monkeypatch): @pytest.mark.parametrize( "weights", [ - ("model-00001-of-00002.safetensors",), # only shard 1 of 2 downloaded - ("model-00002-of-00003.bin",), # a lone middle shard - ( # two shards present but the set claims three + ("model-00001-of-00002.safetensors", "model.safetensors.index.json"), # shard 1 of 2 + ("model-00002-of-00003.bin", "pytorch_model.bin.index.json"), # a lone middle shard + ( # two shards + index but the set claims three "model-00001-of-00003.safetensors", "model-00002-of-00003.safetensors", + "model.safetensors.index.json", + ), + ( # a complete shard set WITHOUT the index map transformers loads it through + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors", ), ], ) def test_marker_rejects_incomplete_shard_set(tmp_path, monkeypatch, weights): # A partially downloaded sharded model must NOT validate: SentenceTransformer's - # Torch backend needs every shard the index names at load time, so accepting an - # incomplete set would pass offline validation and then fail at first indexing -- - # the same validate-then-fail this helper exists to prevent. + # Torch backend needs every shard AND the index map at load time, so accepting an + # incomplete set (or a complete set missing its index) would pass offline + # validation and then fail at first indexing -- the same validate-then-fail this + # helper exists to prevent. _cache_repo_with_files(tmp_path, monkeypatch, *weights) assert mc._embedding_marker_in_hf_cache("org/model") is False @@ -345,13 +351,17 @@ def test_marker_rejects_incomplete_shard_set(tmp_path, monkeypatch, weights): "weights", [ ("pytorch_model.bin",), # torch .bin - ("model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"), # sharded + ( # sharded: every index plus the index map + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors", + "model.safetensors.index.json", + ), ("0_Transformer/model.safetensors",), # weight in a module dir ], ) def test_marker_accepts_recognized_torch_weights(tmp_path, monkeypatch, weights): # The filename recognizer must not over-reject real base-model weights: single - # pytorch_model.bin, sharded model-000NN-of-000NN.safetensors, and weights that + # pytorch_model.bin, a complete sharded set with its index map, and weights that # live inside a module directory all count as loadable. _cache_repo_with_files(tmp_path, monkeypatch, *weights) assert mc._embedding_marker_in_hf_cache("org/model") is True @@ -586,6 +596,36 @@ def _info(model_name, token = None): assert mc.is_embedding_model("org/uncached-embedder") is False # nothing on disk -> not trusted +def test_offline_persisted_verdict_not_trusted_when_snapshot_partial(tmp_path, monkeypatch): + # A persisted verdict is trusted only when the active snapshot carries a COMPLETE + # weight set, not merely that it is materialized. Here the snapshot exists with a + # config but no weights (an interrupted download): _embedding_marker_in_hf_cache + # reads False (materialized, not None), so a bare "materialized" gate would wrongly + # return True and the local_files_only load would then fail. The weight gate must + # reject it. + cache_root = tmp_path / "cache" + snap = cache_root / "models--org--partial" / "snapshots" / "aaa" + snap.mkdir(parents = True) + (snap / "config.json").write_text("{}") # config only -- weights never finished + refs = cache_root / "models--org--partial" / "refs" + refs.mkdir(parents = True) + (refs / "main").write_text("aaa") + monkeypatch.setattr(mc, "_st_cache_roots", lambda: [cache_root]) + + def _info(model_name, token = None): + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + + _fake_hf_model_info(monkeypatch, _info) + assert mc.is_embedding_model("org/partial") is True # online: confirmed + persisted + assert "org/partial" in mc._load_persisted_embedders() + assert mc._embedding_marker_in_hf_cache("org/partial") is False # materialized, not None + + mc._embedding_detection_cache.clear() # restart: memo lost, disk verdict remains + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _fake_hf_model_info(monkeypatch, _no_network) + assert mc.is_embedding_model("org/partial") is False # partial snapshot -> not trusted + + def test_persist_embedder_is_best_effort_when_home_unwritable(tmp_path, monkeypatch): # Persistence is an optimization, never a correctness requirement: if the # Studio home cannot be created/written, the online path must still return its diff --git a/studio/backend/tests/test_rag_embeddings.py b/studio/backend/tests/test_rag_embeddings.py index 28a2f69426c..58c7de25795 100644 --- a/studio/backend/tests/test_rag_embeddings.py +++ b/studio/backend/tests/test_rag_embeddings.py @@ -220,3 +220,44 @@ def test_st_encode_failure_without_llama_binary_reraises(monkeypatch): embeddings._reset_backend() with pytest.raises(RuntimeError, match = "CUDA error during encode"): embeddings.encode(["alpha", "beta"]) + + +def test_get_resolves_default_casing_before_loading(monkeypatch): + # _get() must resolve a repo id to the exact cache casing before constructing + # SentenceTransformer. A configured default whose spelling differs only by case + # from the cache dir is deliberately NOT persist-normalized by /settings (that + # would turn the default into an override), so without resolving here the offline + # local_files_only load would miss the case-sensitive cache dir and fail. Verify + # both the loader and the security gate receive the resolved name. + import sys + import types + + captured = {} + + class _FakeST: + def __init__(self, name, **kwargs): + captured["load_name"] = name + captured["local_files_only"] = kwargs.get("local_files_only") + + fake_st = types.ModuleType("sentence_transformers") + fake_st.SentenceTransformer = _FakeST + monkeypatch.setitem(sys.modules, "sentence_transformers", fake_st) + + monkeypatch.setattr(config, "effective_embedding_model", lambda: "baai/bge-m3") + monkeypatch.setattr( + "utils.models.resolve_st_cached_repo_id_case", + lambda repo_id: "BAAI/bge-m3" if repo_id == "baai/bge-m3" else repo_id, + ) + monkeypatch.setattr(embeddings, "_device", lambda: "cpu") + monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) + guarded = {} + monkeypatch.setattr( + embeddings, "_guard_model_security", lambda name, local_only: guarded.update(name = name) + ) + + monkeypatch.setattr(embeddings, "_model", None, raising = False) + monkeypatch.setattr(embeddings, "_name", None, raising = False) + embeddings._get() + + assert captured["load_name"] == "BAAI/bge-m3" # resolved to the cache casing + assert guarded["name"] == "BAAI/bge-m3" # gate scans the same resolved name diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 605ebc845e9..c2521db9493 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2188,11 +2188,14 @@ def _dir_has_complete_torch_weights(names: set) -> bool: * a single ``model.safetensors`` / ``pytorch_model.bin``; or * a full shard set -- a shard names its own total (``-00001-of-00002``), so - every index ``1..total`` for that (stem, ext, total) must be present. - - A lone shard such as ``model-00001-of-00002.safetensors`` is NOT complete: the - loader needs the missing shard(s) at indexing time, so accepting a partially - downloaded sharded model would validate and then fail under local_files_only. + every index ``1..total`` for that (stem, ext, total) must be present AND its + ``model.safetensors.index.json`` / ``pytorch_model.bin.index.json`` map, + through which transformers discovers and wires the shards, must be present. + + A lone shard such as ``model-00001-of-00002.safetensors``, a full set missing + its index map, is NOT complete: the loader needs every shard and the index at + indexing time, so accepting a partial sharded download would validate and then + fail under local_files_only. """ for stem in ("model", "pytorch_model"): for ext in ("safetensors", "bin"): @@ -2203,26 +2206,26 @@ def _dir_has_complete_torch_weights(names: set) -> bool: m = _ST_SHARD_RE.match(name) if m: shards.setdefault((m.group(1), m.group(4), int(m.group(3))), set()).add(int(m.group(2))) - for (_stem, _ext, total), indices in shards.items(): - if total > 0 and indices == set(range(1, total + 1)): + for (stem, ext, total), indices in shards.items(): + index_map = f"{stem}.{ext}.index.json" + if total > 0 and indices == set(range(1, total + 1)) and index_map in names: return True return False -def _snapshot_is_loadable_st_model(snap: Path) -> bool: - """True when *snap* holds a sentence-transformers model that can actually load. +def _snapshot_has_complete_weights(snap: Path) -> bool: + """True when *snap* is materialized with a config and a COMPLETE Torch weight + set the default SentenceTransformer backend can load -- the ``modules.json`` + sentence-transformers marker aside. - ``modules.json`` alone is not enough: the online security preflight downloads - exactly that one file via ``hf_hub_download``, and a partial download can - leave it behind too, so a snapshot can carry the marker while the weights and - config SentenceTransformer needs are absent. Accepting that offline would - pass validation and then fail on the first RAG load. Requires the marker plus - a config and a COMPLETE recognized Torch base-model weight set - (``_dir_has_complete_torch_weights``) in one directory of the snapshot. + A tag-only feature-extraction embedder confirmed online has exactly this + (config plus weights, no ``modules.json``) and SentenceTransformer's auto-model + fallback loads it, so an online-confirmed positive may be trusted for it + offline. A partial download (config present but weights missing or an + incomplete shard set) does NOT satisfy this, so it must fail validation here + rather than at first indexing -- the validate-then-fail this guards against. """ try: - if not (snap / "modules.json").is_file(): - return False if not any( (snap / name).is_file() for name in ("config.json", "config_sentence_transformers.json") ): @@ -2230,7 +2233,9 @@ def _snapshot_is_loadable_st_model(snap: Path) -> bool: by_dir: dict = {} for path in snap.rglob("*"): try: - if path.is_file() and _ST_WEIGHT_FILE_RE.match(path.name): + if path.is_file() and ( + _ST_WEIGHT_FILE_RE.match(path.name) or path.name.endswith(".index.json") + ): by_dir.setdefault(path.parent, set()).add(path.name) except OSError: continue @@ -2239,6 +2244,55 @@ def _snapshot_is_loadable_st_model(snap: Path) -> bool: return False +def _snapshot_is_loadable_st_model(snap: Path) -> bool: + """True when *snap* holds a sentence-transformers model that can actually load. + + ``modules.json`` alone is not enough: the online security preflight downloads + exactly that one file via ``hf_hub_download``, and a partial download can + leave it behind too, so a snapshot can carry the marker while the weights and + config SentenceTransformer needs are absent. Accepting that offline would + pass validation and then fail on the first RAG load. Requires the marker plus + a config and a COMPLETE recognized Torch base-model weight set + (``_snapshot_has_complete_weights``). + """ + try: + if not (snap / "modules.json").is_file(): + return False + return _snapshot_has_complete_weights(snap) + except OSError: + return False + + +def _active_snapshot_dir(repo_id: str) -> Optional[Path]: + """The materialized snapshot dir the offline ``local_files_only`` load resolves + for *repo_id*, or None. + + Mirrors that resolution: the repo dir :func:`_st_cache_repo_dir` selects (the + casing the settings route persists and the loader opens), its ``refs/main`` + commit (with ``local_files_only`` huggingface_hub resolves the default revision + THROUGH that ref, so a missing / empty / unreadable ref is a cache MISS, never + a reason to scan historical snapshots), and that commit's snapshot dir. None + when any of those is absent -- a cache miss the caller treats as not-cached. + Never raises: a cache mutating underneath (concurrent deletion) reads as None. + """ + try: + repo_dir = _st_cache_repo_dir(repo_id) + if repo_dir is None: + return None + try: + commit = (repo_dir / "refs" / "main").read_text(encoding = "utf-8").strip() + except OSError: + return None # absent or unreadable: the loader cannot resolve it either + if not commit: + return None # empty / whitespace: a partial write, revision unknown + snapshot = repo_dir / "snapshots" / commit + if not snapshot.is_dir(): + return None # ref recorded but not materialized (partial / pruned) + return snapshot + except Exception: + return None + + def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: """Sentence-transformers detection from the local cache, no network call. @@ -2261,22 +2315,10 @@ def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: raises -- a cache mutating underneath (concurrent model deletion) reads as not-cached so callers keep their normal fallback. """ - try: - repo_dir = _st_cache_repo_dir(repo_id) - if repo_dir is None: - return None - try: - commit = (repo_dir / "refs" / "main").read_text(encoding = "utf-8").strip() - except OSError: - return None # absent or unreadable: the loader cannot resolve it either - if not commit: - return None # empty / whitespace: a partial write, revision unknown - snapshot = repo_dir / "snapshots" / commit - if not snapshot.is_dir(): - return None # ref recorded but not materialized (partial / pruned) - return _snapshot_is_loadable_st_model(snapshot) - except Exception: + snapshot = _active_snapshot_dir(repo_id) + if snapshot is None: return None + return _snapshot_is_loadable_st_model(snapshot) def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: @@ -2307,25 +2349,32 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: if _env_offline(): # Offline: the local HF cache is the only source -- a network call cannot - # succeed and would only hang on a DNS error and get retried (#6817). - # Re-probe the marker every call without consulting or populating the memo: - # a cached negative must not stick (a model downloaded later, or a tag-only - # embedder, could not be confirmed here), so a miss is not durable. - marker = _embedding_marker_in_hf_cache(model_name) + # succeed and would only hang on a DNS error and get retried (#6817). Re-probe + # the cache every call without consulting or populating the memo: a cached + # negative must not stick (a model downloaded later, or a tag-only embedder, + # could not be confirmed here), so a miss is not durable. + snapshot = _active_snapshot_dir(model_name) + if snapshot is not None and _snapshot_is_loadable_st_model(snapshot): + return True # a self-describing sentence-transformers snapshot (modules.json) # Retain a positive confirmed online -- this session (memo) or a prior one - # (persisted allowlist) -- ONLY when the active revision is actually - # materialized locally (marker is not None). A recorded True proves - # model_info() tagged the repo an embedder, not that its files are on disk - # (an online /check-embedding call records the verdict without downloading), - # so trusting it for an uncached repo would save a model the local_files_only - # load then fails on. "not None" (not True) still covers a downloaded - # tag-only feature-extraction embedder whose snapshot is present but has no - # modules.json -- recognizable only from that recorded verdict, and the case - # that must survive a restart, not just the _hf_offline_if_dns_dead() mid-load - # flip. - if _known_embedder(model_name, cache_key) and marker is not None: + # (persisted allowlist) -- ONLY when the active snapshot is materialized WITH + # a complete, loadable weight set. A recorded True proves model_info() tagged + # the repo an embedder, not that its files are on disk (an online + # /check-embedding call records the verdict without downloading), so trusting + # it for an uncached or partially downloaded repo would save a model the + # local_files_only load then fails on. The weight gate (not the bare marker) + # is what still covers a downloaded tag-only feature-extraction embedder -- + # weights present, no modules.json, which SentenceTransformer's auto-model + # fallback loads -- while rejecting a config-only or half-sharded snapshot, + # and it is the case that must survive a restart, not just the + # _hf_offline_if_dns_dead() mid-load flip. + if ( + snapshot is not None + and _known_embedder(model_name, cache_key) + and _snapshot_has_complete_weights(snapshot) + ): return True - return marker is True + return False # Online: the Hub is authoritative for the current remote revision. The local # cache marker reflects only the last-downloaded revision, which may lag the From 9b2895c66f0fe5009084a6ee6eb415cc51776fcd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 19 Jul 2026 10:59:39 +0000 Subject: [PATCH 24/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/rag/embeddings.py | 1 - 1 file changed, 1 deletion(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index d65c3910dff..805a2949103 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -207,7 +207,6 @@ def _get(model_name: str | None = None): if not is_local_path(load_name): from utils.models import resolve_st_cached_repo_id_case - load_name = resolve_st_cached_repo_id_case(load_name) logger.info("loading embedding model %s on %s", load_name, device) # Read the offline state ONCE and use that single value for both the From 9753c43b343cad550ba7129de8d17a87f43eb19a Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 19 Jul 2026 11:12:34 +0000 Subject: [PATCH 25/77] studio: require a tokenizer, case-fold verdict ids, and serialize verdict writes Three follow-ups to the offline embedding-model classifier from the latest review: - _snapshot_has_complete_weights now also requires a tokenizer asset. A SentenceTransformer Transformer module builds an AutoTokenizer, so a snapshot with a complete weight set but no tokenizer.json / tokenizer_config.json / vocab still fails the local_files_only load. The check is a permissive union over the common fast-tokenizer, config, and WordPiece/BPE/SentencePiece assets, so an unusual but valid layout is not rejected -- only a genuinely tokenizer-less partial download. - The persisted embedder allowlist is now keyed case-insensitively. model_info() is queried under the requested casing while the settings route saves the cache-resolved casing, so an exact-string lookup missed the persisted positive after a restart (baai/model recorded, BAAI/model looked up) and a loadable tag-only embedder was rejected. Both persist and lookup case-fold the id. - _persist_embedder serializes its read-modify-write under a lock and writes through a per-thread temp file, so concurrent confirmations of different embedders no longer drop each other's entry or collide on the temp path. Cross-process writers stay best-effort (os.replace is atomic; a dropped verdict is only an optimization miss a later online re-confirmation heals). Adds regression tests for the missing-tokenizer reject, alternate tokenizer assets, cross-casing verdict match, and concurrent verdict writes; updates the snapshot test helpers to materialize a tokenizer alongside config and weights. --- .../tests/test_is_embedding_model_offline.py | 75 ++++++++++- studio/backend/utils/models/model_config.py | 117 +++++++++++++----- 2 files changed, 155 insertions(+), 37 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index dd00a408df9..5b7cd4e9477 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -70,8 +70,8 @@ def _repo( """Build a real cache repo dir and point the ST cache root at it. ``snapshots``: (commit, sentence_transformer) tuples. A snapshot is written - fully loadable (config + weights) so only the marker varies. ``main_ref`` - writes refs/main. Returns the snapshot dirs in the given order. + fully loadable (config + tokenizer + weights) so only the marker varies. + ``main_ref`` writes refs/main. Returns the snapshot dirs in the given order. """ cache_root = root if root is not None else tmp_path / "cache" cache_root.mkdir(parents = True, exist_ok = True) @@ -81,6 +81,7 @@ def _repo( d = repo / "snapshots" / name d.mkdir(parents = True, exist_ok = True) (d / "config.json").write_text("{}") + (d / "tokenizer.json").write_text("{}") (d / "model.safetensors").write_bytes(b"\0") if is_st: (d / "modules.json").write_text("[]") @@ -202,6 +203,7 @@ def _st_snapshot( (snap / "modules.json").write_text("[]") if loadable: (snap / "config.json").write_text("{}") + (snap / "tokenizer.json").write_text("{}") (snap / "model.safetensors").write_bytes(b"\0") refs = root / repo_dir / "refs" refs.mkdir(parents = True, exist_ok = True) @@ -276,6 +278,7 @@ def test_marker_onnx_only_snapshot_is_not_loadable(tmp_path, monkeypatch): snap.mkdir(parents = True) (snap / "modules.json").write_text("[]") (snap / "config.json").write_text("{}") + (snap / "tokenizer.json").write_text("{}") # isolate the failure to the weight format (snap / "model.onnx").write_bytes(b"\0") # refs/main so the probe resolves this revision and reaches the weight check # (without it the answer would be None -- a cache miss -- for a different reason). @@ -286,19 +289,41 @@ def test_marker_onnx_only_snapshot_is_not_loadable(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache("org/model") is False +def test_marker_rejects_weights_without_tokenizer(tmp_path, monkeypatch): + # A snapshot with modules.json + config + a complete weight set but NO tokenizer + # asset is not loadable: SentenceTransformer's Transformer module also builds an + # AutoTokenizer, which fails offline when tokenizer.json / tokenizer_config.json / + # a vocab is absent. Accepting it would pass validation and then fail at load. + _cache_repo_with_files(tmp_path, monkeypatch, "model.safetensors", tokenizer = False) + assert mc._embedding_marker_in_hf_cache("org/model") is False + + +@pytest.mark.parametrize("tok_file", ["tokenizer_config.json", "vocab.txt", "spiece.model"]) +def test_marker_accepts_alternate_tokenizer_assets(tmp_path, monkeypatch, tok_file): + # The tokenizer check is a permissive union: any one recognized asset (fast + # tokenizer json, a WordPiece vocab, a SentencePiece model, ...) is enough, so a + # valid non-tokenizer.json layout is not wrongly rejected. + _cache_repo_with_files(tmp_path, monkeypatch, "model.safetensors", tok_file, tokenizer = False) + assert mc._embedding_marker_in_hf_cache("org/model") is True + + def _cache_repo_with_files( tmp_path, monkeypatch, *files, commit = "aaa", + tokenizer = True, ): - """A cache repo whose active snapshot holds modules.json + config + *files*.""" + """A cache repo whose active snapshot holds modules.json + config + a tokenizer + + *files*. ``tokenizer=False`` omits the tokenizer asset (a partial download).""" hf_root = tmp_path / "hf" repo = hf_root / "models--org--model" snap = repo / "snapshots" / commit snap.mkdir(parents = True) (snap / "modules.json").write_text("[]") (snap / "config.json").write_text("{}") + if tokenizer: + (snap / "tokenizer.json").write_text("{}") for name in files: target = snap / name target.parent.mkdir(parents = True, exist_ok = True) @@ -626,6 +651,50 @@ def _info(model_name, token = None): assert mc.is_embedding_model("org/partial") is False # partial snapshot -> not trusted +def test_offline_persisted_verdict_matches_across_casing(tmp_path, monkeypatch): + # The verdict is recorded under the online-request spelling but the settings route + # saves the cache-resolved spelling, so an exact-string lookup would miss it. A + # tag-only embedder verified online as baai/model, then looked up offline as the + # saved BAAI/model, must still be recognized: the allowlist is case-folded. + _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "BAAI/model") + + def _info(model_name, token = None): + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + + _fake_hf_model_info(monkeypatch, _info) + assert mc.is_embedding_model("baai/model") is True # online: confirmed as baai/model + + mc._embedding_detection_cache.clear() # restart: memo lost, disk verdict remains + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _fake_hf_model_info(monkeypatch, _no_network) + # Looked up under the cache-resolved casing the settings route persisted. + assert mc.is_embedding_model("BAAI/model") is True + + +def test_persist_embedder_concurrent_writes_keep_every_verdict(tmp_path, monkeypatch): + # Concurrent online confirmations must not drop each other's entry: the + # read-modify-write is serialized and each writer uses a per-thread temp file, so + # every model persisted from parallel threads survives (a lost verdict would go + # unrecognized offline after a restart). + import threading + + names = [f"org/emb-{i}" for i in range(24)] + barrier = threading.Barrier(len(names)) + + def _writer(name): + barrier.wait() # maximize overlap on the shared file + mc._persist_embedder(name) + + threads = [threading.Thread(target = _writer, args = (n,)) for n in names] + for t in threads: + t.start() + for t in threads: + t.join() + + persisted = mc._load_persisted_embedders() + assert {n.casefold() for n in names} <= persisted + + def test_persist_embedder_is_best_effort_when_home_unwritable(tmp_path, monkeypatch): # Persistence is an optimization, never a correctness requirement: if the # Studio home cannot be created/written, the online path must still return its diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index c2521db9493..80ad3ab2814 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2118,18 +2118,37 @@ def _persisted_embedders_path() -> Path: return _studio_root() / "embedding_verdicts.json" +# Serializes the read-modify-write in _persist_embedder so two threads confirming +# different embedders at once cannot each overwrite the file with a one-entry update +# (a lost verdict) or collide on the temp file. Cross-process writers are inherently +# best-effort -- os.replace stays atomic, and a dropped verdict is only an +# optimization miss a later online re-confirmation heals. +_persist_lock = threading.Lock() + + +def _verdict_key(model_name: str) -> str: + """Case-folded key for the embedder allowlist. The verdict is a property of the + repo, not its spelling: model_info() is queried under the requested casing while + the settings route saves the cache-resolved casing, so keying by the exact string + would miss the persisted positive after a restart (``baai/model`` recorded, + ``BAAI/model`` looked up). Repo ids are case-insensitive, so fold both sides.""" + return model_name.casefold() + + def _load_persisted_embedders() -> set: - """Repo ids recorded as embedders in a prior session; empty set on any error. + """Case-folded repo ids recorded as embedders in a prior session; empty set on any + error. Best-effort and never raises: a missing, empty, corrupt, or non-list file (and any read/decode error) reads as "nothing recorded" so callers fall back - to their normal cache-marker logic. + to their normal cache-marker logic. Folding at load time also normalizes any + mixed-case entries an older build may have written. """ try: with open(_persisted_embedders_path(), encoding = "utf-8") as fh: data = json.load(fh) if isinstance(data, list): - return {name for name in data if isinstance(name, str)} + return {_verdict_key(name) for name in data if isinstance(name, str)} except Exception: pass return set() @@ -2138,24 +2157,28 @@ def _load_persisted_embedders() -> set: def _persist_embedder(model_name: str) -> None: """Record *model_name* as an online-confirmed embedder, best-effort. - Only positive Hub verdicts are written (never a negative or a token): the - file is a durable allowlist the offline branch consults, gated there on the - active revision actually being materialized on disk, so a stale entry for a - since-deleted cache cannot resurrect a False. Any failure (unwritable home, - race) is swallowed -- persistence is an optimization, not a correctness - requirement for the online path that calls it. + Only positive Hub verdicts are written (never a negative or a token), case-folded + so the lookup matches regardless of casing: the file is a durable allowlist the + offline branch consults, gated there on the active revision actually being + materialized on disk, so a stale entry for a since-deleted cache cannot resurrect + a False. The read-modify-write is serialized under ``_persist_lock`` and writes + through a per-thread temp file, so concurrent confirmations cannot drop each + other's entry or collide. Any failure (unwritable home, race) is swallowed -- + persistence is an optimization, not a correctness requirement for the online path. """ try: - current = _load_persisted_embedders() - if model_name in current: - return - current.add(model_name) - path = _persisted_embedders_path() - path.parent.mkdir(parents = True, exist_ok = True) - tmp = path.with_name(path.name + f".{os.getpid()}.tmp") - with open(tmp, "w", encoding = "utf-8") as fh: - json.dump(sorted(current), fh) - os.replace(tmp, path) + key = _verdict_key(model_name) + with _persist_lock: + current = _load_persisted_embedders() + if key in current: + return + current.add(key) + path = _persisted_embedders_path() + path.parent.mkdir(parents = True, exist_ok = True) + tmp = path.with_name(path.name + f".{os.getpid()}.{threading.get_ident()}.tmp") + with open(tmp, "w", encoding = "utf-8") as fh: + json.dump(sorted(current), fh) + os.replace(tmp, path) except Exception as e: logger.debug(f"Could not persist embedder verdict for {model_name}: {e}") @@ -2163,10 +2186,11 @@ def _persist_embedder(model_name: str) -> None: def _known_embedder(model_name: str, cache_key: tuple) -> bool: """True when *model_name* was confirmed an embedder online, this session or a prior one -- the session memo (authoritative, token-scoped) OR the persisted - cross-restart allowlist. Callers still gate this on a materialized snapshot.""" + cross-restart allowlist (matched case-insensitively). Callers still gate this on a + materialized snapshot.""" if _embedding_detection_cache.get(cache_key) is True: return True - return model_name in _load_persisted_embedders() + return _verdict_key(model_name) in _load_persisted_embedders() # The base-model weight files the RAG loader's default backend consumes. _get() @@ -2181,6 +2205,25 @@ def _known_embedder(model_name: str, cache_key: tuple) -> bool: _ST_WEIGHT_FILE_RE = re.compile(r"^(model|pytorch_model)(-\d+-of-\d+)?\.(safetensors|bin)$") _ST_SHARD_RE = re.compile(r"^(model|pytorch_model)-(\d+)-of-(\d+)\.(safetensors|bin)$") +# A SentenceTransformer Transformer module loads its weights AND an AutoTokenizer, +# so a snapshot with a complete weight set but no tokenizer asset still fails the +# local_files_only load. Any ONE of these is enough to load some tokenizer -- kept a +# permissive union (fast tokenizer, config, or a WordPiece/BPE/SentencePiece vocab) +# so an unusual-but-valid layout is not rejected, only a genuinely tokenizer-less +# partial download. +_ST_TOKENIZER_FILES = frozenset( + { + "tokenizer.json", + "tokenizer_config.json", + "tokenizer.model", + "vocab.txt", + "vocab.json", + "merges.txt", + "spiece.model", + "sentencepiece.bpe.model", + } +) + def _dir_has_complete_torch_weights(names: set) -> bool: """True when the filenames in ONE directory form a COMPLETE Torch base-model @@ -2214,16 +2257,17 @@ def _dir_has_complete_torch_weights(names: set) -> bool: def _snapshot_has_complete_weights(snap: Path) -> bool: - """True when *snap* is materialized with a config and a COMPLETE Torch weight - set the default SentenceTransformer backend can load -- the ``modules.json`` - sentence-transformers marker aside. - - A tag-only feature-extraction embedder confirmed online has exactly this - (config plus weights, no ``modules.json``) and SentenceTransformer's auto-model - fallback loads it, so an online-confirmed positive may be trusted for it - offline. A partial download (config present but weights missing or an - incomplete shard set) does NOT satisfy this, so it must fail validation here - rather than at first indexing -- the validate-then-fail this guards against. + """True when *snap* is materialized with a config, a tokenizer asset, and a + COMPLETE Torch weight set the default SentenceTransformer backend can load -- + the ``modules.json`` sentence-transformers marker aside. + + A tag-only feature-extraction embedder confirmed online has exactly this (config, + tokenizer, weights, no ``modules.json``) and SentenceTransformer's auto-model + fallback loads it, so an online-confirmed positive may be trusted for it offline. + A partial download -- config present but weights missing / an incomplete shard + set, or weights present but the tokenizer AutoTokenizer needs absent -- does NOT + satisfy this and must fail validation here rather than at first indexing, the + validate-then-fail this guards against. """ try: if not any( @@ -2231,14 +2275,19 @@ def _snapshot_has_complete_weights(snap: Path) -> bool: ): return False by_dir: dict = {} + has_tokenizer = False for path in snap.rglob("*"): try: - if path.is_file() and ( - _ST_WEIGHT_FILE_RE.match(path.name) or path.name.endswith(".index.json") - ): + if not path.is_file(): + continue + if path.name in _ST_TOKENIZER_FILES: + has_tokenizer = True + if _ST_WEIGHT_FILE_RE.match(path.name) or path.name.endswith(".index.json"): by_dir.setdefault(path.parent, set()).add(path.name) except OSError: continue + if not has_tokenizer: + return False return any(_dir_has_complete_torch_weights(names) for names in by_dir.values()) except OSError: return False From 8f4d3bb66d11be04292c96ad348dcc4f72eed051 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 19 Jul 2026 12:07:59 +0000 Subject: [PATCH 26/77] studio: tighten comments in the offline embedding-model classifier Comment-only pass over the PR's changed files. Collapse the long block comments and docstrings around is_embedding_model, the cache-snapshot and weight-completeness helpers, the embedder-verdict persistence, the offline security gate, and the offline/casing tests to short one- or two-line forms. Preserve the rationale (issue #6817, the local_files_only invariant, the casing and weight-gate reasons) in far fewer words. No code changes. --- studio/backend/core/rag/embeddings.py | 63 ++-- studio/backend/routes/settings.py | 43 ++- .../test_embedding_model_security_gate.py | 24 +- studio/backend/tests/test_hf_offline_flags.py | 48 +-- .../tests/test_is_embedding_model_offline.py | 202 +++++-------- studio/backend/tests/test_rag_embeddings.py | 8 +- studio/backend/utils/models/model_config.py | 276 ++++++------------ .../backend/utils/security/file_security.py | 26 +- studio/backend/utils/utils.py | 11 +- 9 files changed, 255 insertions(+), 446 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 805a2949103..9b8a0264172 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -87,13 +87,11 @@ def _st_module_subdirs(name: str, token: str | None, local_only: bool) -> tuple[ security scan: a flagged pickle directly under one must block. Returns () on any failure (no modules.json, offline, malformed) so the guard never bricks the embedder. - ``local_only`` MUST be the value the caller captured for the load, not a fresh - read of the environment. ``_hf_offline_if_dns_dead()`` flips the process-wide - offline vars and restores them, so re-reading here could force this probe - local-only, return () because modules.json is not cached, and leave the scan - with NO module load roots -- a flagged pickle under ``0_Transformer/`` would - then pass as an unreferenced nested artifact while the loader, using the - captured predicate, still fetched and deserialized it. + ``local_only`` MUST be the value the caller captured for the load, not a fresh env + read: ``_hf_offline_if_dns_dead()`` flips the offline vars and restores them, so + re-reading could force this probe local-only, return () (modules.json not cached), and + leave the scan with no module roots -- a flagged pickle under ``0_Transformer/`` would + then pass while the loader, using the captured predicate, still deserialized it. """ try: import json @@ -113,11 +111,10 @@ def _st_module_subdirs(name: str, token: str | None, local_only: bool) -> tuple[ from huggingface_hub.utils import EntryNotFoundError try: - # local_files_only from the CAPTURED predicate: huggingface_hub - # honors only HF_HUB_OFFLINE natively, so an offline session would - # otherwise block on network timeouts here even though the cached - # snapshot already has modules.json -- but it must be the same - # value the load uses, never a fresh env read (see the docstring). + # local_files_only from the CAPTURED predicate (never a fresh env read; + # see the docstring): huggingface_hub honors only HF_HUB_OFFLINE, so an + # offline session would otherwise block on timeouts despite a cached + # modules.json. local = hf_hub_download( name, "modules.json", @@ -143,13 +140,11 @@ def _guard_model_security(name: str, local_only_load: bool) -> None: /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_load`` MUST be the same value the caller passes to - SentenceTransformer as ``local_files_only`` -- it is what licenses skipping the - Hub scan, so it is taken as an argument rather than re-read from the - environment here. ``_hf_offline_if_dns_dead()`` mutates the process-wide - offline vars and restores them, so two separate reads can disagree: this - could skip the scan while the constructor then loaded with - ``local_files_only=False`` and fetched the unscanned repo. + ``local_only_load`` MUST be the same value the caller passes to SentenceTransformer as + ``local_files_only`` -- it is what licenses skipping the Hub scan, so it is an argument + rather than re-read here. ``_hf_offline_if_dns_dead()`` mutates the offline vars and + restores them, so two reads can disagree: this could skip the scan while the constructor + then loaded with ``local_files_only=False`` and fetched the unscanned repo. """ try: from utils.security import evaluate_file_security, security_load_subdirs @@ -194,14 +189,11 @@ def _get(model_name: str | None = None): from utils.utils import hf_env_offline device = _device() - # Resolve to the exact cache casing the offline local_files_only load - # needs. The /settings route persists that spelling for a custom override - # but deliberately leaves the configured default verbatim (rewriting it - # would turn the default into an override and break later default changes), - # so a default whose casing differs from the cache dir would otherwise - # miss it and fail offline. Resolve here at load time to cover the default - # too; a no-op for a local path or when nothing case-matching is cached, - # and idempotent for an override already normalized at persist time. + # Resolve to the exact cache casing the offline local_files_only load needs. + # The /settings route persists that spelling for a custom override but leaves + # the configured default verbatim (rewriting it would turn the default into an + # override), so resolve here at load time to cover the default too. A no-op for + # a local path or when nothing case-matching is cached. load_name = name from utils.paths import is_local_path @@ -209,18 +201,15 @@ def _get(model_name: str | None = None): from utils.models import resolve_st_cached_repo_id_case load_name = resolve_st_cached_repo_id_case(load_name) logger.info("loading embedding model %s on %s", load_name, device) - # Read the offline state ONCE and use that single value for both the - # security gate and the loader. _hf_offline_if_dns_dead() mutates the - # process-wide offline vars and restores them on exit, so re-reading - # for the constructor could yield False after the guard had already - # skipped the Hub scan on True -- and the load would then fetch and - # deserialize the unscanned repo. + # Read the offline state ONCE for both the security gate and the loader. + # _hf_offline_if_dns_dead() mutates the offline vars and restores them, so + # re-reading for the constructor could yield False after the guard skipped the + # Hub scan on True -- and the load would then fetch the unscanned repo. local_only = hf_env_offline() _guard_model_security(load_name, local_only) - # Propagate the user's offline intent into the loader: SentenceTransformer - # performs its own Hub operations, and huggingface_hub honors only - # HF_HUB_OFFLINE, so a TRANSFORMERS_OFFLINE-only session would otherwise - # still fetch missing repo files over the network. + # Propagate offline intent into the loader: huggingface_hub honors only + # HF_HUB_OFFLINE, so a TRANSFORMERS_OFFLINE-only session would otherwise still + # fetch missing repo files. _model = SentenceTransformer( load_name, device = device, diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index cc46e147716..dbf069fa238 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -413,11 +413,9 @@ def update_embedding_model( scan_st_pickle = ( model != default_embedding_model() and not is_local_gguf and not _llama_backend_active() ) - # Read the offline state ONCE for this request and reuse it: the module probe - # and the scan must agree, and _hf_offline_if_dns_dead() can flip the - # process-wide vars between two reads (a probe forced local-only can return no - # module roots, letting a flagged 0_Transformer/ pickle pass as a nested - # artifact while the scan was skipped on a different value). + # Read the offline state ONCE and reuse it: the module probe and the scan must agree, + # and _hf_offline_if_dns_dead() can flip the vars between two reads (a probe forced + # local-only returns no module roots, letting a flagged 0_Transformer/ pickle pass). from utils.utils import hf_env_offline local_only_load = hf_env_offline() @@ -443,8 +441,7 @@ def update_embedding_model( ) ) # local_only_load: this gate covers the RAG embedder, whose loader pins - # SentenceTransformer to the local cache with the same predicate, so - # offline the scan can only stall on timeouts before failing open. + # SentenceTransformer to the local cache with the same predicate. if evaluate_file_security( model, hf_token = scan_token, @@ -480,26 +477,20 @@ def update_embedding_model( gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error(model, hf_token) if gguf_error: raise HTTPException(status_code = 409, detail = gguf_error) - # Persist the casing the local HF cache actually uses: validation above - # accepts a case-insensitive cache hit, but the offline SentenceTransformer - # load resolves the cache by exact case, so store the cached spelling (a - # no-op when nothing case-matching is cached) to keep the model loadable. - # Resolved against the cache the ST loader itself searches (ST_HOME when set, - # else the Hub cache), so the persisted spelling is one the offline - # exact-case load can actually find. Three cases are deliberately left alone: + # Persist the exact cache casing: validation accepts a case-insensitive hit, but the + # offline ST load resolves the cache by exact case. Resolved against the cache the ST + # loader searches (ST_HOME if set, else the Hub cache). A no-op when nothing + # case-matching is cached. Three cases are left alone: # - # * the default -- rewriting its casing would make the exact-string default - # comparison in set_rag_embedding_model() treat it as a custom override, - # so later changes to the configured default would stop applying; - # * a local path -- a relative directory like "org/model" is loaded from - # disk, and rewriting it to a case-insensitive cache collision - # ("Org/model") would stop resolving to that directory; - # * the llama-server backend -- it does not load through SentenceTransformer - # at all. It derives a GGUF companion via effective_gguf_repo() from this - # saved spelling and fetches it with hf_hub_download, i.e. from the HUB - # cache. Normalizing to an ST_HOME spelling would change the derived repo - # to one _hf_gguf_backend_error() never validated (BAAI/bge-m3-GGUF rather - # than the checked baai/bge-m3-GGUF) and which may be absent offline. + # * the default -- rewriting it would make set_rag_embedding_model()'s exact-string + # default comparison treat it as a custom override, so later default changes stop + # applying; + # * a local path -- a directory like "org/model" is loaded from disk, and rewriting + # to a cache collision ("Org/model") would stop resolving to it; + # * the llama-server backend -- it does not load through SentenceTransformer; it + # derives a GGUF companion via effective_gguf_repo() and fetches it from the HUB + # cache, so an ST_HOME spelling would pick a repo _hf_gguf_backend_error() never + # validated (BAAI/bge-m3-GGUF vs the checked baai/bge-m3-GGUF). if ( model != default_embedding_model() and not is_local_path(model) diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index 5919f0fe981..a2651e55704 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -266,9 +266,8 @@ def test_clean_repo_saves_under_force(client, monkeypatch): def test_custom_model_saved_in_cache_casing(client, monkeypatch): - # A custom repo is persisted in the casing its local HF cache dir uses, so an - # offline exact-case SentenceTransformer load still finds it (baai/bge-m3 -> - # BAAI/bge-m3). force=True bypasses the verification gate to isolate this. + # A custom repo is persisted in its local HF cache casing so the offline exact-case + # ST load finds it (baai/bge-m3 -> BAAI/bge-m3). force=True isolates this from the gate. c, saved = client monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) import utils.models as _models @@ -280,9 +279,8 @@ def test_custom_model_saved_in_cache_casing(client, monkeypatch): def test_local_path_model_is_not_casing_normalized(client, monkeypatch, tmp_path): - # A local directory (here an existing path) is loaded from disk. Rewriting it - # to a case-insensitive HF cache collision would stop resolving to that - # directory and be read as a Hub repo id, so normalization must be skipped. + # A local directory is loaded from disk; casing normalization must be skipped or it + # would resolve to a cache collision instead of the directory. c, saved = client monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) import utils.models as _models @@ -299,11 +297,9 @@ def _must_not_run(m): def test_llama_backend_model_is_not_casing_normalized(monkeypatch): - # On the llama-server backend nothing loads through SentenceTransformer: the - # embedder derives a GGUF companion from this saved spelling via - # effective_gguf_repo() and fetches it with hf_hub_download, i.e. from the HUB - # cache. Rewriting to an ST_HOME spelling would point it at a GGUF repo that - # _hf_gguf_backend_error() never validated and that may be absent offline. + # On the llama-server backend nothing loads through SentenceTransformer: the embedder + # derives a GGUF companion from this spelling and fetches it from the HUB cache, so an + # ST_HOME rewrite would point at a GGUF repo _hf_gguf_backend_error() never validated. saved: dict = {} monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) @@ -330,10 +326,8 @@ def _must_not_run(m): def test_default_model_is_not_casing_normalized(client, monkeypatch): - # Submitting the exact default must NOT be run through cache-casing - # normalization: rewriting it would make set_rag_embedding_model()'s exact - # default comparison treat it as a custom override, so later default changes - # would stop applying. resolve_st_cached_repo_id_case must not be consulted. + # The exact default must NOT be casing-normalized: rewriting it would make + # set_rag_embedding_model()'s exact default comparison treat it as a custom override. c, saved = client import utils.models as _models diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index d96f79162e5..a1eb10decea 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -3,16 +3,11 @@ """Effective-offline handling for the embedding preflight. -``huggingface_hub`` honors only ``HF_HUB_OFFLINE``; ``TRANSFORMERS_OFFLINE`` -expresses the same user intent but does not itself stop a fetch. Studio treats -either as offline (``hf_env_offline``) and makes that real by passing -``local_files_only`` to the loader, which lets the metadata-only Hub security -scan skip straight to its documented fail-open instead of burning both request -timeouts on a session the user declared offline. - -That skip is only sound while the loader really is pinned to the local cache, so -the coupling between the two is pinned here as an explicit invariant rather than -left to a comment. +``huggingface_hub`` honors only ``HF_HUB_OFFLINE``; ``TRANSFORMERS_OFFLINE`` expresses the +same intent but does not stop a fetch. Studio treats either as offline (``hf_env_offline``) +and makes it real by passing ``local_files_only`` to the loader, which lets the Hub security +scan skip to its fail-open instead of burning both timeouts. That skip is sound only while +the loader is pinned to the local cache, so the coupling is pinned here as an invariant. """ from __future__ import annotations @@ -30,8 +25,7 @@ def _maybe_stub(name: str, builder): - # Stub only if the real module is unavailable, so this file never shadows - # real packages for later tests in the same pytest process. + # Stub only if the real module is unavailable, so this file never shadows real packages. try: importlib.import_module(name) except ImportError: @@ -104,10 +98,8 @@ def _model_info(*a, **k): @pytest.mark.parametrize("var", ["HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"]) def test_shared_gate_still_scans_when_offline_by_default(monkeypatch, var): - # THE important one. This gate is shared by every loader (training, MLX, - # export, ...), and most do NOT constrain their loaders to the local cache, - # so an offline-looking env var alone must never disable the malware scan -- - # those paths can still fetch and deserialize an unscanned model. + # THE important one: the gate is shared by loaders that do NOT pin to the local cache, + # so an offline env var alone must never disable the malware scan. import utils.security.file_security as fs calls: list = [] @@ -118,8 +110,7 @@ def test_shared_gate_still_scans_when_offline_by_default(monkeypatch, var): def test_security_scan_short_circuits_for_a_local_only_caller(monkeypatch): - # A caller that guarantees a local-only load gets the Hub round-trip skipped - # instead of burning both request timeouts (10s + 20s) before failing open. + # A local-only caller gets the Hub round-trip skipped instead of burning both timeouts. import utils.security.file_security as fs calls: list = [] @@ -142,15 +133,10 @@ def _read_backend(rel: str) -> str: def test_embedding_loader_forces_local_only_when_offline(): - """The invariant the RAG opt-in rests on. - - core/rag/embeddings.py is allowed to pass local_only_load because its loader - is pinned to the local cache by the SAME value. It must be read ONCE and - shared: _hf_offline_if_dns_dead() mutates the process-wide offline vars and - restores them, so two separate hf_env_offline() reads can disagree and the - guard could skip the scan on True while the constructor then fetched the - unscanned repo with local_files_only=False. Checked at source level because - importing the loader drags in sentence_transformers/torch. + """The invariant the RAG opt-in rests on: embeddings.py may pass local_only_load + because its loader is pinned to the local cache by the SAME value, read ONCE and shared + (two hf_env_offline() reads can disagree since _hf_offline_if_dns_dead() flips the vars). + Checked at source level because importing the loader drags in sentence_transformers. """ src = _read_backend("core/rag/embeddings.py") assert ( @@ -172,11 +158,9 @@ def test_embedding_loader_forces_local_only_when_offline(): def test_only_the_rag_embedding_path_opts_into_the_bypass(): - """No other loader may claim local-only without constraining its loader. - - The MLX/inference, training and export gates call from_pretrained without a - local-only argument, so if one of them started passing local_only_load the - malware gate would be disabled for a path that can still fetch. + """No other loader may claim local-only without constraining its loader: the + MLX/inference, training and export gates call from_pretrained without a local-only + argument, so passing local_only_load there would disable the gate for a fetching path. """ allowed = {"core/rag/embeddings.py", "routes/settings.py"} callers = [ diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 5b7cd4e9477..af91be1982f 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -20,8 +20,7 @@ def _maybe_stub(name: str, builder): - # Stub only if the real module is unavailable, so this file never shadows - # real packages for later tests in the same pytest process. + # Stub only if the real module is unavailable, so this file never shadows real packages. try: importlib.import_module(name) except ImportError: @@ -51,8 +50,7 @@ def _clean_state(tmp_path, monkeypatch): mc._embedding_detection_cache.clear() monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) - # Point the persisted embedder-verdict store at a per-test Studio home so the - # cross-restart allowlist never leaks into another test or the real ~/.unsloth. + # Per-test Studio home so the persisted allowlist never leaks into another test. monkeypatch.delenv("STUDIO_HOME", raising = False) monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "_studio_home")) yield @@ -67,11 +65,9 @@ def _repo( repo_id = "org/model", root = None, ): - """Build a real cache repo dir and point the ST cache root at it. - - ``snapshots``: (commit, sentence_transformer) tuples. A snapshot is written - fully loadable (config + tokenizer + weights) so only the marker varies. - ``main_ref`` writes refs/main. Returns the snapshot dirs in the given order. + """Build a real cache repo dir and point the ST cache root at it. ``snapshots``: + (commit, sentence_transformer) tuples, each written fully loadable so only the marker + varies. ``main_ref`` writes refs/main. Returns the snapshot dirs in order. """ cache_root = root if root is not None else tmp_path / "cache" cache_root.mkdir(parents = True, exist_ok = True) @@ -127,9 +123,8 @@ def test_marker_none_when_not_cached(monkeypatch): def test_marker_judges_the_revision_refs_main_points_at(tmp_path, monkeypatch): - # The repo USED to be a sentence-transformers model (the old snapshot has - # modules.json) but the revision refs/main points at no longer is. The active - # revision must win; scanning any snapshot would wrongly say True. + # The old snapshot has modules.json but the active refs/main revision does not: the + # active revision must win, so scanning any snapshot would wrongly say True. _repo(tmp_path, monkeypatch, ("new", False), ("old", True), main_ref = "new") assert mc._embedding_marker_in_hf_cache("org/model") is False @@ -140,9 +135,8 @@ def test_marker_true_when_refs_main_revision_is_st(tmp_path, monkeypatch): def test_marker_missing_ref_is_a_cache_miss(tmp_path, monkeypatch): - # With local_files_only=True huggingface_hub resolves the default revision - # THROUGH refs/main, so a snapshot dir alone is not discoverable. Accepting - # one here would pass validation and then fail at first indexing. + # local_files_only resolves the default revision THROUGH refs/main, so a snapshot dir + # alone is not discoverable and accepting one would fail at first indexing. _repo(tmp_path, monkeypatch, ("new", False), ("old", True)) # no refs/main assert mc._embedding_marker_in_hf_cache("org/model") is None @@ -172,10 +166,9 @@ def test_marker_empty_ref_is_cache_miss(tmp_path, monkeypatch): def test_marker_scopes_to_the_repo_dir_the_loader_opens(tmp_path, monkeypatch): - # Two cache dirs differing only by case. The loader opens the one the - # settings route persists (exact case first, as resolve_st_cached_repo_id_case - # picks), so a complete model there must validate even when the OTHER variant - # holds a newer, non-ST snapshot -- judging across both would reject it. + # Two cache dirs differing only by case: the loader opens the one persisted (exact case + # first), so a complete model there must validate even when the OTHER variant holds a + # newer, non-ST snapshot -- judging across both would reject it. if not _case_sensitive_fs(tmp_path): pytest.skip("duplicate case variants need a case-sensitive filesystem") root = tmp_path / "cache" @@ -212,9 +205,8 @@ def _st_snapshot( def test_st_probe_uses_sentence_transformers_home(tmp_path, monkeypatch): - # _get() builds SentenceTransformer without cache_folder, so with ST_HOME set - # that is the ONLY cache the load searches. The probe must follow it, or a - # model present there is called uncached and rejected with a 409 offline. + # With ST_HOME set that is the ONLY cache the load searches, so the probe must follow + # it or a model present there is called uncached and 409'd offline. hf_root, st_root = tmp_path / "hf", tmp_path / "st" hf_root.mkdir() _st_snapshot(st_root, "models--org--model") @@ -236,9 +228,8 @@ def test_st_probe_ignores_hub_cache_when_st_home_is_set(tmp_path, monkeypatch): def test_gguf_probe_never_follows_st_home(tmp_path, monkeypatch): - # The GGUF path downloads with hf_hub_download and no cache_dir, so it uses - # the Hub cache. Letting its probe see ST_HOME would select a file the GGUF - # loader cannot find. + # The GGUF path uses the Hub cache (hf_hub_download, no cache_dir), so letting its probe + # see ST_HOME would select a file the GGUF loader cannot find. hf_root, st_root = tmp_path / "hf", tmp_path / "st" hf_root.mkdir() _st_snapshot(st_root, "models--org--model") @@ -256,9 +247,8 @@ def test_st_probe_falls_back_to_hub_cache(tmp_path, monkeypatch): def test_marker_only_snapshot_is_not_loadable(tmp_path, monkeypatch): - # The online security preflight downloads modules.json on its own via - # hf_hub_download, and a partial download leaves it too. Accepting that - # offline passes validation and then fails on the first RAG load. + # The security preflight downloads modules.json on its own, and a partial download + # leaves it too, so accepting it offline would fail on the first RAG load. hf_root = tmp_path / "hf" _st_snapshot(hf_root, "models--org--model", loadable = False) _fake_hf_cache(monkeypatch, hf_root) @@ -267,11 +257,9 @@ def test_marker_only_snapshot_is_not_loadable(tmp_path, monkeypatch): def test_marker_onnx_only_snapshot_is_not_loadable(tmp_path, monkeypatch): - # A snapshot with the marker and config but ONLY an ONNX export cached is not - # loadable: _get() builds SentenceTransformer with the default Torch backend - # (no backend="onnx"), which reads model.safetensors / pytorch_model.bin, so - # accepting the ONNX offline would pass validation and then fail on the first - # RAG load -- the same validate-then-fail as the marker-only case. + # Marker + config but ONLY an ONNX export is not loadable: the default Torch backend + # reads model.safetensors / pytorch_model.bin, so accepting the ONNX offline would fail + # at load. hf_root = tmp_path / "hf" repo = hf_root / "models--org--model" snap = repo / "snapshots" / "aaa" @@ -280,8 +268,7 @@ def test_marker_onnx_only_snapshot_is_not_loadable(tmp_path, monkeypatch): (snap / "config.json").write_text("{}") (snap / "tokenizer.json").write_text("{}") # isolate the failure to the weight format (snap / "model.onnx").write_bytes(b"\0") - # refs/main so the probe resolves this revision and reaches the weight check - # (without it the answer would be None -- a cache miss -- for a different reason). + # refs/main so the probe reaches the weight check (without it the answer is None). (repo / "refs").mkdir(parents = True) (repo / "refs" / "main").write_text("aaa") _fake_hf_cache(monkeypatch, hf_root) @@ -290,18 +277,15 @@ def test_marker_onnx_only_snapshot_is_not_loadable(tmp_path, monkeypatch): def test_marker_rejects_weights_without_tokenizer(tmp_path, monkeypatch): - # A snapshot with modules.json + config + a complete weight set but NO tokenizer - # asset is not loadable: SentenceTransformer's Transformer module also builds an - # AutoTokenizer, which fails offline when tokenizer.json / tokenizer_config.json / - # a vocab is absent. Accepting it would pass validation and then fail at load. + # Marker + config + weights but NO tokenizer asset is not loadable: the Transformer + # module also builds an AutoTokenizer, which fails offline without a tokenizer asset. _cache_repo_with_files(tmp_path, monkeypatch, "model.safetensors", tokenizer = False) assert mc._embedding_marker_in_hf_cache("org/model") is False @pytest.mark.parametrize("tok_file", ["tokenizer_config.json", "vocab.txt", "spiece.model"]) def test_marker_accepts_alternate_tokenizer_assets(tmp_path, monkeypatch, tok_file): - # The tokenizer check is a permissive union: any one recognized asset (fast - # tokenizer json, a WordPiece vocab, a SentencePiece model, ...) is enough, so a + # The tokenizer check is a permissive union: any one recognized asset is enough, so a # valid non-tokenizer.json layout is not wrongly rejected. _cache_repo_with_files(tmp_path, monkeypatch, "model.safetensors", tok_file, tokenizer = False) assert mc._embedding_marker_in_hf_cache("org/model") is True @@ -314,8 +298,8 @@ def _cache_repo_with_files( commit = "aaa", tokenizer = True, ): - """A cache repo whose active snapshot holds modules.json + config + a tokenizer + - *files*. ``tokenizer=False`` omits the tokenizer asset (a partial download).""" + """Cache repo whose active snapshot holds modules.json + config + a tokenizer + *files*. + ``tokenizer=False`` omits the tokenizer asset.""" hf_root = tmp_path / "hf" repo = hf_root / "models--org--model" snap = repo / "snapshots" / commit @@ -335,11 +319,8 @@ def _cache_repo_with_files( def test_marker_rejects_non_base_weight_bins(tmp_path, monkeypatch): - # A partial cache carrying modules.json + config plus only a commonly published - # NON-weight .bin (training_args.bin) is not a loadable base model: the Torch - # backend needs model.safetensors / pytorch_model.bin, so accepting it would - # pass validation and then fail at first indexing. Adapter-only artifacts are - # rejected for the same reason. + # A NON-weight .bin (training_args.bin) or an adapter-only artifact is not a loadable + # base model: the Torch backend needs model.safetensors / pytorch_model.bin. _cache_repo_with_files(tmp_path, monkeypatch, "training_args.bin") assert mc._embedding_marker_in_hf_cache("org/model") is False _cache_repo_with_files(tmp_path / "b", monkeypatch, "adapter_model.safetensors") @@ -363,11 +344,8 @@ def test_marker_rejects_non_base_weight_bins(tmp_path, monkeypatch): ], ) def test_marker_rejects_incomplete_shard_set(tmp_path, monkeypatch, weights): - # A partially downloaded sharded model must NOT validate: SentenceTransformer's - # Torch backend needs every shard AND the index map at load time, so accepting an - # incomplete set (or a complete set missing its index) would pass offline - # validation and then fail at first indexing -- the same validate-then-fail this - # helper exists to prevent. + # A partially downloaded sharded model must NOT validate: the Torch backend needs every + # shard AND the index map, so an incomplete set (or a set missing its index) fails. _cache_repo_with_files(tmp_path, monkeypatch, *weights) assert mc._embedding_marker_in_hf_cache("org/model") is False @@ -385,9 +363,8 @@ def test_marker_rejects_incomplete_shard_set(tmp_path, monkeypatch, weights): ], ) def test_marker_accepts_recognized_torch_weights(tmp_path, monkeypatch, weights): - # The filename recognizer must not over-reject real base-model weights: single - # pytorch_model.bin, a complete sharded set with its index map, and weights that - # live inside a module directory all count as loadable. + # The recognizer must not over-reject real weights: single pytorch_model.bin, a complete + # sharded set with its index, and weights inside a module dir all count as loadable. _cache_repo_with_files(tmp_path, monkeypatch, *weights) assert mc._embedding_marker_in_hf_cache("org/model") is True @@ -399,9 +376,8 @@ def _case_sensitive_fs(tmp_path) -> bool: def test_st_casing_resolves_against_st_home(tmp_path, monkeypatch): - # resolve_cached_repo_id_case scans only the Hub cache, so with ST_HOME set it - # would persist the requested lower-case id while the exact-case offline load - # looks for models--BAAI--bge-m3 in ST_HOME and misses it. + # resolve_cached_repo_id_case scans only the Hub cache, so with ST_HOME set it would + # persist the lower-case id while the exact-case offline load misses it in ST_HOME. if not _case_sensitive_fs(tmp_path): pytest.skip("casing only diverges on a case-sensitive filesystem") hf_root, st_root = tmp_path / "hf", tmp_path / "st" @@ -433,8 +409,7 @@ def test_st_casing_noop_when_uncached(tmp_path, monkeypatch): def test_marker_never_raises_when_cache_mutates(monkeypatch): - # A snapshot vanishing mid-iteration (concurrent cached-model deletion) - # must read as not-cached, not propagate a 500 out of the routes. + # A snapshot vanishing mid-iteration must read as not-cached, not raise a 500. def _exploding_iter(repo): raise FileNotFoundError("snapshot removed underneath") @@ -443,8 +418,7 @@ def _exploding_iter(repo): def test_is_embedding_model_survives_cache_race_online(monkeypatch): - # With the cache probe failing, the online path must still resolve via the - # Hub instead of erroring out. + # With the cache probe failing, the online path must still resolve via the Hub. def _exploding_iter(repo): raise FileNotFoundError("snapshot removed underneath") @@ -462,9 +436,8 @@ def _exploding_iter(repo): def test_offline_cached_st_detected_via_marker_no_network(tmp_path, monkeypatch): - # Offline: a downloaded sentence-transformers repo is classified from its - # modules.json marker with no model_info() network call that would hang on - # DNS retries (#6817). + # Offline: a downloaded sentence-transformers repo is classified from its modules.json + # marker with no model_info() network call that would hang on DNS retries (#6817). _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa", repo_id = "unsloth/bge-small-en-v1.5") monkeypatch.setenv("HF_HUB_OFFLINE", "1") _fake_hf_model_info(monkeypatch, _no_network) @@ -472,10 +445,8 @@ def test_offline_cached_st_detected_via_marker_no_network(tmp_path, monkeypatch) def test_online_defers_to_hub_over_stale_marker(tmp_path, monkeypatch): - # Online: the Hub is authoritative for the current revision. Even with a - # cached modules.json (the repo WAS an embedder), a Hub lookup that no longer - # reports embedding signals wins -- the stale local marker must not - # short-circuit model_info(). + # Online: the Hub is authoritative. Even with a cached modules.json, a Hub lookup that + # no longer reports embedding signals wins -- the stale marker must not short-circuit it. _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa") calls = [] @@ -489,10 +460,8 @@ def _info(model_name, token = None): def test_online_permanent_hub_error_ignores_stale_marker(tmp_path, monkeypatch): - # A permanent Hub error (deleted / gated / typo'd repo) is authoritative: - # even with a cached modules.json, validation must NOT pass on the stale - # marker -- return False so the settings route surfaces its 409, and the - # persisted model can't fail later when the loader refreshes from the Hub. + # A permanent Hub error (deleted / gated / typo'd repo) is authoritative: even with a + # cached modules.json, return False so the settings route surfaces its 409. _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa") class RepositoryNotFoundError(Exception): @@ -506,9 +475,8 @@ def _info(model_name, token = None): def test_online_hub_failure_falls_back_to_marker_uncached(tmp_path, monkeypatch): - # A transient model_info() failure falls back to the local marker WITHOUT - # caching: the degraded result must not become sticky, so a later successful - # Hub lookup can still override it. + # A transient model_info() failure falls back to the local marker WITHOUT caching, so a + # later successful Hub lookup can still override the degraded result. _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa", repo_id = "org/emb") _fake_hf_model_info(monkeypatch, _no_network) # raises -> Hub "unreachable" assert mc.is_embedding_model("org/emb") is True # marker fallback @@ -516,10 +484,9 @@ def test_online_hub_failure_falls_back_to_marker_uncached(tmp_path, monkeypatch) def test_online_negative_does_not_block_later_offline_download(tmp_path, monkeypatch): - # An online Hub lookup authoritatively reports non-embedding and is memoized. - # The repo is then downloaded WITH modules.json and the session goes offline; - # the offline path re-probes the marker (never consulting the online memo), - # so the freshly downloaded embedder is detected instead of the stale False. + # An online lookup reports non-embedding and is memoized. The repo is then downloaded + # with modules.json and the session goes offline; the offline path re-probes the marker + # (never the memo), so the fresh embedder is detected instead of the stale False. _no_cache(monkeypatch) def _info(model_name, token = None): @@ -536,13 +503,10 @@ def _info(model_name, token = None): def test_offline_retains_online_confirmed_positive(tmp_path, monkeypatch): - # A tag-only embedder (feature-extraction, no modules.json) is confirmed online - # and cached True, and its snapshot IS materialized locally. - # _hf_offline_if_dns_dead() then flips the process to offline mid-load; the - # offline path must RETAIN that positive because the files are present, not - # downgrade a model already verified this session. The marker alone reads False - # here (present but no modules.json), so retention rests on the memo plus a - # present active snapshot, not on the marker. + # A tag-only embedder (no modules.json) confirmed online and cached True, with its + # snapshot materialized. Flipped offline mid-load, the offline path must RETAIN the + # positive: the marker reads False, so retention rests on the memo plus a present + # snapshot, not the marker. _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") def _info(model_name, token = None): @@ -557,11 +521,9 @@ def _info(model_name, token = None): def test_offline_metadata_only_positive_not_trusted_without_cache(monkeypatch): - # A positive from an online model_info() call (e.g. /check-embedding) proves - # only that the repo is tagged an embedder, not that any files were downloaded. - # With nothing materialized in the cache, the offline path must NOT trust that - # memo: saving the model would then fail the local_files_only RAG load at first - # indexing. So it returns False despite the cached True. + # An online positive proves only that the repo is tagged an embedder, not that files + # were downloaded. With nothing materialized, the offline path must NOT trust the memo + # (the local_files_only load would fail), so it returns False despite the cached True. _no_cache(monkeypatch) def _info(model_name, token = None): @@ -578,12 +540,9 @@ def _info(model_name, token = None): def test_offline_detects_persisted_tag_only_embedder_after_restart(tmp_path, monkeypatch): - # A tag-only embedder (feature-extraction, no modules.json) is confirmed online - # in one session, which durably records the verdict, and its snapshot is - # materialized on disk. After a RESTART (the session memo is gone) the process - # comes up offline: the marker alone reads False (present but no modules.json), - # so recognition rests on the persisted allowlist plus the present active - # snapshot -- the exact case the memo cannot cover across a restart. + # A tag-only embedder confirmed online durably records the verdict, snapshot + # materialized. After a RESTART (memo gone) the process comes up offline: the marker + # reads False, so recognition rests on the persisted allowlist plus the present snapshot. _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") def _info(model_name, token = None): @@ -600,11 +559,9 @@ def _info(model_name, token = None): def test_offline_persisted_verdict_not_trusted_when_uncached(tmp_path, monkeypatch): - # The persisted allowlist records only that the repo was tagged an embedder, - # not that its files are on disk. After a restart with NOTHING materialized in - # the cache, the offline path must NOT trust the persisted verdict: saving the - # model would then fail the local_files_only RAG load at first indexing. It is - # gated on the active snapshot being present, so an uncached repo returns False. + # The persisted allowlist records only the tag, not on-disk files. After a restart with + # NOTHING materialized, the offline path must NOT trust it (gated on the active snapshot + # being present), so an uncached repo returns False. _no_cache(monkeypatch) def _info(model_name, token = None): @@ -622,12 +579,10 @@ def _info(model_name, token = None): def test_offline_persisted_verdict_not_trusted_when_snapshot_partial(tmp_path, monkeypatch): - # A persisted verdict is trusted only when the active snapshot carries a COMPLETE - # weight set, not merely that it is materialized. Here the snapshot exists with a - # config but no weights (an interrupted download): _embedding_marker_in_hf_cache - # reads False (materialized, not None), so a bare "materialized" gate would wrongly - # return True and the local_files_only load would then fail. The weight gate must - # reject it. + # A persisted verdict is trusted only when the active snapshot carries a COMPLETE weight + # set, not merely that it is materialized. Here config but no weights (interrupted + # download): a bare "materialized" gate would wrongly return True, so the weight gate + # must reject it. cache_root = tmp_path / "cache" snap = cache_root / "models--org--partial" / "snapshots" / "aaa" snap.mkdir(parents = True) @@ -652,10 +607,9 @@ def _info(model_name, token = None): def test_offline_persisted_verdict_matches_across_casing(tmp_path, monkeypatch): - # The verdict is recorded under the online-request spelling but the settings route - # saves the cache-resolved spelling, so an exact-string lookup would miss it. A - # tag-only embedder verified online as baai/model, then looked up offline as the - # saved BAAI/model, must still be recognized: the allowlist is case-folded. + # The verdict is recorded under the request spelling but the settings route saves the + # cache-resolved one, so an exact-string lookup would miss it. Verified as baai/model, + # looked up as the saved BAAI/model, must still match: the allowlist is case-folded. _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "BAAI/model") def _info(model_name, token = None): @@ -672,10 +626,8 @@ def _info(model_name, token = None): def test_persist_embedder_concurrent_writes_keep_every_verdict(tmp_path, monkeypatch): - # Concurrent online confirmations must not drop each other's entry: the - # read-modify-write is serialized and each writer uses a per-thread temp file, so - # every model persisted from parallel threads survives (a lost verdict would go - # unrecognized offline after a restart). + # Concurrent confirmations must not drop each other's entry: serialized read-modify-write + # + per-thread temp files, so every model persisted from parallel threads survives. import threading names = [f"org/emb-{i}" for i in range(24)] @@ -696,9 +648,8 @@ def _writer(name): def test_persist_embedder_is_best_effort_when_home_unwritable(tmp_path, monkeypatch): - # Persistence is an optimization, never a correctness requirement: if the - # Studio home cannot be created/written, the online path must still return its - # verdict rather than raise. Point the home at a path blocked by a file. + # Persistence is an optimization: if the Studio home cannot be written, the online path + # must still return its verdict rather than raise. Point the home at a path blocked by a file. blocker = tmp_path / "blocker" blocker.write_text("not a dir") monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(blocker / "studio")) # parent is a file @@ -727,8 +678,8 @@ def test_offline_not_cached_returns_false_without_network(monkeypatch): def test_online_uncached_still_uses_network(monkeypatch): - # Not offline, not cached: the network model_info path must still run so an - # embedding model that lacks modules.json (feature-extraction tag) is caught. + # Not offline, not cached: the model_info path must still run so a feature-extraction + # embedder lacking modules.json is caught. _no_cache(monkeypatch) calls = [] @@ -742,9 +693,8 @@ def _info(model_name, token = None): def test_offline_negative_is_not_cached_then_online_detects(monkeypatch): - # A tag-only embedder is not identifiable from modules.json. Offline returns - # False WITHOUT caching, so once the env var clears the online model_info - # lookup still runs and detects it -- the negative must not be sticky. + # A tag-only embedder is not identifiable from modules.json. Offline returns False + # WITHOUT caching, so once the env clears the online lookup still detects it. _no_cache(monkeypatch) calls = [] diff --git a/studio/backend/tests/test_rag_embeddings.py b/studio/backend/tests/test_rag_embeddings.py index 58c7de25795..11aa089a2f5 100644 --- a/studio/backend/tests/test_rag_embeddings.py +++ b/studio/backend/tests/test_rag_embeddings.py @@ -224,11 +224,9 @@ def test_st_encode_failure_without_llama_binary_reraises(monkeypatch): def test_get_resolves_default_casing_before_loading(monkeypatch): # _get() must resolve a repo id to the exact cache casing before constructing - # SentenceTransformer. A configured default whose spelling differs only by case - # from the cache dir is deliberately NOT persist-normalized by /settings (that - # would turn the default into an override), so without resolving here the offline - # local_files_only load would miss the case-sensitive cache dir and fail. Verify - # both the loader and the security gate receive the resolved name. + # SentenceTransformer: a default differing only by case is NOT persist-normalized by + # /settings, so without resolving here the offline load would miss the cache dir. + # Verify both the loader and the security gate receive the resolved name. import sys import types diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 80ad3ab2814..1dbb2ac5873 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1653,15 +1653,10 @@ def _hf_hub_cache_root() -> list[Path]: def _st_cache_roots() -> list[Path]: - """The cache root a SentenceTransformer load will ACTUALLY search. - - ``_get()`` constructs SentenceTransformer without ``cache_folder``, so it uses - SENTENCE_TRANSFORMERS_HOME when that is set and the Hub cache otherwise -- - one or the other, never both. Probing the union would let offline validation - pass on a repo cached only in the Hub cache while the loader then searched - ST_HOME and failed. Kept separate from the Hub-cache probe because the GGUF - path downloads with ``hf_hub_download`` and no ``cache_dir``, so it really - does use the Hub cache and must not be told about ST_HOME. + """Cache root a SentenceTransformer load actually searches: SENTENCE_TRANSFORMERS_HOME + if set, else the Hub cache -- one or the other, never both. Probing the union would + pass offline validation on a repo cached only in the other root. Kept separate from + the Hub-cache probe, which the GGUF path needs (it downloads via ``hf_hub_download``). """ st_home = (os.environ.get("SENTENCE_TRANSFORMERS_HOME") or "").strip() if st_home: @@ -1689,16 +1684,10 @@ def _iter_cache_snapshots_in(repo_id: str, roots: list[Path]): def resolve_st_cached_repo_id_case(repo_id: str) -> str: - """*repo_id* in the casing of its dir in the cache the ST loader will search. - - ``resolve_cached_repo_id_case`` scans only the Hub cache, so with - SENTENCE_TRANSFORMERS_HOME set it would leave the requested spelling - unchanged -- and the offline ``local_files_only`` load, which is exact-case on - a case-sensitive filesystem, would then miss the differently cased dir that - detection had just accepted. Normalizes against the same roots detection uses. - Returns *repo_id* unchanged for a local path, a non-repo string, or when - nothing case-matching is cached; prefers an exact match before any variant and - tie-breaks variants deterministically. + """*repo_id* recased to match its dir in the cache the ST loader searches, so the + offline exact-case ``local_files_only`` load finds it. Unchanged for a local path, a + non-repo string, or when nothing case-matching is cached; prefers an exact match, then + a deterministic pick among case variants. """ if is_local_path(repo_id) or "/" not in repo_id: return repo_id @@ -1723,15 +1712,10 @@ def resolve_st_cached_repo_id_case(repo_id: str) -> str: def _st_cache_repo_dir(repo_id: str) -> Optional[Path]: - """The ONE cache repo dir the ST loader will open for *repo_id*, or None. - - Uses the same selection rule as :func:`resolve_st_cached_repo_id_case` -- - exact case first, then a deterministic pick among case variants -- because - that is the spelling the settings route persists and therefore the directory - the loader opens. Scoping to a single directory matters when duplicate case - variants exist: judging the repo by whichever variant happens to hold the - newest snapshot could reject a complete model in the directory that will - actually be loaded (or accept one from a directory that will not be). + """The ONE cache repo dir the ST loader opens for *repo_id*, or None. Same selection + as :func:`resolve_st_cached_repo_id_case` (exact case, then a deterministic pick). + Scoping to one dir matters with duplicate case variants: judging by whichever holds + the newest snapshot could reject the model in the dir actually loaded. """ prefix = "models--" expected = f"{prefix}{repo_id.replace('/', '--')}" @@ -1782,15 +1766,10 @@ def _iter_snapshots_of(repo_dirs: list[Path]): def _iter_hf_cache_snapshots(repo_id: str): """Yield HUB cache snapshot dirs for *repo_id*, newest first. - Deliberately the Hub cache ONLY. Its callers (the GGUF detectors) later - download with ``hf_hub_download`` and no ``cache_dir``, so that is the cache - their load actually uses; adding SENTENCE_TRANSFORMERS_HOME here would let - detection pick a file the GGUF loader cannot find. The Sentence-Transformers - probe is :func:`_st_cache_repo_dir`. - - Empty if the root does not exist, the repo isn't cached, or it has no - snapshots. Repo name match is case-insensitive to handle casing drift between - download time and lookup. + Hub cache ONLY: its GGUF-detector callers download via ``hf_hub_download`` with no + ``cache_dir``, so that is the cache their load uses; the ST probe is + :func:`_st_cache_repo_dir`. Empty if the root or repo isn't cached. Match is + case-insensitive to handle casing drift between download and lookup. """ yield from _iter_cache_snapshots_in(repo_id, _hf_hub_cache_root()) @@ -2107,41 +2086,33 @@ def download_gguf_file( def _persisted_embedders_path() -> Path: - """On-disk store of repo ids confirmed embedders via an online Hub lookup. - - Lives under the resolved Studio home so it survives a restart: the session - memo (:data:`_embedding_detection_cache`) is lost on exit, so without this a - tag-only feature-extraction embedder -- one whose cached snapshot carries no - ``modules.json`` and so cannot be recognized from the cache alone -- would be - misclassified as non-embedding the first offline call after a restart. + """On-disk store of repo ids confirmed embedders online, under the Studio home so it + survives a restart. Without it a tag-only feature-extraction embedder (cached snapshot + has no ``modules.json``, so unrecognizable from the cache alone) would be misclassified + non-embedding the first offline call after the session memo is lost. """ return _studio_root() / "embedding_verdicts.json" -# Serializes the read-modify-write in _persist_embedder so two threads confirming -# different embedders at once cannot each overwrite the file with a one-entry update -# (a lost verdict) or collide on the temp file. Cross-process writers are inherently -# best-effort -- os.replace stays atomic, and a dropped verdict is only an -# optimization miss a later online re-confirmation heals. +# Serializes the read-modify-write in _persist_embedder so concurrent confirmations +# cannot lose a verdict or collide on the temp file. Cross-process writes are +# best-effort: os.replace stays atomic, and a dropped verdict is only an optimization +# miss a later online re-confirmation heals. _persist_lock = threading.Lock() def _verdict_key(model_name: str) -> str: - """Case-folded key for the embedder allowlist. The verdict is a property of the - repo, not its spelling: model_info() is queried under the requested casing while - the settings route saves the cache-resolved casing, so keying by the exact string - would miss the persisted positive after a restart (``baai/model`` recorded, - ``BAAI/model`` looked up). Repo ids are case-insensitive, so fold both sides.""" + """Case-folded allowlist key: the verdict is a property of the repo, not its spelling. + model_info() is queried under the requested casing but the settings route saves the + cache-resolved casing, so an exact-string key would miss the persisted positive after a + restart. Repo ids are case-insensitive, so fold both sides.""" return model_name.casefold() def _load_persisted_embedders() -> set: """Case-folded repo ids recorded as embedders in a prior session; empty set on any - error. - - Best-effort and never raises: a missing, empty, corrupt, or non-list file - (and any read/decode error) reads as "nothing recorded" so callers fall back - to their normal cache-marker logic. Folding at load time also normalizes any + error (missing, empty, corrupt, or non-list file all read as "nothing recorded", so + callers fall back to their cache-marker logic). Never raises. Folding also normalizes mixed-case entries an older build may have written. """ try: @@ -2157,14 +2128,11 @@ def _load_persisted_embedders() -> set: def _persist_embedder(model_name: str) -> None: """Record *model_name* as an online-confirmed embedder, best-effort. - Only positive Hub verdicts are written (never a negative or a token), case-folded - so the lookup matches regardless of casing: the file is a durable allowlist the - offline branch consults, gated there on the active revision actually being - materialized on disk, so a stale entry for a since-deleted cache cannot resurrect - a False. The read-modify-write is serialized under ``_persist_lock`` and writes - through a per-thread temp file, so concurrent confirmations cannot drop each - other's entry or collide. Any failure (unwritable home, race) is swallowed -- - persistence is an optimization, not a correctness requirement for the online path. + Only positive Hub verdicts are written, case-folded. The offline branch gates this + allowlist on the active revision being materialized on disk, so a stale entry cannot + resurrect a False. Serialized under ``_persist_lock`` through a per-thread temp file so + concurrent writers cannot drop an entry or collide. Any failure is swallowed -- + persistence is an optimization, not a correctness requirement. """ try: key = _verdict_key(model_name) @@ -2184,33 +2152,23 @@ def _persist_embedder(model_name: str) -> None: def _known_embedder(model_name: str, cache_key: tuple) -> bool: - """True when *model_name* was confirmed an embedder online, this session or a - prior one -- the session memo (authoritative, token-scoped) OR the persisted - cross-restart allowlist (matched case-insensitively). Callers still gate this on a - materialized snapshot.""" + """True when *model_name* was confirmed an embedder online -- the session memo or the + persisted cross-restart allowlist. Callers still gate this on a materialized snapshot.""" if _embedding_detection_cache.get(cache_key) is True: return True return _verdict_key(model_name) in _load_persisted_embedders() -# The base-model weight files the RAG loader's default backend consumes. _get() -# constructs SentenceTransformer without backend="onnx"/"openvino", so it loads -# through Torch: a single model.safetensors / pytorch_model.bin or their sharded -# model-00001-of-000NN forms, at the snapshot root or inside a module dir -# (0_Transformer/, 2_Dense/). Matched by NAME, not suffix: a bare ".safetensors"/ -# ".bin" match would accept the commonly published training_args.bin / optimizer.bin -# or an adapter-only artifact -- files that carry a weight suffix but are not a -# loadable base model -- and let a partial cache pass offline validation and then -# fail on the first RAG load, the exact validate-then-fail this helper prevents. +# Base-model weight files the default (Torch) SentenceTransformer backend consumes: +# a single model.safetensors / pytorch_model.bin or their sharded -00001-of-000NN forms. +# Matched by NAME, not suffix, so a training_args.bin / optimizer.bin or an adapter-only +# artifact cannot let a partial cache pass offline validation and then fail on first load. _ST_WEIGHT_FILE_RE = re.compile(r"^(model|pytorch_model)(-\d+-of-\d+)?\.(safetensors|bin)$") _ST_SHARD_RE = re.compile(r"^(model|pytorch_model)-(\d+)-of-(\d+)\.(safetensors|bin)$") -# A SentenceTransformer Transformer module loads its weights AND an AutoTokenizer, -# so a snapshot with a complete weight set but no tokenizer asset still fails the -# local_files_only load. Any ONE of these is enough to load some tokenizer -- kept a -# permissive union (fast tokenizer, config, or a WordPiece/BPE/SentencePiece vocab) -# so an unusual-but-valid layout is not rejected, only a genuinely tokenizer-less -# partial download. +# A Transformer module also loads an AutoTokenizer, so weights without a tokenizer asset +# still fail the local_files_only load. Any ONE of these loads some tokenizer -- a +# permissive union so only a genuinely tokenizer-less partial download is rejected. _ST_TOKENIZER_FILES = frozenset( { "tokenizer.json", @@ -2226,19 +2184,14 @@ def _known_embedder(model_name: str, cache_key: tuple) -> bool: def _dir_has_complete_torch_weights(names: set) -> bool: - """True when the filenames in ONE directory form a COMPLETE Torch base-model - weight set the default SentenceTransformer backend can load: + """True when the filenames in ONE dir form a COMPLETE Torch weight set: * a single ``model.safetensors`` / ``pytorch_model.bin``; or - * a full shard set -- a shard names its own total (``-00001-of-00002``), so - every index ``1..total`` for that (stem, ext, total) must be present AND its - ``model.safetensors.index.json`` / ``pytorch_model.bin.index.json`` map, - through which transformers discovers and wires the shards, must be present. - - A lone shard such as ``model-00001-of-00002.safetensors``, a full set missing - its index map, is NOT complete: the loader needs every shard and the index at - indexing time, so accepting a partial sharded download would validate and then - fail under local_files_only. + * a full shard set -- every index ``1..total`` (a shard names its own total, + ``-00001-of-00002``) plus its ``...index.json`` map. + + A lone shard or a shard set missing its index map is NOT complete: accepting a partial + sharded download would validate and then fail under local_files_only. """ for stem in ("model", "pytorch_model"): for ext in ("safetensors", "bin"): @@ -2257,17 +2210,13 @@ def _dir_has_complete_torch_weights(names: set) -> bool: def _snapshot_has_complete_weights(snap: Path) -> bool: - """True when *snap* is materialized with a config, a tokenizer asset, and a - COMPLETE Torch weight set the default SentenceTransformer backend can load -- - the ``modules.json`` sentence-transformers marker aside. - - A tag-only feature-extraction embedder confirmed online has exactly this (config, - tokenizer, weights, no ``modules.json``) and SentenceTransformer's auto-model - fallback loads it, so an online-confirmed positive may be trusted for it offline. - A partial download -- config present but weights missing / an incomplete shard - set, or weights present but the tokenizer AutoTokenizer needs absent -- does NOT - satisfy this and must fail validation here rather than at first indexing, the - validate-then-fail this guards against. + """True when *snap* has a config, a tokenizer asset, and a COMPLETE Torch weight set + (the ``modules.json`` marker aside). + + A tag-only feature-extraction embedder confirmed online has exactly this and ST's + auto-model fallback loads it, so an online-confirmed positive may be trusted offline. + A partial download (weights missing, incomplete shard set, or no tokenizer) fails here + rather than at first indexing. """ try: if not any( @@ -2296,13 +2245,9 @@ def _snapshot_has_complete_weights(snap: Path) -> bool: def _snapshot_is_loadable_st_model(snap: Path) -> bool: """True when *snap* holds a sentence-transformers model that can actually load. - ``modules.json`` alone is not enough: the online security preflight downloads - exactly that one file via ``hf_hub_download``, and a partial download can - leave it behind too, so a snapshot can carry the marker while the weights and - config SentenceTransformer needs are absent. Accepting that offline would - pass validation and then fail on the first RAG load. Requires the marker plus - a config and a COMPLETE recognized Torch base-model weight set - (``_snapshot_has_complete_weights``). + ``modules.json`` alone is not enough: the security preflight downloads exactly that + file, and a partial download can leave it behind while weights/config are absent. + Requires the marker plus ``_snapshot_has_complete_weights``. """ try: if not (snap / "modules.json").is_file(): @@ -2313,16 +2258,11 @@ def _snapshot_is_loadable_st_model(snap: Path) -> bool: def _active_snapshot_dir(repo_id: str) -> Optional[Path]: - """The materialized snapshot dir the offline ``local_files_only`` load resolves - for *repo_id*, or None. - - Mirrors that resolution: the repo dir :func:`_st_cache_repo_dir` selects (the - casing the settings route persists and the loader opens), its ``refs/main`` - commit (with ``local_files_only`` huggingface_hub resolves the default revision - THROUGH that ref, so a missing / empty / unreadable ref is a cache MISS, never - a reason to scan historical snapshots), and that commit's snapshot dir. None - when any of those is absent -- a cache miss the caller treats as not-cached. - Never raises: a cache mutating underneath (concurrent deletion) reads as None. + """The materialized snapshot dir the offline ``local_files_only`` load resolves for + *repo_id*, or None. Mirrors that resolution: the repo dir :func:`_st_cache_repo_dir` + selects, its ``refs/main`` commit (huggingface_hub resolves the default revision THROUGH + that ref, so a missing/empty/unreadable ref is a cache MISS, not a reason to scan + historical snapshots), and that commit's snapshot dir. Never raises. """ try: repo_dir = _st_cache_repo_dir(repo_id) @@ -2343,26 +2283,12 @@ def _active_snapshot_dir(repo_id: str) -> Optional[Path]: def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: - """Sentence-transformers detection from the local cache, no network call. - - Models exactly what an offline ``local_files_only=True`` load resolves, so a - True here means that load can succeed: - - * one repo dir -- the one :func:`_st_cache_repo_dir` selects, i.e. the casing - the settings route persists and the loader opens; - * its ``refs/main``, because with ``local_files_only`` huggingface_hub - resolves the default revision THROUGH that ref. A snapshot directory alone - is not discoverable, so a missing, empty or unreadable ref is a cache MISS - rather than a reason to go scanning historical snapshots -- accepting one - would pass validation and then fail at first indexing; - * that revision's snapshot, which must be materialized and actually loadable - (marker plus config plus weights, not the bare ``modules.json`` the - security preflight fetches on its own). - - True/False when the active revision is / is not a loadable - sentence-transformers snapshot, None when nothing usable is cached. Never - raises -- a cache mutating underneath (concurrent model deletion) reads as - not-cached so callers keep their normal fallback. + """Sentence-transformers detection from the local cache, no network call. Models what + an offline ``local_files_only`` load resolves (via :func:`_active_snapshot_dir`), so a + True means that load can succeed: the active ``refs/main`` snapshot must be materialized + and actually loadable (marker plus config plus weights, not the bare ``modules.json`` + the security preflight fetches on its own). None when nothing usable is cached; never + raises. """ snapshot = _active_snapshot_dir(repo_id) if snapshot is None: @@ -2386,8 +2312,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: """ cache_key = (model_name, hf_token) - # Local paths: check for sentence-transformer marker (modules.json). This is - # authoritative for an explicit path, so memoize it. + # Local paths: the modules.json marker is authoritative for an explicit path, so memoize. if is_local_path(model_name): if cache_key in _embedding_detection_cache: return _embedding_detection_cache[cache_key] @@ -2397,26 +2322,19 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: return is_emb if _env_offline(): - # Offline: the local HF cache is the only source -- a network call cannot - # succeed and would only hang on a DNS error and get retried (#6817). Re-probe - # the cache every call without consulting or populating the memo: a cached - # negative must not stick (a model downloaded later, or a tag-only embedder, - # could not be confirmed here), so a miss is not durable. + # Offline: the local HF cache is the only source -- a network call would only hang + # on DNS and get retried (#6817). Re-probe every call without consulting/populating + # the memo, so a negative never sticks (a model downloaded later must be seen). snapshot = _active_snapshot_dir(model_name) if snapshot is not None and _snapshot_is_loadable_st_model(snapshot): - return True # a self-describing sentence-transformers snapshot (modules.json) - # Retain a positive confirmed online -- this session (memo) or a prior one - # (persisted allowlist) -- ONLY when the active snapshot is materialized WITH - # a complete, loadable weight set. A recorded True proves model_info() tagged - # the repo an embedder, not that its files are on disk (an online - # /check-embedding call records the verdict without downloading), so trusting - # it for an uncached or partially downloaded repo would save a model the - # local_files_only load then fails on. The weight gate (not the bare marker) - # is what still covers a downloaded tag-only feature-extraction embedder -- - # weights present, no modules.json, which SentenceTransformer's auto-model - # fallback loads -- while rejecting a config-only or half-sharded snapshot, - # and it is the case that must survive a restart, not just the - # _hf_offline_if_dns_dead() mid-load flip. + return True # self-describing sentence-transformers snapshot (modules.json) + # Trust a positive confirmed online (memo or persisted allowlist) ONLY when the + # active snapshot is materialized with a complete, loadable weight set. A recorded + # True proves model_info() tagged the repo an embedder, not that its files are on + # disk, so gating on the weights avoids saving a model the local_files_only load + # then fails on. The weight gate (not the bare marker) still covers a downloaded + # tag-only embedder (weights, no modules.json, loaded by ST's auto-model fallback) + # while rejecting a config-only or half-sharded snapshot. if ( snapshot is not None and _known_embedder(model_name, cache_key) @@ -2425,12 +2343,10 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: return True return False - # Online: the Hub is authoritative for the current remote revision. The local - # cache marker reflects only the last-downloaded revision, which may lag the - # Hub (a repo can add or drop modules.json), so model_info() decides and the - # marker is used solely as a fallback when the Hub is unreachable. Only - # Hub-derived results are memoized, so a transient failure never poisons the - # cache and a later fresh download or successful lookup can override it. + # Online: the Hub is authoritative for the current revision; the local marker can lag + # it, so model_info() decides and the marker is only a fallback when the Hub is + # unreachable. Only Hub-derived results are memoized, so a transient failure never + # poisons the cache. if cache_key in _embedding_detection_cache: return _embedding_detection_cache[cache_key] try: @@ -2448,9 +2364,8 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: _embedding_detection_cache[cache_key] = is_emb if is_emb: - # Durably record the positive so a later offline session (after a - # restart, when the memo is gone) can still recognize a downloaded - # tag-only embedder, gated there on its snapshot being materialized. + # Durably record the positive so a later offline session (memo gone) can still + # recognize a downloaded tag-only embedder, gated on its snapshot. _persist_embedder(model_name) logger.info( f"Model {model_name} detected as embedding model: " @@ -2461,12 +2376,10 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: return is_emb except Exception as e: - # A permanent Hub error (deleted / gated / bad revision / typo matching - # stale cache casing) is authoritative: the repo is not a usable - # embedding model, so return False and let the settings route surface its - # documented 409. Falling back to a cached modules.json here would wrongly - # pass validation for a repo the loader can no longer fetch. Matches the - # permanent-vs-transient split the nearby GGUF/vision detectors use. + # A permanent Hub error (deleted / gated / bad revision / typo) is authoritative: + # the repo is not a usable embedding model, so return False and let the settings + # route surface its 409. Falling back to a cached modules.json would wrongly pass a + # repo the loader can no longer fetch. Matches the nearby GGUF/vision detectors. if type(e).__name__ in ( "RepositoryNotFoundError", "GatedRepoError", @@ -2475,8 +2388,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: ): logger.warning(f"Could not determine if {model_name} is embedding model: {e}") return False - # Transient / connectivity / 5xx failure: fall back to the local marker, - # uncached -- a degraded signal a later successful Hub call can override. + # Transient / 5xx failure: fall back to the local marker, uncached. marker = _embedding_marker_in_hf_cache(model_name) if marker is True: logger.info( diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 9c01c793e4e..3b34a809cbe 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -249,13 +249,11 @@ def _fetch_security_status( """ from huggingface_hub import model_info as hf_model_info - # Only when the CALLER guarantees a local-only load. This is deliberately not - # keyed off hf_env_offline() here: this gate is shared by every loader - # (training, MLX/inference, export, ...), and most of them do not pass - # local_files_only, so an offline-looking session can still fetch and - # deserialize an unscanned model through those paths. Skipping the scan for - # them would disable the malware gate outright, so the bypass has to be - # opted into by the callers that actually hold the invariant. + # Only when the CALLER guarantees a local-only load -- deliberately not keyed off + # hf_env_offline(): this gate is shared by every loader (training, MLX/inference, + # export, ...) and most do not pass local_files_only, so an offline-looking session + # can still fetch an unscanned model through those paths. The bypass must be opted + # into by the callers that actually hold the invariant. if local_only_load: logger.debug( "HF security scan skipped for '%s': caller loads local-only; failing open.", @@ -303,15 +301,11 @@ def evaluate_file_security( 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`` lets a caller skip the Hub round-trip when it GUARANTEES the - load it is gating cannot fetch -- e.g. the RAG embedder, which passes - ``local_files_only`` to SentenceTransformer from the same predicate. Pass it only - with that guarantee in hand: the scan exists to stop a poisoned pickle being - downloaded and deserialized, so claiming local-only while the loader can still - fetch disables the gate. Default False, because most callers here (training, - MLX/inference, export) do not constrain their loaders. Even when set, an - already-cached repo still loads unscanned -- the same fail-open this function - documents for an unavailable scan. + ``local_only_load`` lets a caller skip the Hub round-trip when it GUARANTEES the load + cannot fetch -- e.g. the RAG embedder, which passes ``local_files_only`` to + SentenceTransformer from the same predicate. Pass it only with that guarantee: claiming + local-only while the loader can still fetch disables the gate. Default False, since most + callers (training, MLX/inference, export) do not constrain their loaders. """ # 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/. diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index a029191cd45..f70dd111429 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -173,12 +173,9 @@ def _offline_flag(name: str) -> bool: def hf_env_offline() -> bool: """True when either HF offline env var is truthy (strip+lower, on/true/yes/1). - This is the user's *intent* to work offline, which is broader than what - ``huggingface_hub`` enforces: it honors ``HF_HUB_OFFLINE`` natively but - ignores ``TRANSFORMERS_OFFLINE`` entirely. Setting this alone therefore does - NOT stop a fetch -- callers that need that guarantee must pass - ``local_files_only = hf_env_offline()`` to the loader (as - ``core/rag/embeddings.py`` does), and anything that relaxes a check on the - strength of "we are offline" depends on those callers doing so. + The user's *intent* to work offline, broader than what ``huggingface_hub`` enforces + (it honors ``HF_HUB_OFFLINE`` but ignores ``TRANSFORMERS_OFFLINE``). This alone does NOT + stop a fetch -- callers needing that guarantee must pass ``local_files_only = + hf_env_offline()`` to the loader (as ``core/rag/embeddings.py`` does). """ return _offline_flag("HF_HUB_OFFLINE") or _offline_flag("TRANSFORMERS_OFFLINE") From c2614bb96a668c9147c35b13f2b80b212f9336bc Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 19 Jul 2026 13:02:16 +0000 Subject: [PATCH 27/77] studio: drop redundant comments in the offline embedding-model classifier Second comment-reduction pass over the offline embedding-model cache work: delete comments and trailing notes that restate the adjacent code or an assertion, and trim the remaining docstrings and rationale comments to their load-bearing invariants. Comments and docstrings only; no code changes. --- studio/backend/core/rag/embeddings.py | 46 ++--- studio/backend/routes/settings.py | 24 +-- .../test_embedding_model_security_gate.py | 18 +- studio/backend/tests/test_hf_offline_flags.py | 9 +- .../tests/test_is_embedding_model_offline.py | 193 ++++++++---------- studio/backend/tests/test_rag_embeddings.py | 10 +- studio/backend/utils/models/model_config.py | 191 +++++++---------- .../backend/utils/security/file_security.py | 16 +- studio/backend/utils/utils.py | 9 +- 9 files changed, 209 insertions(+), 307 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 9b8a0264172..1fe3987092c 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -87,11 +87,10 @@ def _st_module_subdirs(name: str, token: str | None, local_only: bool) -> tuple[ security scan: a flagged pickle directly under one must block. Returns () on any failure (no modules.json, offline, malformed) so the guard never bricks the embedder. - ``local_only`` MUST be the value the caller captured for the load, not a fresh env - read: ``_hf_offline_if_dns_dead()`` flips the offline vars and restores them, so - re-reading could force this probe local-only, return () (modules.json not cached), and - leave the scan with no module roots -- a flagged pickle under ``0_Transformer/`` would - then pass while the loader, using the captured predicate, still deserialized it. + ``local_only`` MUST be the value the caller captured for the load, not a fresh env read + (``_hf_offline_if_dns_dead()`` flips the offline vars mid-load): re-reading could force + this probe local-only, return (), and leave the scan with no module roots while the + loader still deserializes a flagged pickle under ``0_Transformer/``. """ try: import json @@ -111,10 +110,8 @@ def _st_module_subdirs(name: str, token: str | None, local_only: bool) -> tuple[ from huggingface_hub.utils import EntryNotFoundError try: - # local_files_only from the CAPTURED predicate (never a fresh env read; - # see the docstring): huggingface_hub honors only HF_HUB_OFFLINE, so an - # offline session would otherwise block on timeouts despite a cached - # modules.json. + # CAPTURED predicate, never a fresh env read (see docstring); hf honors only + # HF_HUB_OFFLINE, so offline would otherwise block on timeouts. local = hf_hub_download( name, "modules.json", @@ -140,11 +137,10 @@ def _guard_model_security(name: str, local_only_load: bool) -> None: /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_load`` MUST be the same value the caller passes to SentenceTransformer as - ``local_files_only`` -- it is what licenses skipping the Hub scan, so it is an argument - rather than re-read here. ``_hf_offline_if_dns_dead()`` mutates the offline vars and - restores them, so two reads can disagree: this could skip the scan while the constructor - then loaded with ``local_files_only=False`` and fetched the unscanned repo. + ``local_only_load`` MUST equal the ``local_files_only`` the caller passes to + SentenceTransformer -- it is what licenses skipping the Hub scan. ``_hf_offline_if_dns_dead()`` + mutates then restores the offline vars, so two reads can disagree and the scan could be + skipped while the constructor fetched the unscanned repo. """ try: from utils.security import evaluate_file_security, security_load_subdirs @@ -189,11 +185,10 @@ def _get(model_name: str | None = None): from utils.utils import hf_env_offline device = _device() - # Resolve to the exact cache casing the offline local_files_only load needs. - # The /settings route persists that spelling for a custom override but leaves - # the configured default verbatim (rewriting it would turn the default into an - # override), so resolve here at load time to cover the default too. A no-op for - # a local path or when nothing case-matching is cached. + # Resolve to the exact cache casing the offline local_files_only load needs. The + # /settings route persists that spelling for a custom override but leaves the + # configured default verbatim, so resolve here to cover the default too. No-op + # for a local path or nothing cached. load_name = name from utils.paths import is_local_path @@ -201,15 +196,14 @@ def _get(model_name: str | None = None): from utils.models import resolve_st_cached_repo_id_case load_name = resolve_st_cached_repo_id_case(load_name) logger.info("loading embedding model %s on %s", load_name, device) - # Read the offline state ONCE for both the security gate and the loader. - # _hf_offline_if_dns_dead() mutates the offline vars and restores them, so - # re-reading for the constructor could yield False after the guard skipped the - # Hub scan on True -- and the load would then fetch the unscanned repo. + # Read the offline state ONCE for both the gate and the loader: + # _hf_offline_if_dns_dead() mutates then restores the offline vars, so re-reading + # for the constructor could yield False after the guard skipped the Hub scan on + # True. local_only = hf_env_offline() _guard_model_security(load_name, local_only) - # Propagate offline intent into the loader: huggingface_hub honors only - # HF_HUB_OFFLINE, so a TRANSFORMERS_OFFLINE-only session would otherwise still - # fetch missing repo files. + # huggingface_hub honors only HF_HUB_OFFLINE, so a TRANSFORMERS_OFFLINE-only + # session would otherwise still fetch missing repo files. _model = SentenceTransformer( load_name, device = device, diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index dbf069fa238..60775fbdb0d 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -413,9 +413,8 @@ def update_embedding_model( scan_st_pickle = ( model != default_embedding_model() and not is_local_gguf and not _llama_backend_active() ) - # Read the offline state ONCE and reuse it: the module probe and the scan must agree, - # and _hf_offline_if_dns_dead() can flip the vars between two reads (a probe forced - # local-only returns no module roots, letting a flagged 0_Transformer/ pickle pass). + # Read the offline state ONCE: the module probe and the scan must agree, and + # _hf_offline_if_dns_dead() can flip the vars between two reads. from utils.utils import hf_env_offline local_only_load = hf_env_offline() @@ -477,20 +476,15 @@ def update_embedding_model( gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error(model, hf_token) if gguf_error: raise HTTPException(status_code = 409, detail = gguf_error) - # Persist the exact cache casing: validation accepts a case-insensitive hit, but the - # offline ST load resolves the cache by exact case. Resolved against the cache the ST - # loader searches (ST_HOME if set, else the Hub cache). A no-op when nothing - # case-matching is cached. Three cases are left alone: + # Persist the exact cache casing: validation accepts a case-insensitive hit but the + # offline ST load resolves the cache by exact case (against ST_HOME if set, else the Hub + # cache). No-op when nothing case-matching is cached. Three cases are left alone: # # * the default -- rewriting it would make set_rag_embedding_model()'s exact-string - # default comparison treat it as a custom override, so later default changes stop - # applying; - # * a local path -- a directory like "org/model" is loaded from disk, and rewriting - # to a cache collision ("Org/model") would stop resolving to it; - # * the llama-server backend -- it does not load through SentenceTransformer; it - # derives a GGUF companion via effective_gguf_repo() and fetches it from the HUB - # cache, so an ST_HOME spelling would pick a repo _hf_gguf_backend_error() never - # validated (BAAI/bge-m3-GGUF vs the checked baai/bge-m3-GGUF). + # default comparison treat it as an override, so later default changes stop applying; + # * a local path -- loaded from disk; a cache-collision recasing would stop resolving to it; + # * the llama-server backend -- loads a GGUF companion from the HUB cache, so an ST_HOME + # spelling could pick a repo _hf_gguf_backend_error() never validated. if ( model != default_embedding_model() and not is_local_path(model) diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index a2651e55704..c9aea881c32 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -266,8 +266,7 @@ def test_clean_repo_saves_under_force(client, monkeypatch): def test_custom_model_saved_in_cache_casing(client, monkeypatch): - # A custom repo is persisted in its local HF cache casing so the offline exact-case - # ST load finds it (baai/bge-m3 -> BAAI/bge-m3). force=True isolates this from the gate. + # Persisted in the cache casing so the offline exact-case ST load finds it. c, saved = client monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) import utils.models as _models @@ -279,8 +278,7 @@ def test_custom_model_saved_in_cache_casing(client, monkeypatch): def test_local_path_model_is_not_casing_normalized(client, monkeypatch, tmp_path): - # A local directory is loaded from disk; casing normalization must be skipped or it - # would resolve to a cache collision instead of the directory. + # A local directory is loaded from disk; recasing would resolve to a cache collision. c, saved = client monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) import utils.models as _models @@ -297,9 +295,8 @@ def _must_not_run(m): def test_llama_backend_model_is_not_casing_normalized(monkeypatch): - # On the llama-server backend nothing loads through SentenceTransformer: the embedder - # derives a GGUF companion from this spelling and fetches it from the HUB cache, so an - # ST_HOME rewrite would point at a GGUF repo _hf_gguf_backend_error() never validated. + # The llama backend fetches a GGUF companion from the HUB cache, so an ST_HOME recasing + # could point at a GGUF repo _hf_gguf_backend_error() never validated. saved: dict = {} monkeypatch.setattr(settings, "default_embedding_model", lambda: "unsloth/default-embed") monkeypatch.setattr(settings, "validate_embedding_model", lambda v: v) @@ -322,12 +319,11 @@ def _must_not_run(m): c = TestClient(app, raise_server_exceptions = False) r = c.put("/embedding-model", json = {"embedding_model": "baai/bge-m3", "force": True}) assert r.status_code == 200 - assert saved.get("model") == "baai/bge-m3" # the spelling the GGUF check validated + assert saved.get("model") == "baai/bge-m3" def test_default_model_is_not_casing_normalized(client, monkeypatch): - # The exact default must NOT be casing-normalized: rewriting it would make - # set_rag_embedding_model()'s exact default comparison treat it as a custom override. + # Recasing the default would make set_rag_embedding_model() treat it as a custom override. c, saved = client import utils.models as _models @@ -350,7 +346,7 @@ def test_load_sink_refuses_flagged_model(monkeypatch): def test_load_sink_allows_clean_model(monkeypatch): monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) import core.rag.embeddings as embeddings - embeddings._guard_model_security("acme/clean-embed", False) # no raise + embeddings._guard_model_security("acme/clean-embed", False) def test_sink_threads_ambient_token_into_scan(monkeypatch): diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index a1eb10decea..cf48d3363ee 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -98,19 +98,17 @@ def _model_info(*a, **k): @pytest.mark.parametrize("var", ["HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"]) def test_shared_gate_still_scans_when_offline_by_default(monkeypatch, var): - # THE important one: the gate is shared by loaders that do NOT pin to the local cache, - # so an offline env var alone must never disable the malware scan. + # Shared gate: an offline env var alone must never disable the scan for a fetching loader. import utils.security.file_security as fs calls: list = [] _fake_hub(monkeypatch, calls) monkeypatch.setenv(var, "1") - assert fs._fetch_security_status("org/model", None) is None # fails open on error + assert fs._fetch_security_status("org/model", None) is None assert calls, f"{var} alone must NOT bypass the shared malware gate" def test_security_scan_short_circuits_for_a_local_only_caller(monkeypatch): - # A local-only caller gets the Hub round-trip skipped instead of burning both timeouts. import utils.security.file_security as fs calls: list = [] @@ -124,7 +122,7 @@ def test_security_scan_runs_when_online(monkeypatch): calls: list = [] _fake_hub(monkeypatch, calls) - assert fs._fetch_security_status("org/model", None) is None # fails open on error + assert fs._fetch_security_status("org/model", None) is None assert calls, "online must attempt the Hub" @@ -149,7 +147,6 @@ def test_embedding_loader_forces_local_only_when_offline(): "SentenceTransformer must be pinned with the SAME captured value; a second " "hf_env_offline() read can flip to False and fetch the unscanned repo" ) - # And the guard must not re-derive it internally. guard = src.split("def _guard_model_security", 1)[1].split("\ndef ", 1)[0] assert "hf_env_offline()" not in guard, ( "_guard_model_security must take local_only_load as an argument so it " diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index af91be1982f..24e40a3c733 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -123,8 +123,7 @@ def test_marker_none_when_not_cached(monkeypatch): def test_marker_judges_the_revision_refs_main_points_at(tmp_path, monkeypatch): - # The old snapshot has modules.json but the active refs/main revision does not: the - # active revision must win, so scanning any snapshot would wrongly say True. + # The active refs/main revision wins over an older snapshot that still has modules.json. _repo(tmp_path, monkeypatch, ("new", False), ("old", True), main_ref = "new") assert mc._embedding_marker_in_hf_cache("org/model") is False @@ -135,40 +134,36 @@ def test_marker_true_when_refs_main_revision_is_st(tmp_path, monkeypatch): def test_marker_missing_ref_is_a_cache_miss(tmp_path, monkeypatch): - # local_files_only resolves the default revision THROUGH refs/main, so a snapshot dir - # alone is not discoverable and accepting one would fail at first indexing. - _repo(tmp_path, monkeypatch, ("new", False), ("old", True)) # no refs/main + # local_files_only resolves the default revision through refs/main; a bare snapshot dir + # is not discoverable. + _repo(tmp_path, monkeypatch, ("new", False), ("old", True)) assert mc._embedding_marker_in_hf_cache("org/model") is None def test_marker_ref_points_at_absent_snapshot_is_cache_miss(tmp_path, monkeypatch): - # refs/main names a commit whose snapshot dir is absent (partial download / - # pruning): a miss, never a fall-through to a stale historical snapshot. + # refs/main naming an absent snapshot is a miss, not a fall-through to a stale one. _repo(tmp_path, monkeypatch, ("old", True), main_ref = "missing_commit") assert mc._embedding_marker_in_hf_cache("org/model") is None def test_marker_unreadable_ref_is_cache_miss(tmp_path, monkeypatch): - # refs/main exists but cannot be read (transient I/O / restrictive - # permissions): the loader cannot resolve it either. + # refs/main present but unreadable: the loader cannot resolve it either. dirs = _repo(tmp_path, monkeypatch, ("old", True)) refs_main = dirs[0].parent.parent / "refs" / "main" refs_main.parent.mkdir(parents = True, exist_ok = True) - refs_main.mkdir() # a directory -> read_text raises IsADirectoryError + refs_main.mkdir() # a directory -> read_text raises assert mc._embedding_marker_in_hf_cache("org/model") is None def test_marker_empty_ref_is_cache_miss(tmp_path, monkeypatch): - # refs/main exists but is empty / whitespace (a partial write): the active - # revision is unknown. + # empty / whitespace refs/main (partial write): the active revision is unknown. _repo(tmp_path, monkeypatch, ("old", True), main_ref = " \n") assert mc._embedding_marker_in_hf_cache("org/model") is None def test_marker_scopes_to_the_repo_dir_the_loader_opens(tmp_path, monkeypatch): - # Two cache dirs differing only by case: the loader opens the one persisted (exact case - # first), so a complete model there must validate even when the OTHER variant holds a - # newer, non-ST snapshot -- judging across both would reject it. + # The loader opens the exact-case dir, so a complete model there validates even when the + # other case variant holds a newer non-ST snapshot. if not _case_sensitive_fs(tmp_path): pytest.skip("duplicate case variants need a case-sensitive filesystem") root = tmp_path / "cache" @@ -205,8 +200,7 @@ def _st_snapshot( def test_st_probe_uses_sentence_transformers_home(tmp_path, monkeypatch): - # With ST_HOME set that is the ONLY cache the load searches, so the probe must follow - # it or a model present there is called uncached and 409'd offline. + # ST_HOME is the only cache the load searches, so the probe must follow it. hf_root, st_root = tmp_path / "hf", tmp_path / "st" hf_root.mkdir() _st_snapshot(st_root, "models--org--model") @@ -217,8 +211,7 @@ def test_st_probe_uses_sentence_transformers_home(tmp_path, monkeypatch): def test_st_probe_ignores_hub_cache_when_st_home_is_set(tmp_path, monkeypatch): - # The loader searches ST_HOME only, so a repo cached ONLY in the Hub cache - # must not validate -- otherwise validation passes and the load then fails. + # Loader searches ST_HOME only; a repo cached only in the Hub cache must not validate. hf_root, st_root = tmp_path / "hf", tmp_path / "st" st_root.mkdir() _st_snapshot(hf_root, "models--org--model") @@ -228,8 +221,7 @@ def test_st_probe_ignores_hub_cache_when_st_home_is_set(tmp_path, monkeypatch): def test_gguf_probe_never_follows_st_home(tmp_path, monkeypatch): - # The GGUF path uses the Hub cache (hf_hub_download, no cache_dir), so letting its probe - # see ST_HOME would select a file the GGUF loader cannot find. + # The GGUF path uses the Hub cache (hf_hub_download), so its probe must not follow ST_HOME. hf_root, st_root = tmp_path / "hf", tmp_path / "st" hf_root.mkdir() _st_snapshot(st_root, "models--org--model") @@ -247,8 +239,7 @@ def test_st_probe_falls_back_to_hub_cache(tmp_path, monkeypatch): def test_marker_only_snapshot_is_not_loadable(tmp_path, monkeypatch): - # The security preflight downloads modules.json on its own, and a partial download - # leaves it too, so accepting it offline would fail on the first RAG load. + # The security preflight (and a partial download) leaves a bare modules.json. hf_root = tmp_path / "hf" _st_snapshot(hf_root, "models--org--model", loadable = False) _fake_hf_cache(monkeypatch, hf_root) @@ -257,18 +248,16 @@ def test_marker_only_snapshot_is_not_loadable(tmp_path, monkeypatch): def test_marker_onnx_only_snapshot_is_not_loadable(tmp_path, monkeypatch): - # Marker + config but ONLY an ONNX export is not loadable: the default Torch backend - # reads model.safetensors / pytorch_model.bin, so accepting the ONNX offline would fail - # at load. + # Marker + config but only an ONNX export: the Torch backend needs safetensors/bin. hf_root = tmp_path / "hf" repo = hf_root / "models--org--model" snap = repo / "snapshots" / "aaa" snap.mkdir(parents = True) (snap / "modules.json").write_text("[]") (snap / "config.json").write_text("{}") - (snap / "tokenizer.json").write_text("{}") # isolate the failure to the weight format + (snap / "tokenizer.json").write_text("{}") (snap / "model.onnx").write_bytes(b"\0") - # refs/main so the probe reaches the weight check (without it the answer is None). + # refs/main so the probe reaches the weight check. (repo / "refs").mkdir(parents = True) (repo / "refs" / "main").write_text("aaa") _fake_hf_cache(monkeypatch, hf_root) @@ -277,16 +266,14 @@ def test_marker_onnx_only_snapshot_is_not_loadable(tmp_path, monkeypatch): def test_marker_rejects_weights_without_tokenizer(tmp_path, monkeypatch): - # Marker + config + weights but NO tokenizer asset is not loadable: the Transformer - # module also builds an AutoTokenizer, which fails offline without a tokenizer asset. + # Weights but no tokenizer asset: the Transformer module also builds an AutoTokenizer. _cache_repo_with_files(tmp_path, monkeypatch, "model.safetensors", tokenizer = False) assert mc._embedding_marker_in_hf_cache("org/model") is False @pytest.mark.parametrize("tok_file", ["tokenizer_config.json", "vocab.txt", "spiece.model"]) def test_marker_accepts_alternate_tokenizer_assets(tmp_path, monkeypatch, tok_file): - # The tokenizer check is a permissive union: any one recognized asset is enough, so a - # valid non-tokenizer.json layout is not wrongly rejected. + # Permissive union: any one recognized tokenizer asset suffices. _cache_repo_with_files(tmp_path, monkeypatch, "model.safetensors", tok_file, tokenizer = False) assert mc._embedding_marker_in_hf_cache("org/model") is True @@ -319,8 +306,7 @@ def _cache_repo_with_files( def test_marker_rejects_non_base_weight_bins(tmp_path, monkeypatch): - # A NON-weight .bin (training_args.bin) or an adapter-only artifact is not a loadable - # base model: the Torch backend needs model.safetensors / pytorch_model.bin. + # A non-weight .bin or adapter-only artifact is not a loadable base model. _cache_repo_with_files(tmp_path, monkeypatch, "training_args.bin") assert mc._embedding_marker_in_hf_cache("org/model") is False _cache_repo_with_files(tmp_path / "b", monkeypatch, "adapter_model.safetensors") @@ -344,8 +330,7 @@ def test_marker_rejects_non_base_weight_bins(tmp_path, monkeypatch): ], ) def test_marker_rejects_incomplete_shard_set(tmp_path, monkeypatch, weights): - # A partially downloaded sharded model must NOT validate: the Torch backend needs every - # shard AND the index map, so an incomplete set (or a set missing its index) fails. + # An incomplete shard set (or one missing its index map) must not validate. _cache_repo_with_files(tmp_path, monkeypatch, *weights) assert mc._embedding_marker_in_hf_cache("org/model") is False @@ -363,8 +348,7 @@ def test_marker_rejects_incomplete_shard_set(tmp_path, monkeypatch, weights): ], ) def test_marker_accepts_recognized_torch_weights(tmp_path, monkeypatch, weights): - # The recognizer must not over-reject real weights: single pytorch_model.bin, a complete - # sharded set with its index, and weights inside a module dir all count as loadable. + # Must not over-reject: single bin, complete sharded set, and a weight in a module dir. _cache_repo_with_files(tmp_path, monkeypatch, *weights) assert mc._embedding_marker_in_hf_cache("org/model") is True @@ -376,8 +360,7 @@ def _case_sensitive_fs(tmp_path) -> bool: def test_st_casing_resolves_against_st_home(tmp_path, monkeypatch): - # resolve_cached_repo_id_case scans only the Hub cache, so with ST_HOME set it would - # persist the lower-case id while the exact-case offline load misses it in ST_HOME. + # With ST_HOME set, casing must resolve against it, or the exact-case offline load misses. if not _case_sensitive_fs(tmp_path): pytest.skip("casing only diverges on a case-sensitive filesystem") hf_root, st_root = tmp_path / "hf", tmp_path / "st" @@ -436,8 +419,8 @@ def _exploding_iter(repo): def test_offline_cached_st_detected_via_marker_no_network(tmp_path, monkeypatch): - # Offline: a downloaded sentence-transformers repo is classified from its modules.json - # marker with no model_info() network call that would hang on DNS retries (#6817). + # Offline: classified from the cached modules.json marker, no model_info() call that + # would hang on DNS (#6817). _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa", repo_id = "unsloth/bge-small-en-v1.5") monkeypatch.setenv("HF_HUB_OFFLINE", "1") _fake_hf_model_info(monkeypatch, _no_network) @@ -445,8 +428,8 @@ def test_offline_cached_st_detected_via_marker_no_network(tmp_path, monkeypatch) def test_online_defers_to_hub_over_stale_marker(tmp_path, monkeypatch): - # Online: the Hub is authoritative. Even with a cached modules.json, a Hub lookup that - # no longer reports embedding signals wins -- the stale marker must not short-circuit it. + # Online: the Hub is authoritative; a stale cached marker must not short-circuit a lookup + # that no longer reports embedding signals. _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa") calls = [] @@ -456,12 +439,12 @@ def _info(model_name, token = None): _fake_hf_model_info(monkeypatch, _info) assert mc.is_embedding_model("org/was-embedder") is False - assert calls == ["org/was-embedder"] # Hub consulted, not skipped + assert calls == ["org/was-embedder"] def test_online_permanent_hub_error_ignores_stale_marker(tmp_path, monkeypatch): - # A permanent Hub error (deleted / gated / typo'd repo) is authoritative: even with a - # cached modules.json, return False so the settings route surfaces its 409. + # A permanent Hub error is authoritative: return False even with a cached marker, so the + # settings route surfaces its 409. _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa") class RepositoryNotFoundError(Exception): @@ -475,118 +458,108 @@ def _info(model_name, token = None): def test_online_hub_failure_falls_back_to_marker_uncached(tmp_path, monkeypatch): - # A transient model_info() failure falls back to the local marker WITHOUT caching, so a - # later successful Hub lookup can still override the degraded result. + # A transient failure falls back to the local marker WITHOUT caching, so a later Hub + # lookup can override. _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa", repo_id = "org/emb") - _fake_hf_model_info(monkeypatch, _no_network) # raises -> Hub "unreachable" - assert mc.is_embedding_model("org/emb") is True # marker fallback - assert ("org/emb", None) not in mc._embedding_detection_cache # not poisoned + _fake_hf_model_info(monkeypatch, _no_network) + assert mc.is_embedding_model("org/emb") is True + assert ("org/emb", None) not in mc._embedding_detection_cache def test_online_negative_does_not_block_later_offline_download(tmp_path, monkeypatch): - # An online lookup reports non-embedding and is memoized. The repo is then downloaded - # with modules.json and the session goes offline; the offline path re-probes the marker - # (never the memo), so the fresh embedder is detected instead of the stale False. + # A memoized online negative must not block a later offline detection: the offline path + # re-probes the marker, never the memo. _no_cache(monkeypatch) def _info(model_name, token = None): return types.SimpleNamespace(tags = ["text-generation"], pipeline_tag = "text-generation") _fake_hf_model_info(monkeypatch, _info) - assert mc.is_embedding_model("org/late-embedder") is False # online: Hub says no + assert mc.is_embedding_model("org/late-embedder") is False assert mc._embedding_detection_cache[("org/late-embedder", None)] is False - # Now the model is downloaded (marker appears) and the session goes offline. + # Model now downloaded (marker appears); session goes offline. _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa", repo_id = "org/late-embedder") monkeypatch.setenv("HF_HUB_OFFLINE", "1") - assert mc.is_embedding_model("org/late-embedder") is True # marker re-probed + assert mc.is_embedding_model("org/late-embedder") is True def test_offline_retains_online_confirmed_positive(tmp_path, monkeypatch): - # A tag-only embedder (no modules.json) confirmed online and cached True, with its - # snapshot materialized. Flipped offline mid-load, the offline path must RETAIN the - # positive: the marker reads False, so retention rests on the memo plus a present - # snapshot, not the marker. + # A tag-only embedder (no modules.json) cached True online, snapshot materialized. Offline + # must retain the positive via the memo + present snapshot, since the marker reads False. _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") def _info(model_name, token = None): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) _fake_hf_model_info(monkeypatch, _info) - assert mc.is_embedding_model("org/gte-modernbert") is True # online: cached True + assert mc.is_embedding_model("org/gte-modernbert") is True monkeypatch.setenv("HF_HUB_OFFLINE", "1") - _fake_hf_model_info(monkeypatch, _no_network) # offline must not hit network - assert mc.is_embedding_model("org/gte-modernbert") is True # positive retained + _fake_hf_model_info(monkeypatch, _no_network) + assert mc.is_embedding_model("org/gte-modernbert") is True def test_offline_metadata_only_positive_not_trusted_without_cache(monkeypatch): - # An online positive proves only that the repo is tagged an embedder, not that files - # were downloaded. With nothing materialized, the offline path must NOT trust the memo - # (the local_files_only load would fail), so it returns False despite the cached True. + # An online positive only tags the repo; with nothing materialized the offline load would + # fail, so the memo must NOT be trusted. _no_cache(monkeypatch) def _info(model_name, token = None): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) _fake_hf_model_info(monkeypatch, _info) - assert ( - mc.is_embedding_model("org/uncached-embedder") is True - ) # online: cached True (metadata only) + assert mc.is_embedding_model("org/uncached-embedder") is True monkeypatch.setenv("HF_HUB_OFFLINE", "1") - _fake_hf_model_info(monkeypatch, _no_network) # offline must not hit network - assert mc.is_embedding_model("org/uncached-embedder") is False # not cached -> not trusted + _fake_hf_model_info(monkeypatch, _no_network) + assert mc.is_embedding_model("org/uncached-embedder") is False def test_offline_detects_persisted_tag_only_embedder_after_restart(tmp_path, monkeypatch): - # A tag-only embedder confirmed online durably records the verdict, snapshot - # materialized. After a RESTART (memo gone) the process comes up offline: the marker - # reads False, so recognition rests on the persisted allowlist plus the present snapshot. + # Tag-only embedder persisted online, snapshot materialized. After a restart (memo gone) + # offline recognition rests on the persisted allowlist + present snapshot. _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") def _info(model_name, token = None): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) _fake_hf_model_info(monkeypatch, _info) - assert mc.is_embedding_model("org/gte-modernbert") is True # online: confirmed + persisted - assert "org/gte-modernbert" in mc._load_persisted_embedders() # written to disk + assert mc.is_embedding_model("org/gte-modernbert") is True + assert "org/gte-modernbert" in mc._load_persisted_embedders() - mc._embedding_detection_cache.clear() # simulate a process restart: memo lost + mc._embedding_detection_cache.clear() # simulate a restart: memo lost monkeypatch.setenv("HF_HUB_OFFLINE", "1") - _fake_hf_model_info(monkeypatch, _no_network) # offline must not hit network - assert mc.is_embedding_model("org/gte-modernbert") is True # recovered from disk + _fake_hf_model_info(monkeypatch, _no_network) + assert mc.is_embedding_model("org/gte-modernbert") is True def test_offline_persisted_verdict_not_trusted_when_uncached(tmp_path, monkeypatch): - # The persisted allowlist records only the tag, not on-disk files. After a restart with - # NOTHING materialized, the offline path must NOT trust it (gated on the active snapshot - # being present), so an uncached repo returns False. + # The persisted allowlist records only the tag; after a restart with nothing materialized + # it must NOT be trusted. _no_cache(monkeypatch) def _info(model_name, token = None): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) _fake_hf_model_info(monkeypatch, _info) - assert mc.is_embedding_model("org/uncached-embedder") is True # online: confirmed + persisted + assert mc.is_embedding_model("org/uncached-embedder") is True assert "org/uncached-embedder" in mc._load_persisted_embedders() mc._embedding_detection_cache.clear() # restart: memo lost, disk verdict remains _no_cache(monkeypatch) monkeypatch.setenv("HF_HUB_OFFLINE", "1") _fake_hf_model_info(monkeypatch, _no_network) - assert mc.is_embedding_model("org/uncached-embedder") is False # nothing on disk -> not trusted + assert mc.is_embedding_model("org/uncached-embedder") is False def test_offline_persisted_verdict_not_trusted_when_snapshot_partial(tmp_path, monkeypatch): - # A persisted verdict is trusted only when the active snapshot carries a COMPLETE weight - # set, not merely that it is materialized. Here config but no weights (interrupted - # download): a bare "materialized" gate would wrongly return True, so the weight gate - # must reject it. + # A persisted verdict is trusted only with a COMPLETE weight set, not mere materialization. + # Here config but no weights (interrupted download) must be rejected. cache_root = tmp_path / "cache" snap = cache_root / "models--org--partial" / "snapshots" / "aaa" snap.mkdir(parents = True) - (snap / "config.json").write_text("{}") # config only -- weights never finished + (snap / "config.json").write_text("{}") refs = cache_root / "models--org--partial" / "refs" refs.mkdir(parents = True) (refs / "main").write_text("aaa") @@ -596,45 +569,42 @@ def _info(model_name, token = None): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) _fake_hf_model_info(monkeypatch, _info) - assert mc.is_embedding_model("org/partial") is True # online: confirmed + persisted + assert mc.is_embedding_model("org/partial") is True assert "org/partial" in mc._load_persisted_embedders() assert mc._embedding_marker_in_hf_cache("org/partial") is False # materialized, not None mc._embedding_detection_cache.clear() # restart: memo lost, disk verdict remains monkeypatch.setenv("HF_HUB_OFFLINE", "1") _fake_hf_model_info(monkeypatch, _no_network) - assert mc.is_embedding_model("org/partial") is False # partial snapshot -> not trusted + assert mc.is_embedding_model("org/partial") is False def test_offline_persisted_verdict_matches_across_casing(tmp_path, monkeypatch): - # The verdict is recorded under the request spelling but the settings route saves the - # cache-resolved one, so an exact-string lookup would miss it. Verified as baai/model, - # looked up as the saved BAAI/model, must still match: the allowlist is case-folded. + # The verdict is saved under the cache-resolved casing; the case-folded allowlist must + # still match a differently-cased lookup. _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "BAAI/model") def _info(model_name, token = None): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) _fake_hf_model_info(monkeypatch, _info) - assert mc.is_embedding_model("baai/model") is True # online: confirmed as baai/model + assert mc.is_embedding_model("baai/model") is True mc._embedding_detection_cache.clear() # restart: memo lost, disk verdict remains monkeypatch.setenv("HF_HUB_OFFLINE", "1") _fake_hf_model_info(monkeypatch, _no_network) - # Looked up under the cache-resolved casing the settings route persisted. assert mc.is_embedding_model("BAAI/model") is True def test_persist_embedder_concurrent_writes_keep_every_verdict(tmp_path, monkeypatch): - # Concurrent confirmations must not drop each other's entry: serialized read-modify-write - # + per-thread temp files, so every model persisted from parallel threads survives. + # Concurrent confirmations must not drop entries: serialized writes + per-thread temp files. import threading names = [f"org/emb-{i}" for i in range(24)] barrier = threading.Barrier(len(names)) def _writer(name): - barrier.wait() # maximize overlap on the shared file + barrier.wait() mc._persist_embedder(name) threads = [threading.Thread(target = _writer, args = (n,)) for n in names] @@ -648,19 +618,19 @@ def _writer(name): def test_persist_embedder_is_best_effort_when_home_unwritable(tmp_path, monkeypatch): - # Persistence is an optimization: if the Studio home cannot be written, the online path - # must still return its verdict rather than raise. Point the home at a path blocked by a file. + # Persistence is best-effort: an unwritable Studio home must not raise; the online verdict + # still returns. blocker = tmp_path / "blocker" blocker.write_text("not a dir") - monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(blocker / "studio")) # parent is a file + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(blocker / "studio")) _no_cache(monkeypatch) def _info(model_name, token = None): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) _fake_hf_model_info(monkeypatch, _info) - assert mc.is_embedding_model("org/emb") is True # no exception despite unwritable home - assert mc._load_persisted_embedders() == set() # nothing recorded, silently + assert mc.is_embedding_model("org/emb") is True + assert mc._load_persisted_embedders() == set() def test_offline_cached_non_st_returns_false_without_network(tmp_path, monkeypatch): @@ -678,8 +648,7 @@ def test_offline_not_cached_returns_false_without_network(monkeypatch): def test_online_uncached_still_uses_network(monkeypatch): - # Not offline, not cached: the model_info path must still run so a feature-extraction - # embedder lacking modules.json is caught. + # Not offline, not cached: model_info still runs so a tag-only embedder is caught. _no_cache(monkeypatch) calls = [] @@ -693,8 +662,8 @@ def _info(model_name, token = None): def test_offline_negative_is_not_cached_then_online_detects(monkeypatch): - # A tag-only embedder is not identifiable from modules.json. Offline returns False - # WITHOUT caching, so once the env clears the online lookup still detects it. + # A tag-only embedder isn't identifiable from modules.json; offline returns False WITHOUT + # caching so online can still detect it later. _no_cache(monkeypatch) calls = [] @@ -706,9 +675,9 @@ def _info(model_name, token = None): monkeypatch.setenv("HF_HUB_OFFLINE", "1") assert mc.is_embedding_model("org/gte-modernbert") is False - assert calls == [] # offline: no network + assert calls == [] assert ("org/gte-modernbert", None) not in mc._embedding_detection_cache monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) - assert mc.is_embedding_model("org/gte-modernbert") is True # now detected online + assert mc.is_embedding_model("org/gte-modernbert") is True assert calls == ["org/gte-modernbert"] diff --git a/studio/backend/tests/test_rag_embeddings.py b/studio/backend/tests/test_rag_embeddings.py index 11aa089a2f5..4c03053c775 100644 --- a/studio/backend/tests/test_rag_embeddings.py +++ b/studio/backend/tests/test_rag_embeddings.py @@ -223,10 +223,8 @@ def test_st_encode_failure_without_llama_binary_reraises(monkeypatch): def test_get_resolves_default_casing_before_loading(monkeypatch): - # _get() must resolve a repo id to the exact cache casing before constructing - # SentenceTransformer: a default differing only by case is NOT persist-normalized by - # /settings, so without resolving here the offline load would miss the cache dir. - # Verify both the loader and the security gate receive the resolved name. + # _get() resolves the cache casing before loading (a default is NOT persist-normalized + # by /settings), and both the loader and the gate get the resolved name. import sys import types @@ -257,5 +255,5 @@ def __init__(self, name, **kwargs): monkeypatch.setattr(embeddings, "_name", None, raising = False) embeddings._get() - assert captured["load_name"] == "BAAI/bge-m3" # resolved to the cache casing - assert guarded["name"] == "BAAI/bge-m3" # gate scans the same resolved name + assert captured["load_name"] == "BAAI/bge-m3" + assert guarded["name"] == "BAAI/bge-m3" diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 1dbb2ac5873..5b1de808564 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1644,7 +1644,7 @@ def _local_gguf_companion_search_root(selected_path: str, gguf_file: str) -> str def _hf_hub_cache_root() -> list[Path]: - """The Hub cache root, i.e. what ``hf_hub_download`` uses with no ``cache_dir``.""" + """The Hub cache root -- what ``hf_hub_download`` uses with no ``cache_dir``.""" try: from huggingface_hub import constants as hf_constants return [Path(hf_constants.HF_HUB_CACHE)] @@ -1653,11 +1653,9 @@ def _hf_hub_cache_root() -> list[Path]: def _st_cache_roots() -> list[Path]: - """Cache root a SentenceTransformer load actually searches: SENTENCE_TRANSFORMERS_HOME - if set, else the Hub cache -- one or the other, never both. Probing the union would - pass offline validation on a repo cached only in the other root. Kept separate from - the Hub-cache probe, which the GGUF path needs (it downloads via ``hf_hub_download``). - """ + """Cache root a SentenceTransformer load searches: SENTENCE_TRANSFORMERS_HOME if set, + else the Hub cache -- never the union (that would pass validation on a repo cached only + in the other root).""" st_home = (os.environ.get("SENTENCE_TRANSFORMERS_HOME") or "").strip() if st_home: try: @@ -1668,7 +1666,7 @@ def _st_cache_roots() -> list[Path]: def _iter_cache_snapshots_in(repo_id: str, roots: list[Path]): - """Snapshot dirs for *repo_id* under *roots*, newest first (shared machinery).""" + """Snapshot dirs for *repo_id* under *roots*, newest first.""" target = f"models--{repo_id.replace('/', '--')}".lower() repo_dirs: list[Path] = [] for cache_dir in roots: @@ -1684,11 +1682,9 @@ def _iter_cache_snapshots_in(repo_id: str, roots: list[Path]): def resolve_st_cached_repo_id_case(repo_id: str) -> str: - """*repo_id* recased to match its dir in the cache the ST loader searches, so the - offline exact-case ``local_files_only`` load finds it. Unchanged for a local path, a - non-repo string, or when nothing case-matching is cached; prefers an exact match, then - a deterministic pick among case variants. - """ + """*repo_id* recased to match its cache dir so the offline exact-case + ``local_files_only`` load finds it; unchanged for a local path, non-repo, or nothing + cached (exact match, then a deterministic pick among case variants).""" if is_local_path(repo_id) or "/" not in repo_id: return repo_id prefix = "models--" @@ -1700,7 +1696,7 @@ def resolve_st_cached_repo_id_case(repo_id: str) -> str: if not cache_dir.is_dir(): continue if (cache_dir / expected).is_dir(): - return repo_id # exact case already cached: keep the request + return repo_id for entry in cache_dir.iterdir(): if entry.is_dir() and entry.name.lower() == target: variants.append(entry.name) @@ -1712,11 +1708,9 @@ def resolve_st_cached_repo_id_case(repo_id: str) -> str: def _st_cache_repo_dir(repo_id: str) -> Optional[Path]: - """The ONE cache repo dir the ST loader opens for *repo_id*, or None. Same selection - as :func:`resolve_st_cached_repo_id_case` (exact case, then a deterministic pick). - Scoping to one dir matters with duplicate case variants: judging by whichever holds - the newest snapshot could reject the model in the dir actually loaded. - """ + """The ONE cache repo dir the ST loader opens for *repo_id*, or None (exact case, then + a deterministic pick). Scoping to one dir avoids judging a duplicate case variant that + is not the one loaded.""" prefix = "models--" expected = f"{prefix}{repo_id.replace('/', '--')}" target = expected.lower() @@ -1764,13 +1758,9 @@ def _iter_snapshots_of(repo_dirs: list[Path]): def _iter_hf_cache_snapshots(repo_id: str): - """Yield HUB cache snapshot dirs for *repo_id*, newest first. - - Hub cache ONLY: its GGUF-detector callers download via ``hf_hub_download`` with no - ``cache_dir``, so that is the cache their load uses; the ST probe is - :func:`_st_cache_repo_dir`. Empty if the root or repo isn't cached. Match is - case-insensitive to handle casing drift between download and lookup. - """ + """Hub cache snapshot dirs for *repo_id*, newest first. Hub cache ONLY -- its GGUF + callers download via ``hf_hub_download`` (the ST probe is :func:`_st_cache_repo_dir`); + case-insensitive match for casing drift between download and lookup.""" yield from _iter_cache_snapshots_in(repo_id, _hf_hub_cache_root()) @@ -2086,35 +2076,28 @@ def download_gguf_file( def _persisted_embedders_path() -> Path: - """On-disk store of repo ids confirmed embedders online, under the Studio home so it - survives a restart. Without it a tag-only feature-extraction embedder (cached snapshot - has no ``modules.json``, so unrecognizable from the cache alone) would be misclassified - non-embedding the first offline call after the session memo is lost. - """ + """On-disk store (under Studio home) of repo ids confirmed embedders online, so it + survives a restart -- without it a tag-only embedder (no ``modules.json`` in cache) is + misclassified the first offline call after the session memo is lost.""" return _studio_root() / "embedding_verdicts.json" -# Serializes the read-modify-write in _persist_embedder so concurrent confirmations -# cannot lose a verdict or collide on the temp file. Cross-process writes are -# best-effort: os.replace stays atomic, and a dropped verdict is only an optimization -# miss a later online re-confirmation heals. +# Serializes the read-modify-write in _persist_embedder. Cross-process writes are +# best-effort: os.replace is atomic and a dropped verdict is only a missed optimization a +# later online re-confirmation heals. _persist_lock = threading.Lock() def _verdict_key(model_name: str) -> str: - """Case-folded allowlist key: the verdict is a property of the repo, not its spelling. - model_info() is queried under the requested casing but the settings route saves the - cache-resolved casing, so an exact-string key would miss the persisted positive after a - restart. Repo ids are case-insensitive, so fold both sides.""" + """Case-folded allowlist key: model_info() runs under the requested casing but the + settings route saves the cache-resolved casing, so fold both sides (repo ids are + case-insensitive).""" return model_name.casefold() def _load_persisted_embedders() -> set: """Case-folded repo ids recorded as embedders in a prior session; empty set on any - error (missing, empty, corrupt, or non-list file all read as "nothing recorded", so - callers fall back to their cache-marker logic). Never raises. Folding also normalizes - mixed-case entries an older build may have written. - """ + error (missing / empty / corrupt read as "nothing recorded"). Never raises.""" try: with open(_persisted_embedders_path(), encoding = "utf-8") as fh: data = json.load(fh) @@ -2126,14 +2109,10 @@ def _load_persisted_embedders() -> set: def _persist_embedder(model_name: str) -> None: - """Record *model_name* as an online-confirmed embedder, best-effort. - - Only positive Hub verdicts are written, case-folded. The offline branch gates this - allowlist on the active revision being materialized on disk, so a stale entry cannot - resurrect a False. Serialized under ``_persist_lock`` through a per-thread temp file so - concurrent writers cannot drop an entry or collide. Any failure is swallowed -- - persistence is an optimization, not a correctness requirement. - """ + """Record *model_name* as an online-confirmed embedder, best-effort. Positive Hub + verdicts only, case-folded, serialized under ``_persist_lock`` via a per-thread temp + file. Any failure is swallowed -- persistence is an optimization, not a correctness + requirement.""" try: key = _verdict_key(model_name) with _persist_lock: @@ -2152,23 +2131,22 @@ def _persist_embedder(model_name: str) -> None: def _known_embedder(model_name: str, cache_key: tuple) -> bool: - """True when *model_name* was confirmed an embedder online -- the session memo or the - persisted cross-restart allowlist. Callers still gate this on a materialized snapshot.""" + """True when *model_name* was confirmed an embedder online (session memo or persisted + allowlist). Callers still gate this on a materialized snapshot.""" if _embedding_detection_cache.get(cache_key) is True: return True return _verdict_key(model_name) in _load_persisted_embedders() -# Base-model weight files the default (Torch) SentenceTransformer backend consumes: -# a single model.safetensors / pytorch_model.bin or their sharded -00001-of-000NN forms. -# Matched by NAME, not suffix, so a training_args.bin / optimizer.bin or an adapter-only -# artifact cannot let a partial cache pass offline validation and then fail on first load. +# Base-model weight files the default Torch SentenceTransformer backend consumes: +# model/pytorch_model .safetensors/.bin, plain or sharded. Matched by NAME not suffix so a +# training_args.bin / adapter-only artifact can't pass offline validation then fail on load. _ST_WEIGHT_FILE_RE = re.compile(r"^(model|pytorch_model)(-\d+-of-\d+)?\.(safetensors|bin)$") _ST_SHARD_RE = re.compile(r"^(model|pytorch_model)-(\d+)-of-(\d+)\.(safetensors|bin)$") -# A Transformer module also loads an AutoTokenizer, so weights without a tokenizer asset -# still fail the local_files_only load. Any ONE of these loads some tokenizer -- a -# permissive union so only a genuinely tokenizer-less partial download is rejected. +# A Transformer module also loads an AutoTokenizer, so weights without any tokenizer asset +# fail the local_files_only load. Any ONE of these suffices -- a permissive union, so only +# a genuinely tokenizer-less partial download is rejected. _ST_TOKENIZER_FILES = frozenset( { "tokenizer.json", @@ -2184,15 +2162,10 @@ def _known_embedder(model_name: str, cache_key: tuple) -> bool: def _dir_has_complete_torch_weights(names: set) -> bool: - """True when the filenames in ONE dir form a COMPLETE Torch weight set: - - * a single ``model.safetensors`` / ``pytorch_model.bin``; or - * a full shard set -- every index ``1..total`` (a shard names its own total, - ``-00001-of-00002``) plus its ``...index.json`` map. - - A lone shard or a shard set missing its index map is NOT complete: accepting a partial - sharded download would validate and then fail under local_files_only. - """ + """True when *names* form a COMPLETE Torch weight set: a single ``model.safetensors`` / + ``pytorch_model.bin``, or a full shard set (every index ``1..total`` plus its + ``.index.json`` map). A lone shard or a missing index map is incomplete -- it would + validate then fail under local_files_only.""" for stem in ("model", "pytorch_model"): for ext in ("safetensors", "bin"): if f"{stem}.{ext}" in names: @@ -2211,13 +2184,9 @@ def _dir_has_complete_torch_weights(names: set) -> bool: def _snapshot_has_complete_weights(snap: Path) -> bool: """True when *snap* has a config, a tokenizer asset, and a COMPLETE Torch weight set - (the ``modules.json`` marker aside). - - A tag-only feature-extraction embedder confirmed online has exactly this and ST's - auto-model fallback loads it, so an online-confirmed positive may be trusted offline. - A partial download (weights missing, incomplete shard set, or no tokenizer) fails here - rather than at first indexing. - """ + (``modules.json`` aside). A tag-only embedder confirmed online has exactly this and ST's + auto-model fallback loads it; a partial download fails here rather than at first + indexing.""" try: if not any( (snap / name).is_file() for name in ("config.json", "config_sentence_transformers.json") @@ -2243,12 +2212,9 @@ def _snapshot_has_complete_weights(snap: Path) -> bool: def _snapshot_is_loadable_st_model(snap: Path) -> bool: - """True when *snap* holds a sentence-transformers model that can actually load. - - ``modules.json`` alone is not enough: the security preflight downloads exactly that - file, and a partial download can leave it behind while weights/config are absent. - Requires the marker plus ``_snapshot_has_complete_weights``. - """ + """True when *snap* holds a loadable sentence-transformers model. ``modules.json`` alone + is insufficient -- the security preflight downloads exactly that file, so require the + marker plus ``_snapshot_has_complete_weights``.""" try: if not (snap / "modules.json").is_file(): return False @@ -2259,11 +2225,10 @@ def _snapshot_is_loadable_st_model(snap: Path) -> bool: def _active_snapshot_dir(repo_id: str) -> Optional[Path]: """The materialized snapshot dir the offline ``local_files_only`` load resolves for - *repo_id*, or None. Mirrors that resolution: the repo dir :func:`_st_cache_repo_dir` - selects, its ``refs/main`` commit (huggingface_hub resolves the default revision THROUGH - that ref, so a missing/empty/unreadable ref is a cache MISS, not a reason to scan - historical snapshots), and that commit's snapshot dir. Never raises. - """ + *repo_id*, or None. Mirrors that resolution: the :func:`_st_cache_repo_dir` repo dir, its + ``refs/main`` commit (the loader resolves the default revision through that ref, so a + missing/empty/unreadable ref is a cache MISS, not a reason to scan historical + snapshots), and that commit's snapshot dir. Never raises.""" try: repo_dir = _st_cache_repo_dir(repo_id) if repo_dir is None: @@ -2271,25 +2236,23 @@ def _active_snapshot_dir(repo_id: str) -> Optional[Path]: try: commit = (repo_dir / "refs" / "main").read_text(encoding = "utf-8").strip() except OSError: - return None # absent or unreadable: the loader cannot resolve it either + return None if not commit: - return None # empty / whitespace: a partial write, revision unknown + return None snapshot = repo_dir / "snapshots" / commit if not snapshot.is_dir(): - return None # ref recorded but not materialized (partial / pruned) + return None return snapshot except Exception: return None def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: - """Sentence-transformers detection from the local cache, no network call. Models what - an offline ``local_files_only`` load resolves (via :func:`_active_snapshot_dir`), so a - True means that load can succeed: the active ``refs/main`` snapshot must be materialized - and actually loadable (marker plus config plus weights, not the bare ``modules.json`` - the security preflight fetches on its own). None when nothing usable is cached; never - raises. - """ + """Sentence-transformers detection from the local cache, no network. Models what an + offline ``local_files_only`` load resolves (via :func:`_active_snapshot_dir`): the active + ``refs/main`` snapshot must be materialized and loadable (marker plus config plus + weights, not the bare ``modules.json`` the preflight fetches). None when nothing usable + is cached; never raises.""" snapshot = _active_snapshot_dir(repo_id) if snapshot is None: return None @@ -2312,7 +2275,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: """ cache_key = (model_name, hf_token) - # Local paths: the modules.json marker is authoritative for an explicit path, so memoize. + # Local path: modules.json marker is authoritative, so memoize. if is_local_path(model_name): if cache_key in _embedding_detection_cache: return _embedding_detection_cache[cache_key] @@ -2322,19 +2285,16 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: return is_emb if _env_offline(): - # Offline: the local HF cache is the only source -- a network call would only hang - # on DNS and get retried (#6817). Re-probe every call without consulting/populating - # the memo, so a negative never sticks (a model downloaded later must be seen). + # Offline: the local cache is the only source; a network call would hang on DNS + # (#6817). Re-probe every call without the memo so a negative never sticks. snapshot = _active_snapshot_dir(model_name) if snapshot is not None and _snapshot_is_loadable_st_model(snapshot): - return True # self-describing sentence-transformers snapshot (modules.json) - # Trust a positive confirmed online (memo or persisted allowlist) ONLY when the - # active snapshot is materialized with a complete, loadable weight set. A recorded - # True proves model_info() tagged the repo an embedder, not that its files are on - # disk, so gating on the weights avoids saving a model the local_files_only load - # then fails on. The weight gate (not the bare marker) still covers a downloaded - # tag-only embedder (weights, no modules.json, loaded by ST's auto-model fallback) - # while rejecting a config-only or half-sharded snapshot. + return True + # Trust an online-confirmed positive (memo or persisted allowlist) only when the + # active snapshot has a complete weight set: a recorded True proves the Hub tagged it + # an embedder, not that its files are on disk. The weight gate covers a downloaded + # tag-only embedder (no modules.json) while rejecting a config-only or half-sharded + # snapshot. if ( snapshot is not None and _known_embedder(model_name, cache_key) @@ -2343,10 +2303,8 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: return True return False - # Online: the Hub is authoritative for the current revision; the local marker can lag - # it, so model_info() decides and the marker is only a fallback when the Hub is - # unreachable. Only Hub-derived results are memoized, so a transient failure never - # poisons the cache. + # Online: the Hub is authoritative for the current revision (the local marker can lag). + # Only Hub-derived results are memoized, so a transient failure never poisons the cache. if cache_key in _embedding_detection_cache: return _embedding_detection_cache[cache_key] try: @@ -2364,8 +2322,8 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: _embedding_detection_cache[cache_key] = is_emb if is_emb: - # Durably record the positive so a later offline session (memo gone) can still - # recognize a downloaded tag-only embedder, gated on its snapshot. + # Durably record the positive so a later offline session can recognize a + # downloaded tag-only embedder. _persist_embedder(model_name) logger.info( f"Model {model_name} detected as embedding model: " @@ -2377,9 +2335,8 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: except Exception as e: # A permanent Hub error (deleted / gated / bad revision / typo) is authoritative: - # the repo is not a usable embedding model, so return False and let the settings - # route surface its 409. Falling back to a cached modules.json would wrongly pass a - # repo the loader can no longer fetch. Matches the nearby GGUF/vision detectors. + # return False and let the settings route surface its 409, rather than passing a + # repo the loader can no longer fetch. if type(e).__name__ in ( "RepositoryNotFoundError", "GatedRepoError", diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 3b34a809cbe..dc5039d7662 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -249,11 +249,10 @@ def _fetch_security_status( """ from huggingface_hub import model_info as hf_model_info - # Only when the CALLER guarantees a local-only load -- deliberately not keyed off - # hf_env_offline(): this gate is shared by every loader (training, MLX/inference, - # export, ...) and most do not pass local_files_only, so an offline-looking session - # can still fetch an unscanned model through those paths. The bypass must be opted - # into by the callers that actually hold the invariant. + # Skip only when the CALLER guarantees a local-only load -- deliberately not keyed off + # hf_env_offline(): this gate is shared by every loader (training, MLX, export, ...) and + # most do not pass local_files_only, so an offline-looking session could still fetch an + # unscanned model. The bypass must be opted into. if local_only_load: logger.debug( "HF security scan skipped for '%s': caller loads local-only; failing open.", @@ -302,10 +301,9 @@ def evaluate_file_security( is root-level there and blocks, and an index inside it is honored when scoping shards. ``local_only_load`` lets a caller skip the Hub round-trip when it GUARANTEES the load - cannot fetch -- e.g. the RAG embedder, which passes ``local_files_only`` to - SentenceTransformer from the same predicate. Pass it only with that guarantee: claiming - local-only while the loader can still fetch disables the gate. Default False, since most - callers (training, MLX/inference, export) do not constrain their loaders. + cannot fetch (e.g. the RAG embedder, which passes ``local_files_only`` to + SentenceTransformer from the same predicate). Pass it only with that guarantee; claiming + local-only while the loader can still fetch disables the gate. Default False. """ # 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/. diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index f70dd111429..c9da0a563a5 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -171,11 +171,10 @@ def _offline_flag(name: str) -> bool: def hf_env_offline() -> bool: - """True when either HF offline env var is truthy (strip+lower, on/true/yes/1). + """True when HF_HUB_OFFLINE or TRANSFORMERS_OFFLINE is truthy (strip+lower, on/true/yes/1). - The user's *intent* to work offline, broader than what ``huggingface_hub`` enforces - (it honors ``HF_HUB_OFFLINE`` but ignores ``TRANSFORMERS_OFFLINE``). This alone does NOT - stop a fetch -- callers needing that guarantee must pass ``local_files_only = - hf_env_offline()`` to the loader (as ``core/rag/embeddings.py`` does). + The user's intent to work offline, broader than ``huggingface_hub`` (which ignores + ``TRANSFORMERS_OFFLINE``). Does NOT itself stop a fetch -- callers needing that must pass + ``local_files_only = hf_env_offline()`` to the loader. """ return _offline_flag("HF_HUB_OFFLINE") or _offline_flag("TRANSFORMERS_OFFLINE") From b83c15d4b3d41cebed8698e23de1f751e40e7471 Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sun, 19 Jul 2026 21:40:40 +0300 Subject: [PATCH 28/77] studio: pin embedder verdicts to a revision, canonicalize default aliases - A persisted verdict recorded that the Hub tagged ONE revision an embedder, but was stored per repo. Once refs/main advanced to a complete but non-embedding Transformer snapshot, the offline path still returned True: the settings route accepted the updated model without force and RAG could silently load it as an embedder. Verdicts now carry the commit they were confirmed at and are trusted only while the active revision matches. One confirmed before the repo was cached has no revision to compare, so the first revision observed afterwards is pinned then -- which is what lets a later advance be caught. The persisted file gains a {id: commit} form and still reads the previous list format. - tokenizer_config.json no longer counts as a tokenizer asset. It only DESCRIBES a tokenizer, so a snapshot with config, weights and just that file passed validation and then failed AutoTokenizer.from_pretrained(local_files_only=True) at first indexing for common BERT/GPT-style models. - A casing-only alias of the default is canonicalized to the default up front. Repo ids are case-insensitive but every gate here compares exact strings, so saving "Unsloth/bge-m3" against a default of "unsloth/bge-m3" ran the verification and scan for a custom model and then persisted an override -- after which later changes to the configured default stopped applying. - verify_import_hoist.py replays __all__ assignments in order instead of unioning them. Only the final value exports anything, so a later plain "=" that drops a name must leave its import counted as unused; "+=" still extends, and an unreadable rebind keeps the earlier names rather than flagging real re-exports. --- scripts/verify_import_hoist.py | 53 ++++++++-- studio/backend/routes/settings.py | 26 +++-- .../test_embedding_model_security_gate.py | 18 ++++ .../tests/test_is_embedding_model_offline.py | 76 ++++++++++++- studio/backend/utils/models/model_config.py | 100 ++++++++++++++---- 5 files changed, 228 insertions(+), 45 deletions(-) diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index d46633c066a..817fcc82b32 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -448,30 +448,48 @@ def _legb_chain(scope: Scope) -> list[Scope]: def _collect_dunder_all(tree: ast.Module) -> set[str]: - """Module-level ``__all__`` string entries (``= [...]``, ``+= [...]``, or an - annotated assign). + """The FINAL module-level ``__all__`` string entries (``= [...]``, ``+= [...]``, or + an annotated assign). A name listed in ``__all__`` is a public re-export, which is a real use of the import that binds it. ``__all__`` entries are string constants, not ``Name`` loads, so the load-based use scan never sees them -- without this a package ``__init__`` that adds ``from .x import y`` purely to re-export ``y`` looks like an unused hoist and trips ``HOISTED-IMPORT-UNUSED``. + + Only the final value counts, so assignments are replayed in order: a plain ``=`` + REPLACES the list and ``+=`` extends it. Unioning every assignment instead would + let ``__all__ = ["y"]`` followed by ``__all__ = []`` still mark ``y`` used, so a + genuinely unused hoist would slip through. """ names: set[str] = set() for node in tree.body: if isinstance(node, ast.Assign): - targets = node.targets - elif isinstance(node, (ast.AnnAssign, ast.AugAssign)): - targets = [node.target] + targets, replaces = node.targets, True + elif isinstance(node, ast.AnnAssign): + targets, replaces = [node.target], True + elif isinstance(node, ast.AugAssign): + targets, replaces = [node.target], False else: continue if not any(isinstance(t, ast.Name) and t.id == "__all__" for t in targets): continue value = node.value - if isinstance(value, (ast.List, ast.Tuple, ast.Set)): + entries: set[str] = set() + literal = isinstance(value, (ast.List, ast.Tuple, ast.Set)) + if literal: for elt in value.elts: if isinstance(elt, ast.Constant) and isinstance(elt.value, str): - names.add(elt.value) + entries.add(elt.value) + if replaces: + if not literal: + # Rebound to something we cannot read statically (a call, a name, a + # comprehension). Its contents are unknown, so keep what we had rather + # than claim the list is empty and flag real re-exports as unused. + continue + names = entries + else: + names |= entries return names @@ -758,6 +776,27 @@ def used_targets(analysis) -> set[str]: 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n', "BLOCKER", ), + "reassigned_all_drops_the_reexport": ( + # only the FINAL __all__ exports anything: a later plain "=" REPLACES the list, + # so b is not re-exported and its import is a genuine unused hoist. Unioning + # every assignment would have marked it used and let this through. + 'from pkg import a\n__all__ = ["a"]\n', + 'from pkg import a\nfrom pkg import b\n__all__ = ["a", "b"]\n__all__ = ["a"]\n', + "BLOCKER", + ), + "augmented_all_extends_the_reexport": ( + # "+=" extends rather than replaces, so b IS re-exported -> must not block + 'from pkg import a\n__all__ = ["a"]\n', + 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n__all__ += ["b"]\n', + None, + ), + "unreadable_all_keeps_earlier_reexports": ( + # rebound to something not statically readable: contents unknown, so keep what + # we had rather than flag a real re-export as unused + 'from pkg import a\n__all__ = ["a"]\n', + 'from pkg import a\nfrom pkg import b\n__all__ = ["a", "b"]\n__all__ = sorted(__all__)\n', + None, + ), } diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 7ebd345382e..85e0864c10f 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -406,6 +406,15 @@ def update_embedding_model( log = logger, ) from exc hf_token = (payload.hf_token or "").strip() or None + # Repo ids are case-insensitive, so a casing-only alias of the default IS the + # default. Canonicalize before any comparison: set_rag_embedding_model() and the + # gates below all compare exact strings, so leaving the alias in place would run + # the verification/scan for a custom model AND persist an override, after which + # later changes to the configured default would stop applying and the UI would + # report a custom selection. + _default_model = default_embedding_model() + if not is_local_path(model) and model.casefold() == _default_model.casefold(): + model = _default_model # 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. @@ -415,9 +424,7 @@ def update_embedding_model( # (inert) from effective_gguf_repo(), so scanning the ST repo's pickle here would # wrongly reject a custom repo whose GGUF companion is clean; the GGUF availability # checks below cover that path instead. - scan_st_pickle = ( - model != default_embedding_model() and not is_local_gguf and not _llama_backend_active() - ) + scan_st_pickle = model != _default_model and not is_local_gguf and not _llama_backend_active() # Read the offline state ONCE: the module probe and the scan must agree, and # _hf_offline_if_dns_dead() can flip the vars between two reads. from utils.utils import hf_env_offline @@ -461,7 +468,7 @@ def update_embedding_model( "cannot be used as the embedding model." ), ) - if model != default_embedding_model() and not payload.force and not is_local_gguf: + if model != _default_model and not payload.force and not is_local_gguf: from core.rag import config as rag_config # A GGUF-named repo on the llama-server backend is loaded from its .gguf @@ -485,16 +492,13 @@ def update_embedding_model( # offline ST load resolves the cache by exact case (against ST_HOME if set, else the Hub # cache). No-op when nothing case-matching is cached. Three cases are left alone: # - # * the default -- rewriting it would make set_rag_embedding_model()'s exact-string - # default comparison treat it as an override, so later default changes stop applying; + # * the default, including a casing-only alias of it -- rewriting it would make + # set_rag_embedding_model()'s exact-string default comparison treat it as an + # override, so later default changes stop applying; # * a local path -- loaded from disk; a cache-collision recasing would stop resolving to it; # * the llama-server backend -- loads a GGUF companion from the HUB cache, so an ST_HOME # spelling could pick a repo _hf_gguf_backend_error() never validated. - if ( - model != default_embedding_model() - and not is_local_path(model) - and not _llama_backend_active() - ): + if model != _default_model and not is_local_path(model) and not _llama_backend_active(): model = resolve_st_cached_repo_id_case(model) set_rag_embedding_model(model) logger.info( diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index 77d5652cb52..9a40df63d96 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -311,6 +311,24 @@ def _must_not_run(m): assert saved.get("model") == str(local_dir) +def test_casing_alias_of_the_default_is_stored_as_the_default(client, monkeypatch): + # Repo ids are case-insensitive, but set_rag_embedding_model() compares exact + # strings: persisting "Unsloth/default-embed" against a default of + # "unsloth/default-embed" would store a custom override, so later changes to the + # configured default would stop applying and the UI would show a custom selection. + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) + import utils.models as _models + + def _must_not_run(m): + raise AssertionError("a casing alias of the default must not be ST-normalized") + + monkeypatch.setattr(_models, "resolve_st_cached_repo_id_case", _must_not_run) + r = c.put("/embedding-model", json = {"embedding_model": "Unsloth/Default-Embed"}) + assert r.status_code == 200 + assert saved.get("model") == "unsloth/default-embed" # the canonical default + + def test_llama_backend_model_is_not_casing_normalized(monkeypatch): # The llama backend fetches a GGUF companion from the HUB cache, so an ST_HOME recasing # could point at a GGUF repo _hf_gguf_backend_error() never validated. diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 24e40a3c733..42f25b2c120 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -271,13 +271,24 @@ def test_marker_rejects_weights_without_tokenizer(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache("org/model") is False -@pytest.mark.parametrize("tok_file", ["tokenizer_config.json", "vocab.txt", "spiece.model"]) +@pytest.mark.parametrize("tok_file", ["tokenizer.json", "vocab.txt", "spiece.model"]) def test_marker_accepts_alternate_tokenizer_assets(tmp_path, monkeypatch, tok_file): - # Permissive union: any one recognized tokenizer asset suffices. + # Permissive union: any one recognized tokenizer ARTIFACT suffices. _cache_repo_with_files(tmp_path, monkeypatch, "model.safetensors", tok_file, tokenizer = False) assert mc._embedding_marker_in_hf_cache("org/model") is True +def test_marker_rejects_tokenizer_config_without_a_vocabulary(tmp_path, monkeypatch): + # tokenizer_config.json only DESCRIBES a tokenizer. Without vocab.txt / + # vocab.json+merges.txt / tokenizer.json, AutoTokenizer.from_pretrained( + # local_files_only=True) fails for common BERT/GPT-style models -- so accepting + # it would validate the save and fail at first indexing. + _cache_repo_with_files( + tmp_path, monkeypatch, "model.safetensors", "tokenizer_config.json", tokenizer = False + ) + assert mc._embedding_marker_in_hf_cache("org/model") is False + + def _cache_repo_with_files( tmp_path, monkeypatch, @@ -596,6 +607,61 @@ def _info(model_name, token = None): assert mc.is_embedding_model("BAAI/model") is True +def _tag_only_repo( + tmp_path, + monkeypatch, + commit, + repo_id = "org/emb", +): + """A complete, loadable snapshot with NO modules.json -- the tag-only embedder shape, + which the marker cannot recognize, so detection falls through to the recorded verdict.""" + hf_root = tmp_path / "hf" + repo = hf_root / f"models--{repo_id.replace('/', '--')}" + snap = repo / "snapshots" / commit + snap.mkdir(parents = True) + (snap / "config.json").write_text("{}") + (snap / "tokenizer.json").write_text("{}") + (snap / "model.safetensors").write_bytes(b"\0") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + + +def test_online_verdict_still_applies_at_the_confirmed_revision(tmp_path, monkeypatch): + # Control for the test below: same shape, matching revision -> still trusted. + monkeypatch.setattr(mc, "_load_persisted_embedders", lambda: {"org/emb": "commit_a"}) + monkeypatch.setattr(mc, "_persist_embedder", lambda name, commit: None) + _tag_only_repo(tmp_path, monkeypatch, "commit_a") + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + assert mc.is_embedding_model("org/emb") is True + + +def test_online_verdict_stops_applying_when_the_revision_advances(tmp_path, monkeypatch): + # A verdict records that the Hub tagged ONE revision an embedder. Once refs/main + # advances -- e.g. to a complete but non-embedding Transformer snapshot -- the old + # positive says nothing about what would now load, so the offline path must not + # accept it without force and let RAG silently load a non-embedder. + monkeypatch.setattr(mc, "_load_persisted_embedders", lambda: {"org/emb": "commit_a"}) + monkeypatch.setattr(mc, "_persist_embedder", lambda name, commit: None) + _tag_only_repo(tmp_path, monkeypatch, "commit_b") # refs/main moved on + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + assert mc.is_embedding_model("org/emb") is False + + +def test_verdict_without_a_revision_is_pinned_to_the_one_it_meets(tmp_path, monkeypatch): + # Confirmed online before the repo was cached, so there was no revision to record. + # The first revision observed afterwards is the one it was about: trust it, and pin + # it so a LATER advance is caught. + pinned: dict = {} + monkeypatch.setattr(mc, "_load_persisted_embedders", lambda: {"org/emb": None}) + monkeypatch.setattr(mc, "_persist_embedder", lambda name, commit: pinned.update({name: commit})) + _tag_only_repo(tmp_path, monkeypatch, "first_commit") + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + assert mc.is_embedding_model("org/emb") is True + assert pinned == {"org/emb": "first_commit"} + + def test_persist_embedder_concurrent_writes_keep_every_verdict(tmp_path, monkeypatch): # Concurrent confirmations must not drop entries: serialized writes + per-thread temp files. import threading @@ -605,7 +671,7 @@ def test_persist_embedder_concurrent_writes_keep_every_verdict(tmp_path, monkeyp def _writer(name): barrier.wait() - mc._persist_embedder(name) + mc._persist_embedder(name, "aaa") threads = [threading.Thread(target = _writer, args = (n,)) for n in names] for t in threads: @@ -614,7 +680,7 @@ def _writer(name): t.join() persisted = mc._load_persisted_embedders() - assert {n.casefold() for n in names} <= persisted + assert {n.casefold() for n in names} <= set(persisted) def test_persist_embedder_is_best_effort_when_home_unwritable(tmp_path, monkeypatch): @@ -630,7 +696,7 @@ def _info(model_name, token = None): _fake_hf_model_info(monkeypatch, _info) assert mc.is_embedding_model("org/emb") is True - assert mc._load_persisted_embedders() == set() + assert mc._load_persisted_embedders() == {} def test_offline_cached_non_st_returns_false_without_network(tmp_path, monkeypatch): diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 5b1de808564..9a7f09c24ca 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2073,6 +2073,10 @@ def download_gguf_file( # Cache embedding detection per session to avoid repeated HF API calls _embedding_detection_cache: Dict[tuple, bool] = {} +# refs/main commit a positive verdict was confirmed AT (None when the repo was not +# cached then). A verdict says the Hub tagged THAT revision an embedder, so it must +# not be trusted once the active revision moves on. +_embedding_verdict_commit: Dict[tuple, Optional[str]] = {} def _persisted_embedders_path() -> Path: @@ -2095,47 +2099,77 @@ def _verdict_key(model_name: str) -> str: return model_name.casefold() -def _load_persisted_embedders() -> set: - """Case-folded repo ids recorded as embedders in a prior session; empty set on any - error (missing / empty / corrupt read as "nothing recorded"). Never raises.""" +def _load_persisted_embedders() -> dict: + """Case-folded repo id -> the refs/main commit the verdict was confirmed at (None when + unknown). Empty on any error (missing / empty / corrupt read as "nothing recorded"). + A legacy list file is read as ids with an unknown revision. Never raises.""" try: with open(_persisted_embedders_path(), encoding = "utf-8") as fh: data = json.load(fh) - if isinstance(data, list): - return {_verdict_key(name) for name in data if isinstance(name, str)} + if isinstance(data, dict): + return { + _verdict_key(name): (commit if isinstance(commit, str) and commit else None) + for name, commit in data.items() + if isinstance(name, str) + } + if isinstance(data, list): # pre-revision-pinning format + return {_verdict_key(name): None for name in data if isinstance(name, str)} except Exception: pass - return set() + return {} -def _persist_embedder(model_name: str) -> None: - """Record *model_name* as an online-confirmed embedder, best-effort. Positive Hub - verdicts only, case-folded, serialized under ``_persist_lock`` via a per-thread temp - file. Any failure is swallowed -- persistence is an optimization, not a correctness - requirement.""" +def _persist_embedder(model_name: str, commit: Optional[str]) -> None: + """Record *model_name* as an online-confirmed embedder AT *commit*, best-effort. + Positive Hub verdicts only, case-folded, serialized under ``_persist_lock`` via a + per-thread temp file. Any failure is swallowed -- persistence is an optimization, not a + correctness requirement.""" try: key = _verdict_key(model_name) with _persist_lock: current = _load_persisted_embedders() - if key in current: + if current.get(key) == commit and key in current: return - current.add(key) + current[key] = commit path = _persisted_embedders_path() path.parent.mkdir(parents = True, exist_ok = True) tmp = path.with_name(path.name + f".{os.getpid()}.{threading.get_ident()}.tmp") with open(tmp, "w", encoding = "utf-8") as fh: - json.dump(sorted(current), fh) + json.dump(current, fh, sort_keys = True) os.replace(tmp, path) except Exception as e: logger.debug(f"Could not persist embedder verdict for {model_name}: {e}") -def _known_embedder(model_name: str, cache_key: tuple) -> bool: - """True when *model_name* was confirmed an embedder online (session memo or persisted - allowlist). Callers still gate this on a materialized snapshot.""" +def _known_embedder(model_name: str, cache_key: tuple, active_commit: Optional[str]) -> bool: + """True when *model_name* was confirmed an embedder online FOR THE ACTIVE REVISION. + + A verdict records that the Hub tagged one particular revision an embedder. If + ``refs/main`` has since advanced -- say to a complete but non-embedding Transformer + snapshot -- the recorded positive says nothing about the revision that would now load, + so it is not trusted and the caller falls through to re-verification. + + A verdict confirmed before the repo was cached has no revision to compare; the first + revision observed afterwards is the one it was about, so it is trusted and PINNED here, + which is what lets a later advance be caught. Callers still gate on a materialized + snapshot. + """ + known, recorded = False, None if _embedding_detection_cache.get(cache_key) is True: + known, recorded = True, _embedding_verdict_commit.get(cache_key) + else: + persisted = _load_persisted_embedders() + key = _verdict_key(model_name) + if key in persisted: + known, recorded = True, persisted[key] + if not known: + return False + if recorded is None: + if active_commit is not None: + _embedding_verdict_commit[cache_key] = active_commit + _persist_embedder(model_name, active_commit) return True - return _verdict_key(model_name) in _load_persisted_embedders() + return recorded == active_commit # Base-model weight files the default Torch SentenceTransformer backend consumes: @@ -2147,10 +2181,14 @@ def _known_embedder(model_name: str, cache_key: tuple) -> bool: # A Transformer module also loads an AutoTokenizer, so weights without any tokenizer asset # fail the local_files_only load. Any ONE of these suffices -- a permissive union, so only # a genuinely tokenizer-less partial download is rejected. +# Real tokenizer ARTIFACTS -- a serialized tokenizer or its vocabulary. Note that +# tokenizer_config.json is deliberately absent: it only DESCRIBES a tokenizer, so a +# snapshot carrying it without vocab.txt / vocab.json+merges.txt / tokenizer.json +# still fails AutoTokenizer.from_pretrained(local_files_only=True) for common +# BERT/GPT-style models -- validating here and failing at first indexing. _ST_TOKENIZER_FILES = frozenset( { "tokenizer.json", - "tokenizer_config.json", "tokenizer.model", "vocab.txt", "vocab.json", @@ -2223,6 +2261,20 @@ def _snapshot_is_loadable_st_model(snap: Path) -> bool: return False +def _active_commit(repo_id: str) -> Optional[str]: + """The commit ``refs/main`` resolves to for *repo_id*, or None when not cached / + unreadable. This is the revision an offline load would use, so it is what an + online verdict has to be pinned to.""" + try: + repo_dir = _st_cache_repo_dir(repo_id) + if repo_dir is None: + return None + commit = (repo_dir / "refs" / "main").read_text(encoding = "utf-8").strip() + return commit or None + except Exception: + return None + + def _active_snapshot_dir(repo_id: str) -> Optional[Path]: """The materialized snapshot dir the offline ``local_files_only`` load resolves for *repo_id*, or None. Mirrors that resolution: the :func:`_st_cache_repo_dir` repo dir, its @@ -2290,6 +2342,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: snapshot = _active_snapshot_dir(model_name) if snapshot is not None and _snapshot_is_loadable_st_model(snapshot): return True + active_commit = _active_commit(model_name) # Trust an online-confirmed positive (memo or persisted allowlist) only when the # active snapshot has a complete weight set: a recorded True proves the Hub tagged it # an embedder, not that its files are on disk. The weight gate covers a downloaded @@ -2297,7 +2350,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: # snapshot. if ( snapshot is not None - and _known_embedder(model_name, cache_key) + and _known_embedder(model_name, cache_key, active_commit) and _snapshot_has_complete_weights(snapshot) ): return True @@ -2323,8 +2376,11 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: _embedding_detection_cache[cache_key] = is_emb if is_emb: # Durably record the positive so a later offline session can recognize a - # downloaded tag-only embedder. - _persist_embedder(model_name) + # downloaded tag-only embedder -- pinned to the revision it was confirmed at, + # so it stops applying once refs/main advances. + confirmed_at = _active_commit(model_name) + _embedding_verdict_commit[cache_key] = confirmed_at + _persist_embedder(model_name, confirmed_at) logger.info( f"Model {model_name} detected as embedding model: " f"pipeline_tag={pipeline_tag}, " From 35113e68ce0da0d4b8e64cd39aa782ef8563086c Mon Sep 17 00:00:00 2001 From: Hakan Baysal Date: Sun, 19 Jul 2026 22:10:30 +0300 Subject: [PATCH 29/77] studio: validate the real ST load root, and pin verdicts to the Hub revision Four ways the offline probe still disagreed with what the loader does: - Verdicts were pinned to the LOCAL refs/main, but model_info() describes the current HUB revision. With a stale cache the two differ, so an older snapshot nobody verified was allowlisted. The pin is now info.sha, taken from the ModelInfo that produced the positive. A verdict carrying no revision (a legacy entry) is no longer trusted at all -- trusting it meant pinning whatever happened to be cached, which is the same bug; the next online check re-records it properly. - config, tokenizer and weights had to exist somewhere in the snapshot, not together. modules.json can send SentenceTransformer at 0_Transformer/, which is loaded FROM that directory, so a cache with the config at the root and only 0_Transformer/model.safetensors passed and then failed the local-only load. Each directory is now checked as a complete load root, which covers both the plain HF layout and the ST module layout. - vocab.json and merges.txt counted independently, but BPE needs the pair unless a serialized tokenizer.json is present, so half a pair validated and then failed AutoTokenizer.from_pretrained(local_files_only=True). - A slashless short name like all-MiniLM-L6-v2 is a supported ST alias that the loader resolves through the sentence-transformers/ organization, so its snapshot is cached under that full id. Probing only the bare name reported a miss and 409'd a model that was cached and loadable; the bare id is still tried first, matching the loader's own order. --- .../tests/test_is_embedding_model_offline.py | 110 ++++++++++++++--- studio/backend/utils/models/model_config.py | 116 +++++++++++------- 2 files changed, 164 insertions(+), 62 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 42f25b2c120..48b473e55d8 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -355,15 +355,63 @@ def test_marker_rejects_incomplete_shard_set(tmp_path, monkeypatch, weights): "model-00002-of-00002.safetensors", "model.safetensors.index.json", ), - ("0_Transformer/model.safetensors",), # weight in a module dir + ( # ST module dir that is itself a complete load root + "0_Transformer/model.safetensors", + "0_Transformer/config.json", + "0_Transformer/tokenizer.json", + ), ], ) def test_marker_accepts_recognized_torch_weights(tmp_path, monkeypatch, weights): - # Must not over-reject: single bin, complete sharded set, and a weight in a module dir. + # Must not over-reject: single bin, complete sharded set, and a complete module dir. _cache_repo_with_files(tmp_path, monkeypatch, *weights) assert mc._embedding_marker_in_hf_cache("org/model") is True +def test_marker_rejects_weights_split_from_their_config(tmp_path, monkeypatch): + # modules.json sends SentenceTransformer at 0_Transformer/, which is loaded FROM that + # directory. Weights there with the config and tokenizer only at the root is a partial + # cache: a scattered any-of check would pass it and the local-only load would fail. + _cache_repo_with_files(tmp_path, monkeypatch, "0_Transformer/model.safetensors") + assert mc._embedding_marker_in_hf_cache("org/model") is False + + +@pytest.mark.parametrize( + ("tok_files", "expected"), + [ + (("vocab.json", "merges.txt"), True), # the complete BPE pair + (("vocab.json",), False), # half a pair: AutoTokenizer still fails + (("merges.txt",), False), + (("vocab.json", "tokenizer.json"), True), # serialized tokenizer is self-contained + ], +) +def test_marker_requires_both_bpe_tokenizer_files(tmp_path, monkeypatch, tok_files, expected): + # BPE needs vocab.json AND merges.txt; either alone would validate here and then fail + # AutoTokenizer.from_pretrained(local_files_only=True) at first indexing. + _cache_repo_with_files(tmp_path, monkeypatch, "model.safetensors", *tok_files, tokenizer = False) + assert mc._embedding_marker_in_hf_cache("org/model") is expected + + +def test_short_name_resolves_through_the_sentence_transformers_org(tmp_path, monkeypatch): + # "all-MiniLM-L6-v2" is a supported ST alias: the loader falls back to the + # sentence-transformers/ org, so the snapshot is cached under that full id. Probing only + # the bare name would report a miss and 409 a model that is cached and loadable. + hf_root = tmp_path / "hf" + _st_snapshot(hf_root, "models--sentence-transformers--all-MiniLM-L6-v2") + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + assert mc._embedding_marker_in_hf_cache("all-MiniLM-L6-v2") is True + + +def test_short_name_prefers_a_bare_cache_dir_when_present(tmp_path, monkeypatch): + # The bare id is tried first, matching the loader's own resolution order. + hf_root = tmp_path / "hf" + _st_snapshot(hf_root, "models--all-MiniLM-L6-v2") + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + assert mc._st_cache_repo_dir("all-MiniLM-L6-v2") == hf_root / "models--all-MiniLM-L6-v2" + + def _case_sensitive_fs(tmp_path) -> bool: probe = tmp_path / "_CaseProbe" probe.mkdir() @@ -501,7 +549,7 @@ def test_offline_retains_online_confirmed_positive(tmp_path, monkeypatch): _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") def _info(model_name, token = None): - return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) assert mc.is_embedding_model("org/gte-modernbert") is True @@ -517,7 +565,7 @@ def test_offline_metadata_only_positive_not_trusted_without_cache(monkeypatch): _no_cache(monkeypatch) def _info(model_name, token = None): - return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) assert mc.is_embedding_model("org/uncached-embedder") is True @@ -533,7 +581,7 @@ def test_offline_detects_persisted_tag_only_embedder_after_restart(tmp_path, mon _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") def _info(model_name, token = None): - return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) assert mc.is_embedding_model("org/gte-modernbert") is True @@ -551,7 +599,7 @@ def test_offline_persisted_verdict_not_trusted_when_uncached(tmp_path, monkeypat _no_cache(monkeypatch) def _info(model_name, token = None): - return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) assert mc.is_embedding_model("org/uncached-embedder") is True @@ -577,7 +625,7 @@ def test_offline_persisted_verdict_not_trusted_when_snapshot_partial(tmp_path, m monkeypatch.setattr(mc, "_st_cache_roots", lambda: [cache_root]) def _info(model_name, token = None): - return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) assert mc.is_embedding_model("org/partial") is True @@ -596,7 +644,7 @@ def test_offline_persisted_verdict_matches_across_casing(tmp_path, monkeypatch): _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "BAAI/model") def _info(model_name, token = None): - return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) assert mc.is_embedding_model("baai/model") is True @@ -649,17 +697,41 @@ def test_online_verdict_stops_applying_when_the_revision_advances(tmp_path, monk assert mc.is_embedding_model("org/emb") is False -def test_verdict_without_a_revision_is_pinned_to_the_one_it_meets(tmp_path, monkeypatch): - # Confirmed online before the repo was cached, so there was no revision to record. - # The first revision observed afterwards is the one it was about: trust it, and pin - # it so a LATER advance is caught. - pinned: dict = {} +def test_verdict_without_a_revision_is_not_trusted(tmp_path, monkeypatch): + # A legacy entry from before revision pinning proves nothing about what would load + # now, so it must not be trusted -- trusting it and pinning whatever happens to be in + # the local cache would allowlist a revision nobody verified. The next online check + # re-records it against info.sha. monkeypatch.setattr(mc, "_load_persisted_embedders", lambda: {"org/emb": None}) - monkeypatch.setattr(mc, "_persist_embedder", lambda name, commit: pinned.update({name: commit})) - _tag_only_repo(tmp_path, monkeypatch, "first_commit") + _tag_only_repo(tmp_path, monkeypatch, "some_commit") monkeypatch.setenv("HF_HUB_OFFLINE", "1") + assert mc.is_embedding_model("org/emb") is False + + +def test_verdict_pins_the_hub_revision_not_the_cached_one(tmp_path, monkeypatch): + # model_info() describes the CURRENT Hub revision. With a stale local cache the two + # differ, and pinning refs/main would allowlist the older cached snapshot even though + # its metadata was never verified. + recorded: dict = {} + monkeypatch.setattr( + mc, "_persist_embedder", lambda name, commit: recorded.update({name: commit}) + ) + _tag_only_repo(tmp_path, monkeypatch, "stale_local_commit") + + def _info(model_name, token = None): + return types.SimpleNamespace( + tags = ["feature-extraction"], pipeline_tag = None, sha = "hub_head_commit" + ) + + _fake_hf_model_info(monkeypatch, _info) assert mc.is_embedding_model("org/emb") is True - assert pinned == {"org/emb": "first_commit"} + assert recorded == {"org/emb": "hub_head_commit"} + + # ...and offline that verdict does NOT vouch for the stale cached revision. + mc._embedding_detection_cache.clear() + monkeypatch.setattr(mc, "_load_persisted_embedders", lambda: {"org/emb": "hub_head_commit"}) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + assert mc.is_embedding_model("org/emb") is False def test_persist_embedder_concurrent_writes_keep_every_verdict(tmp_path, monkeypatch): @@ -692,7 +764,7 @@ def test_persist_embedder_is_best_effort_when_home_unwritable(tmp_path, monkeypa _no_cache(monkeypatch) def _info(model_name, token = None): - return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) assert mc.is_embedding_model("org/emb") is True @@ -720,7 +792,7 @@ def test_online_uncached_still_uses_network(monkeypatch): def _info(model_name, token = None): calls.append(model_name) - return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) assert mc.is_embedding_model("org/gte-modernbert") is True @@ -735,7 +807,7 @@ def test_offline_negative_is_not_cached_then_online_detects(monkeypatch): def _info(model_name, token = None): calls.append(model_name) - return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None) + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 9a7f09c24ca..b3789acbe29 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1707,10 +1707,33 @@ def resolve_st_cached_repo_id_case(repo_id: str) -> str: return repo_id +def _st_repo_id_candidates(repo_id: str) -> list[str]: + """Repo ids the ST loader could resolve *repo_id* to, in the order it tries them. + + A slashless short name like ``all-MiniLM-L6-v2`` is a supported Sentence Transformers + alias: the loader falls back to the ``sentence-transformers/`` organization, so the + snapshot is cached under ``models--sentence-transformers--all-MiniLM-L6-v2``. Probing + only the bare name would report a cache miss -- a 409 offline -- for a model that is + cached and loadable.""" + if "/" in repo_id: + return [repo_id] + return [repo_id, f"sentence-transformers/{repo_id}"] + + def _st_cache_repo_dir(repo_id: str) -> Optional[Path]: - """The ONE cache repo dir the ST loader opens for *repo_id*, or None (exact case, then - a deterministic pick). Scoping to one dir avoids judging a duplicate case variant that - is not the one loaded.""" + """The ONE cache repo dir the ST loader opens for *repo_id*, or None -- following the + same short-name alias fallback the loader uses.""" + for candidate in _st_repo_id_candidates(repo_id): + found = _repo_dir_for_exact_id(candidate) + if found is not None: + return found + return None + + +def _repo_dir_for_exact_id(repo_id: str) -> Optional[Path]: + """The cache dir for exactly *repo_id* (exact case, then a deterministic case variant), + or None. Scoping to one dir avoids judging a duplicate case variant that is not the one + loaded.""" prefix = "models--" expected = f"{prefix}{repo_id.replace('/', '--')}" target = expected.lower() @@ -2149,10 +2172,9 @@ def _known_embedder(model_name: str, cache_key: tuple, active_commit: Optional[s snapshot -- the recorded positive says nothing about the revision that would now load, so it is not trusted and the caller falls through to re-verification. - A verdict confirmed before the repo was cached has no revision to compare; the first - revision observed afterwards is the one it was about, so it is trusted and PINNED here, - which is what lets a later advance be caught. Callers still gate on a materialized - snapshot. + A verdict with no recorded revision (a legacy entry from before pinning) proves nothing + about what would load now, so it is NOT trusted -- the next online check re-records it + against info.sha. Callers still gate on a materialized snapshot. """ known, recorded = False, None if _embedding_detection_cache.get(cache_key) is True: @@ -2162,13 +2184,8 @@ def _known_embedder(model_name: str, cache_key: tuple, active_commit: Optional[s key = _verdict_key(model_name) if key in persisted: known, recorded = True, persisted[key] - if not known: + if not known or recorded is None or active_commit is None: return False - if recorded is None: - if active_commit is not None: - _embedding_verdict_commit[cache_key] = active_commit - _persist_embedder(model_name, active_commit) - return True return recorded == active_commit @@ -2182,21 +2199,28 @@ def _known_embedder(model_name: str, cache_key: tuple, active_commit: Optional[s # fail the local_files_only load. Any ONE of these suffices -- a permissive union, so only # a genuinely tokenizer-less partial download is rejected. # Real tokenizer ARTIFACTS -- a serialized tokenizer or its vocabulary. Note that -# tokenizer_config.json is deliberately absent: it only DESCRIBES a tokenizer, so a -# snapshot carrying it without vocab.txt / vocab.json+merges.txt / tokenizer.json -# still fails AutoTokenizer.from_pretrained(local_files_only=True) for common -# BERT/GPT-style models -- validating here and failing at first indexing. -_ST_TOKENIZER_FILES = frozenset( +# A tokenizer counts only as a COMPLETE artifact set. Deliberately absent: +# tokenizer_config.json (only DESCRIBES a tokenizer, supplies no vocabulary) and half a +# BPE pair. Either would validate here and then fail at first indexing. +_ST_TOKENIZER_SOLO_FILES = frozenset( { - "tokenizer.json", - "tokenizer.model", - "vocab.txt", - "vocab.json", - "merges.txt", + "tokenizer.json", # serialized fast tokenizer: self-contained + "tokenizer.model", # sentencepiece "spiece.model", "sentencepiece.bpe.model", + "vocab.txt", # wordpiece (BERT-style) } ) +# BPE needs BOTH halves: vocab.json or merges.txt alone still fails +# AutoTokenizer.from_pretrained(local_files_only=True) unless tokenizer.json is there. +_ST_TOKENIZER_BPE_PAIR = ("vocab.json", "merges.txt") + + +def _names_have_tokenizer(names: set) -> bool: + """True when *names* (one directory's files) supply a usable tokenizer.""" + if names & _ST_TOKENIZER_SOLO_FILES: + return True + return all(part in names for part in _ST_TOKENIZER_BPE_PAIR) def _dir_has_complete_torch_weights(names: set) -> bool: @@ -2220,31 +2244,35 @@ def _dir_has_complete_torch_weights(names: set) -> bool: return False +_ST_CONFIG_FILES = ("config.json", "config_sentence_transformers.json") + + def _snapshot_has_complete_weights(snap: Path) -> bool: - """True when *snap* has a config, a tokenizer asset, and a COMPLETE Torch weight set - (``modules.json`` aside). A tag-only embedder confirmed online has exactly this and ST's - auto-model fallback loads it; a partial download fails here rather than at first - indexing.""" + """True when some LOAD ROOT inside *snap* is complete: a config, a tokenizer and a + COMPLETE Torch weight set, all in the SAME directory. + + Co-location is the point. ``modules.json`` can send SentenceTransformer at a module + subdirectory (``0_Transformer/``), loaded from that directory -- so a snapshot with the + config at the root, a tokenizer elsewhere and ``0_Transformer/model.safetensors`` but no + ``0_Transformer/config.json`` would pass a scattered any-of check and then fail the + local-only load. Per-directory covers both layouts: a plain HF model has all three at + the root, an ST model inside its Transformer module dir.""" try: - if not any( - (snap / name).is_file() for name in ("config.json", "config_sentence_transformers.json") - ): - return False by_dir: dict = {} - has_tokenizer = False for path in snap.rglob("*"): try: - if not path.is_file(): - continue - if path.name in _ST_TOKENIZER_FILES: - has_tokenizer = True - if _ST_WEIGHT_FILE_RE.match(path.name) or path.name.endswith(".index.json"): + if path.is_file(): by_dir.setdefault(path.parent, set()).add(path.name) except OSError: continue - if not has_tokenizer: - return False - return any(_dir_has_complete_torch_weights(names) for names in by_dir.values()) + for names in by_dir.values(): + if not any(cfg in names for cfg in _ST_CONFIG_FILES): + continue + if not _names_have_tokenizer(names): + continue + if _dir_has_complete_torch_weights(names): + return True + return False except OSError: return False @@ -2376,9 +2404,11 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: _embedding_detection_cache[cache_key] = is_emb if is_emb: # Durably record the positive so a later offline session can recognize a - # downloaded tag-only embedder -- pinned to the revision it was confirmed at, - # so it stops applying once refs/main advances. - confirmed_at = _active_commit(model_name) + # downloaded tag-only embedder -- pinned to the revision whose metadata was + # actually verified. That is info.sha, the Hub revision model_info() just + # described, NOT the local cache's refs/main: with a stale cache the two differ + # and pinning the local one would allowlist a snapshot nobody verified. + confirmed_at = getattr(info, "sha", None) or None _embedding_verdict_commit[cache_key] = confirmed_at _persist_embedder(model_name, confirmed_at) logger.info( From 8407e0ae29fcf08f275d4428670e14232a5d833d Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 03:16:38 +0000 Subject: [PATCH 30/77] studio: fail closed for an offline security scan instead of failing open A local_only (offline) load cannot fetch Hugging Face's malware scan, and the previous behaviour skipped the scan and failed OPEN, so a cached repo with a poisoned pickle weight could deserialize under SentenceTransformer(local_files_only=True). Evaluate it fail-CLOSED against the cached files instead: block a base-model pickle weight the load would deserialize (pytorch_model.bin and its shards, in a directory with no safetensors alternative) and allow a pickle-free (safetensors / gguf are inert) cache. A cached pickle model must be reloaded online once to be scanned, or shipped as safetensors. Nothing cached is not a security event. _fetch_security_status no longer needs the local_only_load skip (the offline branch is handled in evaluate_file_security). Adds a regression test covering the safetensors-allow and pickle-block paths with no Hub call. --- studio/backend/tests/test_hf_offline_flags.py | 23 +++- .../backend/utils/security/file_security.py | 121 ++++++++++++++---- 2 files changed, 115 insertions(+), 29 deletions(-) diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index cf48d3363ee..d5bc83b42b9 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -108,13 +108,30 @@ def test_shared_gate_still_scans_when_offline_by_default(monkeypatch, var): assert calls, f"{var} alone must NOT bypass the shared malware gate" -def test_security_scan_short_circuits_for_a_local_only_caller(monkeypatch): +def test_local_only_load_fails_closed_offline_never_hitting_the_hub(monkeypatch, tmp_path): + # A local-only (offline) load never hits the Hub; it is evaluated fail-CLOSED against the + # cached files: a pickle-free (safetensors) cache is allowed, a cached pickle is blocked. + import utils.models.model_config as mc import utils.security.file_security as fs calls: list = [] _fake_hub(monkeypatch, calls) - assert fs._fetch_security_status("org/model", None, True) is None - assert calls == [], "a local-only caller must not hit the Hub" + + safe = tmp_path / "safe" / "aaa" + safe.mkdir(parents = True) + (safe / "model.safetensors").write_bytes(b"\0") + monkeypatch.setattr(mc, "_active_snapshot_dir", lambda name: safe) + allowed = fs.evaluate_file_security("org/model", None, local_only_load = True) + assert calls == [], "a local-only load must not hit the Hub" + assert allowed.blocked is False # safetensors cache is inert + + bad = tmp_path / "bad" / "aaa" + bad.mkdir(parents = True) + (bad / "pytorch_model.bin").write_bytes(b"\0") + monkeypatch.setattr(mc, "_active_snapshot_dir", lambda name: bad) + blocked = fs.evaluate_file_security("org/model", None, local_only_load = True) + assert calls == [], "still no Hub call when blocking" + assert blocked.blocked is True # unscanned cached pickle -> fail closed def test_security_scan_runs_when_online(monkeypatch): diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index dc5039d7662..a8666407cb4 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -15,8 +15,10 @@ * Block whenever ``filesWithIssues`` lists a non-``safe`` level, regardless of ``scansDone`` (often false even for clean repos). Unknown/future levels fail CLOSED (block) so Hub schema drift cannot silently allow a bad verdict; only a - small allowlist of clean / not-yet-scanned levels is non-blocking. The sole - fail-open path is an unavailable status (missing field / offline / error). + small allowlist of clean / not-yet-scanned levels is non-blocking. An + unavailable status (missing field / error) fails open, but an explicit + local-only (offline) load fails CLOSED against the cached files instead: an + unscanned cached pickle is blocked, a pickle-free (safetensors) cache allowed. * Scope to the load-path RCE vector: a root-level (or load-subdir-level), code-executing file. Inert formats (safetensors / gguf / config / text) and subdirectory pickles that no root weight-index references are NOT loaded, so @@ -29,6 +31,7 @@ scanned so a repo cannot dodge the gate by suffixing its name. """ +import re from dataclasses import dataclass, field from typing import Optional @@ -236,29 +239,88 @@ def _load_scan_target(model_name: str, load_subdirs: tuple) -> tuple: return model_name, load_subdirs -def _fetch_security_status( - model_name: str, - hf_token: Optional[str], - local_only_load: bool = False, -): - """``security_repo_status`` (a dict) or None if unavailable. Hub metadata only; - retries once on a transient error, then returns None so the caller fails open. +# Pickle weight formats a from_pretrained load deserializes (the RCE vector); safetensors +# and gguf are inert. Matched by base-model NAME so training_args.bin / optimizer.pt (not +# loaded as model weights) do not trip the offline gate. +_PICKLE_WEIGHT_RE = re.compile( + r"^(model|pytorch_model)(-\d+-of-\d+)?\.(bin|pt|pth|ckpt|pkl|pickle)$" +) - ``local_only_load`` is an explicit promise from the caller that the load it is - gating cannot reach the Hub. See :func:`evaluate_file_security`. + +def _cached_pickle_weight_files(snap) -> list: + """Base-model pickle weight files in the snapshot's module dirs that have NO safetensors + alternative -- i.e. the pickles a from_pretrained load actually deserializes. A dir + holding ``model.safetensors`` loads that (inert) and ignores any pickle sibling.""" + by_dir_pickle: dict = {} + dirs_with_safetensors: set = set() + try: + for path in snap.rglob("*"): + try: + if not path.is_file(): + continue + low = path.name.lower() + if low.endswith(".safetensors"): + dirs_with_safetensors.add(path.parent) + elif _PICKLE_WEIGHT_RE.match(low): + by_dir_pickle.setdefault(path.parent, []).append(path.name) + except OSError: + continue + except OSError: + return [] + hits: set = set() + for directory, names in by_dir_pickle.items(): + if directory not in dirs_with_safetensors: + hits.update(names) + return sorted(hits) + + +def _evaluate_local_only(model_name: str) -> "FileSecurityDecision": + """Fail-CLOSED security decision for an offline (local_files_only) load. + + The Hub scan cannot be fetched offline, so instead of failing OPEN we inspect the cached + files we already have: block the actual RCE vector -- a pickle weight the load + deserializes -- and allow only a pickle-free cache (safetensors / gguf are inert). A + previously-scanned pickle model must be reloaded online once to pass, or shipped as + safetensors. Nothing cached means there is nothing to deserialize, so it is not blocked + (the load fails downstream on its own, which is not a security event). """ - from huggingface_hub import model_info as hf_model_info + try: + from utils.models.model_config import _active_snapshot_dir + snap = _active_snapshot_dir(model_name) + except Exception: + snap = None + if snap is None: + return FileSecurityDecision(model_name, False, reason = "offline; nothing cached to scan") - # Skip only when the CALLER guarantees a local-only load -- deliberately not keyed off - # hf_env_offline(): this gate is shared by every loader (training, MLX, export, ...) and - # most do not pass local_files_only, so an offline-looking session could still fetch an - # unscanned model. The bypass must be opted into. - if local_only_load: - logger.debug( - "HF security scan skipped for '%s': caller loads local-only; failing open.", - model_name, + pickles = _cached_pickle_weight_files(snap) + if not pickles: + return FileSecurityDecision( + model_name, False, reason = "offline; cached weights are pickle-free (inert)" ) - return None + + names = ", ".join(pickles) + logger.warning( + "Blocking offline load of '%s': cached pickle weights cannot be security-scanned " + "offline (%s). Reconnect once to scan, or use safetensors weights.", + model_name, + names, + ) + return FileSecurityDecision( + model_name, + True, + unsafe_files = [{"path": p, "level": "unscanned"} for p in pickles], + reason = ( + "offline: cached pickle weights are unscanned and cannot be verified; " + f"reconnect once to scan, or use safetensors weights ({names})" + ), + ) + + +def _fetch_security_status(model_name: str, hf_token: Optional[str]): + """``security_repo_status`` (a dict) or None if unavailable. Hub metadata only; + retries once on a transient error, then returns None so the caller fails open. + """ + from huggingface_hub import model_info as hf_model_info token_arg = hf_token if hf_token else False last_exc = None @@ -300,10 +362,12 @@ def evaluate_file_security( 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`` lets a caller skip the Hub round-trip when it GUARANTEES the load - cannot fetch (e.g. the RAG embedder, which passes ``local_files_only`` to - SentenceTransformer from the same predicate). Pass it only with that guarantee; claiming - local-only while the loader can still fetch disables the gate. Default False. + ``local_only_load`` marks a load the caller GUARANTEES cannot fetch (e.g. the RAG + embedder, which passes ``local_files_only`` to SentenceTransformer from the same + predicate). It cannot reach the Hub scan, so it is evaluated fail-CLOSED against the + cached files (:func:`_evaluate_local_only`): a cached pickle weight is blocked, a + pickle-free (safetensors) cache is allowed. Pass it only with that guarantee; claiming + local-only while the loader can still fetch changes gate semantics. Default False. """ # 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/. @@ -319,7 +383,12 @@ 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") - status = _fetch_security_status(model_name, hf_token, local_only_load) + # An offline (local-only) load cannot fetch the Hub scan: fail CLOSED against the cache + # instead of skipping the gate, so an unscanned cached pickle cannot deserialize. + 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( model_name, False, reason = "scan unavailable; allowed (fail-open)" From 4d4909659167f8d61d52cd1e76b07832fce082c0 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 04:18:02 +0000 Subject: [PATCH 31/77] studio: only suppress an offline pickle when a loadable safetensors weight exists The offline security gate treated any .safetensors in a directory as covering a pickle weight, so a cache with pytorch_model.bin beside a bare adapter_model.safetensors (or an orphan shard with no index) passed the fail-closed check even though from_pretrained still selects and deserializes the pickle. Require a genuinely loadable safetensors weight -- an unsharded base file or a complete indexed shard set -- before treating the pickle as covered. Also make the import-hoist analyzer preserve uncertainty when __all__ is extended by a value it cannot read statically (__all__ += dynamic()), matching how it already handles an unreadable rebind, so a dynamically-supplied re-export is not flagged HOISTED-IMPORT-UNUSED. --- scripts/verify_import_hoist.py | 61 +++++++++++++------ studio/backend/tests/test_hf_offline_flags.py | 56 ++++++++++++----- .../backend/utils/security/file_security.py | 51 +++++++++++++--- 3 files changed, 127 insertions(+), 41 deletions(-) diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index 817fcc82b32..b5116396df0 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -447,9 +447,9 @@ def _legb_chain(scope: Scope) -> list[Scope]: # ---------------------------------------------------------------- analysis -def _collect_dunder_all(tree: ast.Module) -> set[str]: - """The FINAL module-level ``__all__`` string entries (``= [...]``, ``+= [...]``, or - an annotated assign). +def _collect_dunder_all(tree: ast.Module) -> tuple[set[str], bool]: + """The FINAL module-level ``__all__`` string entries plus an ``opaque`` flag + (``= [...]``, ``+= [...]``, or an annotated assign). A name listed in ``__all__`` is a public re-export, which is a real use of the import that binds it. ``__all__`` entries are string constants, not ``Name`` @@ -461,8 +461,16 @@ def _collect_dunder_all(tree: ast.Module) -> set[str]: REPLACES the list and ``+=`` extends it. Unioning every assignment instead would let ``__all__ = ["y"]`` followed by ``__all__ = []`` still mark ``y`` used, so a genuinely unused hoist would slip through. + + ``opaque`` is True when ``__all__`` is rebound OR extended by a value we cannot read + statically (a call, a name, a comprehension, a spread). The static entry set is then + not known to be exhaustive, so the caller must NOT treat a name's absence from it as + proof the import is unexported -- a re-export supplied dynamically + (``__all__ += _exports()``) would otherwise trip ``HOISTED-IMPORT-UNUSED``. An + unreadable ``=`` and an unreadable ``+=`` set it alike, so the two stay consistent. """ names: set[str] = set() + opaque = False for node in tree.body: if isinstance(node, ast.Assign): targets, replaces = node.targets, True @@ -476,21 +484,25 @@ def _collect_dunder_all(tree: ast.Module) -> set[str]: continue value = node.value entries: set[str] = set() - literal = isinstance(value, (ast.List, ast.Tuple, ast.Set)) - if literal: + readable = True # every element is a statically-known string constant + if isinstance(value, (ast.List, ast.Tuple, ast.Set)): for elt in value.elts: if isinstance(elt, ast.Constant) and isinstance(elt.value, str): entries.add(elt.value) + else: + readable = False # a spread / computed element: unknown extra exports + else: + readable = False # a call, a name, a comprehension: contents unknown + if not readable: + # Contents partly/wholly unknown. Keep the names we can see, but mark the set + # non-exhaustive so the caller preserves uncertainty rather than flagging a + # dynamically-supplied re-export as unused. Applies to both `=` and `+=`. + opaque = True if replaces: - if not literal: - # Rebound to something we cannot read statically (a call, a name, a - # comprehension). Its contents are unknown, so keep what we had rather - # than claim the list is empty and flag real re-exports as unused. - continue names = entries else: names |= entries - return names + return names, opaque def _analyze(src: str): @@ -525,13 +537,18 @@ def _analyze(src: str): # Re-exports count as uses: a name listed in module-level __all__ that is bound # by a module import is deliberately exported, not a dangling hoist. Fold its # targets into the used set so HOISTED-IMPORT-UNUSED does not fire on a - # legitimately-added `from .x import y` in a package __init__. - for _n in _collect_dunder_all(tree): - _bs = module_imports.get(_n) - if _bs: - targets_by_scope.setdefault(module.qualname, set()).update( - x.target for x in _bs if x.target - ) + # legitimately-added `from .x import y` in a package __init__. When __all__ is + # opaque (a dynamically-computed part), the static list is not known to be + # exhaustive, so credit EVERY module import rather than risk flagging a + # dynamically-exported hoist as unused. + _all_names, _all_opaque = _collect_dunder_all(tree) + _reexported = module_imports.values() if _all_opaque else ( + _bs for _n in _all_names if (_bs := module_imports.get(_n)) + ) + for _bs in _reexported: + targets_by_scope.setdefault(module.qualname, set()).update( + x.target for x in _bs if x.target + ) module_dup = { n for n, bs in module.bindings.items() @@ -797,6 +814,14 @@ def used_targets(analysis) -> set[str]: 'from pkg import a\nfrom pkg import b\n__all__ = ["a", "b"]\n__all__ = sorted(__all__)\n', None, ), + "unreadable_augmented_all_keeps_the_reexport": ( + # "+=" extended by a value we cannot read statically: the augmented contents are + # unknown, so a dynamically-supplied re-export must not be flagged as unused + # (mirrors the unreadable "=" case, which also preserves uncertainty) + 'from pkg import a\n__all__ = ["a"]\n', + 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n__all__ += sorted(["b"])\n', + None, + ), } diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index d5bc83b42b9..65e882a1a52 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -117,21 +117,47 @@ def test_local_only_load_fails_closed_offline_never_hitting_the_hub(monkeypatch, calls: list = [] _fake_hub(monkeypatch, calls) - safe = tmp_path / "safe" / "aaa" - safe.mkdir(parents = True) - (safe / "model.safetensors").write_bytes(b"\0") - monkeypatch.setattr(mc, "_active_snapshot_dir", lambda name: safe) - allowed = fs.evaluate_file_security("org/model", None, local_only_load = True) - assert calls == [], "a local-only load must not hit the Hub" - assert allowed.blocked is False # safetensors cache is inert - - bad = tmp_path / "bad" / "aaa" - bad.mkdir(parents = True) - (bad / "pytorch_model.bin").write_bytes(b"\0") - monkeypatch.setattr(mc, "_active_snapshot_dir", lambda name: bad) - blocked = fs.evaluate_file_security("org/model", None, local_only_load = True) - assert calls == [], "still no Hub call when blocking" - assert blocked.blocked is True # unscanned cached pickle -> fail closed + def _evaluate(dir_path): + monkeypatch.setattr(mc, "_active_snapshot_dir", lambda name: dir_path) + decision = fs.evaluate_file_security("org/model", None, local_only_load = True) + assert calls == [], "a local-only load must not hit the Hub" + return decision.blocked + + def _snap(name, files: dict): + d = tmp_path / name / "aaa" + d.mkdir(parents = True) + for fname, body in files.items(): + (d / fname).write_bytes(body if isinstance(body, bytes) else body.encode()) + return d + + # A genuinely inert (unsharded safetensors) cache is allowed. + assert _evaluate(_snap("safe", {"model.safetensors": b"\0"})) is False + + # A cached pickle weight with no loadable safetensors alternative is blocked. + assert _evaluate(_snap("bad", {"pytorch_model.bin": b"\0"})) is True + + # A pickle beside a bare ADAPTER (not a base weight) still loads the pickle -> blocked: + # from_pretrained cannot use adapter_model.safetensors as the base checkpoint. + assert _evaluate(_snap("adapter", { + "pytorch_model.bin": b"\0", "adapter_model.safetensors": b"\0", + })) is True + + # A pickle beside a lone ORPHAN shard (no index) still loads the pickle -> blocked: + # a sharded safetensors load needs the index the loader reads to locate every shard. + assert _evaluate(_snap("orphan", { + "pytorch_model.bin": b"\0", "model-00001-of-00002.safetensors": b"\0", + })) is True + + # A COMPLETE indexed safetensors shard set is what the loader picks instead of the + # pickle -> allowed. + index = '{"weight_map": {"a": "model-00001-of-00002.safetensors", ' \ + '"b": "model-00002-of-00002.safetensors"}}' + assert _evaluate(_snap("sharded", { + "pytorch_model.bin": b"\0", + "model-00001-of-00002.safetensors": b"\0", + "model-00002-of-00002.safetensors": b"\0", + "model.safetensors.index.json": index, + })) is False def test_security_scan_runs_when_online(monkeypatch): diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index a8666407cb4..83fc0d15343 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -247,21 +247,56 @@ def _load_scan_target(model_name: str, load_subdirs: tuple) -> tuple: ) +# A from_pretrained load prefers safetensors over a pickle ONLY when the directory holds a +# safetensors weight it can actually load in its place: an unsharded base file, or an index +# whose every referenced shard is present. A bare adapter (``adapter_model.safetensors``) or +# an orphan shard with no index is NOT a loadable base weight -- the loader falls back to and +# deserializes the pickle, which therefore stays the live RCE vector. +_SAFETENSORS_BASE_UNSHARDED = ("model.safetensors", "pytorch_model.safetensors") +_SAFETENSORS_BASE_INDEX = ("model.safetensors.index.json", "pytorch_model.safetensors.index.json") + + +def _safetensors_index_complete(index_path, present_lower: set) -> bool: + """True when every shard the safetensors index maps is present in the same directory.""" + import json + + try: + weight_map = (json.loads(index_path.read_text(encoding = "utf-8")) or {}).get("weight_map") or {} + except (OSError, ValueError): + return False # unreadable index -> not a usable safetensors set -> keep the pickle blocked + shards = {str(shard).rsplit("/", 1)[-1].lower() for shard in weight_map.values()} + return bool(shards) and shards <= present_lower + + +def _dir_has_loadable_safetensors(files: dict) -> bool: + """True when *files* (lower-name -> Path for one directory) hold a safetensors weight a + from_pretrained load will read INSTEAD of a pickle sibling: an unsharded base file, or a + complete indexed shard set. A bare adapter or an orphan shard does not qualify.""" + if any(name in files for name in _SAFETENSORS_BASE_UNSHARDED): + return True + present_lower = set(files) + for index_name in _SAFETENSORS_BASE_INDEX: + index_path = files.get(index_name) + if index_path is not None and _safetensors_index_complete(index_path, present_lower): + return True + return False + + def _cached_pickle_weight_files(snap) -> list: - """Base-model pickle weight files in the snapshot's module dirs that have NO safetensors - alternative -- i.e. the pickles a from_pretrained load actually deserializes. A dir - holding ``model.safetensors`` loads that (inert) and ignores any pickle sibling.""" + """Base-model pickle weight files in the snapshot's module dirs with NO loadable safetensors + alternative -- i.e. the pickles a from_pretrained load actually deserializes. A dir is + covered only by a safetensors weight the loader would pick instead (unsharded base file or + a complete indexed shard set); a bare adapter or an orphan shard leaves the pickle live.""" by_dir_pickle: dict = {} - dirs_with_safetensors: set = set() + by_dir_files: dict = {} # directory -> {lower-name: Path} try: for path in snap.rglob("*"): try: if not path.is_file(): continue low = path.name.lower() - if low.endswith(".safetensors"): - dirs_with_safetensors.add(path.parent) - elif _PICKLE_WEIGHT_RE.match(low): + by_dir_files.setdefault(path.parent, {})[low] = path + if _PICKLE_WEIGHT_RE.match(low): by_dir_pickle.setdefault(path.parent, []).append(path.name) except OSError: continue @@ -269,7 +304,7 @@ def _cached_pickle_weight_files(snap) -> list: return [] hits: set = set() for directory, names in by_dir_pickle.items(): - if directory not in dirs_with_safetensors: + if not _dir_has_loadable_safetensors(by_dir_files.get(directory, {})): hits.update(names) return sorted(hits) From ba98c68804970d4e7e291bb38b0b2d635e16b6d7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 04:19:09 +0000 Subject: [PATCH 32/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- scripts/verify_import_hoist.py | 6 +- studio/backend/tests/test_hf_offline_flags.py | 56 ++++++++++++++----- .../backend/utils/security/file_security.py | 4 +- 3 files changed, 49 insertions(+), 17 deletions(-) diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index b5116396df0..48beb94bd36 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -542,8 +542,10 @@ def _analyze(src: str): # exhaustive, so credit EVERY module import rather than risk flagging a # dynamically-exported hoist as unused. _all_names, _all_opaque = _collect_dunder_all(tree) - _reexported = module_imports.values() if _all_opaque else ( - _bs for _n in _all_names if (_bs := module_imports.get(_n)) + _reexported = ( + module_imports.values() + if _all_opaque + else (_bs for _n in _all_names if (_bs := module_imports.get(_n))) ) for _bs in _reexported: targets_by_scope.setdefault(module.qualname, set()).update( diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index 65e882a1a52..88354cf0c27 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -138,26 +138,54 @@ def _snap(name, files: dict): # A pickle beside a bare ADAPTER (not a base weight) still loads the pickle -> blocked: # from_pretrained cannot use adapter_model.safetensors as the base checkpoint. - assert _evaluate(_snap("adapter", { - "pytorch_model.bin": b"\0", "adapter_model.safetensors": b"\0", - })) is True + assert ( + _evaluate( + _snap( + "adapter", + { + "pytorch_model.bin": b"\0", + "adapter_model.safetensors": b"\0", + }, + ) + ) + is True + ) # A pickle beside a lone ORPHAN shard (no index) still loads the pickle -> blocked: # a sharded safetensors load needs the index the loader reads to locate every shard. - assert _evaluate(_snap("orphan", { - "pytorch_model.bin": b"\0", "model-00001-of-00002.safetensors": b"\0", - })) is True + assert ( + _evaluate( + _snap( + "orphan", + { + "pytorch_model.bin": b"\0", + "model-00001-of-00002.safetensors": b"\0", + }, + ) + ) + is True + ) # A COMPLETE indexed safetensors shard set is what the loader picks instead of the # pickle -> allowed. - index = '{"weight_map": {"a": "model-00001-of-00002.safetensors", ' \ - '"b": "model-00002-of-00002.safetensors"}}' - assert _evaluate(_snap("sharded", { - "pytorch_model.bin": b"\0", - "model-00001-of-00002.safetensors": b"\0", - "model-00002-of-00002.safetensors": b"\0", - "model.safetensors.index.json": index, - })) is False + index = ( + '{"weight_map": {"a": "model-00001-of-00002.safetensors", ' + '"b": "model-00002-of-00002.safetensors"}}' + ) + assert ( + _evaluate( + _snap( + "sharded", + { + "pytorch_model.bin": b"\0", + "model-00001-of-00002.safetensors": b"\0", + "model-00002-of-00002.safetensors": b"\0", + "model.safetensors.index.json": index, + }, + ) + ) + is False + ) def test_security_scan_runs_when_online(monkeypatch): diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 83fc0d15343..32e64ed4372 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -261,7 +261,9 @@ def _safetensors_index_complete(index_path, present_lower: set) -> bool: import json try: - weight_map = (json.loads(index_path.read_text(encoding = "utf-8")) or {}).get("weight_map") or {} + weight_map = (json.loads(index_path.read_text(encoding = "utf-8")) or {}).get( + "weight_map" + ) or {} except (OSError, ValueError): return False # unreadable index -> not a usable safetensors set -> keep the pickle blocked shards = {str(shard).rsplit("/", 1)[-1].lower() for shard in weight_map.values()} From fbf9977e677452b98c62c62544dc16b903660e8a Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 07:02:43 +0000 Subject: [PATCH 33/77] studio: scope the offline pickle scan to load paths; reset __all__ opacity on rebind Address three review follow-ups on the offline security gate and the import-hoist analyzer: - The offline pickle scan walked the whole snapshot, so a stray pickle in a non-load subdirectory (archive/, nemo/) that SentenceTransformer never deserializes was blocked. Scope it to real from_pretrained load roots -- the snapshot root, or a subdir that holds its own config.json -- matching the online scan's load-path scoping. - _collect_dunder_all kept a sticky opaque flag: a readable replacing assignment after an unreadable extend (__all__ += dynamic(); __all__ = []) still credited every import, so a genuinely unused hoist went unreported. A replacing assignment now resets opacity. - A bare __all__: list[str] annotation has no runtime value; it was treated as an unreadable assignment and marked the export set opaque. Skip annotation-only declarations. --- scripts/verify_import_hoist.py | 41 ++++++++++++++----- studio/backend/tests/test_hf_offline_flags.py | 16 ++++++++ .../backend/utils/security/file_security.py | 18 +++++--- 3 files changed, 59 insertions(+), 16 deletions(-) diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index 48beb94bd36..e21a30d8599 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -462,12 +462,15 @@ def _collect_dunder_all(tree: ast.Module) -> tuple[set[str], bool]: let ``__all__ = ["y"]`` followed by ``__all__ = []`` still mark ``y`` used, so a genuinely unused hoist would slip through. - ``opaque`` is True when ``__all__`` is rebound OR extended by a value we cannot read - statically (a call, a name, a comprehension, a spread). The static entry set is then - not known to be exhaustive, so the caller must NOT treat a name's absence from it as + ``opaque`` is True when the FINAL ``__all__`` value is not statically exhaustive: + it is rebound or extended by a value we cannot read (a call, name, comprehension, + spread). The caller must then NOT treat a name's absence from the static set as proof the import is unexported -- a re-export supplied dynamically - (``__all__ += _exports()``) would otherwise trip ``HOISTED-IMPORT-UNUSED``. An - unreadable ``=`` and an unreadable ``+=`` set it alike, so the two stay consistent. + (``__all__ += _exports()``) would otherwise trip ``HOISTED-IMPORT-UNUSED``. A + replacing ``=`` resets opacity (it discards the prior list AND any prior dynamic + part), so a later readable ``__all__ = [...]`` clears an earlier opaque ``+=``; an + ``+=`` extends, so prior opacity persists. A bare ``__all__: list[str]`` annotation + has no runtime value and is skipped entirely. """ names: set[str] = set() opaque = False @@ -475,6 +478,8 @@ def _collect_dunder_all(tree: ast.Module) -> tuple[set[str], bool]: if isinstance(node, ast.Assign): targets, replaces = node.targets, True elif isinstance(node, ast.AnnAssign): + if node.value is None: + continue # bare `__all__: list[str]` annotation: no runtime value targets, replaces = [node.target], True elif isinstance(node, ast.AugAssign): targets, replaces = [node.target], False @@ -493,15 +498,15 @@ def _collect_dunder_all(tree: ast.Module) -> tuple[set[str], bool]: readable = False # a spread / computed element: unknown extra exports else: readable = False # a call, a name, a comprehension: contents unknown - if not readable: - # Contents partly/wholly unknown. Keep the names we can see, but mark the set - # non-exhaustive so the caller preserves uncertainty rather than flagging a - # dynamically-supplied re-export as unused. Applies to both `=` and `+=`. - opaque = True if replaces: - names = entries + # A replace discards the prior list and its opacity: an unreadable rebind + # is opaque, a readable one is fully known again. + names, opaque = entries, not readable else: + # An extend keeps the prior list, so prior opacity persists; an unreadable + # extend adds names we cannot see. names |= entries + opaque = opaque or not readable return names, opaque @@ -824,6 +829,20 @@ def used_targets(analysis) -> set[str]: 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n__all__ += sorted(["b"])\n', None, ), + "readable_reassign_resets_opacity": ( + # an opaque "+=" then a readable replacing "=": the final __all__ is empty, so a + # hoisted unused import is a genuine bad hoist -- the reset must un-stick opacity + 'from pkg import a\n__all__ = ["a"]\n', + 'from pkg import a\nfrom pkg import b\n__all__ += sorted(["x"])\n__all__ = []\n', + "BLOCKER", + ), + "annotation_only_all_is_not_opaque": ( + # a bare `__all__: list[str]` annotation has no runtime value: it must be skipped, + # not treated as an unreadable assignment, so a new unused hoist is still caught + "from pkg import a\n__all__: list[str]\n", + "from pkg import a\nfrom pkg import b\n__all__: list[str]\n", + "BLOCKER", + ), } diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index 88354cf0c27..091edba2abe 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -187,6 +187,22 @@ def _snap(name, files: dict): is False ) + # A stray pickle in a NON-load subdir (no config.json) is never deserialized by + # from_pretrained, so it must not block -- matching the online scan's load-path scope. + stray = tmp_path / "stray" / "aaa" + (stray / "archive").mkdir(parents = True) + (stray / "model.safetensors").write_bytes(b"\0") + (stray / "archive" / "pytorch_model.bin").write_bytes(b"\0") + assert _evaluate(stray) is False + + # A pickle in a real MODULE load root (a subdir with config.json) and no safetensors + # there is a live vector -> blocked. + modroot = tmp_path / "modroot" / "aaa" + (modroot / "0_Transformer").mkdir(parents = True) + (modroot / "0_Transformer" / "config.json").write_bytes(b"{}") + (modroot / "0_Transformer" / "pytorch_model.bin").write_bytes(b"\0") + assert _evaluate(modroot) is True + def test_security_scan_runs_when_online(monkeypatch): import utils.security.file_security as fs diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 32e64ed4372..8678567052d 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -285,10 +285,13 @@ def _dir_has_loadable_safetensors(files: dict) -> bool: def _cached_pickle_weight_files(snap) -> list: - """Base-model pickle weight files in the snapshot's module dirs with NO loadable safetensors - alternative -- i.e. the pickles a from_pretrained load actually deserializes. A dir is - covered only by a safetensors weight the loader would pick instead (unsharded base file or - a complete indexed shard set); a bare adapter or an orphan shard leaves the pickle live.""" + """Base-model pickle weight files a from_pretrained load actually deserializes: at a real + load root (the snapshot root, or a subdir that is itself a load root -- it holds a + ``config.json``) and with NO loadable safetensors alternative there. A stray pickle in a + non-load subdir (``archive/``, ``nemo/``) that no load opens is not a vector, matching the + online scan's load-path scoping. A dir is covered by a safetensors weight the loader would + pick instead (unsharded base file or a complete indexed shard set); a bare adapter or an + orphan shard leaves the pickle live.""" by_dir_pickle: dict = {} by_dir_files: dict = {} # directory -> {lower-name: Path} try: @@ -306,7 +309,12 @@ def _cached_pickle_weight_files(snap) -> list: return [] hits: set = set() for directory, names in by_dir_pickle.items(): - if not _dir_has_loadable_safetensors(by_dir_files.get(directory, {})): + files = by_dir_files.get(directory, {}) + # from_pretrained deserializes a pickle only at a load root: the snapshot root, or a + # subdir that is itself a load root (has config.json). A pickle elsewhere is unread. + if directory != snap and "config.json" not in files: + continue + if not _dir_has_loadable_safetensors(files): hits.update(names) return sorted(hits) From c96decf3fd7cb790cf70f1e5531ee8cc2728e71d Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 07:22:20 +0000 Subject: [PATCH 34/77] studio: recase slashless ST aliases and accept a pinned embedder after a transient failure Two offline-detection gaps on well-formed input: - resolve_st_cached_repo_id_case bailed on every slashless name, so a differently-cased short alias (all-minilm-l6-v2) validated case-insensitively but was loaded verbatim; the SentenceTransformer loader rewrites it to sentence-transformers/all-minilm-l6-v2 and looks it up case-sensitively, missing the canonical sentence-transformers/all-MiniLM-L6-v2 cache dir. Resolve through _st_cache_repo_dir, which follows the same org alias, and hand back the on-disk casing. - On a transient (non-permanent) Hub failure, is_embedding_model only accepted a cached modules.json marker, so a downloaded tag-only embedder (no modules.json) with a verdict pinned to the active revision was rejected even though the offline branch accepts the identical cache. Mirror the offline branch's pinned-verdict acceptance. --- .../tests/test_is_embedding_model_offline.py | 34 ++++++++++++ studio/backend/utils/models/model_config.py | 55 +++++++++++-------- 2 files changed, 67 insertions(+), 22 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 48b473e55d8..3a63a7e2d3b 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -450,6 +450,22 @@ def test_st_casing_noop_when_uncached(tmp_path, monkeypatch): assert mc.resolve_st_cached_repo_id_case("org/not-cached") == "org/not-cached" +def test_st_casing_recases_a_slashless_alias_to_the_org_dir(tmp_path, monkeypatch): + # The loader rewrites a slashless short name to sentence-transformers/ and looks it + # up CASE-SENSITIVELY, so a differently-cased alias must resolve to the canonical org cache + # dir; returning it verbatim makes the exact-case offline load miss a genuinely cached model. + if not _case_sensitive_fs(tmp_path): + pytest.skip("casing only diverges on a case-sensitive filesystem") + hf_root = tmp_path / "hf" + (hf_root / "models--sentence-transformers--all-MiniLM-L6-v2").mkdir(parents = True) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + assert ( + mc.resolve_st_cached_repo_id_case("all-minilm-l6-v2") + == "sentence-transformers/all-MiniLM-L6-v2" + ) + + def test_marker_never_raises_when_cache_mutates(monkeypatch): # A snapshot vanishing mid-iteration must read as not-cached, not raise a 500. def _exploding_iter(repo): @@ -593,6 +609,24 @@ def _info(model_name, token = None): assert mc.is_embedding_model("org/gte-modernbert") is True +def test_transient_hub_failure_accepts_a_pinned_tag_only_embedder(tmp_path, monkeypatch): + # Symmetry with the offline branch: on a TRANSIENT Hub failure (offline env NOT set) a + # downloaded tag-only embedder (no modules.json) whose verdict is pinned to the active + # revision must still be recognized, just as the offline branch recognizes it -- otherwise + # a flaky network reclassifies the same cache the offline path accepts. + _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") + + def _info(model_name, token = None): + return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") + + _fake_hf_model_info(monkeypatch, _info) + assert mc.is_embedding_model("org/gte-modernbert") is True # online confirm + persist + + mc._embedding_detection_cache.clear() # force the next call past the memo + _fake_hf_model_info(monkeypatch, _no_network) # a non-permanent error -> transient branch + assert mc.is_embedding_model("org/gte-modernbert") is True + + def test_offline_persisted_verdict_not_trusted_when_uncached(tmp_path, monkeypatch): # The persisted allowlist records only the tag; after a restart with nothing materialized # it must NOT be trusted. diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index b3789acbe29..97e0ec43ced 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1682,29 +1682,26 @@ def _iter_cache_snapshots_in(repo_id: str, roots: list[Path]): def resolve_st_cached_repo_id_case(repo_id: str) -> str: - """*repo_id* recased to match its cache dir so the offline exact-case - ``local_files_only`` load finds it; unchanged for a local path, non-repo, or nothing - cached (exact match, then a deterministic pick among case variants).""" - if is_local_path(repo_id) or "/" not in repo_id: + """*repo_id* recased (and org-qualified) to match the cache dir the ST loader actually + opens, so the offline exact-case ``local_files_only`` load finds it; unchanged for a local + path or nothing cached. + + Resolved through :func:`_st_cache_repo_dir`, which already follows the loader's short-name + alias: a slashless name like ``all-minilm-l6-v2`` is rewritten to + ``sentence-transformers/`` and looked up CASE-SENSITIVELY, so it must map to the + canonical ``sentence-transformers/all-MiniLM-L6-v2`` cache dir rather than being persisted + verbatim and missing that lookup. A slashed name recases the same way (exact-case dir + preferred, else a deterministic case variant).""" + if is_local_path(repo_id): + return repo_id + repo_dir = _st_cache_repo_dir(repo_id) + if repo_dir is None: return repo_id prefix = "models--" - expected = f"{prefix}{repo_id.replace('/', '--')}" - target = expected.lower() - variants: list[str] = [] - for cache_dir in _st_cache_roots(): - try: - if not cache_dir.is_dir(): - continue - if (cache_dir / expected).is_dir(): - return repo_id - for entry in cache_dir.iterdir(): - if entry.is_dir() and entry.name.lower() == target: - variants.append(entry.name) - except OSError: - continue - if variants: - return sorted(variants)[0][len(prefix) :].replace("--", "/") - return repo_id + name = repo_dir.name + if not name.startswith(prefix): + return repo_id + return name[len(prefix) :].replace("--", "/") def _st_repo_id_candidates(repo_id: str) -> list[str]: @@ -2431,7 +2428,10 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: ): logger.warning(f"Could not determine if {model_name} is embedding model: {e}") return False - # Transient / 5xx failure: fall back to the local marker, uncached. + # Transient / 5xx failure: fall back to the local cache, uncached. Mirror the offline + # branch so the two agree on a downloaded tag-only embedder (no modules.json): accept a + # complete cached snapshot whose embedder verdict is pinned to the active revision, not + # only one that carries a modules.json marker. marker = _embedding_marker_in_hf_cache(model_name) if marker is True: logger.info( @@ -2439,6 +2439,17 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: f"(modules.json) after Hub lookup failed: {e}" ) return True + snapshot = _active_snapshot_dir(model_name) + if ( + snapshot is not None + and _known_embedder(model_name, cache_key, _active_commit(model_name)) + and _snapshot_has_complete_weights(snapshot) + ): + logger.info( + f"Model {model_name} recognized as a previously-verified embedding model " + f"from the local cache after Hub lookup failed: {e}" + ) + return True logger.warning(f"Could not determine if {model_name} is embedding model: {e}") return False From 56201d01b025a965ef79ec6db30763c2ecd33f40 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 07:38:30 +0000 Subject: [PATCH 35/77] studio: scan modules.json-declared module roots in the offline pickle gate The offline pickle scan treated only the snapshot root and config.json-bearing subdirs as load roots, so a pickle in a non-Transformer SentenceTransformer module directory that has no config.json (e.g. a 0_WordEmbeddings/ module: wordembedding_config.json + pytorch_model.bin) was skipped even though the loader deserializes it. Parse modules.json (and thread through load_subdirs) to treat every declared module directory as a load root, so such a pickle is scanned and fail-closed offline. --- studio/backend/tests/test_hf_offline_flags.py | 13 +++++ .../backend/utils/security/file_security.py | 52 ++++++++++++++----- 2 files changed, 52 insertions(+), 13 deletions(-) diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index 091edba2abe..b5ac39c7f58 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -203,6 +203,19 @@ def _snap(name, files: dict): (modroot / "0_Transformer" / "pytorch_model.bin").write_bytes(b"\0") assert _evaluate(modroot) is True + # A pickle in a modules.json-declared module dir WITHOUT a config.json (a WordEmbeddings + # module: wordembedding_config.json + pytorch_model.bin) is still deserialized by the ST + # loader, so it must be scanned -> blocked, not skipped for lacking config.json. + we = tmp_path / "wordemb" / "aaa" + (we / "0_WordEmbeddings").mkdir(parents = True) + (we / "modules.json").write_text( + '[{"idx": 0, "name": "0_WordEmbeddings", "path": "0_WordEmbeddings", ' + '"type": "sentence_transformers.models.WordEmbeddings"}]' + ) + (we / "0_WordEmbeddings" / "wordembedding_config.json").write_bytes(b"{}") + (we / "0_WordEmbeddings" / "pytorch_model.bin").write_bytes(b"\0") + assert _evaluate(we) is True + def test_security_scan_runs_when_online(monkeypatch): import utils.security.file_security as fs diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 8678567052d..9de7e5e6169 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -284,14 +284,40 @@ def _dir_has_loadable_safetensors(files: dict) -> bool: return False -def _cached_pickle_weight_files(snap) -> list: +def _st_load_roots(snap, load_subdirs = ()) -> set: + """Directories a load opens ``from_pretrained`` on: the snapshot root, every module path + ``modules.json`` declares, and each passed-in ``load_subdirs`` entry. A SentenceTransformer + module can load from a directory without ``config.json`` -- e.g. a ``0_WordEmbeddings/`` + module with ``wordembedding_config.json`` + ``pytorch_model.bin`` -- so a pickle there is + still deserialized and must be treated as a load root.""" + roots = {snap} + for subdir in load_subdirs or (): + rel = _normalize_repo_path(str(subdir)).strip("/") + if rel: + roots.add(snap / rel) + try: + import json + modules = json.loads((snap / "modules.json").read_text(encoding = "utf-8")) + except (OSError, ValueError): + return roots + if isinstance(modules, list): + for module in modules: + if isinstance(module, dict): + rel = _normalize_repo_path(str(module.get("path") or "")).strip("/") + if rel: + roots.add(snap / rel) + return roots + + +def _cached_pickle_weight_files(snap, load_subdirs = ()) -> list: """Base-model pickle weight files a from_pretrained load actually deserializes: at a real - load root (the snapshot root, or a subdir that is itself a load root -- it holds a - ``config.json``) and with NO loadable safetensors alternative there. A stray pickle in a - non-load subdir (``archive/``, ``nemo/``) that no load opens is not a vector, matching the - online scan's load-path scoping. A dir is covered by a safetensors weight the loader would - pick instead (unsharded base file or a complete indexed shard set); a bare adapter or an - orphan shard leaves the pickle live.""" + load root and with NO loadable safetensors alternative there. A load root is the snapshot + root, a directory ``modules.json`` / ``load_subdirs`` declares, or a plain from_pretrained + root (holds ``config.json``). A stray pickle in a non-load subdir (``archive/``, ``nemo/``) + that no load opens is not a vector, matching the online scan's load-path scoping. A dir is + covered by a safetensors weight the loader would pick instead (unsharded base file or a + complete indexed shard set); a bare adapter or an orphan shard leaves the pickle live.""" + roots = _st_load_roots(snap, load_subdirs) by_dir_pickle: dict = {} by_dir_files: dict = {} # directory -> {lower-name: Path} try: @@ -310,16 +336,16 @@ def _cached_pickle_weight_files(snap) -> list: hits: set = set() for directory, names in by_dir_pickle.items(): files = by_dir_files.get(directory, {}) - # from_pretrained deserializes a pickle only at a load root: the snapshot root, or a - # subdir that is itself a load root (has config.json). A pickle elsewhere is unread. - if directory != snap and "config.json" not in files: + # A pickle is deserialized only at a load root (a declared module dir, load subdir, the + # snapshot root, or a plain from_pretrained root with config.json). Elsewhere it is unread. + if directory not in roots and "config.json" not in files: continue if not _dir_has_loadable_safetensors(files): hits.update(names) return sorted(hits) -def _evaluate_local_only(model_name: str) -> "FileSecurityDecision": +def _evaluate_local_only(model_name: str, load_subdirs = ()) -> "FileSecurityDecision": """Fail-CLOSED security decision for an offline (local_files_only) load. The Hub scan cannot be fetched offline, so instead of failing OPEN we inspect the cached @@ -337,7 +363,7 @@ def _evaluate_local_only(model_name: str) -> "FileSecurityDecision": if snap is None: return FileSecurityDecision(model_name, False, reason = "offline; nothing cached to scan") - pickles = _cached_pickle_weight_files(snap) + pickles = _cached_pickle_weight_files(snap, load_subdirs) if not pickles: return FileSecurityDecision( model_name, False, reason = "offline; cached weights are pickle-free (inert)" @@ -431,7 +457,7 @@ def evaluate_file_security( # An offline (local-only) load cannot fetch the Hub scan: fail CLOSED against the cache # instead of skipping the gate, so an unscanned cached pickle cannot deserialize. if local_only_load: - return _evaluate_local_only(model_name) + return _evaluate_local_only(model_name, load_subdirs) status = _fetch_security_status(model_name, hf_token) if not isinstance(status, dict): From 947beb1f0fbb1d6a955b6ddd7d879e5932af0b91 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 07:52:12 +0000 Subject: [PATCH 36/77] studio: classify cached non-Transformer SentenceTransformer models offline _snapshot_has_complete_weights recognized only a Transformer-shaped load root (config + tokenizer + weights co-located), so a fully-cached model built from a non-Transformer module (0_WordEmbeddings uses wordembedding_config.json + embedding weights and its own tokenizer, no HF config.json; BoW keeps its vocab in config.json) was classified non-embedding offline and the settings endpoint returned 409. Add _snapshot_modules_all_loadable, which parses modules.json and accepts a snapshot when every declared module's path directory carries the files that module class's own load() reads (a Transformer/root module still needs the full HF load root; a WordEmbeddings module needs its config plus a complete weight set; other modules need their *_config.json), and at least one embedding-producing module is present. It is OR-ed after the Transformer check, so it only ever accepts more and cannot regress the existing path or reject a pruned cache. --- .../tests/test_is_embedding_model_offline.py | 136 ++++++++++++++++++ studio/backend/utils/models/model_config.py | 107 +++++++++++++- 2 files changed, 236 insertions(+), 7 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 3a63a7e2d3b..ce285fbcd1e 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -8,6 +8,7 @@ from __future__ import annotations import importlib +import json import sys import types from pathlib import Path @@ -392,6 +393,141 @@ def test_marker_requires_both_bpe_tokenizer_files(tmp_path, monkeypatch, tok_fil assert mc._embedding_marker_in_hf_cache("org/model") is expected +# ── non-Transformer ST models (WordEmbeddings / BoW) built from modules.json ── + + +def _modules_json(*modules): + """A modules.json body from (name, path, type) triples, in declared order.""" + return json.dumps( + [ + {"idx": i, "name": name, "path": path, "type": mtype} + for i, (name, path, mtype) in enumerate(modules) + ] + ) + + +def _wordembeddings_repo( + tmp_path, + monkeypatch, + *, + weight_file = "model.safetensors", + include_weights = True, + commit = "aaa", + repo_id = "sentence-transformers/average_word_embeddings_glove.6B.300d", +): + """A cached SentenceTransformer WordEmbeddings model mirroring the real on-disk layout of + ``sentence-transformers/average_word_embeddings_glove.6B.300d``: a root ``modules.json`` + + ``config_sentence_transformers.json``, a ``0_WordEmbeddings`` module dir holding + ``wordembedding_config.json`` (+ a whitespace-tokenizer config + embedding weights) and a + ``1_Pooling`` module dir holding ``config.json``. It carries NO HF ``config.json`` / + tokenizer, so the Transformer-shaped weight check alone misclassifies it as non-embedding + and the settings route 409s it (#7218).""" + hf_root = tmp_path / "hf" + repo = hf_root / f"models--{repo_id.replace('/', '--')}" + snap = repo / "snapshots" / commit + (snap / "0_WordEmbeddings").mkdir(parents = True) + (snap / "modules.json").write_text( + _modules_json( + ("0", "0_WordEmbeddings", "sentence_transformers.models.WordEmbeddings"), + ("1", "1_Pooling", "sentence_transformers.models.Pooling"), + ) + ) + (snap / "config_sentence_transformers.json").write_text("{}") + (snap / "0_WordEmbeddings" / "wordembedding_config.json").write_text("{}") + (snap / "0_WordEmbeddings" / "whitespacetokenizer_config.json").write_text("{}") + if include_weights: + (snap / "0_WordEmbeddings" / weight_file).write_bytes(b"\0") + (snap / "1_Pooling").mkdir(parents = True) + (snap / "1_Pooling" / "config.json").write_text("{}") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + return snap + + +_GLOVE = "sentence-transformers/average_word_embeddings_glove.6B.300d" + + +@pytest.mark.parametrize("weight_file", ["model.safetensors", "pytorch_model.bin"]) +def test_marker_accepts_complete_wordembeddings_model(tmp_path, monkeypatch, weight_file): + # A fully-cached WordEmbeddings model (glove-style: no HF config.json / tokenizer) loads + # offline via modules.json -> WordEmbeddings.load(), so it must be recognized as an + # embedding model even though it carries none of the Transformer-shaped assets (#7218). + _wordembeddings_repo(tmp_path, monkeypatch, weight_file = weight_file) + assert mc._embedding_marker_in_hf_cache(_GLOVE) is True + + +def test_offline_wordembeddings_is_embedding_without_network(tmp_path, monkeypatch): + # End to end: the offline settings path must not 409 a cached WordEmbeddings model, with no + # model_info() network call. + _wordembeddings_repo(tmp_path, monkeypatch) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _fake_hf_model_info(monkeypatch, _no_network) + assert mc.is_embedding_model(_GLOVE) is True + + +def test_marker_rejects_wordembeddings_module_dir_empty(tmp_path, monkeypatch): + # Control: modules.json declares 0_WordEmbeddings but the dir holds no files (never + # materialized / pruned). The offline load would fail, so this must NOT validate. + snap = _wordembeddings_repo(tmp_path, monkeypatch) + for child in (snap / "0_WordEmbeddings").iterdir(): + child.unlink() + assert mc._embedding_marker_in_hf_cache(_GLOVE) is False + + +def test_marker_rejects_wordembeddings_without_weights(tmp_path, monkeypatch): + # WordEmbeddings.load() hard-loads model.safetensors / pytorch_model.bin (no fallback), so a + # config-only module dir (weights pruned) would fail the load and must NOT validate. + _wordembeddings_repo(tmp_path, monkeypatch, include_weights = False) + assert mc._embedding_marker_in_hf_cache(_GLOVE) is False + + +def test_marker_accepts_complete_bow_model(tmp_path, monkeypatch): + # A BoW module keeps its vocab in config.json and writes NO weight file; a complete cache is + # still loadable via BoW.load(config.json), so it must validate. + hf_root = tmp_path / "hf" + repo = hf_root / "models--org--bow" + snap = repo / "snapshots" / "aaa" + (snap / "0_BoW").mkdir(parents = True) + (snap / "modules.json").write_text( + _modules_json( + ("0", "0_BoW", "sentence_transformers.models.BoW"), + ("1", "1_Pooling", "sentence_transformers.models.Pooling"), + ) + ) + (snap / "config_sentence_transformers.json").write_text("{}") + (snap / "0_BoW" / "config.json").write_text('{"vocab": ["a", "b"]}') + (snap / "1_Pooling").mkdir(parents = True) + (snap / "1_Pooling" / "config.json").write_text("{}") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text("aaa") + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + assert mc._embedding_marker_in_hf_cache("org/bow") is True + + +def test_marker_rejects_structural_only_module_list(tmp_path, monkeypatch): + # A degenerate modules.json with only structural modules (Pooling / Normalize) has no source + # of embeddings and must NOT validate on the non-Transformer path. + hf_root = tmp_path / "hf" + repo = hf_root / "models--org--degenerate" + snap = repo / "snapshots" / "aaa" + (snap / "1_Pooling").mkdir(parents = True) + (snap / "modules.json").write_text( + _modules_json( + ("0", "1_Pooling", "sentence_transformers.models.Pooling"), + ("1", "2_Normalize", "sentence_transformers.models.Normalize"), + ) + ) + (snap / "1_Pooling" / "config.json").write_text("{}") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text("aaa") + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + assert mc._embedding_marker_in_hf_cache("org/degenerate") is False + + def test_short_name_resolves_through_the_sentence_transformers_org(tmp_path, monkeypatch): # "all-MiniLM-L6-v2" is a supported ST alias: the loader falls back to the # sentence-transformers/ org, so the snapshot is cached under that full id. Probing only diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 97e0ec43ced..d784fea3b08 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2244,6 +2244,18 @@ def _dir_has_complete_torch_weights(names: set) -> bool: _ST_CONFIG_FILES = ("config.json", "config_sentence_transformers.json") +def _dir_is_transformer_load_root(names: set) -> bool: + """True when one directory's files (*names*) form a COMPLETE Transformer / plain-HF load + root: an ST/HF config, a usable tokenizer and a complete Torch weight set, all co-located. + This is the shape ``modules.json`` sends a Transformer module (or a bare from_pretrained) + at -- the per-directory unit the snapshot checks share.""" + if not any(cfg in names for cfg in _ST_CONFIG_FILES): + return False + if not _names_have_tokenizer(names): + return False + return _dir_has_complete_torch_weights(names) + + def _snapshot_has_complete_weights(snap: Path) -> bool: """True when some LOAD ROOT inside *snap* is complete: a config, a tokenizer and a COMPLETE Torch weight set, all in the SAME directory. @@ -2263,25 +2275,106 @@ def _snapshot_has_complete_weights(snap: Path) -> bool: except OSError: continue for names in by_dir.values(): - if not any(cfg in names for cfg in _ST_CONFIG_FILES): - continue - if not _names_have_tokenizer(names): - continue - if _dir_has_complete_torch_weights(names): + if _dir_is_transformer_load_root(names): return True return False except OSError: return False +# Sentence-transformers modules whose ST ``load()`` reads no embedding payload of their own -- +# a ``modules.json`` listing ONLY these has no source of embeddings, so it is not a loadable +# model on its own and must not validate the non-Transformer acceptance path below. +_ST_STRUCTURAL_MODULE_NAMES = frozenset({"pooling", "normalize"}) + + +def _dir_file_names(dir_path: Path) -> set: + """The file basenames directly in *dir_path* (not recursive), or an empty set when the + directory is absent / unreadable / not a directory. Never raises.""" + try: + return {entry.name for entry in dir_path.iterdir() if entry.is_file()} + except OSError: + return set() + + +def _module_dir_is_loadable(cls: str, is_root: bool, dir_path: Path) -> bool: + """True when *dir_path* carries the files the sentence-transformers module class *cls* + reads in its own ``load()`` (see sentence_transformers/models/*.py): + + * a root / ``Transformer`` module is a full HF load root (config + tokenizer + weights); + * ``Normalize`` reads nothing; + * ``WordEmbeddings`` hard-loads ``model.safetensors`` / ``pytorch_model.bin`` (no fallback), + so it needs its ``wordembedding_config.json`` AND a complete weight set; + * every other module (``BoW``, ``Pooling``, ``Dense``, ``CNN``, ``LSTM`` ...) reads a + mandatory ``config.json`` / ``*_config.json`` from its directory (its weights, when used, + are optional there), so a present module config is the load requirement.""" + if is_root or "transformer" in cls: + return _dir_is_transformer_load_root(_dir_file_names(dir_path)) + if cls == "normalize": + return True + names = _dir_file_names(dir_path) + if not any(name == "config.json" or name.endswith("_config.json") for name in names): + return False + if cls == "wordembeddings": + return _dir_has_complete_torch_weights(names) + return True + + +def _snapshot_modules_all_loadable(snap: Path) -> bool: + """True when *snap* is a COMPLETE non-Transformer sentence-transformers model that the + Transformer-shaped :func:`_snapshot_has_complete_weights` misses -- e.g. one built from a + ``0_WordEmbeddings`` module (``wordembedding_config.json`` + embedding weights, no HF + ``config.json`` / tokenizer) or ``BoW`` (vocab in ``config.json``). The offline + ``local_files_only`` load builds these from ``modules.json`` by calling + ``module_class.load(path)`` per declared module. + + Additive and conservative: it only ever ACCEPTS (the caller OR-s it with the Transformer + check) and never rejects, so it cannot regress the Transformer path. It validates only when + ``modules.json`` parses as a non-empty list AND every declared module's ``path`` directory + carries the files that module's own ``load()`` reads -- so a bare / pruned cache (empty + ``modules.json``, a missing module directory, a WordEmbeddings module with its weights + stripped) does not validate. It also requires at least one embedding-producing module, so a + degenerate Pooling/Normalize-only list is rejected. Never raises.""" + try: + modules = json.loads((snap / "modules.json").read_text(encoding = "utf-8")) + except (OSError, ValueError): + return False + if not isinstance(modules, list) or not modules: + return False + saw_content_module = False + for module in modules: + if not isinstance(module, dict): + return False + cls = str(module.get("type") or "").rsplit(".", 1)[-1].strip().lower() + if not cls: + return False + raw_path = str(module.get("path") or "").strip() + is_root = raw_path in ("", ".") + if is_root: + dir_path = snap + else: + rel = raw_path.strip("/") + if not rel or ".." in Path(rel).parts: + return False # never resolve a module path outside the snapshot + dir_path = snap / rel + if not _module_dir_is_loadable(cls, is_root, dir_path): + return False + if is_root or cls not in _ST_STRUCTURAL_MODULE_NAMES: + saw_content_module = True + return saw_content_module + + def _snapshot_is_loadable_st_model(snap: Path) -> bool: """True when *snap* holds a loadable sentence-transformers model. ``modules.json`` alone is insufficient -- the security preflight downloads exactly that file, so require the - marker plus ``_snapshot_has_complete_weights``.""" + marker plus a complete load: a Transformer-shaped weight set + (:func:`_snapshot_has_complete_weights`), OR a complete non-Transformer module set such as + ``0_WordEmbeddings`` / ``BoW`` (:func:`_snapshot_modules_all_loadable`), which carries none + of the Transformer config/tokenizer/weights the first check requires.""" try: if not (snap / "modules.json").is_file(): return False - return _snapshot_has_complete_weights(snap) + return _snapshot_has_complete_weights(snap) or _snapshot_modules_all_loadable(snap) except OSError: return False From f56b7127f96dbfec27c241961855b492e544b647 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 08:05:41 +0000 Subject: [PATCH 37/77] studio: scan PEFT adapter pickle weights in the offline security gate from_pretrained auto-detects an adapter_config.json in the load root and deserializes the adapter weights on top of the base model, so adapter_model.bin is a separate pickle RCE vector that a safetensors base weight does not cover. The offline scan matched only base-model pickle names, so an offline local-only load with safetensors base weights plus a cached adapter_model.bin was allowed despite the live adapter pickle. Scan adapter pickles too, scoped to a load root where adapter_config.json is present and no adapter_model.safetensors exists. --- studio/backend/tests/test_hf_offline_flags.py | 20 +++++++++++ .../backend/utils/security/file_security.py | 35 +++++++++++++------ 2 files changed, 45 insertions(+), 10 deletions(-) diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index b5ac39c7f58..291c09e69ff 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -216,6 +216,26 @@ def _snap(name, files: dict): (we / "0_WordEmbeddings" / "pytorch_model.bin").write_bytes(b"\0") assert _evaluate(we) is True + # A PEFT adapter pickle is a SEPARATE vector: from_pretrained auto-detects adapter_config.json + # and deserializes adapter_model.bin ON TOP of the base, so an inert safetensors base does not + # cover it -> blocked. + adbin = tmp_path / "adapterbin" / "aaa" + adbin.mkdir(parents = True) + (adbin / "config.json").write_bytes(b"{}") + (adbin / "model.safetensors").write_bytes(b"\0") + (adbin / "adapter_config.json").write_bytes(b"{}") + (adbin / "adapter_model.bin").write_bytes(b"\0") + assert _evaluate(adbin) is True + + # The same adapter shipped as safetensors is inert -> allowed. + adsafe = tmp_path / "adaptersafe" / "aaa" + adsafe.mkdir(parents = True) + (adsafe / "config.json").write_bytes(b"{}") + (adsafe / "model.safetensors").write_bytes(b"\0") + (adsafe / "adapter_config.json").write_bytes(b"{}") + (adsafe / "adapter_model.safetensors").write_bytes(b"\0") + assert _evaluate(adsafe) is False + def test_security_scan_runs_when_online(monkeypatch): import utils.security.file_security as fs diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 9de7e5e6169..a5e13bf13a5 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -246,6 +246,12 @@ def _load_scan_target(model_name: str, load_subdirs: tuple) -> tuple: r"^(model|pytorch_model)(-\d+-of-\d+)?\.(bin|pt|pth|ckpt|pkl|pickle)$" ) +# PEFT adapter pickle weights. from_pretrained auto-detects an adapter_config.json in the load +# root and deserializes the adapter weights on top of the base model, so an adapter pickle is a +# SEPARATE RCE vector from the base weights -- a safetensors base does not cover it. It has its +# own safetensors preference: adapter_model.safetensors is loaded in place of the pickle. +_ADAPTER_PICKLE_RE = re.compile(r"^adapter_model\.(bin|pt|pth|ckpt|pkl|pickle)$") + # A from_pretrained load prefers safetensors over a pickle ONLY when the directory holds a # safetensors weight it can actually load in its place: an unsharded base file, or an index @@ -310,13 +316,16 @@ def _st_load_roots(snap, load_subdirs = ()) -> set: def _cached_pickle_weight_files(snap, load_subdirs = ()) -> list: - """Base-model pickle weight files a from_pretrained load actually deserializes: at a real - load root and with NO loadable safetensors alternative there. A load root is the snapshot - root, a directory ``modules.json`` / ``load_subdirs`` declares, or a plain from_pretrained - root (holds ``config.json``). A stray pickle in a non-load subdir (``archive/``, ``nemo/``) - that no load opens is not a vector, matching the online scan's load-path scoping. A dir is - covered by a safetensors weight the loader would pick instead (unsharded base file or a - complete indexed shard set); a bare adapter or an orphan shard leaves the pickle live.""" + """Pickle weight files a from_pretrained load actually deserializes: at a real load root and + with NO loadable safetensors alternative there. A load root is the snapshot root, a directory + ``modules.json`` / ``load_subdirs`` declares, or a plain from_pretrained root (holds + ``config.json``). A stray pickle in a non-load subdir (``archive/``, ``nemo/``) that no load + opens is not a vector, matching the online scan's load-path scoping. Two weight classes are + scoped independently: a BASE pickle (``pytorch_model.bin`` ...) is covered by a loadable base + safetensors the loader picks instead; a PEFT ADAPTER pickle (``adapter_model.bin``), which + from_pretrained auto-loads when ``adapter_config.json`` is present, is covered only by + ``adapter_model.safetensors`` -- a safetensors base does NOT cover it. A bare adapter or an + orphan shard leaves its pickle live.""" roots = _st_load_roots(snap, load_subdirs) by_dir_pickle: dict = {} by_dir_files: dict = {} # directory -> {lower-name: Path} @@ -327,7 +336,7 @@ def _cached_pickle_weight_files(snap, load_subdirs = ()) -> list: continue low = path.name.lower() by_dir_files.setdefault(path.parent, {})[low] = path - if _PICKLE_WEIGHT_RE.match(low): + if _PICKLE_WEIGHT_RE.match(low) or _ADAPTER_PICKLE_RE.match(low): by_dir_pickle.setdefault(path.parent, []).append(path.name) except OSError: continue @@ -340,8 +349,14 @@ def _cached_pickle_weight_files(snap, load_subdirs = ()) -> list: # snapshot root, or a plain from_pretrained root with config.json). Elsewhere it is unread. if directory not in roots and "config.json" not in files: continue - if not _dir_has_loadable_safetensors(files): - hits.update(names) + base = [n for n in names if _PICKLE_WEIGHT_RE.match(n.lower())] + if base and not _dir_has_loadable_safetensors(files): + hits.update(base) + # An adapter pickle is deserialized only when from_pretrained auto-detects the adapter + # (adapter_config.json present) and there is no adapter_model.safetensors to load instead. + adapter = [n for n in names if _ADAPTER_PICKLE_RE.match(n.lower())] + if adapter and "adapter_config.json" in files and "adapter_model.safetensors" not in files: + hits.update(adapter) return sorted(hits) From 4c04b8786b0941a71c2c1c855ec156cfa6e1ff6e Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 08:17:49 +0000 Subject: [PATCH 38/77] studio: require weights for Dense/CNN/LSTM SentenceTransformer modules offline _module_dir_is_loadable accepted a Dense, CNN, or LSTM module dir with only its config, but those modules' load() hard-load model.safetensors else pytorch_model.bin (verified against sentence-transformers source: no fallback, raises if neither exists) -- exactly like WordEmbeddings. A cache with such a module's config but no weights would validate and then fail the local_files_only load. Require a complete weight set for every weighted module, not just WordEmbeddings. --- .../tests/test_is_embedding_model_offline.py | 37 +++++++++++++++++++ studio/backend/utils/models/model_config.py | 20 +++++++--- 2 files changed, 51 insertions(+), 6 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index ce285fbcd1e..5ac7bc54033 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -483,6 +483,43 @@ def test_marker_rejects_wordembeddings_without_weights(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache(_GLOVE) is False +def _we_dense_repo(tmp_path, monkeypatch, commit = "aaa"): + # WordEmbeddings + Pooling + a Dense projection module, with the Dense weights ABSENT. + hf_root = tmp_path / "hf" + repo = hf_root / "models--org--we-dense" + snap = repo / "snapshots" / commit + (snap / "0_WordEmbeddings").mkdir(parents = True) + (snap / "modules.json").write_text( + _modules_json( + ("0", "0_WordEmbeddings", "sentence_transformers.models.WordEmbeddings"), + ("1", "1_Pooling", "sentence_transformers.models.Pooling"), + ("2", "2_Dense", "sentence_transformers.models.Dense"), + ) + ) + (snap / "config_sentence_transformers.json").write_text("{}") + (snap / "0_WordEmbeddings" / "wordembedding_config.json").write_text("{}") + (snap / "0_WordEmbeddings" / "model.safetensors").write_bytes(b"\0") + (snap / "1_Pooling").mkdir(parents = True) + (snap / "1_Pooling" / "config.json").write_text("{}") + (snap / "2_Dense").mkdir(parents = True) + (snap / "2_Dense" / "config.json").write_text("{}") # config present, weights absent + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + return snap + + +def test_marker_requires_weights_for_a_dense_module(tmp_path, monkeypatch): + # Dense / CNN / LSTM.load() hard-load model.safetensors / pytorch_model.bin (no fallback), so + # a Dense module dir with only its config would fail the offline load and must NOT validate; + # adding the weights makes the model loadable. + snap = _we_dense_repo(tmp_path, monkeypatch) + assert mc._embedding_marker_in_hf_cache("org/we-dense") is False + (snap / "2_Dense" / "model.safetensors").write_bytes(b"\0") + assert mc._embedding_marker_in_hf_cache("org/we-dense") is True + + def test_marker_accepts_complete_bow_model(tmp_path, monkeypatch): # A BoW module keeps its vocab in config.json and writes NO weight file; a complete cache is # still loadable via BoW.load(config.json), so it must validate. diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index d784fea3b08..28f68bdf81b 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2297,17 +2297,25 @@ def _dir_file_names(dir_path: Path) -> set: return set() +# Non-Transformer modules whose ST ``load()`` reads a weight file. Each does +# ``load_safetensors(model.safetensors)`` else ``torch.load(pytorch_model.bin)`` with no +# fallback, so it needs its module config AND a complete Torch weight set -- a config alone +# would validate here and then raise FileNotFoundError at load. ``Pooling`` / ``BoW`` / +# ``Normalize`` read no weights (BoW keeps its data in config.json), so a config suffices. +_ST_WEIGHTED_MODULE_NAMES = frozenset({"wordembeddings", "dense", "cnn", "lstm"}) + + def _module_dir_is_loadable(cls: str, is_root: bool, dir_path: Path) -> bool: """True when *dir_path* carries the files the sentence-transformers module class *cls* reads in its own ``load()`` (see sentence_transformers/models/*.py): * a root / ``Transformer`` module is a full HF load root (config + tokenizer + weights); * ``Normalize`` reads nothing; - * ``WordEmbeddings`` hard-loads ``model.safetensors`` / ``pytorch_model.bin`` (no fallback), - so it needs its ``wordembedding_config.json`` AND a complete weight set; - * every other module (``BoW``, ``Pooling``, ``Dense``, ``CNN``, ``LSTM`` ...) reads a - mandatory ``config.json`` / ``*_config.json`` from its directory (its weights, when used, - are optional there), so a present module config is the load requirement.""" + * a WEIGHTED module (``WordEmbeddings`` / ``Dense`` / ``CNN`` / ``LSTM``) needs its module + config AND a complete Torch weight set (its ``load()`` hard-loads ``model.safetensors`` / + ``pytorch_model.bin``); + * every other module (``BoW``, ``Pooling`` ...) reads only a ``config.json`` / ``*_config``, + so a present module config is the load requirement.""" if is_root or "transformer" in cls: return _dir_is_transformer_load_root(_dir_file_names(dir_path)) if cls == "normalize": @@ -2315,7 +2323,7 @@ def _module_dir_is_loadable(cls: str, is_root: bool, dir_path: Path) -> bool: names = _dir_file_names(dir_path) if not any(name == "config.json" or name.endswith("_config.json") for name in names): return False - if cls == "wordembeddings": + if cls in _ST_WEIGHTED_MODULE_NAMES: return _dir_has_complete_torch_weights(names) return True From ac0a89cb5a25d3f82591424dcf296d3edf2f557d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 08:19:07 +0000 Subject: [PATCH 39/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_is_embedding_model_offline.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 5ac7bc54033..0801d7cb12e 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -483,7 +483,11 @@ def test_marker_rejects_wordembeddings_without_weights(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache(_GLOVE) is False -def _we_dense_repo(tmp_path, monkeypatch, commit = "aaa"): +def _we_dense_repo( + tmp_path, + monkeypatch, + commit = "aaa", +): # WordEmbeddings + Pooling + a Dense projection module, with the Dense weights ABSENT. hf_root = tmp_path / "hf" repo = hf_root / "models--org--we-dense" From fbfbe0722fbd166f6b4a98cc9fba6c88951fee2a Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 08:46:37 +0000 Subject: [PATCH 40/77] studio: scan root-index subdir pickle shards offline; handle __all__.append/.extend - The offline pickle scan followed only load-root directories, so a shard mapped by a root pytorch_model.bin.index.json into a non-root subdirectory was skipped even though from_pretrained follows the index weight_map and deserializes it (a layout an attacker can craft to evade the scanner). Read the local index and scan its referenced pickle shards, covered by a loadable base safetensors at the index root -- mirroring the online scan. - The import-hoist analyzer ignored __all__.append("X") / __all__.extend([...]) runtime re-export mutators, so an import added solely for one tripped HOISTED-IMPORT-UNUSED. Read their string args like +=, and treat any other __all__ method call as opaque. --- scripts/verify_import_hoist.py | 59 ++++++++++++++++++- studio/backend/tests/test_hf_offline_flags.py | 25 ++++++++ .../backend/utils/security/file_security.py | 40 +++++++++++++ 3 files changed, 123 insertions(+), 1 deletion(-) diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index e21a30d8599..5387ff69efb 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -447,6 +447,29 @@ def _legb_chain(scope: Scope) -> list[Scope]: # ---------------------------------------------------------------- analysis +def _all_mutator_entries(method: str, args: list) -> tuple[set[str], bool]: + """String entries a ``__all__.(args)`` call adds, and whether it was fully + readable. ``append(x)`` adds one string constant; ``extend([...])`` adds a list literal's + string constants. A dynamic arg, or any other mutator (``insert``/``remove``/...), is not + statically readable -> caller marks the export set opaque.""" + entries: set[str] = set() + if method == "append" and len(args) == 1: + arg = args[0] + if isinstance(arg, ast.Constant) and isinstance(arg.value, str): + entries.add(arg.value) + return entries, True + return entries, False + if method == "extend" and len(args) == 1 and isinstance(args[0], (ast.List, ast.Tuple, ast.Set)): + readable = True + for elt in args[0].elts: + if isinstance(elt, ast.Constant) and isinstance(elt.value, str): + entries.add(elt.value) + else: + readable = False + return entries, readable + return entries, False # unknown mutator or dynamic arg -> opaque + + def _collect_dunder_all(tree: ast.Module) -> tuple[set[str], bool]: """The FINAL module-level ``__all__`` string entries plus an ``opaque`` flag (``= [...]``, ``+= [...]``, or an annotated assign). @@ -470,11 +493,27 @@ def _collect_dunder_all(tree: ast.Module) -> tuple[set[str], bool]: replacing ``=`` resets opacity (it discards the prior list AND any prior dynamic part), so a later readable ``__all__ = [...]`` clears an earlier opaque ``+=``; an ``+=`` extends, so prior opacity persists. A bare ``__all__: list[str]`` annotation - has no runtime value and is skipped entirely. + has no runtime value and is skipped entirely. ``__all__.append("X")`` / + ``.extend([...])`` are runtime extenders read the same way as ``+=``; any other + ``__all__`` method call (``insert``/``remove``/...) is opaque. """ names: set[str] = set() opaque = False for node in tree.body: + if isinstance(node, ast.Expr): + # __all__.append("X") / .extend([...]) / other mutators: a runtime EXTEND of the + # export list, never a replace. Read append/extend string args; anything else opaque. + call = node.value + if ( + isinstance(call, ast.Call) + and isinstance(call.func, ast.Attribute) + and isinstance(call.func.value, ast.Name) + and call.func.value.id == "__all__" + ): + entries, readable = _all_mutator_entries(call.func.attr, call.args) + names |= entries + opaque = opaque or not readable + continue if isinstance(node, ast.Assign): targets, replaces = node.targets, True elif isinstance(node, ast.AnnAssign): @@ -843,6 +882,24 @@ def used_targets(analysis) -> set[str]: "from pkg import a\nfrom pkg import b\n__all__: list[str]\n", "BLOCKER", ), + "all_append_reexport_is_used": ( + # __all__.append("b") re-exports b -> must not flag the hoisted import as unused + 'from pkg import a\n__all__ = ["a"]\n', + 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n__all__.append("b")\n', + None, + ), + "all_extend_reexport_is_used": ( + # __all__.extend(["b"]) re-exports b as well + 'from pkg import a\n__all__ = ["a"]\n', + 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n__all__.extend(["b"])\n', + None, + ), + "all_append_unrelated_still_blocks": ( + # a readable append of a DIFFERENT name does not credit an unrelated unused hoist + 'from pkg import a\n__all__ = ["a"]\n', + 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n__all__.append("a")\n', + "BLOCKER", + ), } diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index 291c09e69ff..b6728cd6b88 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -236,6 +236,31 @@ def _snap(name, files: dict): (adsafe / "adapter_model.safetensors").write_bytes(b"\0") assert _evaluate(adsafe) is False + # A root pickle index whose weight_map points shards into a SUBDIR: from_pretrained follows + # the index and deserializes them, so they must be scanned even though the subdir is not a + # load root of its own (no config.json). + idxsub = tmp_path / "idxsub" / "aaa" + (idxsub / "sharded").mkdir(parents = True) + (idxsub / "config.json").write_bytes(b"{}") + (idxsub / "pytorch_model.bin.index.json").write_text( + '{"weight_map": {"a": "sharded/pytorch_model-00001-of-00002.bin", ' + '"b": "sharded/pytorch_model-00002-of-00002.bin"}}' + ) + (idxsub / "sharded" / "pytorch_model-00001-of-00002.bin").write_bytes(b"\0") + (idxsub / "sharded" / "pytorch_model-00002-of-00002.bin").write_bytes(b"\0") + assert _evaluate(idxsub) is True + + # Those subdir shards are covered when a loadable base safetensors sits at the index root. + idxsafe = tmp_path / "idxsafe" / "aaa" + (idxsafe / "sharded").mkdir(parents = True) + (idxsafe / "config.json").write_bytes(b"{}") + (idxsafe / "model.safetensors").write_bytes(b"\0") + (idxsafe / "pytorch_model.bin.index.json").write_text( + '{"weight_map": {"a": "sharded/pytorch_model-00001-of-00002.bin"}}' + ) + (idxsafe / "sharded" / "pytorch_model-00001-of-00002.bin").write_bytes(b"\0") + assert _evaluate(idxsafe) is False + def test_security_scan_runs_when_online(monkeypatch): import utils.security.file_security as fs diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index a5e13bf13a5..0045e960540 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -252,6 +252,28 @@ def _load_scan_target(model_name: str, load_subdirs: tuple) -> tuple: # own safetensors preference: adapter_model.safetensors is loaded in place of the pickle. _ADAPTER_PICKLE_RE = re.compile(r"^adapter_model\.(bin|pt|pth|ckpt|pkl|pickle)$") +# Suffixes a from_pretrained load deserializes as a pickle (the RCE vector). Used for shards a +# weight-index maps, which are model weights regardless of their exact filename. +_PICKLE_SUFFIXES = frozenset({".bin", ".pt", ".pth", ".ckpt", ".pkl", ".pickle"}) + +# Root pickle weight-index. Its weight_map can point shards into SUBDIRECTORIES; from_pretrained +# follows the map and deserializes them wherever they live. +_PICKLE_INDEX_FILE = "pytorch_model.bin.index.json" + + +def _index_weight_map_values(index_path) -> set: + """Repo-relative shard paths a weight-index maps (its ``weight_map`` values), or empty on + an unreadable / malformed index.""" + import json + + try: + weight_map = (json.loads(index_path.read_text(encoding = "utf-8")) or {}).get( + "weight_map" + ) or {} + except (OSError, ValueError): + return set() + return {_normalize_repo_path(str(shard)) for shard in weight_map.values()} + # A from_pretrained load prefers safetensors over a pickle ONLY when the directory holds a # safetensors weight it can actually load in its place: an unsharded base file, or an index @@ -357,6 +379,24 @@ def _cached_pickle_weight_files(snap, load_subdirs = ()) -> list: adapter = [n for n in names if _ADAPTER_PICKLE_RE.match(n.lower())] if adapter and "adapter_config.json" in files and "adapter_model.safetensors" not in files: hits.update(adapter) + # A load-root pickle index (pytorch_model.bin.index.json) can map shards into SUBDIRECTORIES + # that are not themselves load roots; from_pretrained follows the map and deserializes them, + # so include those referenced pickle shards (unless a loadable base safetensors at the index + # root covers the base weights). Mirrors the online _indexed_shard_paths scan, read from disk. + for root_dir in roots: + files = by_dir_files.get(root_dir, {}) + index_path = files.get(_PICKLE_INDEX_FILE) + if index_path is None or _dir_has_loadable_safetensors(files): + continue + for shard_rel in _index_weight_map_values(index_path): + if _file_suffix(shard_rel) not in _PICKLE_SUFFIXES: + continue + shard = root_dir.joinpath(*shard_rel.split("/")) + try: + if shard.is_file(): + hits.add(shard_rel) + except OSError: + continue return sorted(hits) From 8e9a0e4215ee7239386af98c9ea038ed550b96a8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 08:48:37 +0000 Subject: [PATCH 41/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- scripts/verify_import_hoist.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index 5387ff69efb..b35b61613cb 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -459,7 +459,11 @@ def _all_mutator_entries(method: str, args: list) -> tuple[set[str], bool]: entries.add(arg.value) return entries, True return entries, False - if method == "extend" and len(args) == 1 and isinstance(args[0], (ast.List, ast.Tuple, ast.Set)): + if ( + method == "extend" + and len(args) == 1 + and isinstance(args[0], (ast.List, ast.Tuple, ast.Set)) + ): readable = True for elt in args[0].elts: if isinstance(elt, ast.Constant) and isinstance(elt.value, str): From b85a7712da6fc971ec3245b003a3c65519ff4c3e Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 10:10:22 +0000 Subject: [PATCH 42/77] studio: classify StaticEmbedding offline, require WordEmbeddings tokenizer, bound model_info - A StaticEmbedding module (e.g. sentence-transformers/static-retrieval-mrl-en-v1's 0_StaticEmbedding/) holds tokenizer.json + weights and NO config, so the config-gated non-Transformer path 409'd it offline. Recognize it by what StaticEmbedding.load() reads: a tokenizer.json plus a complete Torch weight set. - WordEmbeddings.load() rebuilds its tokenizer via the configured tokenizer_class.load() from the module dir, so a WordEmbeddings module now also requires a tokenizer artifact (whitespacetokenizer_config.json / phrasetokenizer_config.json, or a shared HF tokenizer asset), not just its config + weights. - With neither offline env var set, an unbounded model_info() could hang on connect/DNS retries for networkless users (the #6817 symptom). Bound it with a 15s timeout so a dead network fails fast and the existing transient-failure cache fallback resolves a cached model, while a reachable Hub still wins. (Documented caveat: a stalled DNS getaddrinfo may exceed this.) --- .../tests/test_is_embedding_model_offline.py | 153 ++++++++++++++++-- studio/backend/utils/models/model_config.py | 61 ++++++- 2 files changed, 195 insertions(+), 19 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 0801d7cb12e..72a35cba7e5 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -412,6 +412,7 @@ def _wordembeddings_repo( *, weight_file = "model.safetensors", include_weights = True, + include_tokenizer = True, commit = "aaa", repo_id = "sentence-transformers/average_word_embeddings_glove.6B.300d", ): @@ -434,7 +435,10 @@ def _wordembeddings_repo( ) (snap / "config_sentence_transformers.json").write_text("{}") (snap / "0_WordEmbeddings" / "wordembedding_config.json").write_text("{}") - (snap / "0_WordEmbeddings" / "whitespacetokenizer_config.json").write_text("{}") + if include_tokenizer: + # WhitespaceTokenizer.load() reads this from the module dir (WordEmbeddings.load -> + # tokenizer_class.load(dir)); see sentence_transformers/models/tokenizer/*.py. + (snap / "0_WordEmbeddings" / "whitespacetokenizer_config.json").write_text("{}") if include_weights: (snap / "0_WordEmbeddings" / weight_file).write_bytes(b"\0") (snap / "1_Pooling").mkdir(parents = True) @@ -483,6 +487,101 @@ def test_marker_rejects_wordembeddings_without_weights(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache(_GLOVE) is False +@pytest.mark.parametrize( + "tok_file", + [ + "whitespacetokenizer_config.json", # WhitespaceTokenizer.load reads this + "phrasetokenizer_config.json", # PhraseTokenizer.load reads this + "tokenizer.json", # transformers-backed wrapper: shared HF tokenizer asset + "vocab.txt", + ], +) +def test_marker_requires_a_tokenizer_for_wordembeddings(tmp_path, monkeypatch, tok_file): + # WordEmbeddings.load() calls its configured tokenizer_class.load(dir) from the module dir, so + # a module with wordembedding_config.json + weights but NO tokenizer artifact would validate + # on the config+weights path and then fail at tokenizer_class.load() -- it must NOT validate. + # Restoring any recognized tokenizer artifact makes it loadable (#7218). + snap = _wordembeddings_repo(tmp_path, monkeypatch, include_tokenizer = False) + assert mc._embedding_marker_in_hf_cache(_GLOVE) is False + (snap / "0_WordEmbeddings" / tok_file).write_text("{}") + assert mc._embedding_marker_in_hf_cache(_GLOVE) is True + + +# ── StaticEmbedding models (model2vec / static-retrieval) built from modules.json ── + + +_STATIC = "sentence-transformers/static-retrieval-mrl-en-v1" + + +def _staticembedding_repo( + tmp_path, + monkeypatch, + *, + weight_file = "model.safetensors", + include_weights = True, + include_tokenizer = True, + commit = "aaa", + repo_id = _STATIC, +): + """A cached SentenceTransformer StaticEmbedding model mirroring the real on-disk layout of + ``sentence-transformers/static-retrieval-mrl-en-v1``: a root ``modules.json`` + + ``config_sentence_transformers.json`` and a ``0_StaticEmbedding`` module dir holding just + ``tokenizer.json`` + ``model.safetensors`` -- exactly what ``StaticEmbedding.load()`` reads, + which writes NO config. It carries NO HF/module ``config.json``, so the config-gated + non-Transformer path alone misclassifies it and the settings route 409s it offline (#7218).""" + hf_root = tmp_path / "hf" + repo = hf_root / f"models--{repo_id.replace('/', '--')}" + snap = repo / "snapshots" / commit + (snap / "0_StaticEmbedding").mkdir(parents = True) + (snap / "modules.json").write_text( + _modules_json( + ("0", "0_StaticEmbedding", "sentence_transformers.models.StaticEmbedding"), + ) + ) + (snap / "config_sentence_transformers.json").write_text("{}") + if include_tokenizer: + (snap / "0_StaticEmbedding" / "tokenizer.json").write_text("{}") + if include_weights: + (snap / "0_StaticEmbedding" / weight_file).write_bytes(b"\0") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + return snap + + +@pytest.mark.parametrize("weight_file", ["model.safetensors", "pytorch_model.bin"]) +def test_marker_accepts_complete_staticembedding_model(tmp_path, monkeypatch, weight_file): + # A fully-cached StaticEmbedding model (0_StaticEmbedding/ with tokenizer.json + weights, NO + # config) loads offline via modules.json -> StaticEmbedding.load(), so it must be recognized + # even though it carries none of the Transformer-shaped config the other path requires (#7218). + _staticembedding_repo(tmp_path, monkeypatch, weight_file = weight_file) + assert mc._embedding_marker_in_hf_cache(_STATIC) is True + + +def test_marker_rejects_staticembedding_without_weights(tmp_path, monkeypatch): + # StaticEmbedding.load() hard-reads model.safetensors / pytorch_model.bin, so a tokenizer-only + # module dir (weights pruned) would fail the offline load and must NOT validate. + _staticembedding_repo(tmp_path, monkeypatch, include_weights = False) + assert mc._embedding_marker_in_hf_cache(_STATIC) is False + + +def test_marker_rejects_staticembedding_without_tokenizer(tmp_path, monkeypatch): + # StaticEmbedding.load() hard-reads tokenizer.json (Tokenizer.from_file), so a weights-only + # module dir would fail the offline load and must NOT validate. + _staticembedding_repo(tmp_path, monkeypatch, include_tokenizer = False) + assert mc._embedding_marker_in_hf_cache(_STATIC) is False + + +def test_offline_staticembedding_is_embedding_without_network(tmp_path, monkeypatch): + # End to end: the offline settings path must not 409 a cached StaticEmbedding model, with no + # model_info() network call. + _staticembedding_repo(tmp_path, monkeypatch) + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + _fake_hf_model_info(monkeypatch, _no_network) + assert mc.is_embedding_model(_STATIC) is True + + def _we_dense_repo( tmp_path, monkeypatch, @@ -502,6 +601,9 @@ def _we_dense_repo( ) (snap / "config_sentence_transformers.json").write_text("{}") (snap / "0_WordEmbeddings" / "wordembedding_config.json").write_text("{}") + # WordEmbeddings.load() rebuilds its tokenizer from the module dir, so a complete module + # carries this too; the variable under test here is the Dense module's weights. + (snap / "0_WordEmbeddings" / "whitespacetokenizer_config.json").write_text("{}") (snap / "0_WordEmbeddings" / "model.safetensors").write_bytes(b"\0") (snap / "1_Pooling").mkdir(parents = True) (snap / "1_Pooling" / "config.json").write_text("{}") @@ -660,7 +762,7 @@ def _exploding_iter(repo): monkeypatch.setattr(mc, "_st_cache_repo_dir", _exploding_iter) _fake_hf_model_info( monkeypatch, - lambda name, token = None: types.SimpleNamespace( + lambda name, token = None, **kwargs: types.SimpleNamespace( tags = ["sentence-transformers"], pipeline_tag = None ), ) @@ -685,7 +787,7 @@ def test_online_defers_to_hub_over_stale_marker(tmp_path, monkeypatch): _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa") calls = [] - def _info(model_name, token = None): + def _info(model_name, token = None, **kwargs): calls.append(model_name) return types.SimpleNamespace(tags = ["text-generation"], pipeline_tag = "text-generation") @@ -702,7 +804,7 @@ def test_online_permanent_hub_error_ignores_stale_marker(tmp_path, monkeypatch): class RepositoryNotFoundError(Exception): pass - def _info(model_name, token = None): + def _info(model_name, token = None, **kwargs): raise RepositoryNotFoundError("404 not found") _fake_hf_model_info(monkeypatch, _info) @@ -718,12 +820,31 @@ def test_online_hub_failure_falls_back_to_marker_uncached(tmp_path, monkeypatch) assert ("org/emb", None) not in mc._embedding_detection_cache +def test_online_model_info_is_bounded_and_falls_back_on_dns_death(tmp_path, monkeypatch): + # #6817: with NEITHER offline env set, a dead / renamed network must not hang model_info() + # indefinitely. The call is made with a BOUNDED timeout and, when it raises a network error, + # a CACHED embedder still resolves via the transient fallback -- no real network needed. + _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa", repo_id = "org/emb") + calls = [] + + def _info(model_name, token = None, **kwargs): + calls.append(kwargs.get("timeout")) + raise OSError("Temporary failure in name resolution") # getaddrinfo / DNS-dead + + _fake_hf_model_info(monkeypatch, _info) + assert mc.is_embedding_model("org/emb") is True # transient fallback -> cached marker + assert len(calls) == 1 # model_info WAS attempted, exactly once (no retry storm) + timeout = calls[0] + # A bounded, positive timeout was passed so a dead network fails fast instead of hanging. + assert isinstance(timeout, (int, float)) and 0 < timeout <= 60 + + def test_online_negative_does_not_block_later_offline_download(tmp_path, monkeypatch): # A memoized online negative must not block a later offline detection: the offline path # re-probes the marker, never the memo. _no_cache(monkeypatch) - def _info(model_name, token = None): + def _info(model_name, token = None, **kwargs): return types.SimpleNamespace(tags = ["text-generation"], pipeline_tag = "text-generation") _fake_hf_model_info(monkeypatch, _info) @@ -741,7 +862,7 @@ def test_offline_retains_online_confirmed_positive(tmp_path, monkeypatch): # must retain the positive via the memo + present snapshot, since the marker reads False. _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") - def _info(model_name, token = None): + def _info(model_name, token = None, **kwargs): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -757,7 +878,7 @@ def test_offline_metadata_only_positive_not_trusted_without_cache(monkeypatch): # fail, so the memo must NOT be trusted. _no_cache(monkeypatch) - def _info(model_name, token = None): + def _info(model_name, token = None, **kwargs): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -773,7 +894,7 @@ def test_offline_detects_persisted_tag_only_embedder_after_restart(tmp_path, mon # offline recognition rests on the persisted allowlist + present snapshot. _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") - def _info(model_name, token = None): + def _info(model_name, token = None, **kwargs): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -793,7 +914,7 @@ def test_transient_hub_failure_accepts_a_pinned_tag_only_embedder(tmp_path, monk # a flaky network reclassifies the same cache the offline path accepts. _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") - def _info(model_name, token = None): + def _info(model_name, token = None, **kwargs): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -809,7 +930,7 @@ def test_offline_persisted_verdict_not_trusted_when_uncached(tmp_path, monkeypat # it must NOT be trusted. _no_cache(monkeypatch) - def _info(model_name, token = None): + def _info(model_name, token = None, **kwargs): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -835,7 +956,7 @@ def test_offline_persisted_verdict_not_trusted_when_snapshot_partial(tmp_path, m (refs / "main").write_text("aaa") monkeypatch.setattr(mc, "_st_cache_roots", lambda: [cache_root]) - def _info(model_name, token = None): + def _info(model_name, token = None, **kwargs): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -854,7 +975,7 @@ def test_offline_persisted_verdict_matches_across_casing(tmp_path, monkeypatch): # still match a differently-cased lookup. _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "BAAI/model") - def _info(model_name, token = None): + def _info(model_name, token = None, **kwargs): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -929,7 +1050,7 @@ def test_verdict_pins_the_hub_revision_not_the_cached_one(tmp_path, monkeypatch) ) _tag_only_repo(tmp_path, monkeypatch, "stale_local_commit") - def _info(model_name, token = None): + def _info(model_name, token = None, **kwargs): return types.SimpleNamespace( tags = ["feature-extraction"], pipeline_tag = None, sha = "hub_head_commit" ) @@ -974,7 +1095,7 @@ def test_persist_embedder_is_best_effort_when_home_unwritable(tmp_path, monkeypa monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(blocker / "studio")) _no_cache(monkeypatch) - def _info(model_name, token = None): + def _info(model_name, token = None, **kwargs): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -1001,7 +1122,7 @@ def test_online_uncached_still_uses_network(monkeypatch): _no_cache(monkeypatch) calls = [] - def _info(model_name, token = None): + def _info(model_name, token = None, **kwargs): calls.append(model_name) return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") @@ -1016,7 +1137,7 @@ def test_offline_negative_is_not_cached_then_online_detects(monkeypatch): _no_cache(monkeypatch) calls = [] - def _info(model_name, token = None): + def _info(model_name, token = None, **kwargs): calls.append(model_name) return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 28f68bdf81b..02f2411f66c 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2304,6 +2304,35 @@ def _dir_file_names(dir_path: Path) -> set: # ``Normalize`` read no weights (BoW keeps its data in config.json), so a config suffices. _ST_WEIGHTED_MODULE_NAMES = frozenset({"wordembeddings", "dense", "cnn", "lstm"}) +# ``StaticEmbedding.load()`` (sentence_transformers/models/StaticEmbedding.py) reads a +# ``tokenizer.json`` via ``Tokenizer.from_file`` and then ``model.safetensors`` else +# ``pytorch_model.bin`` -- ``save()`` emits ONLY those two files and NO config. So a +# ``StaticEmbedding`` module placed in its own subdir (e.g. +# ``sentence-transformers/static-retrieval-mrl-en-v1`` -> ``0_StaticEmbedding/`` holding just +# ``tokenizer.json`` + ``model.safetensors``) is a WEIGHTED content module recognized by that +# tokenizer plus a complete Torch weight set, NOT by a module config. (A model2vec-style +# StaticEmbedding at the root path ``.`` ships a ``config.json`` and already validates via the +# root / Transformer load-root path.) +_ST_STATIC_EMBEDDING_TOKENIZER_FILE = "tokenizer.json" + +# ``WordEmbeddings.load()`` rebuilds its tokenizer by calling the configured +# ``tokenizer_class.load(dir)`` from the SAME module dir: ``WhitespaceTokenizer`` / +# ``PhraseTokenizer`` read a ``_config.json`` (see sentence_transformers/models/ +# tokenizer/*.py); a transformers-backed wrapper (newer sentence-transformers) reads the shared +# HF tokenizer assets instead. So a WordEmbeddings dir needs a tokenizer artifact ON TOP of its +# ``wordembedding_config.json`` + weights, else it validates on the config+weights path and then +# fails at ``tokenizer_class.load()``. ``Dense`` / ``CNN`` / ``LSTM`` read no tokenizer. +_ST_WORD_TOKENIZER_CONFIG_FILES = frozenset( + {"whitespacetokenizer_config.json", "phrasetokenizer_config.json"} +) + + +def _names_have_word_embeddings_tokenizer(names: set) -> bool: + """True when *names* supply the tokenizer ``WordEmbeddings.load()`` rebuilds from the module + dir: a ``WhitespaceTokenizer`` / ``PhraseTokenizer`` ``*_config.json`` (its own ``load()`` + reads it), or a shared HF tokenizer asset for a transformers-backed wrapper.""" + return bool(names & _ST_WORD_TOKENIZER_CONFIG_FILES) or _names_have_tokenizer(names) + def _module_dir_is_loadable(cls: str, is_root: bool, dir_path: Path) -> bool: """True when *dir_path* carries the files the sentence-transformers module class *cls* @@ -2311,9 +2340,11 @@ def _module_dir_is_loadable(cls: str, is_root: bool, dir_path: Path) -> bool: * a root / ``Transformer`` module is a full HF load root (config + tokenizer + weights); * ``Normalize`` reads nothing; + * ``StaticEmbedding`` reads a ``tokenizer.json`` + a complete Torch weight set and NO config; * a WEIGHTED module (``WordEmbeddings`` / ``Dense`` / ``CNN`` / ``LSTM``) needs its module config AND a complete Torch weight set (its ``load()`` hard-loads ``model.safetensors`` / - ``pytorch_model.bin``); + ``pytorch_model.bin``); ``WordEmbeddings`` additionally rebuilds a tokenizer from the dir, + so it also needs a tokenizer artifact; * every other module (``BoW``, ``Pooling`` ...) reads only a ``config.json`` / ``*_config``, so a present module config is the load requirement.""" if is_root or "transformer" in cls: @@ -2321,10 +2352,20 @@ def _module_dir_is_loadable(cls: str, is_root: bool, dir_path: Path) -> bool: if cls == "normalize": return True names = _dir_file_names(dir_path) + if cls == "staticembedding": + # No config on disk: recognized by the tokenizer + weights its load() actually reads. + return ( + _ST_STATIC_EMBEDDING_TOKENIZER_FILE in names + and _dir_has_complete_torch_weights(names) + ) if not any(name == "config.json" or name.endswith("_config.json") for name in names): return False if cls in _ST_WEIGHTED_MODULE_NAMES: - return _dir_has_complete_torch_weights(names) + if not _dir_has_complete_torch_weights(names): + return False + if cls == "wordembeddings": + return _names_have_word_embeddings_tokenizer(names) + return True return True @@ -2437,6 +2478,20 @@ def _embedding_marker_in_hf_cache(repo_id: str) -> Optional[bool]: return _snapshot_is_loadable_st_model(snapshot) +# Bounded metadata fetch for the ONLINE branch (#6817). With NEITHER HF_HUB_OFFLINE nor +# TRANSFORMERS_OFFLINE set, an unbounded model_info() can hang on connect / DNS retries when the +# network is dead or the Hub is unreachable -- the #6817 symptom for users who never set an +# offline flag. Cap the call (mirrors utils/security/file_security's bounded model_info fetch, +# _REQUEST_TIMEOUT=10 / _RETRY_TIMEOUT=20) so an unreachable Hub raises FAST and the transient- +# failure handler below resolves a cached model, WITHOUT surrendering the Hub-authoritative +# verdict when the Hub IS reachable (metadata answers well inside this window, so a reachable Hub +# still wins). CAVEAT: huggingface_hub forwards timeout to requests' get(..., timeout=...), which +# bounds the connect + read; a stalled DNS getaddrinfo() can still block past it on some +# resolvers, so this NARROWS -- but does not fully close -- the hang window. HF_HUB_OFFLINE +# remains the only hard guarantee. +_HUB_MODEL_INFO_TIMEOUT = 15.0 + + def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: """Detect embedding/sentence-transformer models via HF metadata. @@ -2489,7 +2544,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 "" From 5847763538b20204af690cc4835b02e430854156 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 10:11:46 +0000 Subject: [PATCH 43/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../tests/test_is_embedding_model_offline.py | 90 +++++++++++++++---- studio/backend/utils/models/model_config.py | 5 +- 2 files changed, 77 insertions(+), 18 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 72a35cba7e5..aa32e8342d6 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -787,7 +787,11 @@ def test_online_defers_to_hub_over_stale_marker(tmp_path, monkeypatch): _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa") calls = [] - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): calls.append(model_name) return types.SimpleNamespace(tags = ["text-generation"], pipeline_tag = "text-generation") @@ -804,7 +808,11 @@ def test_online_permanent_hub_error_ignores_stale_marker(tmp_path, monkeypatch): class RepositoryNotFoundError(Exception): pass - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): raise RepositoryNotFoundError("404 not found") _fake_hf_model_info(monkeypatch, _info) @@ -827,7 +835,11 @@ def test_online_model_info_is_bounded_and_falls_back_on_dns_death(tmp_path, monk _repo(tmp_path, monkeypatch, ("aaa", True), main_ref = "aaa", repo_id = "org/emb") calls = [] - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): calls.append(kwargs.get("timeout")) raise OSError("Temporary failure in name resolution") # getaddrinfo / DNS-dead @@ -844,7 +856,11 @@ def test_online_negative_does_not_block_later_offline_download(tmp_path, monkeyp # re-probes the marker, never the memo. _no_cache(monkeypatch) - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): return types.SimpleNamespace(tags = ["text-generation"], pipeline_tag = "text-generation") _fake_hf_model_info(monkeypatch, _info) @@ -862,7 +878,11 @@ def test_offline_retains_online_confirmed_positive(tmp_path, monkeypatch): # must retain the positive via the memo + present snapshot, since the marker reads False. _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -878,7 +898,11 @@ def test_offline_metadata_only_positive_not_trusted_without_cache(monkeypatch): # fail, so the memo must NOT be trusted. _no_cache(monkeypatch) - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -894,7 +918,11 @@ def test_offline_detects_persisted_tag_only_embedder_after_restart(tmp_path, mon # offline recognition rests on the persisted allowlist + present snapshot. _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -914,7 +942,11 @@ def test_transient_hub_failure_accepts_a_pinned_tag_only_embedder(tmp_path, monk # a flaky network reclassifies the same cache the offline path accepts. _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "org/gte-modernbert") - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -930,7 +962,11 @@ def test_offline_persisted_verdict_not_trusted_when_uncached(tmp_path, monkeypat # it must NOT be trusted. _no_cache(monkeypatch) - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -956,7 +992,11 @@ def test_offline_persisted_verdict_not_trusted_when_snapshot_partial(tmp_path, m (refs / "main").write_text("aaa") monkeypatch.setattr(mc, "_st_cache_roots", lambda: [cache_root]) - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -975,7 +1015,11 @@ def test_offline_persisted_verdict_matches_across_casing(tmp_path, monkeypatch): # still match a differently-cased lookup. _repo(tmp_path, monkeypatch, ("aaa", False), main_ref = "aaa", repo_id = "BAAI/model") - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -1050,7 +1094,11 @@ def test_verdict_pins_the_hub_revision_not_the_cached_one(tmp_path, monkeypatch) ) _tag_only_repo(tmp_path, monkeypatch, "stale_local_commit") - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): return types.SimpleNamespace( tags = ["feature-extraction"], pipeline_tag = None, sha = "hub_head_commit" ) @@ -1095,7 +1143,11 @@ def test_persist_embedder_is_best_effort_when_home_unwritable(tmp_path, monkeypa monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(blocker / "studio")) _no_cache(monkeypatch) - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") _fake_hf_model_info(monkeypatch, _info) @@ -1122,7 +1174,11 @@ def test_online_uncached_still_uses_network(monkeypatch): _no_cache(monkeypatch) calls = [] - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): calls.append(model_name) return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") @@ -1137,7 +1193,11 @@ def test_offline_negative_is_not_cached_then_online_detects(monkeypatch): _no_cache(monkeypatch) calls = [] - def _info(model_name, token = None, **kwargs): + def _info( + model_name, + token = None, + **kwargs, + ): calls.append(model_name) return types.SimpleNamespace(tags = ["feature-extraction"], pipeline_tag = None, sha = "aaa") diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 02f2411f66c..d9259786544 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2354,9 +2354,8 @@ def _module_dir_is_loadable(cls: str, is_root: bool, dir_path: Path) -> bool: names = _dir_file_names(dir_path) if cls == "staticembedding": # No config on disk: recognized by the tokenizer + weights its load() actually reads. - return ( - _ST_STATIC_EMBEDDING_TOKENIZER_FILE in names - and _dir_has_complete_torch_weights(names) + return _ST_STATIC_EMBEDDING_TOKENIZER_FILE in names and _dir_has_complete_torch_weights( + names ) if not any(name == "config.json" or name.endswith("_config.json") for name in names): return False From 95af81bbd8217749adbe6bfff885539a4eed23b2 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 10:30:39 +0000 Subject: [PATCH 44/77] Resolve indexed safetensors shards relative to their index _safetensors_index_complete compared shard basenames against the flat set of files in the index directory, so an index whose weight_map names shards in a subdirectory was treated as incomplete whenever a legacy pytorch_model.bin sat beside it. That falsely blocked a snapshot whose pickle weights are fully covered by a complete, loadable safetensors shard set. Resolve each shard path relative to the index directory instead, and add a regression test for the subdir-mapped shard case. --- studio/backend/tests/test_hf_offline_flags.py | 15 ++++++++++++ .../backend/utils/security/file_security.py | 23 ++++++++++++++----- 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index b6728cd6b88..935b2299f44 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -261,6 +261,21 @@ def _snap(name, files: dict): (idxsafe / "sharded" / "pytorch_model-00001-of-00002.bin").write_bytes(b"\0") assert _evaluate(idxsafe) is False + # A complete safetensors shard set that a safetensors index names in a SUBDIR covers a legacy + # pytorch_model.bin: the loader reads the inert safetensors, so this must be allowed (the + # index shards resolve relative to the index dir, not by basename). + stsub = tmp_path / "stsub" / "aaa" + (stsub / "weights").mkdir(parents = True) + (stsub / "config.json").write_bytes(b"{}") + (stsub / "pytorch_model.bin").write_bytes(b"\0") + (stsub / "model.safetensors.index.json").write_text( + '{"weight_map": {"a": "weights/model-00001-of-00002.safetensors", ' + '"b": "weights/model-00002-of-00002.safetensors"}}' + ) + (stsub / "weights" / "model-00001-of-00002.safetensors").write_bytes(b"\0") + (stsub / "weights" / "model-00002-of-00002.safetensors").write_bytes(b"\0") + assert _evaluate(stsub) is False + def test_security_scan_runs_when_online(monkeypatch): import utils.security.file_security as fs diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 0045e960540..9783ae08c98 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -284,8 +284,11 @@ def _index_weight_map_values(index_path) -> set: _SAFETENSORS_BASE_INDEX = ("model.safetensors.index.json", "pytorch_model.safetensors.index.json") -def _safetensors_index_complete(index_path, present_lower: set) -> bool: - """True when every shard the safetensors index maps is present in the same directory.""" +def _safetensors_index_complete(index_path) -> bool: + """True when every shard the safetensors index maps is present, resolved RELATIVE TO the + index's own directory -- a ``weight_map`` value may name a subdirectory + (``weights/model-00001-of-00002.safetensors``), so comparing basenames alone would miss a + complete set and wrongly treat the inert safetensors as absent.""" import json try: @@ -294,8 +297,17 @@ def _safetensors_index_complete(index_path, present_lower: set) -> bool: ) or {} except (OSError, ValueError): return False # unreadable index -> not a usable safetensors set -> keep the pickle blocked - shards = {str(shard).rsplit("/", 1)[-1].lower() for shard in weight_map.values()} - return bool(shards) and shards <= present_lower + shards = {_normalize_repo_path(str(shard)) for shard in weight_map.values()} + if not shards: + return False + base = index_path.parent + for shard_rel in shards: + try: + if not base.joinpath(*shard_rel.split("/")).is_file(): + return False + except OSError: + return False + return True def _dir_has_loadable_safetensors(files: dict) -> bool: @@ -304,10 +316,9 @@ def _dir_has_loadable_safetensors(files: dict) -> bool: complete indexed shard set. A bare adapter or an orphan shard does not qualify.""" if any(name in files for name in _SAFETENSORS_BASE_UNSHARDED): return True - present_lower = set(files) for index_name in _SAFETENSORS_BASE_INDEX: index_path = files.get(index_name) - if index_path is not None and _safetensors_index_complete(index_path, present_lower): + if index_path is not None and _safetensors_index_complete(index_path): return True return False From 67ac0efaf1e71657d25d4fc06c28b79328bdc4f5 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 10:36:05 +0000 Subject: [PATCH 45/77] Restrict offline weight-completeness check to declared load roots _snapshot_has_complete_weights scanned every directory in a snapshot and accepted it when ANY directory was a complete Transformer load root. When modules.json is present a SentenceTransformer load only opens the declared module paths, so a snapshot whose declared modules are incomplete but which happens to contain an unrelated complete directory was accepted offline and then failed at the first local_files_only load. Restrict the candidate directories to the roots a load actually opens: the snapshot root plus each modules.json module path. For a well-formed snapshot the verdict is unchanged; only a complete directory at an undeclared path no longer vouches for an otherwise-incomplete snapshot. --- .../tests/test_is_embedding_model_offline.py | 73 ++++++++++++++++++- studio/backend/utils/models/model_config.py | 49 ++++++++++++- 2 files changed, 117 insertions(+), 5 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index aa32e8342d6..4239fb57915 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -298,12 +298,18 @@ def _cache_repo_with_files( tokenizer = True, ): """Cache repo whose active snapshot holds modules.json + config + a tokenizer + *files*. - ``tokenizer=False`` omits the tokenizer asset.""" + ``tokenizer=False`` omits the tokenizer asset. + + When a file lives under a subdirectory, ``modules.json`` DECLARES that subdir as a + Transformer module path -- the real SentenceTransformer layout, where each module is loaded + FROM its declared path. A complete directory at an UNDECLARED path is never opened by the + loader, so an ``[]`` modules.json next to files in ``0_Transformer/`` would be an + unrealistic snapshot (#7218). Root-level files need no declaration: a plain load reads the + snapshot root, which is always a candidate.""" hf_root = tmp_path / "hf" repo = hf_root / "models--org--model" snap = repo / "snapshots" / commit snap.mkdir(parents = True) - (snap / "modules.json").write_text("[]") (snap / "config.json").write_text("{}") if tokenizer: (snap / "tokenizer.json").write_text("{}") @@ -311,6 +317,22 @@ def _cache_repo_with_files( target = snap / name target.parent.mkdir(parents = True, exist_ok = True) target.write_bytes(b"\0") + module_dirs = [] + for name in files: + parent = Path(name).parent.as_posix() + if parent not in ("", ".") and parent not in module_dirs: + module_dirs.append(parent) + if module_dirs: + (snap / "modules.json").write_text( + _modules_json( + *( + (str(i), path, "sentence_transformers.models.Transformer") + for i, path in enumerate(module_dirs) + ) + ) + ) + else: + (snap / "modules.json").write_text("[]") (repo / "refs").mkdir(parents = True) (repo / "refs" / "main").write_text(commit) _fake_hf_cache(monkeypatch, hf_root) @@ -377,6 +399,53 @@ def test_marker_rejects_weights_split_from_their_config(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache("org/model") is False +def test_marker_rejects_complete_dir_at_an_undeclared_path(tmp_path, monkeypatch): + # modules.json declares 0_Transformer, so the ST load opens THAT directory. Here it is + # incomplete (config only) while a complete config+tokenizer+weights load root sits at an + # UNDECLARED sibling dir (stray_complete/) the loader never opens. Judging any complete + # directory would accept this snapshot and then 409 at the first local_files_only load, so + # it must be restricted to the declared load roots and read as NOT loadable (#7218 P2). + hf_root = tmp_path / "hf" + repo = hf_root / "models--org--model" + snap = repo / "snapshots" / "aaa" + (snap / "0_Transformer").mkdir(parents = True) + (snap / "modules.json").write_text( + _modules_json(("0", "0_Transformer", "sentence_transformers.models.Transformer")) + ) + (snap / "0_Transformer" / "config.json").write_text("{}") # declared root, incomplete + stray = snap / "stray_complete" # complete Transformer at an UNDECLARED path + stray.mkdir(parents = True) + (stray / "config.json").write_text("{}") + (stray / "tokenizer.json").write_text("{}") + (stray / "model.safetensors").write_bytes(b"\0") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text("aaa") + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + assert mc._embedding_marker_in_hf_cache("org/model") is False + + +def test_marker_accepts_complete_dir_at_the_declared_module_path(tmp_path, monkeypatch): + # Companion to the rejection above: when the complete config+tokenizer+weights sit at the + # DECLARED module path (0_Transformer/), the loader opens exactly that directory, so a + # normal Transformer model must still be recognized (#7218 P2). + hf_root = tmp_path / "hf" + repo = hf_root / "models--org--model" + snap = repo / "snapshots" / "aaa" + (snap / "0_Transformer").mkdir(parents = True) + (snap / "modules.json").write_text( + _modules_json(("0", "0_Transformer", "sentence_transformers.models.Transformer")) + ) + (snap / "0_Transformer" / "config.json").write_text("{}") + (snap / "0_Transformer" / "tokenizer.json").write_text("{}") + (snap / "0_Transformer" / "model.safetensors").write_bytes(b"\0") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text("aaa") + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + assert mc._embedding_marker_in_hf_cache("org/model") is True + + @pytest.mark.parametrize( ("tok_files", "expected"), [ diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index d9259786544..037b416452d 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2256,8 +2256,43 @@ def _dir_is_transformer_load_root(names: set) -> bool: return _dir_has_complete_torch_weights(names) +def _declared_load_root_dirs(snap: Path) -> set: + """The directories a load actually opens as a Transformer/plain-HF load root for *snap*: + the snapshot ROOT plus every module ``path`` declared in ``modules.json``. + + ``modules.json`` is authoritative about WHERE a SentenceTransformer load reads from -- it + calls ``module_class.load(path)`` for each declared module, NOT for an arbitrary complete + directory that happens to sit in the snapshot. So the Transformer-shaped completeness check + must only judge these roots: a complete directory at an UNDECLARED path is never opened by + the loader and must not vouch for a snapshot whose declared modules are incomplete. + + When ``modules.json`` is absent / unreadable / not a list (a plain ``from_pretrained`` + model, or the tag-only embedder shape whose weights live at the root), the load reads the + snapshot root, so return just ``{snap}``. Mirrors the path resolution in + :func:`_snapshot_modules_all_loadable` (empty / ``.`` -> the root; a ``..`` component is + rejected as it would resolve outside the snapshot). Never raises.""" + roots = {snap} + try: + modules = json.loads((snap / "modules.json").read_text(encoding = "utf-8")) + except (OSError, ValueError): + return roots + if not isinstance(modules, list): + return roots + for module in modules: + if not isinstance(module, dict): + continue + raw_path = str(module.get("path") or "").strip() + if raw_path in ("", "."): + continue # the root, already included + rel = raw_path.strip("/") + if not rel or ".." in Path(rel).parts: + continue # never resolve a module path outside the snapshot + roots.add(snap / rel) + return roots + + def _snapshot_has_complete_weights(snap: Path) -> bool: - """True when some LOAD ROOT inside *snap* is complete: a config, a tokenizer and a + """True when some DECLARED LOAD ROOT inside *snap* is complete: a config, a tokenizer and a COMPLETE Torch weight set, all in the SAME directory. Co-location is the point. ``modules.json`` can send SentenceTransformer at a module @@ -2265,12 +2300,20 @@ def _snapshot_has_complete_weights(snap: Path) -> bool: config at the root, a tokenizer elsewhere and ``0_Transformer/model.safetensors`` but no ``0_Transformer/config.json`` would pass a scattered any-of check and then fail the local-only load. Per-directory covers both layouts: a plain HF model has all three at - the root, an ST model inside its Transformer module dir.""" + the root, an ST model inside its Transformer module dir. + + The candidate directories are restricted to the roots a load actually opens + (:func:`_declared_load_root_dirs`): the snapshot root plus each ``modules.json`` module + ``path``. A complete Transformer sitting at some UNDECLARED path (e.g. a stray copy the + loader never reads) must NOT validate a snapshot whose declared modules are incomplete. + For a well-formed snapshot -- a normal download / ``save_pretrained`` places its complete + files only at the root or a declared module path -- this yields the identical verdict.""" try: + candidate_dirs = _declared_load_root_dirs(snap) by_dir: dict = {} for path in snap.rglob("*"): try: - if path.is_file(): + if path.is_file() and path.parent in candidate_dirs: by_dir.setdefault(path.parent, set()).add(path.name) except OSError: continue From 22506970af771b9dac858ef90c811a14b066c961 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 10:46:05 +0000 Subject: [PATCH 46/77] Scan SentenceTransformer Router child module weights offline A Router (legacy Asym) snapshot declares its child sub-modules only in router_config.json, not the top-level modules.json, and Router.load() deserializes each child's weights from its own subdir. A config.json-less child such as query_0_WordEmbeddings (wordembedding_config.json plus a pickle pytorch_model.bin loaded via torch.load) was therefore neither a modules.json-declared load root nor a config.json-bearing dir, so the offline gate skipped its pickle even though the loader deserializes it. Parse router_config.json at each load root and treat every declared child subdir as a load root (bounded BFS, so nested routers are covered), so those child pickles are scanned. Add Router regression tests: a pickle child blocks, a safetensors child is allowed, and a Router in a declared subfolder is followed. --- studio/backend/tests/test_hf_offline_flags.py | 61 +++++++++++++++++++ .../backend/utils/security/file_security.py | 46 +++++++++++++- 2 files changed, 106 insertions(+), 1 deletion(-) diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index 935b2299f44..6502a105c0a 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -276,6 +276,67 @@ def _snap(name, files: dict): (stsub / "weights" / "model-00002-of-00002.safetensors").write_bytes(b"\0") assert _evaluate(stsub) is False + # A SentenceTransformer Router (legacy Asym) declares its child sub-modules ONLY in + # router_config.json (never the top-level modules.json), and Router.load() deserializes each + # child's weights from its own subdir. A config.json-less child (query_0_WordEmbeddings: + # wordembedding_config.json + pytorch_model.bin, no safetensors) is a live pickle vector and + # must be scanned -> blocked, not skipped for lacking a config.json or a modules.json entry. + router = tmp_path / "router" / "aaa" + (router / "query_0_WordEmbeddings").mkdir(parents = True) + (router / "document_0_Transformer").mkdir(parents = True) + (router / "modules.json").write_text( + '[{"idx": 0, "name": "0", "path": "", ' + '"type": "sentence_transformers.models.Router"}]' + ) + (router / "router_config.json").write_text( + '{"types": {"query_0_WordEmbeddings": ' + '"sentence_transformers.models.WordEmbeddings", ' + '"document_0_Transformer": "sentence_transformers.models.Transformer"}, ' + '"structure": {"query": ["query_0_WordEmbeddings"], ' + '"document": ["document_0_Transformer"]}, "parameters": {}}' + ) + (router / "query_0_WordEmbeddings" / "wordembedding_config.json").write_bytes(b"{}") + (router / "query_0_WordEmbeddings" / "pytorch_model.bin").write_bytes(b"\0") + (router / "document_0_Transformer" / "config.json").write_bytes(b"{}") + (router / "document_0_Transformer" / "model.safetensors").write_bytes(b"\0") + assert _evaluate(router) is True + + # The same Router whose child ships model.safetensors beside the pickle is inert: the child is + # scoped as a load root, but the loader reads the safetensors it prefers -> allowed (the fix + # scopes the child without over-blocking a genuinely safe one). + routersafe = tmp_path / "routersafe" / "aaa" + (routersafe / "query_0_WordEmbeddings").mkdir(parents = True) + (routersafe / "modules.json").write_text( + '[{"idx": 0, "name": "0", "path": "", ' + '"type": "sentence_transformers.models.Router"}]' + ) + (routersafe / "router_config.json").write_text( + '{"types": {"query_0_WordEmbeddings": ' + '"sentence_transformers.models.WordEmbeddings"}}' + ) + (routersafe / "query_0_WordEmbeddings" / "wordembedding_config.json").write_bytes(b"{}") + (routersafe / "query_0_WordEmbeddings" / "pytorch_model.bin").write_bytes(b"\0") + (routersafe / "query_0_WordEmbeddings" / "model.safetensors").write_bytes(b"\0") + assert _evaluate(routersafe) is False + + # The Router can itself sit in a modules.json-declared subfolder; the child scan must follow + # router_config.json from THAT dir, not only the snapshot root. + routersub = tmp_path / "routersub" / "aaa" + (routersub / "1_Router" / "query_0_WordEmbeddings").mkdir(parents = True) + (routersub / "modules.json").write_text( + '[{"idx": 0, "name": "1_Router", "path": "1_Router", ' + '"type": "sentence_transformers.models.Router"}]' + ) + (routersub / "1_Router" / "router_config.json").write_text( + '{"types": {"query_0_WordEmbeddings": ' + '"sentence_transformers.models.WordEmbeddings"}}' + ) + (routersub / "1_Router" / "query_0_WordEmbeddings" / "wordembedding_config.json").write_bytes( + b"{}" + ) + (routersub / "1_Router" / "query_0_WordEmbeddings" / "pytorch_model.bin").write_bytes(b"\0") + assert _evaluate(routersub) is True + def test_security_scan_runs_when_online(monkeypatch): import utils.security.file_security as fs diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 9783ae08c98..5e3a1a856e0 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -260,6 +260,12 @@ def _load_scan_target(model_name: str, load_subdirs: tuple) -> tuple: # follows the map and deserializes them wherever they live. _PICKLE_INDEX_FILE = "pytorch_model.bin.index.json" +# SentenceTransformer Router (a.k.a. the legacy Asym) config. It declares the router's child +# sub-modules in its ``types`` map ({route}_{idx}_{ClassName} -> module class) and is the ONLY +# place they are declared -- the top-level modules.json lists just the Router. Router.load() +# deserializes each child from its own subdir, so those subdirs are load roots too. +_ROUTER_CONFIG_FILE = "router_config.json" + def _index_weight_map_values(index_path) -> set: """Repo-relative shard paths a weight-index maps (its ``weight_map`` values), or empty on @@ -323,6 +329,32 @@ def _dir_has_loadable_safetensors(files: dict) -> bool: return False +def _router_child_dirs(root) -> set: + """Child sub-module directories a SentenceTransformer ``Router`` (legacy ``Asym``) at *root* + deserializes. A Router declares its children only in ``router_config.json`` (``types`` maps + ``{route}_{idx}_{ClassName}`` -> module class), NOT in the top-level ``modules.json``, and + ``Router.load()`` calls ``module_class.load(subfolder=model_id)`` on each. A child such as + ``query_0_WordEmbeddings/`` holds ``wordembedding_config.json`` + ``pytorch_model.bin`` and no + ``config.json``, so its pickle is still deserialized and its dir must be a load root. Returns + an empty set when there is no readable router config.""" + import json + + try: + config = json.loads((root / _ROUTER_CONFIG_FILE).read_text(encoding = "utf-8")) + except (OSError, ValueError): + return set() + if not isinstance(config, dict): + return set() + types = config.get("types") + children: set = set() + if isinstance(types, dict): + for model_id in types: + rel = _normalize_repo_path(str(model_id)).strip("/") + if rel and ".." not in rel.split("/"): + children.add(root / rel) + return children + + def _st_load_roots(snap, load_subdirs = ()) -> set: """Directories a load opens ``from_pretrained`` on: the snapshot root, every module path ``modules.json`` declares, and each passed-in ``load_subdirs`` entry. A SentenceTransformer @@ -338,13 +370,25 @@ def _st_load_roots(snap, load_subdirs = ()) -> set: import json modules = json.loads((snap / "modules.json").read_text(encoding = "utf-8")) except (OSError, ValueError): - return roots + modules = None if isinstance(modules, list): for module in modules: if isinstance(module, dict): rel = _normalize_repo_path(str(module.get("path") or "")).strip("/") if rel: roots.add(snap / rel) + # A Router/Asym module declares its child sub-modules in router_config.json, not modules.json, + # and Router.load() deserializes each child's weights from its own subdir. Treat those child + # dirs as load roots too so a pickle in a config.json-less child (e.g. + # query_0_WordEmbeddings/pytorch_model.bin) is scanned. Bounded BFS: every child is a strict + # subpath and a visited set stops any cycle; scanning extra dirs only tightens the gate. + pending = list(roots) + while pending: + current = pending.pop() + for child in _router_child_dirs(current): + if child not in roots: + roots.add(child) + pending.append(child) return roots From ae15078c25109dfd4bd654fdf4f40e6aaf96b008 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 10:47:20 +0000 Subject: [PATCH 47/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_hf_offline_flags.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index 6502a105c0a..0909b306aa2 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -285,8 +285,7 @@ def _snap(name, files: dict): (router / "query_0_WordEmbeddings").mkdir(parents = True) (router / "document_0_Transformer").mkdir(parents = True) (router / "modules.json").write_text( - '[{"idx": 0, "name": "0", "path": "", ' - '"type": "sentence_transformers.models.Router"}]' + '[{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Router"}]' ) (router / "router_config.json").write_text( '{"types": {"query_0_WordEmbeddings": ' @@ -307,12 +306,10 @@ def _snap(name, files: dict): routersafe = tmp_path / "routersafe" / "aaa" (routersafe / "query_0_WordEmbeddings").mkdir(parents = True) (routersafe / "modules.json").write_text( - '[{"idx": 0, "name": "0", "path": "", ' - '"type": "sentence_transformers.models.Router"}]' + '[{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Router"}]' ) (routersafe / "router_config.json").write_text( - '{"types": {"query_0_WordEmbeddings": ' - '"sentence_transformers.models.WordEmbeddings"}}' + '{"types": {"query_0_WordEmbeddings": "sentence_transformers.models.WordEmbeddings"}}' ) (routersafe / "query_0_WordEmbeddings" / "wordembedding_config.json").write_bytes(b"{}") (routersafe / "query_0_WordEmbeddings" / "pytorch_model.bin").write_bytes(b"\0") @@ -328,8 +325,7 @@ def _snap(name, files: dict): '"type": "sentence_transformers.models.Router"}]' ) (routersub / "1_Router" / "router_config.json").write_text( - '{"types": {"query_0_WordEmbeddings": ' - '"sentence_transformers.models.WordEmbeddings"}}' + '{"types": {"query_0_WordEmbeddings": "sentence_transformers.models.WordEmbeddings"}}' ) (routersub / "1_Router" / "query_0_WordEmbeddings" / "wordembedding_config.json").write_bytes( b"{}" From 9fd2b134b136a5d62a5a85865cb83cf76545fb28 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 11:10:39 +0000 Subject: [PATCH 48/77] Do not treat an unreferenced config subdir as an offline load root The offline pickle gate skipped a directory only when it was neither a declared load root nor held a config.json. Because _st_load_roots already resolves every real load root (snapshot root, modules.json / load_subdirs dirs, Router children), the config.json fallback only ever promoted an UNREFERENCED subdir -- a nested checkpoint-500/ or archive/ that ships its own config.json + pytorch_model.bin -- to a load root. from_pretrained never descends into such a subdir and the online scan ignores the same unindexed pickle, so offline mode wrongly blocked a model the loader reads from a clean safetensors root. Scope the pickle to directory in roots only, and add a regression test (a stray checkpoint-500/ no longer blocks; a modules.json-declared module dir still does). --- studio/backend/tests/test_hf_offline_flags.py | 21 +++++++++++++++++-- .../backend/utils/security/file_security.py | 10 ++++++--- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index 0909b306aa2..d0776b90548 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -195,14 +195,31 @@ def _snap(name, files: dict): (stray / "archive" / "pytorch_model.bin").write_bytes(b"\0") assert _evaluate(stray) is False - # A pickle in a real MODULE load root (a subdir with config.json) and no safetensors - # there is a live vector -> blocked. + # A pickle in a real MODULE load root (a modules.json-declared Transformer subdir) and no + # safetensors there is a live vector -> blocked. modroot = tmp_path / "modroot" / "aaa" (modroot / "0_Transformer").mkdir(parents = True) + (modroot / "modules.json").write_text( + '[{"idx": 0, "name": "0_Transformer", "path": "0_Transformer", ' + '"type": "sentence_transformers.models.Transformer"}]' + ) (modroot / "0_Transformer" / "config.json").write_bytes(b"{}") (modroot / "0_Transformer" / "pytorch_model.bin").write_bytes(b"\0") assert _evaluate(modroot) is True + # An UNREFERENCED nested checkpoint (its own config.json + pickle) is NOT a load root: no + # modules.json / router_config declares it and from_pretrained never descends into it, so it + # must not block a model the loader reads from a clean safetensors root -- matching the online + # scan, which ignores the same unindexed subdir pickle. A stray config.json does not make a + # subdir a load root. + unref = tmp_path / "unref" / "aaa" + (unref / "checkpoint-500").mkdir(parents = True) + (unref / "config.json").write_bytes(b"{}") + (unref / "model.safetensors").write_bytes(b"\0") + (unref / "checkpoint-500" / "config.json").write_bytes(b"{}") + (unref / "checkpoint-500" / "pytorch_model.bin").write_bytes(b"\0") + assert _evaluate(unref) is False + # A pickle in a modules.json-declared module dir WITHOUT a config.json (a WordEmbeddings # module: wordembedding_config.json + pytorch_model.bin) is still deserialized by the ST # loader, so it must be scanned -> blocked, not skipped for lacking config.json. diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 5e3a1a856e0..cf2ecbf6e0f 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -422,9 +422,13 @@ def _cached_pickle_weight_files(snap, load_subdirs = ()) -> list: hits: set = set() for directory, names in by_dir_pickle.items(): files = by_dir_files.get(directory, {}) - # A pickle is deserialized only at a load root (a declared module dir, load subdir, the - # snapshot root, or a plain from_pretrained root with config.json). Elsewhere it is unread. - if directory not in roots and "config.json" not in files: + # A pickle is deserialized only at an actual load root: the snapshot root, a declared + # modules.json / load_subdirs dir, or a Router child (all resolved by _st_load_roots). + # A stray config.json in an UNREFERENCED subdir (a nested checkpoint-500/ or archive/) + # does NOT make it a load root -- from_pretrained never descends into it and the ST load + # opens only declared modules -- so it must not block, matching the online scan which + # ignores the same unindexed subdir pickle. + if directory not in roots: continue base = [n for n in names if _PICKLE_WEIGHT_RE.match(n.lower())] if base and not _dir_has_loadable_safetensors(files): From be9a5810c3478403189d0e312800b86135cca0df Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 11:10:39 +0000 Subject: [PATCH 49/77] Classify a root Router (Asym) model as loadable offline _module_dir_is_loadable applied Transformer root requirements (config + tokenizer + weights) to every root module, so a Router saved at the snapshot root -- which carries only modules.json + router_config.json and loads its weights from child subdirs -- was classified not loadable offline, and is_embedding_model missed a cached Router embedder. Dispatch on the module class before the root Transformer fallback: a Router/Asym dir is loadable when router_config.json parses and every declared child subdir is loadable (validated recursively through _module_dir_is_loadable, so nested routers and every child type are covered) with at least one embedding-producing child. This also tightens a non-root Router, which previously validated on the mere presence of router_config.json without checking its children. Add Router regression tests (root and declared subfolder, complete and incomplete-child). --- .../tests/test_is_embedding_model_offline.py | 87 +++++++++++++++++++ studio/backend/utils/models/model_config.py | 49 ++++++++++- 2 files changed, 134 insertions(+), 2 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 4239fb57915..3b13198b9b6 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -576,6 +576,93 @@ def test_marker_requires_a_tokenizer_for_wordembeddings(tmp_path, monkeypatch, t assert mc._embedding_marker_in_hf_cache(_GLOVE) is True +# ── Router / Asym models (asymmetric query/document, sparse encoders) built from modules.json ── + + +_ROUTER_REPO = "sentence-transformers/router-asym-demo" + + +def _router_repo( + tmp_path, + monkeypatch, + *, + root_path = "", + complete = True, + commit = "aaa", + repo_id = _ROUTER_REPO, +): + """A cached SentenceTransformer Router (legacy Asym) model. The top-level modules.json declares + ONLY the Router (at ``root_path``); ``router_config.json`` declares the child sub-modules in its + ``types`` map and each child lives in its own subdir. The Router carries NO config / tokenizer / + weights of its own -- ``Router.load()`` reads router_config.json and loads each child from its + subdir -- so the Transformer-shaped root check alone misclassifies it (#7218). ``complete=False`` + strips the Transformer child's weights so the load would fail.""" + hf_root = tmp_path / "hf" + repo = hf_root / f"models--{repo_id.replace('/', '--')}" + snap = repo / "snapshots" / commit + snap.mkdir(parents = True) + router_dir = snap if root_path in ("", ".") else snap / root_path + router_dir.mkdir(parents = True, exist_ok = True) + (snap / "modules.json").write_text( + _modules_json(("0", root_path, "sentence_transformers.models.Router")) + ) + (router_dir / "router_config.json").write_text( + json.dumps( + { + "types": { + "query_0_Transformer": "sentence_transformers.models.Transformer", + "query_1_Pooling": "sentence_transformers.models.Pooling", + }, + "structure": {"query": ["query_0_Transformer", "query_1_Pooling"]}, + "parameters": {}, + } + ) + ) + transformer = router_dir / "query_0_Transformer" + transformer.mkdir(parents = True) + (transformer / "config.json").write_text("{}") + (transformer / "tokenizer.json").write_text("{}") + if complete: + (transformer / "model.safetensors").write_bytes(b"\0") + pooling = router_dir / "query_1_Pooling" + pooling.mkdir(parents = True) + (pooling / "config.json").write_text("{}") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + return snap + + +def test_marker_accepts_root_router_model(tmp_path, monkeypatch): + # A root Router (Asym): modules.json declares only the Router at the empty root path, and it + # loads via router_config.json + child subdirs with no root weights of its own. It must be + # recognized offline instead of being held to Transformer root requirements (#7218). + _router_repo(tmp_path, monkeypatch) + assert mc._embedding_marker_in_hf_cache(_ROUTER_REPO) is True + + +def test_marker_rejects_root_router_with_incomplete_child(tmp_path, monkeypatch): + # Control: the Router's Transformer child is missing its weights, so Router.load() would fail + # -> must NOT validate. + _router_repo(tmp_path, monkeypatch, complete = False) + assert mc._embedding_marker_in_hf_cache(_ROUTER_REPO) is False + + +def test_marker_accepts_router_in_declared_subfolder(tmp_path, monkeypatch): + # The Router can sit in a declared subfolder; validation must follow router_config.json from + # that dir, not the snapshot root. + _router_repo(tmp_path, monkeypatch, root_path = "2_Router") + assert mc._embedding_marker_in_hf_cache(_ROUTER_REPO) is True + + +def test_marker_rejects_subfolder_router_with_incomplete_child(tmp_path, monkeypatch): + # A Router (in a declared subfolder) whose child weights are missing must not validate: the + # check validates the router's declared children, not merely that router_config.json exists. + _router_repo(tmp_path, monkeypatch, root_path = "2_Router", complete = False) + assert mc._embedding_marker_in_hf_cache(_ROUTER_REPO) is False + + # ── StaticEmbedding models (model2vec / static-retrieval) built from modules.json ── diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 037b416452d..0eb22833014 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2377,11 +2377,22 @@ def _names_have_word_embeddings_tokenizer(names: set) -> bool: return bool(names & _ST_WORD_TOKENIZER_CONFIG_FILES) or _names_have_tokenizer(names) +# ``Router.load()`` (a.k.a. the legacy ``Asym``) reads ``router_config.json`` and loads each child +# sub-module from its OWN subdir named there; it carries no Transformer weights of its own, so a +# root Router without root config/tokenizer/weights is still loadable when its children are. The +# children are declared only in this config (never the top-level modules.json). +_ST_ROUTER_MODULE_NAMES = frozenset({"router", "asym"}) +_ST_ROUTER_CONFIG_FILE = "router_config.json" + + def _module_dir_is_loadable(cls: str, is_root: bool, dir_path: Path) -> bool: """True when *dir_path* carries the files the sentence-transformers module class *cls* reads in its own ``load()`` (see sentence_transformers/models/*.py): - * a root / ``Transformer`` module is a full HF load root (config + tokenizer + weights); + * a ``Transformer`` module (or a plain root with no recognized ST module class) is a full HF + load root (config + tokenizer + weights); + * a ``Router`` / ``Asym`` reads ``router_config.json`` and loads its declared children, so it + is loadable when those children are -- even at the root, where it has no weights of its own; * ``Normalize`` reads nothing; * ``StaticEmbedding`` reads a ``tokenizer.json`` + a complete Torch weight set and NO config; * a WEIGHTED module (``WordEmbeddings`` / ``Dense`` / ``CNN`` / ``LSTM``) needs its module @@ -2390,7 +2401,11 @@ def _module_dir_is_loadable(cls: str, is_root: bool, dir_path: Path) -> bool: so it also needs a tokenizer artifact; * every other module (``BoW``, ``Pooling`` ...) reads only a ``config.json`` / ``*_config``, so a present module config is the load requirement.""" - if is_root or "transformer" in cls: + if "transformer" in cls: + return _dir_is_transformer_load_root(_dir_file_names(dir_path)) + if cls in _ST_ROUTER_MODULE_NAMES: + return _router_dir_is_loadable(dir_path) + if is_root: return _dir_is_transformer_load_root(_dir_file_names(dir_path)) if cls == "normalize": return True @@ -2411,6 +2426,36 @@ def _module_dir_is_loadable(cls: str, is_root: bool, dir_path: Path) -> bool: return True +def _router_dir_is_loadable(dir_path: Path) -> bool: + """True when *dir_path* holds a loadable sentence-transformers ``Router`` (legacy ``Asym``): + a readable ``router_config.json`` whose ``types`` maps each child ``{route}_{idx}_{ClassName}`` + to a module class, with every child subdir carrying the files that child's own ``load()`` reads + (validated through :func:`_module_dir_is_loadable`, so nested routers and every child type are + covered) and at least one embedding-producing child. ``Router.load()`` reads this config and + loads each child from its subdir, so a Router directory needs no Transformer weights of its own. + Never raises.""" + try: + config = json.loads((dir_path / _ST_ROUTER_CONFIG_FILE).read_text(encoding = "utf-8")) + except (OSError, ValueError): + return False + if not isinstance(config, dict): + return False + types = config.get("types") + if not isinstance(types, dict) or not types: + return False + saw_content_module = False + for model_id, module_type in types.items(): + child_cls = str(module_type or "").rsplit(".", 1)[-1].strip().lower() + rel = str(model_id or "").strip().strip("/") + if not child_cls or not rel or ".." in Path(rel).parts: + return False # malformed / traversing child reference + if not _module_dir_is_loadable(child_cls, False, dir_path / rel): + return False + if child_cls not in _ST_STRUCTURAL_MODULE_NAMES: + saw_content_module = True + return saw_content_module + + def _snapshot_modules_all_loadable(snap: Path) -> bool: """True when *snap* is a COMPLETE non-Transformer sentence-transformers model that the Transformer-shaped :func:`_snapshot_has_complete_weights` misses -- e.g. one built from a From 9c77fd78f1542db62e692f65a4ea63c6b486f698 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 12:07:11 +0000 Subject: [PATCH 50/77] Require every declared module before accepting an offline cache _snapshot_is_loadable_st_model returned has_complete_weights OR modules_all_loadable, so a complete 0_Transformer short-circuited the or and vouched for the whole snapshot even when a declared sibling module was missing its serialized weights; SentenceTransformer builds every module in modules.json, so that snapshot passed offline validation and then failed the local-only load. When modules.json declares a non-empty list it is now authoritative (modules_all_loadable validates every declared module); has_complete_weights stays the fallback only for an empty/non-list modules.json (the plain from_pretrained root). Also add the weight-bearing modules whose load() hard-loads via load_torch_weights and previously fell to the config-only path -- LayerNorm, WeightedLayerPooling, SparseAutoEncoder -- to _ST_WEIGHTED_MODULE_NAMES, with source citations and the deliberate exclusions (Pooling/Normalize/BoW/WordWeights read no weights on load). Add parametrized regression tests over LayerNorm/WeightedLayerPooling/Dense (a weightless sibling rejects, a complete sibling accepts). --- .../tests/test_is_embedding_model_offline.py | 94 ++++++++++++++++++ studio/backend/utils/models/model_config.py | 99 +++++++++++++++---- 2 files changed, 174 insertions(+), 19 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 3b13198b9b6..1cfec6d3fe7 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -782,6 +782,100 @@ def test_marker_requires_weights_for_a_dense_module(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache("org/we-dense") is True +# ── a complete Transformer must not vouch for a weightless declared sibling module (#7218 P2) ── + + +_TF_PLUS = "org/transformer-plus-module" + + +def _transformer_plus_module_repo( + tmp_path, + monkeypatch, + *, + module_dir, + module_type, + include_module_weights, + commit = "aaa", + repo_id = _TF_PLUS, +): + """A cached SentenceTransformer whose modules.json declares a COMPLETE ``0_Transformer`` load + root (config + tokenizer + weights) PLUS a second weight-bearing module (``LayerNorm`` / + ``WeightedLayerPooling`` / ``Dense``) that always carries its ``config.json`` and, only when + ``include_module_weights``, its ``model.safetensors``. A SentenceTransformer load builds EVERY + declared module -- each of these modules' ``load()`` ends in ``load_torch_weights``, which + RAISES without a weight file -- so a weightless sibling must fail offline validation even + though the Transformer alone is complete.""" + hf_root = tmp_path / "hf" + repo = hf_root / f"models--{repo_id.replace('/', '--')}" + snap = repo / "snapshots" / commit + transformer = snap / "0_Transformer" + transformer.mkdir(parents = True) + (transformer / "config.json").write_text("{}") + (transformer / "tokenizer.json").write_text("{}") + (transformer / "model.safetensors").write_bytes(b"\0") + module = snap / module_dir + module.mkdir(parents = True) + (module / "config.json").write_text("{}") + if include_module_weights: + (module / "model.safetensors").write_bytes(b"\0") + (snap / "modules.json").write_text( + _modules_json( + ("0", "0_Transformer", "sentence_transformers.models.Transformer"), + ("1", module_dir, module_type), + ) + ) + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + return snap + + +_TF_SIBLING_MODULES = [ + ("1_LayerNorm", "sentence_transformers.models.LayerNorm"), + ("1_WeightedLayerPooling", "sentence_transformers.models.WeightedLayerPooling"), + ("1_Dense", "sentence_transformers.models.Dense"), +] + + +@pytest.mark.parametrize(("module_dir", "module_type"), _TF_SIBLING_MODULES) +def test_marker_rejects_complete_transformer_with_weightless_sibling( + tmp_path, monkeypatch, module_dir, module_type +): + # #7218 P2: a snapshot with a complete 0_Transformer AND a declared sibling module whose + # weights are missing must NOT validate -- the old `_snapshot_has_complete_weights(...) OR + # _snapshot_modules_all_loadable(...)` accepted it on the Transformer alone and then 409'd at + # the local-only load. LayerNorm / WeightedLayerPooling / Dense.load() each end in + # load_torch_weights (raises ValueError without model.safetensors/pytorch_model.bin). Dense + # was ALREADY a recognized weight-bearing module, so its case proves the or-short-circuit fix + # independently of the newly-added module names. + _transformer_plus_module_repo( + tmp_path, + monkeypatch, + module_dir = module_dir, + module_type = module_type, + include_module_weights = False, + ) + assert mc._embedding_marker_in_hf_cache(_TF_PLUS) is False + + +@pytest.mark.parametrize(("module_dir", "module_type"), _TF_SIBLING_MODULES) +def test_marker_accepts_complete_transformer_with_complete_sibling( + tmp_path, monkeypatch, module_dir, module_type +): + # Companion to the rejection above: once the sibling module carries its weights the whole + # declared set is loadable, so a WELL-FORMED complete snapshot must still validate -- the fix + # only stops accepting a genuinely incomplete declared module, it does not over-reject. + _transformer_plus_module_repo( + tmp_path, + monkeypatch, + module_dir = module_dir, + module_type = module_type, + include_module_weights = True, + ) + assert mc._embedding_marker_in_hf_cache(_TF_PLUS) is True + + def test_marker_accepts_complete_bow_model(tmp_path, monkeypatch): # A BoW module keeps its vocab in config.json and writes NO weight file; a complete cache is # still loadable via BoW.load(config.json), so it must validate. diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 0eb22833014..47a9e7e16a0 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2340,12 +2340,41 @@ def _dir_file_names(dir_path: Path) -> set: return set() -# Non-Transformer modules whose ST ``load()`` reads a weight file. Each does -# ``load_safetensors(model.safetensors)`` else ``torch.load(pytorch_model.bin)`` with no -# fallback, so it needs its module config AND a complete Torch weight set -- a config alone -# would validate here and then raise FileNotFoundError at load. ``Pooling`` / ``BoW`` / -# ``Normalize`` read no weights (BoW keeps its data in config.json), so a config suffices. -_ST_WEIGHTED_MODULE_NAMES = frozenset({"wordembeddings", "dense", "cnn", "lstm"}) +# Non-Transformer sentence-transformers modules whose ``load()`` HARD-loads a serialized Torch +# weight set: each ends in ``Module.load_torch_weights`` (sentence_transformers/base/modules/ +# module.py), which loads ``model.safetensors`` else ``pytorch_model.bin`` and RAISES +# ``ValueError`` ("Could not find 'model.safetensors' or 'pytorch_model.bin' ...") when NEITHER +# exists. So each needs its module config AND a complete Torch weight set -- a config alone would +# pass offline validation here and then 409 at the local-only load (#7218 P2). Enumerated from the +# sentence-transformers v5 source; each cite is the ``file:line`` of the ``load_torch_weights`` +# call in that class's ``load()``: +# * WordEmbeddings sentence_transformer/modules/word_embeddings.py:141 (+ tokenizer, below) +# * Dense base/modules/dense.py:124 +# * CNN sentence_transformer/modules/cnn.py:90 (cnn_config.json) +# * LSTM sentence_transformer/modules/lstm.py:96 (lstm_config.json) +# * LayerNorm sentence_transformer/modules/layer_norm.py:52 +# * WeightedLayerPooling sentence_transformer/modules/weighted_layer_pooling.py:74 +# * SparseAutoEncoder sparse_encoder/modules/sparse_auto_encoder.py:226 +# All ship a ``config.json`` / ``*_config.json`` (LayerNorm / WeightedLayerPooling / +# SparseAutoEncoder inherit the default ``config_file_name = "config.json"``), so they pass the +# config-presence gate above. Deliberately EXCLUDED (their ``load()`` reads no weight file): +# ``Pooling`` / ``SpladePooling`` / ``Normalize`` / ``Dropout`` read nothing; ``BoW`` / +# ``WordWeights`` keep their data in ``config.json`` and reconstruct it on the config-only base +# ``Module.load()`` (their ``save()`` writes weights, but ``load()`` never reads them). +# ``StaticEmbedding`` (no config: tokenizer + weights) is handled by its own branch above, and +# ``SparseStaticEmbedding`` is only CONDITIONALLY weighted (its ``load()`` skips the weight file +# when ``config.json`` names an ``idf.json`` ``path``), so neither is listed here. +_ST_WEIGHTED_MODULE_NAMES = frozenset( + { + "wordembeddings", + "dense", + "cnn", + "lstm", + "layernorm", + "weightedlayerpooling", + "sparseautoencoder", + } +) # ``StaticEmbedding.load()`` (sentence_transformers/models/StaticEmbedding.py) reads a # ``tokenizer.json`` via ``Tokenizer.from_file`` and then ``model.safetensors`` else @@ -2464,13 +2493,18 @@ def _snapshot_modules_all_loadable(snap: Path) -> bool: ``local_files_only`` load builds these from ``modules.json`` by calling ``module_class.load(path)`` per declared module. - Additive and conservative: it only ever ACCEPTS (the caller OR-s it with the Transformer - check) and never rejects, so it cannot regress the Transformer path. It validates only when - ``modules.json`` parses as a non-empty list AND every declared module's ``path`` directory - carries the files that module's own ``load()`` reads -- so a bare / pruned cache (empty - ``modules.json``, a missing module directory, a WordEmbeddings module with its weights - stripped) does not validate. It also requires at least one embedding-producing module, so a - degenerate Pooling/Normalize-only list is rejected. Never raises.""" + When ``modules.json`` declares a non-empty list this is the AUTHORITATIVE snapshot check: + :func:`_snapshot_is_loadable_st_model` calls it INSTEAD of the Transformer-shaped + :func:`_snapshot_has_complete_weights`, because the load builds EVERY declared module and one + complete ``0_Transformer`` must NOT vouch for a sibling module (e.g. ``LayerNorm`` / + ``WeightedLayerPooling`` / ``Dense``) whose serialized weights are missing (#7218 P2). It + validates only when ``modules.json`` parses as a non-empty list AND every declared module's + ``path`` directory carries the files that module's own ``load()`` reads -- so a bare / pruned + cache (empty ``modules.json``, a missing module directory, a weight-bearing module with its + weights stripped) does not validate. It also requires at least one embedding-producing module, + so a degenerate Pooling/Normalize-only list is rejected. Behaviour-preserving for a WELL-FORMED + complete snapshot (every declared module has its files, so acceptance is unchanged). Never + raises.""" try: modules = json.loads((snap / "modules.json").read_text(encoding = "utf-8")) except (OSError, ValueError): @@ -2500,17 +2534,44 @@ def _snapshot_modules_all_loadable(snap: Path) -> bool: return saw_content_module +def _modules_json_declares_modules(snap: Path) -> bool: + """True when *snap*'s ``modules.json`` parses as a NON-EMPTY list of module entries. + + ``modules.json`` is authoritative about what a SentenceTransformer load builds: when it + declares one or more modules the loader calls ``module_class.load(path)`` for EVERY one, so + the whole declared set must be loadable. An empty list / a non-list / an unreadable file + declares nothing usable -- the load then reads the snapshot root as a plain Transformer -- so + the caller falls back to the Transformer-shaped weight check. Never raises.""" + try: + modules = json.loads((snap / "modules.json").read_text(encoding = "utf-8")) + except (OSError, ValueError): + return False + return isinstance(modules, list) and len(modules) > 0 + + def _snapshot_is_loadable_st_model(snap: Path) -> bool: """True when *snap* holds a loadable sentence-transformers model. ``modules.json`` alone - is insufficient -- the security preflight downloads exactly that file, so require the - marker plus a complete load: a Transformer-shaped weight set - (:func:`_snapshot_has_complete_weights`), OR a complete non-Transformer module set such as - ``0_WordEmbeddings`` / ``BoW`` (:func:`_snapshot_modules_all_loadable`), which carries none - of the Transformer config/tokenizer/weights the first check requires.""" + is insufficient -- the security preflight downloads exactly that file, so require the marker + plus a complete load. + + ``modules.json`` is authoritative about what the loader builds. When it declares a non-empty + list of modules, a SentenceTransformer load builds EVERY one, so the whole declared set must + be loadable (:func:`_snapshot_modules_all_loadable`): a complete ``0_Transformer`` must NOT + vouch for a sibling module (e.g. ``LayerNorm`` / ``WeightedLayerPooling`` / ``Dense``) whose + serialized weights are missing -- that snapshot passes offline validation and then 409s at + the local-only load (#7218 P2). Only when ``modules.json`` declares nothing usable (an empty + list / not a list / unreadable) does the load read the snapshot ROOT as a plain Transformer, + so fall back to the Transformer-shaped weight set (:func:`_snapshot_has_complete_weights`), + which also covers the plain ``from_pretrained`` root layout. Behaviour-preserving for a + WELL-FORMED complete snapshot -- a normal download has every declared module's weights, so + acceptance is unchanged; only a snapshot with a genuinely incomplete declared module is now + rejected.""" try: if not (snap / "modules.json").is_file(): return False - return _snapshot_has_complete_weights(snap) or _snapshot_modules_all_loadable(snap) + if _modules_json_declares_modules(snap): + return _snapshot_modules_all_loadable(snap) + return _snapshot_has_complete_weights(snap) except OSError: return False From cf8ece8b6e8befc1a65d1ba420d4eb965c1093d1 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 12:17:39 +0000 Subject: [PATCH 51/77] Reject self-referential Router children instead of recursing forever _router_dir_is_loadable validates each router_config.json child through _module_dir_is_loadable, which re-enters _router_dir_is_loadable for a Router child. A malformed types entry naming the router's own directory (a key of ".", which normalizes to the same dir) made that recursion never descend, so it looped until RecursionError -- breaking the documented never-raises contract and turning a crafted/corrupted cached model into a 500 from is_embedding_model instead of a graceful unverifiable result. A real child reference is a subdir and always resolves deeper, so reject any child whose resolved path is the router dir itself. Add a regression test (a router_config naming "." as a Router child returns False without raising). --- .../tests/test_is_embedding_model_offline.py | 22 +++++++++++++++++++ studio/backend/utils/models/model_config.py | 9 ++++++-- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 1cfec6d3fe7..dc16d37e462 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -663,6 +663,28 @@ def test_marker_rejects_subfolder_router_with_incomplete_child(tmp_path, monkeyp assert mc._embedding_marker_in_hf_cache(_ROUTER_REPO) is False +def test_marker_rejects_self_referential_router_without_recursing(tmp_path, monkeypatch): + # A malformed router_config.json whose types names "." (the Router's own dir) as a Router child + # must be rejected, not recursed into forever: the child path resolves back to the same dir, so + # without the self-reference guard _router_dir_is_loadable would re-enter until RecursionError + # and turn is_embedding_model into a 500. It must instead return False gracefully (never raise). + hf_root = tmp_path / "hf" + repo = hf_root / "models--org--self-router" + snap = repo / "snapshots" / "aaa" + snap.mkdir(parents = True) + (snap / "modules.json").write_text( + _modules_json(("0", "", "sentence_transformers.models.Router")) + ) + (snap / "router_config.json").write_text( + json.dumps({"types": {".": "sentence_transformers.models.Router"}}) + ) + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text("aaa") + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + assert mc._embedding_marker_in_hf_cache("org/self-router") is False + + # ── StaticEmbedding models (model2vec / static-retrieval) built from modules.json ── diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 47a9e7e16a0..8a0cf7011b4 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2462,7 +2462,9 @@ def _router_dir_is_loadable(dir_path: Path) -> bool: (validated through :func:`_module_dir_is_loadable`, so nested routers and every child type are covered) and at least one embedding-producing child. ``Router.load()`` reads this config and loads each child from its subdir, so a Router directory needs no Transformer weights of its own. - Never raises.""" + A child reference is a subdir name, so a normalized child path always resolves DEEPER; a + self-referential entry (a ``types`` key of ``"."`` -> the same dir) is rejected so a malformed + ``router_config.json`` cannot recurse forever -- this function never raises.""" try: config = json.loads((dir_path / _ST_ROUTER_CONFIG_FILE).read_text(encoding = "utf-8")) except (OSError, ValueError): @@ -2478,7 +2480,10 @@ def _router_dir_is_loadable(dir_path: Path) -> bool: rel = str(model_id or "").strip().strip("/") if not child_cls or not rel or ".." in Path(rel).parts: return False # malformed / traversing child reference - if not _module_dir_is_loadable(child_cls, False, dir_path / rel): + child_dir = dir_path / rel + if child_dir == dir_path: + return False # self-referential child (e.g. "."): would recurse until RecursionError + if not _module_dir_is_loadable(child_cls, False, child_dir): return False if child_cls not in _ST_STRUCTURAL_MODULE_NAMES: saw_content_module = True From e46c7efc2b5bf00f64c9bf3a140e05681ff2bc6b Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 12:28:03 +0000 Subject: [PATCH 52/77] Treat a destructuring __all__ assignment as opaque _collect_dunder_all detected __all__ only as a direct ast.Name assignment target, so a binding through a destructuring target (__all__, meta = [...], v -> an ast.Tuple) was skipped entirely, leaving an empty, non-opaque export set. A newly hoisted import re-exported only through that assignment was then falsely flagged HOISTED-IMPORT-UNUSED. Its value cannot be mapped statically, so mark the export set opaque when __all__ is reached only through a destructuring / item / attr target, matching how the collector already handles other unreadable __all__ forms. Add a self-test case. --- scripts/verify_import_hoist.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index b35b61613cb..b4d05aec4f6 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -499,7 +499,9 @@ def _collect_dunder_all(tree: ast.Module) -> tuple[set[str], bool]: ``+=`` extends, so prior opacity persists. A bare ``__all__: list[str]`` annotation has no runtime value and is skipped entirely. ``__all__.append("X")`` / ``.extend([...])`` are runtime extenders read the same way as ``+=``; any other - ``__all__`` method call (``insert``/``remove``/...) is opaque. + ``__all__`` method call (``insert``/``remove``/...) is opaque. ``__all__`` bound only + through a destructuring target (``__all__, meta = [...], v``) or an item/attr target + is opaque too -- its value cannot be mapped statically, so the set is not exhaustive. """ names: set[str] = set() opaque = False @@ -528,7 +530,17 @@ def _collect_dunder_all(tree: ast.Module) -> tuple[set[str], bool]: targets, replaces = [node.target], False else: continue + # A direct `__all__ = [...]` / `+= [...]` target lets us read the value below. `__all__` + # reached only through a destructuring target (`__all__, meta = [...], v` -> an ast.Tuple) + # or an item/attr target cannot be mapped to its value statically, so mark the export set + # opaque instead of silently missing the binding and false-positiving HOISTED-IMPORT-UNUSED. if not any(isinstance(t, ast.Name) and t.id == "__all__" for t in targets): + if any( + isinstance(sub, ast.Name) and sub.id == "__all__" + for t in targets + for sub in ast.walk(t) + ): + opaque = True continue value = node.value entries: set[str] = set() @@ -872,6 +884,14 @@ def used_targets(analysis) -> set[str]: 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n__all__ += sorted(["b"])\n', None, ), + "destructured_all_is_opaque_keeps_reexport": ( + # __all__ bound only through a tuple-unpacking target cannot be mapped to its value + # statically, so it is opaque: a new import re-exported only through it must not be + # flagged as an unused hoist (mirrors the unreadable "+=" case) + 'from pkg import a\n__all__, meta = ["a"], 1\n', + 'from pkg import a\nfrom pkg import b\n__all__, meta = ["a", "b"], 1\n', + None, + ), "readable_reassign_resets_opacity": ( # an opaque "+=" then a readable replacing "=": the final __all__ is empty, so a # hoisted unused import is a genuine bad hoist -- the reset must un-stick opacity From f84c7e17b7ed41afd12e26c3f8dfc3855becd9cd Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 12:39:41 +0000 Subject: [PATCH 53/77] Canonicalize declared module paths before scoping the offline pickle gate A repo could declare a traversing module path such as 0/../evil in modules.json (or a router_config child), which SentenceTransformer resolves to evil/ and deserializes evil/pytorch_model.bin. _st_load_roots recorded the raw snap/"0/../evil", which never equals the snap/evil that rglob yields, so the offline pickle gate skipped that directory and a malicious repo slipped a pickle past the newly added gate. Add _canonical_load_dir to collapse ./ and ../ components lexically and reject an upward escape, and route the modules.json paths, load_subdirs and router children through it so the gate scopes the same normalized directory the loader opens. Add regression tests for a traversing modules.json path and router child. --- studio/backend/tests/test_hf_offline_flags.py | 28 +++++++++++++++ .../backend/utils/security/file_security.py | 36 ++++++++++++++----- 2 files changed, 55 insertions(+), 9 deletions(-) diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index d0776b90548..da3362e3a22 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -350,6 +350,34 @@ def _snap(name, files: dict): (routersub / "1_Router" / "query_0_WordEmbeddings" / "pytorch_model.bin").write_bytes(b"\0") assert _evaluate(routersub) is True + # A modules.json declaring a TRAVERSING module path ("0/../evil") is resolved by the ST loader + # to evil/, which it deserializes; the gate must scope the NORMALIZED dir and block + # evil/pytorch_model.bin -- recording the raw snap/"0/../evil" (which never equals the real + # snap/evil rglob yields) would let a malicious repo slip a pickle past the offline gate. + trav = tmp_path / "trav" / "aaa" + (trav / "evil").mkdir(parents = True) + (trav / "modules.json").write_text( + '[{"idx": 0, "name": "m", "path": "0/../evil", ' + '"type": "sentence_transformers.models.WordEmbeddings"}]' + ) + (trav / "evil" / "wordembedding_config.json").write_bytes(b"{}") + (trav / "evil" / "pytorch_model.bin").write_bytes(b"\0") + assert _evaluate(trav) is True + + # The same traversal through a Router child path ("0/../evil") must also be normalized and + # blocked, not skipped for containing "..". + rtrav = tmp_path / "rtrav" / "aaa" + (rtrav / "evil").mkdir(parents = True) + (rtrav / "modules.json").write_text( + '[{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Router"}]' + ) + (rtrav / "router_config.json").write_text( + '{"types": {"0/../evil": "sentence_transformers.models.WordEmbeddings"}}' + ) + (rtrav / "evil" / "wordembedding_config.json").write_bytes(b"{}") + (rtrav / "evil" / "pytorch_model.bin").write_bytes(b"\0") + assert _evaluate(rtrav) is True + def test_security_scan_runs_when_online(monkeypatch): import utils.security.file_security as fs diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index cf2ecbf6e0f..e600e4f000b 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -93,6 +93,24 @@ def _normalize_repo_path(path: str) -> str: return p +def _canonical_load_dir(base, rel: str): + """``base`` / ``rel`` with ``.`` / ``..`` components collapsed lexically (no filesystem + access), or None when ``rel`` is empty, the base itself, or escapes ``base``. + + A ``modules.json`` / ``router_config.json`` / load-subdir path is repo-controlled. + SentenceTransformer resolves an entry such as ``0/../evil`` to ``evil/`` and deserializes + ``evil/pytorch_model.bin``, so the offline gate must scope that SAME normalized directory: + the raw ``base / "0/../evil"`` never equals the ``base / "evil"`` that ``rglob`` yields, so a + pickle there would slip the gate. A path that escapes ``base`` (a leading ``..`` after + normalization) is rejected -- a legitimate declared module path never traverses upward.""" + import posixpath + + norm = posixpath.normpath(_normalize_repo_path(rel).strip("/")) + if norm in ("", ".") or norm == ".." or norm.startswith("../"): + return None + return base / norm + + def _file_suffix(path: str) -> str: """Lowercase ``.ext`` of the basename, or ``""`` if none.""" base = _normalize_repo_path(path).rsplit("/", 1)[-1] @@ -349,9 +367,9 @@ def _router_child_dirs(root) -> set: children: set = set() if isinstance(types, dict): for model_id in types: - rel = _normalize_repo_path(str(model_id)).strip("/") - if rel and ".." not in rel.split("/"): - children.add(root / rel) + child = _canonical_load_dir(root, str(model_id)) + if child is not None: + children.add(child) return children @@ -363,9 +381,9 @@ def _st_load_roots(snap, load_subdirs = ()) -> set: still deserialized and must be treated as a load root.""" roots = {snap} for subdir in load_subdirs or (): - rel = _normalize_repo_path(str(subdir)).strip("/") - if rel: - roots.add(snap / rel) + root = _canonical_load_dir(snap, str(subdir)) + if root is not None: + roots.add(root) try: import json modules = json.loads((snap / "modules.json").read_text(encoding = "utf-8")) @@ -374,9 +392,9 @@ def _st_load_roots(snap, load_subdirs = ()) -> set: if isinstance(modules, list): for module in modules: if isinstance(module, dict): - rel = _normalize_repo_path(str(module.get("path") or "")).strip("/") - if rel: - roots.add(snap / rel) + root = _canonical_load_dir(snap, str(module.get("path") or "")) + if root is not None: + roots.add(root) # A Router/Asym module declares its child sub-modules in router_config.json, not modules.json, # and Router.load() deserializes each child's weights from its own subdir. Treat those child # dirs as load roots too so a pickle in a config.json-less child (e.g. From a1323be2c500c54b1ba4d28ebc151d7cbb175cbf Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 13:36:48 +0000 Subject: [PATCH 54/77] Close offline embedding-classification completeness gaps Five real offline misclassifications, each a false negative (the #6817 hang recurs) or false positive (accepted then 409s at the local_files_only load). Dispatch _module_dir_is_loadable on the module class before the root Transformer fallback. A module with save_in_root=True (every InputModule: WordEmbeddings, StaticEmbedding, SparseStaticEmbedding, Transformer, Router) is saved at the snapshot root, so a root WordEmbeddings was wrongly held to Transformer requirements (an HF tokenizer it never writes) and classified not loadable. CLIPModel is Transformer-shaped: CLIPModel.load() reads AutoModel weights plus AutoProcessor, so a config-only CLIP dir must not validate. SparseStaticEmbedding needs a tokenizer plus either idf.json or a complete torch weight set (conditionally weight-bearing); a config alone is not enough. A present but empty or malformed modules.json is not loadable and does not fall back to a root Transformer: with modules.json present the loader never takes the plain-Transformer path (base/model.py _load_config_modules). The tag-only no-modules.json embedder is classified separately via _snapshot_has_complete_weights. Validate a sharded weight index against its weight_map (every mapped shard present, resolved relative to the index dir) instead of trusting the index file's mere existence, mirroring the security-side check. Add regression tests for all five. --- .../tests/test_is_embedding_model_offline.py | 372 +++++++++++++++++- studio/backend/utils/models/model_config.py | 228 +++++++++-- 2 files changed, 553 insertions(+), 47 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index dc16d37e462..6a1cddc7385 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -46,6 +46,19 @@ def _build_structlog_stub(): import utils.models.model_config as mc # noqa: E402 +# A realistic modules.json for a snapshot whose config + tokenizer + weights sit at the ROOT: a +# single Transformer declared at path "" -- exactly how SentenceTransformer saves the first module +# when its save_in_root is True (base/model.py:661-666). SentenceTransformer treats the PRESENCE of +# modules.json as the ST marker but builds ZERO modules from an empty "[]" list and NEVER falls back +# to a root Transformer while modules.json is present (base/model.py _load_config_modules), so an +# "[]" body is an UNREALISTIC, unloadable snapshot (#7218 P4). Declaring the real root Transformer +# keeps these fixtures asserting a genuinely loadable model while the root files remain the single +# gating load root the tests exercise. +_ROOT_TRANSFORMER_MODULES = json.dumps( + [{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Transformer"}] +) + + @pytest.fixture(autouse = True) def _clean_state(tmp_path, monkeypatch): mc._embedding_detection_cache.clear() @@ -81,7 +94,7 @@ def _repo( (d / "tokenizer.json").write_text("{}") (d / "model.safetensors").write_bytes(b"\0") if is_st: - (d / "modules.json").write_text("[]") + (d / "modules.json").write_text(_ROOT_TRANSFORMER_MODULES) dirs.append(d) if main_ref is not None: refs = repo / "refs" @@ -189,7 +202,7 @@ def _st_snapshot( ): snap = root / repo_dir / "snapshots" / commit snap.mkdir(parents = True) - (snap / "modules.json").write_text("[]") + (snap / "modules.json").write_text(_ROOT_TRANSFORMER_MODULES) if loadable: (snap / "config.json").write_text("{}") (snap / "tokenizer.json").write_text("{}") @@ -304,8 +317,9 @@ def _cache_repo_with_files( Transformer module path -- the real SentenceTransformer layout, where each module is loaded FROM its declared path. A complete directory at an UNDECLARED path is never opened by the loader, so an ``[]`` modules.json next to files in ``0_Transformer/`` would be an - unrealistic snapshot (#7218). Root-level files need no declaration: a plain load reads the - snapshot root, which is always a candidate.""" + unrealistic snapshot (#7218). Root-level files declare a single Transformer at path ``""`` -- + the save_in_root layout a plain load reads from the snapshot root -- rather than an empty + ``[]`` list, which builds ZERO modules and is not a loadable ST model (#7218 P4).""" hf_root = tmp_path / "hf" repo = hf_root / "models--org--model" snap = repo / "snapshots" / commit @@ -332,7 +346,7 @@ def _cache_repo_with_files( ) ) else: - (snap / "modules.json").write_text("[]") + (snap / "modules.json").write_text(_ROOT_TRANSFORMER_MODULES) (repo / "refs").mkdir(parents = True) (repo / "refs" / "main").write_text(commit) _fake_hf_cache(monkeypatch, hf_root) @@ -1483,3 +1497,351 @@ def _info( monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) assert mc.is_embedding_model("org/gte-modernbert") is True assert calls == ["org/gte-modernbert"] + + +# ── #7218 P1: a save_in_root=True module declared at the snapshot root path "" ── + + +_ROOT_WE = "org/root-wordembeddings" + + +def _root_wordembeddings_repo( + tmp_path, + monkeypatch, + *, + include_tokenizer = True, + commit = "aaa", + repo_id = _ROOT_WE, +): + """A WordEmbeddings model whose WordEmbeddings module is the FIRST module, saved at the ROOT + path "" -- the save_in_root=True layout (WordEmbeddings inherits InputModule.save_in_root=True; + SentenceTransformer.save writes the first such module at the root, base/model.py:661-666). The + root holds config_sentence_transformers.json + wordembedding_config.json + a whitespace-tokenizer + config + embedding weights, with a 1_Pooling subdir and NO HF config.json / tokenizer.json. + WordEmbeddings.load() reads exactly those files, but the Transformer-shaped root check fails on + the tokenizer, so the pristine `if is_root: _dir_is_transformer_load_root(...)` returned False and + the offline settings route 409'd it (#7218 P1). include_tokenizer=False prunes the tokenizer.""" + hf_root = tmp_path / "hf" + repo = hf_root / f"models--{repo_id.replace('/', '--')}" + snap = repo / "snapshots" / commit + (snap / "1_Pooling").mkdir(parents = True) + (snap / "modules.json").write_text( + _modules_json( + ("0", "", "sentence_transformers.models.WordEmbeddings"), + ("1", "1_Pooling", "sentence_transformers.models.Pooling"), + ) + ) + (snap / "config_sentence_transformers.json").write_text("{}") + (snap / "wordembedding_config.json").write_text("{}") + if include_tokenizer: + (snap / "whitespacetokenizer_config.json").write_text("{}") + (snap / "model.safetensors").write_bytes(b"\0") + (snap / "1_Pooling" / "config.json").write_text("{}") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + return snap + + +def test_marker_accepts_root_wordembeddings_model(tmp_path, monkeypatch): + # A modern WordEmbeddings model saves its first module at the root path "" (save_in_root=True), + # so the WE files (wordembedding_config.json + whitespacetokenizer_config.json + weights) live at + # the root with NO HF config.json / tokenizer.json. Its WordEmbeddings.load() reads exactly those + # files, so it must be recognized -- the pristine code held the root to Transformer requirements + # (a usable HF tokenizer) it never saves and returned False, so #6817 recurred (#7218 P1). + _root_wordembeddings_repo(tmp_path, monkeypatch) + assert mc._embedding_marker_in_hf_cache(_ROOT_WE) is True + + +def test_marker_rejects_root_wordembeddings_without_tokenizer(tmp_path, monkeypatch): + # Control: strip the whitespace-tokenizer config -> WordEmbeddings.load() would fail at + # tokenizer_class.load(dir), so the root module is not loadable. + _root_wordembeddings_repo(tmp_path, monkeypatch, include_tokenizer = False) + assert mc._embedding_marker_in_hf_cache(_ROOT_WE) is False + + +_ROOT_STATIC = "org/root-staticembedding" + + +def _root_staticembedding_repo(tmp_path, monkeypatch, *, include_weights = True, commit = "aaa"): + """A StaticEmbedding model whose StaticEmbedding module is saved at the ROOT path "" (save_in_root + =True). The root holds config_sentence_transformers.json + tokenizer.json + weights and NO module + config. StaticEmbedding already passes the pristine is_root Transformer check because it ships a + tokenizer.json + config_sentence_transformers.json + weights; dispatching StaticEmbedding through + its own branch BEFORE the root fallback must keep it loadable (#7218 P1).""" + hf_root = tmp_path / "hf" + repo = hf_root / f"models--{_ROOT_STATIC.replace('/', '--')}" + snap = repo / "snapshots" / commit + snap.mkdir(parents = True) + (snap / "modules.json").write_text( + _modules_json(("0", "", "sentence_transformers.models.StaticEmbedding")) + ) + (snap / "config_sentence_transformers.json").write_text("{}") + (snap / "tokenizer.json").write_text("{}") + if include_weights: + (snap / "model.safetensors").write_bytes(b"\0") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + return snap + + +def test_marker_accepts_root_staticembedding_model(tmp_path, monkeypatch): + # A root StaticEmbedding (tokenizer.json + weights, NO config) must stay loadable once dispatch + # is on the class before the is_root Transformer fallback (#7218 P1). + _root_staticembedding_repo(tmp_path, monkeypatch) + assert mc._embedding_marker_in_hf_cache(_ROOT_STATIC) is True + + +def test_marker_rejects_root_staticembedding_without_weights(tmp_path, monkeypatch): + # The class branch still enforces StaticEmbedding.load()'s requirements: a root StaticEmbedding + # with its weights pruned (tokenizer + config_sentence_transformers.json only) is not loadable. + _root_staticembedding_repo(tmp_path, monkeypatch, include_weights = False) + assert mc._embedding_marker_in_hf_cache(_ROOT_STATIC) is False + + +# ── #7218 P3: CLIPModel is a full HF model + processor, not config-only ── + + +_CLIP = "org/clip" + + +def _clip_repo(tmp_path, monkeypatch, *, complete, commit = "aaa"): + """A SentenceTransformer whose modules.json declares a CLIPModel at 0_CLIPModel. CLIPModel(load) + -> Transformer.__init__ reads AutoConfig + AutoModel.from_pretrained (weights) + AutoProcessor + (processor/tokenizer), so a config-only dir is NOT loadable. complete=False writes only + config.json (no processor/tokenizer asset, no weights).""" + hf_root = tmp_path / "hf" + repo = hf_root / f"models--{_CLIP.replace('/', '--')}" + snap = repo / "snapshots" / commit + (snap / "0_CLIPModel").mkdir(parents = True) + (snap / "modules.json").write_text( + _modules_json(("0", "0_CLIPModel", "sentence_transformers.models.CLIPModel")) + ) + (snap / "config_sentence_transformers.json").write_text("{}") + (snap / "0_CLIPModel" / "config.json").write_text("{}") + if complete: + (snap / "0_CLIPModel" / "tokenizer.json").write_text("{}") # AutoProcessor asset + (snap / "0_CLIPModel" / "model.safetensors").write_bytes(b"\0") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + + +def test_marker_rejects_config_only_clip_module(tmp_path, monkeypatch): + # #7218 P3: CLIPModel.load() loads a full HF model + processor, so a config-only CLIP module (no + # processor/tokenizer asset, no weights) is NOT loadable. Pristine routed clipmodel to the + # generic config-only `return True` and accepted it, then 409'd at the local-only load. + _clip_repo(tmp_path, monkeypatch, complete = False) + assert mc._embedding_marker_in_hf_cache(_CLIP) is False + + +def test_marker_accepts_complete_clip_module(tmp_path, monkeypatch): + # Companion: config + processor/tokenizer + weights make the CLIP module a loadable HF load root. + _clip_repo(tmp_path, monkeypatch, complete = True) + assert mc._embedding_marker_in_hf_cache(_CLIP) is True + + +# ── #7218 P3: SparseStaticEmbedding is CONDITIONALLY weight-bearing (idf.json OR weights) ── + + +_SPARSE_STATIC = "org/sparse-static" + + +def _sparse_static_repo( + tmp_path, + monkeypatch, + *, + payload, + include_tokenizer = True, + commit = "aaa", +): + """A SparseEncoder whose modules.json declares a SparseStaticEmbedding at 0_SparseStaticEmbedding. + Its load() reads a tokenizer (AutoTokenizer.from_pretrained) + EITHER an idf.json (from_json + branch, NO weights) OR a complete Torch weight set (load_torch_weights). ``payload`` selects what + backs it: "idf" (idf.json), "weights" (model.safetensors) or "none" (neither).""" + hf_root = tmp_path / "hf" + repo = hf_root / f"models--{_SPARSE_STATIC.replace('/', '--')}" + snap = repo / "snapshots" / commit + mod = snap / "0_SparseStaticEmbedding" + mod.mkdir(parents = True) + (snap / "modules.json").write_text( + _modules_json( + ("0", "0_SparseStaticEmbedding", "sentence_transformers.sparse_encoder.models.SparseStaticEmbedding") + ) + ) + (snap / "config_sentence_transformers.json").write_text("{}") + (mod / "config.json").write_text("{}") + if include_tokenizer: + (mod / "tokenizer.json").write_text("{}") + if payload == "idf": + (mod / "idf.json").write_text("{}") + elif payload == "weights": + (mod / "model.safetensors").write_bytes(b"\0") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + + +@pytest.mark.parametrize("payload", ["idf", "weights"]) +def test_marker_accepts_sparse_static_embedding_variants(tmp_path, monkeypatch, payload): + # Both complete variants are loadable: the idf.json (from_json) variant has NO model weights, so + # requiring weights would wrongly reject it; the weights variant loads via load_torch_weights. + _sparse_static_repo(tmp_path, monkeypatch, payload = payload) + assert mc._embedding_marker_in_hf_cache(_SPARSE_STATIC) is True + + +def test_marker_rejects_config_only_sparse_static_embedding(tmp_path, monkeypatch): + # #7218 P3: a SparseStaticEmbedding with a tokenizer but NEITHER idf.json NOR weights has no + # payload (load_torch_weights raises), so it is not loadable. Pristine hit the generic config-only + # `return True` and accepted it. + _sparse_static_repo(tmp_path, monkeypatch, payload = "none") + assert mc._embedding_marker_in_hf_cache(_SPARSE_STATIC) is False + + +def test_marker_rejects_sparse_static_embedding_without_tokenizer(tmp_path, monkeypatch): + # AutoTokenizer.from_pretrained is unconditional in SparseStaticEmbedding.load(), so an idf.json + + # config WITHOUT any tokenizer asset is not loadable -- pristine accepted it on the config alone. + _sparse_static_repo(tmp_path, monkeypatch, payload = "idf", include_tokenizer = False) + assert mc._embedding_marker_in_hf_cache(_SPARSE_STATIC) is False + + +# ── #7218 P4: a PRESENT but empty / malformed modules.json is not a loadable ST model ── + + +def _root_weights_snapshot(tmp_path, monkeypatch, modules_body, *, repo_id = "org/present-modules"): + """A snapshot with a COMPLETE root weight set (config + tokenizer + weights) and a modules.json + whose raw body is ``modules_body``. Lets a test assert that a present-but-empty / malformed + modules.json is not loadable even though the root weights are complete.""" + hf_root = tmp_path / "hf" + repo = hf_root / f"models--{repo_id.replace('/', '--')}" + snap = repo / "snapshots" / "aaa" + snap.mkdir(parents = True) + (snap / "modules.json").write_text(modules_body) + (snap / "config.json").write_text("{}") + (snap / "tokenizer.json").write_text("{}") + (snap / "model.safetensors").write_bytes(b"\0") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text("aaa") + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + return repo_id + + +def test_marker_empty_modules_json_is_not_loadable(tmp_path, monkeypatch): + # #7218 P4: an empty "[]" modules.json builds ZERO modules (base/model.py _load_config_modules + # iterates the list) and NEVER falls back to a root Transformer while modules.json is present, so + # a complete root weight set does NOT make it loadable. Pristine fell back to + # _snapshot_has_complete_weights and accepted it, then 409'd at the local-only load. + repo_id = _root_weights_snapshot(tmp_path, monkeypatch, "[]") + assert mc._embedding_marker_in_hf_cache(repo_id) is False + + +def test_marker_malformed_modules_json_is_not_loadable(tmp_path, monkeypatch): + # A malformed modules.json RAISES in json.load during _load_config_modules, so the load fails; it + # must NOT fall back to the root Transformer weight set (#7218 P4). Pristine accepted it via the + # _snapshot_has_complete_weights fall-through. + repo_id = _root_weights_snapshot(tmp_path, monkeypatch, "{ not valid json") + assert mc._embedding_marker_in_hf_cache(repo_id) is False + + +def test_marker_non_list_modules_json_is_not_loadable(tmp_path, monkeypatch): + # A modules.json that parses but is not a list declares nothing the loader can build. + repo_id = _root_weights_snapshot(tmp_path, monkeypatch, '{"modules": []}') + assert mc._embedding_marker_in_hf_cache(repo_id) is False + + +# ── #7218 P5: a sharded weight-index must have every mapped shard present ── + + +_SHARDED = "org/sharded" + + +def _sharded_transformer_repo( + tmp_path, + monkeypatch, + *, + weight_map, + present_shards, + module_path = "0_Transformer", + commit = "aaa", +): + """A SentenceTransformer with a Transformer declared at ``module_path`` holding config + + tokenizer + a model.safetensors.index.json whose ``weight_map`` is the given dict, plus + ``present_shards`` written to disk (paths relative to the module dir). Lets a test reference a + shard the index maps but leave it absent, or place shards in a subdirectory.""" + hf_root = tmp_path / "hf" + repo = hf_root / f"models--{_SHARDED.replace('/', '--')}" + snap = repo / "snapshots" / commit + mod = snap / module_path + mod.mkdir(parents = True) + (snap / "modules.json").write_text( + _modules_json(("0", module_path, "sentence_transformers.models.Transformer")) + ) + (snap / "config_sentence_transformers.json").write_text("{}") + (mod / "config.json").write_text("{}") + (mod / "tokenizer.json").write_text("{}") + (mod / "model.safetensors.index.json").write_text(json.dumps({"weight_map": weight_map})) + for shard_rel in present_shards: + target = mod / Path(shard_rel) + target.parent.mkdir(parents = True, exist_ok = True) + target.write_bytes(b"\0") + (repo / "refs").mkdir(parents = True) + (repo / "refs" / "main").write_text(commit) + _fake_hf_cache(monkeypatch, hf_root) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + + +def test_marker_rejects_sharded_index_with_a_missing_mapped_shard(tmp_path, monkeypatch): + # #7218 P5: the index's weight_map is authoritative. It maps two numbered shards (present) AND a + # third shard that is MISSING, so the shard set is incomplete. The pristine filename-numbering + # heuristic saw the complete {1,2} set + the index FILE and accepted it; parsing the weight_map + # rejects the incomplete set, mirroring file_security._safetensors_index_complete. + _sharded_transformer_repo( + tmp_path, + monkeypatch, + weight_map = { + "a": "model-00001-of-00002.safetensors", + "b": "model-00002-of-00002.safetensors", + "c": "model-extra.safetensors", + }, + present_shards = ["model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"], + ) + assert mc._embedding_marker_in_hf_cache(_SHARDED) is False + + +def test_marker_accepts_sharded_index_with_subdir_shards(tmp_path, monkeypatch): + # #7218 P5: a weight_map value may name a subdirectory. Every mapped shard present under weights/ + # -> loadable. Pristine inspected only the module dir's basenames (the shards live in a subdir, so + # it found no numbered shards) and wrongly rejected the complete set (a false negative -> #6817). + _sharded_transformer_repo( + tmp_path, + monkeypatch, + weight_map = { + "a": "weights/model-00001-of-00002.safetensors", + "b": "weights/model-00002-of-00002.safetensors", + }, + present_shards = [ + "weights/model-00001-of-00002.safetensors", + "weights/model-00002-of-00002.safetensors", + ], + ) + assert mc._embedding_marker_in_hf_cache(_SHARDED) is True + + +def test_marker_accepts_sharded_index_with_all_mapped_shards(tmp_path, monkeypatch): + # Companion: a complete uniform shard set the weight_map maps -> loadable. + _sharded_transformer_repo( + tmp_path, + monkeypatch, + weight_map = { + "a": "model-00001-of-00002.safetensors", + "b": "model-00002-of-00002.safetensors", + }, + present_shards = ["model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"], + ) + assert mc._embedding_marker_in_hf_cache(_SHARDED) is True diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 8a0cf7011b4..9828ba61eca 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2241,19 +2241,99 @@ def _dir_has_complete_torch_weights(names: set) -> bool: return False +# Weight-index files a ``from_pretrained`` / sentence-transformers load follows to locate sharded +# weights. Its ``weight_map`` is AUTHORITATIVE about which shards the load reads, so every mapped +# shard must be present -- the ``model-*-of-*`` FILENAME numbering above is only a proxy for it. +# All (stem, ext) combinations of the base weight set are listed so a ``pytorch_model.bin`` shard +# set is validated the same way as a ``model.safetensors`` one. +_ST_WEIGHT_INDEX_FILES = ( + "model.safetensors.index.json", + "model.bin.index.json", + "pytorch_model.safetensors.index.json", + "pytorch_model.bin.index.json", +) + + +def _index_weight_set_complete(index_path: Path) -> Optional[bool]: + """Parse a weight-index and confirm every shard its ``weight_map`` references is present, + resolved RELATIVE TO the index's own directory (a ``weight_map`` value may name a + subdirectory, e.g. ``weights/model-00001-of-00002.safetensors``). Mirrors + ``utils/security/file_security._safetensors_index_complete`` on the classification side (#7218 + P5) so a snapshot with a sharded index but a MISSING shard is rejected instead of validated + then 409'd at the local_files_only load. + + Returns True when every mapped shard is present, False when the index parses but its shard set + is empty or incomplete, and None when the index is unreadable / not valid JSON / has no + ``weight_map`` (a stub or partial write) so the caller can fall back to the filename-numbering + heuristic. Never raises.""" + try: + data = json.loads(index_path.read_text(encoding = "utf-8")) + except (OSError, ValueError): + return None + weight_map = data.get("weight_map") if isinstance(data, dict) else None + if not isinstance(weight_map, dict): + return None + shards = set() + for shard in weight_map.values(): + rel = str(shard).strip().replace("\\", "/") + while rel.startswith("./"): + rel = rel[2:] + rel = rel.lstrip("/") + if rel: + shards.add(rel) + if not shards: + return None + base = index_path.parent + for shard_rel in shards: + try: + if not base.joinpath(*shard_rel.split("/")).is_file(): + return False + except OSError: + return False + return True + + +def _dir_weight_set_is_complete(dir_path: Path, names: set) -> bool: + """True when *dir_path* (whose direct file basenames are *names*) holds a COMPLETE Torch weight + set the way a load reads it: a single unsharded ``model.safetensors`` / ``pytorch_model.bin``; + or, when a weight-index is present AND parseable, every shard its ``weight_map`` references + (authoritative, resolved relative to *dir_path*, subdirectory values included); else the + filename-numbering shard heuristic (:func:`_dir_has_complete_torch_weights`) as a fallback for + an unreadable / stub index. Never raises.""" + # 1. An unsharded base weight is always sufficient. + for stem in ("model", "pytorch_model"): + for ext in ("safetensors", "bin"): + if f"{stem}.{ext}" in names: + return True + # 2. A present, parseable weight-index is authoritative: require every mapped shard present. + saw_parseable_index = False + for index_name in _ST_WEIGHT_INDEX_FILES: + if index_name in names: + verdict = _index_weight_set_complete(dir_path / index_name) + if verdict is True: + return True + if verdict is False: + saw_parseable_index = True + if saw_parseable_index: + return False # a mapped shard is missing -> the local_files_only load would fail + # 3. Fallback: numbering heuristic on basenames (unreadable / stub index, or no index at all). + return _dir_has_complete_torch_weights(names) + + _ST_CONFIG_FILES = ("config.json", "config_sentence_transformers.json") -def _dir_is_transformer_load_root(names: set) -> bool: - """True when one directory's files (*names*) form a COMPLETE Transformer / plain-HF load - root: an ST/HF config, a usable tokenizer and a complete Torch weight set, all co-located. - This is the shape ``modules.json`` sends a Transformer module (or a bare from_pretrained) - at -- the per-directory unit the snapshot checks share.""" +def _dir_is_transformer_load_root(dir_path: Path, names: set) -> bool: + """True when one directory's files (*names*, taken from *dir_path*) form a COMPLETE Transformer + / plain-HF load root: an ST/HF config, a usable tokenizer and a complete Torch weight set, all + co-located. This is the shape ``modules.json`` sends a Transformer module (or a bare + from_pretrained) at -- the per-directory unit the snapshot checks share. *dir_path* is threaded + through so a sharded weight-index is validated against its ``weight_map`` (#7218 P5).""" if not any(cfg in names for cfg in _ST_CONFIG_FILES): return False if not _names_have_tokenizer(names): return False - return _dir_has_complete_torch_weights(names) + return _dir_weight_set_is_complete(dir_path, names) def _declared_load_root_dirs(snap: Path) -> set: @@ -2317,8 +2397,8 @@ def _snapshot_has_complete_weights(snap: Path) -> bool: by_dir.setdefault(path.parent, set()).add(path.name) except OSError: continue - for names in by_dir.values(): - if _dir_is_transformer_load_root(names): + for dir_path, names in by_dir.items(): + if _dir_is_transformer_load_root(dir_path, names): return True return False except OSError: @@ -2356,8 +2436,8 @@ def _dir_file_names(dir_path: Path) -> set: # * WeightedLayerPooling sentence_transformer/modules/weighted_layer_pooling.py:74 # * SparseAutoEncoder sparse_encoder/modules/sparse_auto_encoder.py:226 # All ship a ``config.json`` / ``*_config.json`` (LayerNorm / WeightedLayerPooling / -# SparseAutoEncoder inherit the default ``config_file_name = "config.json"``), so they pass the -# config-presence gate above. Deliberately EXCLUDED (their ``load()`` reads no weight file): +# SparseAutoEncoder inherit the default ``config_file_name = "config.json"``), so they satisfy the +# module-config gate in their branch. Deliberately EXCLUDED (their ``load()`` reads no weight file): # ``Pooling`` / ``SpladePooling`` / ``Normalize`` / ``Dropout`` read nothing; ``BoW`` / # ``WordWeights`` keep their data in ``config.json`` and reconstruct it on the config-only base # ``Module.load()`` (their ``save()`` writes weights, but ``load()`` never reads them). @@ -2413,46 +2493,102 @@ def _names_have_word_embeddings_tokenizer(names: set) -> bool: _ST_ROUTER_MODULE_NAMES = frozenset({"router", "asym"}) _ST_ROUTER_CONFIG_FILE = "router_config.json" +# Transformer-SHAPED module classes: sentence-transformers loads each as a full HF model with a +# tokenizer/processor, so a directory needs config + tokenizer/processor + weights (a Transformer +# load root), not a config alone. ``CLIPModel(Transformer)`` (sentence_transformer/modules/ +# clip_model.py:11) -- its ``load()`` calls ``Transformer._load_init_kwargs`` then +# ``Transformer.__init__``, which reads ``AutoConfig`` + ``AutoModel.from_pretrained`` (weights) + +# ``AutoProcessor.from_pretrained`` (processor/tokenizer) (base/modules/transformer.py:654-674), so +# a config-only CLIP dir validates offline then 409s at the local-only load (#7218 P3). Classes +# whose name already contains ``transformer`` (``Transformer`` / ``MLMTransformer``) route through +# the same check without being listed here. +_ST_TRANSFORMER_SHAPED_MODULE_NAMES = frozenset({"clipmodel"}) + +# ``SparseStaticEmbedding.load()`` (sparse_encoder/modules/sparse_static_embedding.py:157-226) reads +# a tokenizer (``AutoTokenizer.from_pretrained``, line 193) and then EITHER, when ``config.json`` +# names a ``.json`` ``path``, an IDF json via ``from_json`` -> ``load_file_path(..., "idf.json")`` +# (line 131-141, NO weight file) OR a complete Torch weight set via ``load_torch_weights`` +# (line 217, RAISES without ``model.safetensors`` / ``pytorch_model.bin``). So it is CONDITIONALLY +# weight-bearing -- adding it to ``_ST_WEIGHTED_MODULE_NAMES`` would wrongly require weights for the +# idf.json variant -- and gets its own branch requiring tokenizer + (idf data OR weights) (#7218 P3). +_ST_SPARSE_IDF_DATA_FILE = "idf.json" + +# Sentence-transformers module classes whose ``load()`` reads ONLY a config file (no weights, no +# tokenizer): a present module config is their sole load requirement. From the v5 source: +# ``Pooling`` / ``SpladePooling`` / ``Dropout`` read config only; ``BoW`` / ``WordWeights`` keep +# their vocab/weights INSIDE ``config.json`` and reconstruct it on the config-only base +# ``Module.load()`` (their ``save()`` writes a weight file, but ``load()`` never reads it). +_ST_CONFIG_ONLY_MODULE_NAMES = frozenset( + {"pooling", "spladepooling", "dropout", "bow", "wordweights"} +) + + +def _dir_has_module_config(names: set) -> bool: + """True when *names* hold a module config (``config.json`` or any ``*_config.json``).""" + return any(name == "config.json" or name.endswith("_config.json") for name in names) + def _module_dir_is_loadable(cls: str, is_root: bool, dir_path: Path) -> bool: """True when *dir_path* carries the files the sentence-transformers module class *cls* - reads in its own ``load()`` (see sentence_transformers/models/*.py): - - * a ``Transformer`` module (or a plain root with no recognized ST module class) is a full HF - load root (config + tokenizer + weights); + reads in its own ``load()`` (see sentence_transformers/{base,sentence_transformer,sparse_encoder} + /modules/*.py). Dispatch is on *cls* FIRST, BEFORE any root-Transformer fallback: a module with + ``save_in_root = True`` (every ``InputModule`` -- ``WordEmbeddings`` / ``StaticEmbedding`` / + ``SparseStaticEmbedding`` / ``Transformer`` / ``Router`` ...) is saved at the snapshot ROOT + (path ``""``; base/model.py:661-666), so a declared non-Transformer root module must be judged by + ITS OWN ``load()``, not held to Transformer root requirements it never writes (#7218 P1): + + * a ``Transformer`` / ``MLMTransformer`` / ``CLIPModel`` module is a full HF load root + (config + tokenizer/processor + weights), at the root or in a subdir; * a ``Router`` / ``Asym`` reads ``router_config.json`` and loads its declared children, so it is loadable when those children are -- even at the root, where it has no weights of its own; * ``Normalize`` reads nothing; * ``StaticEmbedding`` reads a ``tokenizer.json`` + a complete Torch weight set and NO config; - * a WEIGHTED module (``WordEmbeddings`` / ``Dense`` / ``CNN`` / ``LSTM``) needs its module - config AND a complete Torch weight set (its ``load()`` hard-loads ``model.safetensors`` / - ``pytorch_model.bin``); ``WordEmbeddings`` additionally rebuilds a tokenizer from the dir, - so it also needs a tokenizer artifact; - * every other module (``BoW``, ``Pooling`` ...) reads only a ``config.json`` / ``*_config``, - so a present module config is the load requirement.""" - if "transformer" in cls: - return _dir_is_transformer_load_root(_dir_file_names(dir_path)) + * ``SparseStaticEmbedding`` reads a tokenizer + EITHER an ``idf.json`` OR a complete Torch weight + set (conditionally weight-bearing); + * a WEIGHTED module (``WordEmbeddings`` / ``Dense`` / ``CNN`` / ``LSTM`` / ``LayerNorm`` / + ``WeightedLayerPooling`` / ``SparseAutoEncoder``) needs its module config AND a complete Torch + weight set (its ``load()`` hard-loads ``model.safetensors`` / ``pytorch_model.bin``); + ``WordEmbeddings`` additionally rebuilds a tokenizer from the dir, so it also needs a tokenizer; + * a recognized config-only module (``BoW`` / ``Pooling`` / ``Dropout`` / ``WordWeights`` / + ``SpladePooling``) reads only a config, so a present module config is the load requirement; + * an UNRECOGNIZED / absent class falls back to the plain-``from_pretrained`` load root at the + snapshot root (config + tokenizer + weights), else the lenient config-only requirement in a + declared subdir (preserves prior behaviour for module classes not enumerated here).""" + names = _dir_file_names(dir_path) + # Transformer-shaped: a full HF load root (config + tokenizer/processor + weights), root or subdir. + if "transformer" in cls or cls in _ST_TRANSFORMER_SHAPED_MODULE_NAMES: + return _dir_is_transformer_load_root(dir_path, names) + # Router / Asym: loadable when its router_config.json children are (no weights of its own). if cls in _ST_ROUTER_MODULE_NAMES: return _router_dir_is_loadable(dir_path) - if is_root: - return _dir_is_transformer_load_root(_dir_file_names(dir_path)) + # Normalize reads nothing. if cls == "normalize": return True - names = _dir_file_names(dir_path) + # StaticEmbedding: a tokenizer.json + a complete Torch weight set, and NO config file. if cls == "staticembedding": - # No config on disk: recognized by the tokenizer + weights its load() actually reads. - return _ST_STATIC_EMBEDDING_TOKENIZER_FILE in names and _dir_has_complete_torch_weights( - names + return _ST_STATIC_EMBEDDING_TOKENIZER_FILE in names and _dir_weight_set_is_complete( + dir_path, names ) - if not any(name == "config.json" or name.endswith("_config.json") for name in names): - return False + # SparseStaticEmbedding: a tokenizer + (an idf.json OR a complete Torch weight set). + if cls == "sparsestaticembedding": + if not _names_have_tokenizer(names): + return False + return _ST_SPARSE_IDF_DATA_FILE in names or _dir_weight_set_is_complete(dir_path, names) + # A weighted non-Transformer module: its module config AND a complete Torch weight set. if cls in _ST_WEIGHTED_MODULE_NAMES: - if not _dir_has_complete_torch_weights(names): + if not _dir_has_module_config(names) or not _dir_weight_set_is_complete(dir_path, names): return False if cls == "wordembeddings": return _names_have_word_embeddings_tokenizer(names) return True - return True + # A recognized config-only module: a present module config is the load requirement. + if cls in _ST_CONFIG_ONLY_MODULE_NAMES: + return _dir_has_module_config(names) + # Unrecognized / absent class: the plain from_pretrained load root at the snapshot root, else the + # lenient config-only requirement in a declared subdir. + if is_root: + return _dir_is_transformer_load_root(dir_path, names) + return _dir_has_module_config(names) def _router_dir_is_loadable(dir_path: Path) -> bool: @@ -2545,8 +2681,9 @@ def _modules_json_declares_modules(snap: Path) -> bool: ``modules.json`` is authoritative about what a SentenceTransformer load builds: when it declares one or more modules the loader calls ``module_class.load(path)`` for EVERY one, so the whole declared set must be loadable. An empty list / a non-list / an unreadable file - declares nothing usable -- the load then reads the snapshot root as a plain Transformer -- so - the caller falls back to the Transformer-shaped weight check. Never raises.""" + declares nothing usable, and (because ``modules.json`` is PRESENT) the load does NOT fall back + to a root Transformer -- an empty list builds zero modules and a malformed file raises -- so the + caller reads such a snapshot as not loadable. Never raises.""" try: modules = json.loads((snap / "modules.json").read_text(encoding = "utf-8")) except (OSError, ValueError): @@ -2564,19 +2701,26 @@ def _snapshot_is_loadable_st_model(snap: Path) -> bool: be loadable (:func:`_snapshot_modules_all_loadable`): a complete ``0_Transformer`` must NOT vouch for a sibling module (e.g. ``LayerNorm`` / ``WeightedLayerPooling`` / ``Dense``) whose serialized weights are missing -- that snapshot passes offline validation and then 409s at - the local-only load (#7218 P2). Only when ``modules.json`` declares nothing usable (an empty - list / not a list / unreadable) does the load read the snapshot ROOT as a plain Transformer, - so fall back to the Transformer-shaped weight set (:func:`_snapshot_has_complete_weights`), - which also covers the plain ``from_pretrained`` root layout. Behaviour-preserving for a - WELL-FORMED complete snapshot -- a normal download has every declared module's weights, so - acceptance is unchanged; only a snapshot with a genuinely incomplete declared module is now - rejected.""" + the local-only load (#7218 P2). + + A PRESENT ``modules.json`` that declares nothing usable (an empty list / not a list) or is + malformed is NOT a loadable model and does NOT fall back to a root Transformer: with + ``modules.json`` present the loader takes ``base/model.py`` ``_load_config_modules``, which + ``json.load``s the file (a malformed body RAISES) and iterates the list (an empty list builds + ZERO modules); the plain-Transformer fallback (``_load_default_modules``) is reached ONLY when + ``modules.json`` is ABSENT (``_load_modules`` -> ``modules_json_path is None``). So the earlier + fall-through to :func:`_snapshot_has_complete_weights` here was a false positive -- a snapshot + with a stub/empty ``modules.json`` but complete root weights was accepted and then 409'd at the + local-only load (#7218 P4). The plain ``from_pretrained`` root (a tag-only embedder with NO + ``modules.json``) is classified elsewhere in :func:`is_embedding_model` via the recorded verdict + plus :func:`_snapshot_has_complete_weights`, so restricting this marker check to a real declared + module set does not regress it.""" try: if not (snap / "modules.json").is_file(): return False if _modules_json_declares_modules(snap): return _snapshot_modules_all_loadable(snap) - return _snapshot_has_complete_weights(snap) + return False except OSError: return False From 47d4cd3374cb31716d0e6aa979a57c1398991408 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 13:37:10 +0000 Subject: [PATCH 55/77] Close case-folding and online-traversal holes in the offline pickle gate Two gate bypasses where the security scan credited or scoped a path differently from what the loader actually resolves: The safetensors credit was case-folded. _cached_pickle_weight_files lowercases every filename, and the loadable-safetensors and adapter checks tested those folded keys against the exact-lowercase names. On a case-sensitive filesystem (Linux, the Studio default) a crafted repo shipping Model.SafeTensors plus a malicious pytorch_model.bin makes transformers and sentence-transformers miss the exact-name model.safetensors and deserialize the pickle, while the gate credited an inert safetensors and did not block. Credit safetensors case-sensitively against real filenames, and drop pytorch_model.safetensors from the credit set (transformers loads only model.safetensors, never that name). Pickle matching stays case-insensitive (over-blocking a mis-cased pickle the loader would not load is the safe direction). The online scan did not canonicalize traversing paths while the offline gate did. A repo-controlled modules.json path (threaded into the online scan via the RAG guard) or a weight_map shard entry like 0/../evil / ../evil was compared verbatim, so a flagged evil/pytorch_model.bin never matched and evaded the online scan though the loader resolves and deserializes it. Canonicalize the repo-controlled load-subdir prefixes and weight_map shards the same way the offline gate does, so offline and online agree. Add regression tests for both bypasses. --- studio/backend/tests/test_hf_offline_flags.py | 183 ++++++++++++++++++ .../backend/utils/security/file_security.py | 95 +++++++-- 2 files changed, 257 insertions(+), 21 deletions(-) diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index da3362e3a22..993c01444b2 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -435,3 +435,186 @@ def test_only_the_rag_embedding_path_opts_into_the_bypass(): f"{rel} passes local_only_load but does not pin its loader to the " "local cache; that would disable the malware gate for a fetching path" ) + + +# ── FIX 1: the safetensors CREDIT is case-SENSITIVE and only credits names the loader loads ── +# +# The offline gate skips a pickle only when a safetensors weight the loader reads INSTEAD is +# present. transformers (SAFE_WEIGHTS_NAME) and sentence-transformers (Module.load_torch_weights) +# both look up the EXACT name ``model.safetensors`` (then its index); on a case-sensitive +# filesystem a mixed-case decoy is NOT that file, so the loader falls back to the pickle. Folding +# the decoy's name into the credit (or crediting a name the loader never loads) fails OPEN. + + +def _snap_dir(tmp_path, name: str, files: dict): + d = tmp_path / name / "aaa" + d.mkdir(parents = True) + for fname, body in files.items(): + (d / fname).write_bytes(body if isinstance(body, bytes) else body.encode()) + return d + + +def _offline_blocked(monkeypatch, snap) -> bool: + """Run the offline (local-only) gate against a prebuilt cache snapshot, asserting it never + reaches the Hub, and return whether the load is blocked.""" + import utils.models.model_config as mc + import utils.security.file_security as fs + + calls: list = [] + _fake_hub(monkeypatch, calls) + monkeypatch.setattr(mc, "_active_snapshot_dir", lambda name: snap) + decision = fs.evaluate_file_security("org/model", None, local_only_load = True) + assert calls == [], "a local-only load must not hit the Hub" + return decision.blocked + + +def test_offline_mixed_case_base_safetensors_does_not_credit_pickle(monkeypatch, tmp_path): + # ``Model.SafeTensors`` is not the exact ``model.safetensors`` the loader reads, so it cannot + # vouch for a live ``pytorch_model.bin`` -- the loader deserializes the pickle -> BLOCKED. + snap = _snap_dir( + tmp_path, "mixedbase", {"Model.SafeTensors": b"\0", "pytorch_model.bin": b"\0"} + ) + assert _offline_blocked(monkeypatch, snap) is True + + +def test_offline_exact_case_base_safetensors_credits_pickle(monkeypatch, tmp_path): + # Control: the exact ``model.safetensors`` is what the loader reads instead of the pickle. + snap = _snap_dir( + tmp_path, "exactbase", {"model.safetensors": b"\0", "pytorch_model.bin": b"\0"} + ) + assert _offline_blocked(monkeypatch, snap) is False + + +def test_offline_pytorch_model_safetensors_decoy_does_not_credit(monkeypatch, tmp_path): + # transformers/ST never look up ``pytorch_model.safetensors``; an inert file of that name + # must not vouch for the live ``pytorch_model.bin`` -> BLOCKED. + snap = _snap_dir( + tmp_path, "decoy", {"pytorch_model.safetensors": b"\0", "pytorch_model.bin": b"\0"} + ) + assert _offline_blocked(monkeypatch, snap) is True + + +def test_offline_mixed_case_adapter_safetensors_does_not_credit(monkeypatch, tmp_path): + # ``Adapter_Model.SafeTensors`` is not the exact ``adapter_model.safetensors`` PEFT loads, + # so the live ``adapter_model.bin`` stays a separate RCE vector -> BLOCKED (an inert + # safetensors base does not cover the adapter pickle). + snap = _snap_dir( + tmp_path, + "mixedadapter", + { + "config.json": b"{}", + "model.safetensors": b"\0", + "adapter_config.json": b"{}", + "adapter_model.bin": b"\0", + "Adapter_Model.SafeTensors": b"\0", + }, + ) + assert _offline_blocked(monkeypatch, snap) is True + + +def test_offline_exact_case_adapter_safetensors_credits(monkeypatch, tmp_path): + # Control: the exact ``adapter_model.safetensors`` covers the adapter pickle -> allowed. + snap = _snap_dir( + tmp_path, + "exactadapter", + { + "config.json": b"{}", + "model.safetensors": b"\0", + "adapter_config.json": b"{}", + "adapter_model.bin": b"\0", + "adapter_model.safetensors": b"\0", + }, + ) + assert _offline_blocked(monkeypatch, snap) is False + + +# ── FIX 2: the ONLINE scan canonicalizes repo-controlled ``..`` paths, matching the offline gate ── +# +# The offline gate resolves a traversing declared path / weight_map shard (``0/../evil`` -> +# ``evil``) so a pickle the loader deserializes there is blocked. The online scan compares the +# Hub's CANONICAL flagged path against repo-controlled load-subdirs (a ``modules.json`` module +# ``path`` the RAG guard unions in) and ``weight_map`` shards; without collapsing ``..`` the raw +# ``0/../evil`` never matches ``evil/...`` and the flagged pickle slips through. + + +def test_online_load_subdir_prefixes_are_canonicalized(): + import utils.security.file_security as fs + + assert fs._index_prefixes(("0/../evil",)) == ("", "evil/") + assert fs._load_relative_path("evil/pytorch_model.bin", ("0/../evil",)) == "pytorch_model.bin" + # A non-traversing subdir is unchanged. + assert fs._load_relative_path("LLM/pytorch_model.bin", ("LLM",)) == "pytorch_model.bin" + # An escaping subdir (leading ..) is dropped -- it can never name a repo file. + assert fs._load_relative_path("evil/x.bin", ("../evil",)) == "evil/x.bin" + + +def test_online_scan_blocks_flagged_pickle_under_traversing_module_subdir(monkeypatch): + # A repo-controlled ``modules.json`` path ``0/../evil`` reaches the online scan (unioned in + # by the RAG guard). The loader resolves it to ``evil/`` and deserializes + # ``evil/pytorch_model.bin``; the gate must scope the canonical dir and block it root-level. + import utils.security.file_security as fs + + monkeypatch.setattr( + fs, + "_fetch_security_status", + lambda name, token: { + "filesWithIssues": [{"path": "evil/pytorch_model.bin", "level": "unsafe"}] + }, + ) + # Definitive "no weight index" -> a missed shard is SKIPPED (allowed) on the pristine gate; + # the fix instead treats evil/ as a load root so the pickle is root-level and blocks. + monkeypatch.setattr(fs, "_indexed_shard_paths", lambda *a, **k: set()) + decision = fs.evaluate_file_security( + "org/model", None, load_subdirs = ("0/../evil",), local_only_load = False + ) + assert decision.blocked is True + + +def test_online_scan_blocks_flagged_pickle_under_plain_module_subdir(monkeypatch): + # Control: a non-traversing declared subdir still blocks a flagged pickle directly under it. + import utils.security.file_security as fs + + monkeypatch.setattr( + fs, + "_fetch_security_status", + lambda name, token: { + "filesWithIssues": [{"path": "0_Transformer/pytorch_model.bin", "level": "unsafe"}] + }, + ) + monkeypatch.setattr(fs, "_indexed_shard_paths", lambda *a, **k: set()) + decision = fs.evaluate_file_security( + "org/model", None, load_subdirs = ("0_Transformer",), local_only_load = False + ) + assert decision.blocked is True + + +def test_online_indexed_shard_paths_canonicalize_weight_map_traversal(monkeypatch, tmp_path): + # A repo-controlled ``weight_map`` shard that traverses (``sub/../evil/pytorch_model.bin``) + # must be recorded at the canonical ``evil/pytorch_model.bin`` the Hub reports for the + # flagged file, so the online maybe-shard check matches it and blocks the load. + import utils.security.file_security as fs + + fake = types.ModuleType("huggingface_hub") + fake_utils = types.ModuleType("huggingface_hub.utils") + + class _EntryNotFound(Exception): + pass + + fake_utils.EntryNotFoundError = _EntryNotFound + fake.utils = fake_utils + + index_file = tmp_path / "pytorch_model.bin.index.json" + index_file.write_text('{"weight_map": {"w": "sub/../evil/pytorch_model.bin"}}') + + def _download(model_name, filename, token = None): + if filename == "pytorch_model.bin.index.json": + return str(index_file) + raise _EntryNotFound() + + fake.hf_hub_download = _download + monkeypatch.setitem(sys.modules, "huggingface_hub", fake) + monkeypatch.setitem(sys.modules, "huggingface_hub.utils", fake_utils) + + paths = fs._indexed_shard_paths("org/model", None) + assert "evil/pytorch_model.bin" in paths + assert "sub/../evil/pytorch_model.bin" not in paths diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index e600e4f000b..0e3e876cac4 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -93,22 +93,31 @@ def _normalize_repo_path(path: str) -> str: return p -def _canonical_load_dir(base, rel: str): - """``base`` / ``rel`` with ``.`` / ``..`` components collapsed lexically (no filesystem - access), or None when ``rel`` is empty, the base itself, or escapes ``base``. - - A ``modules.json`` / ``router_config.json`` / load-subdir path is repo-controlled. - SentenceTransformer resolves an entry such as ``0/../evil`` to ``evil/`` and deserializes - ``evil/pytorch_model.bin``, so the offline gate must scope that SAME normalized directory: - the raw ``base / "0/../evil"`` never equals the ``base / "evil"`` that ``rglob`` yields, so a - pickle there would slip the gate. A path that escapes ``base`` (a leading ``..`` after - normalization) is rejected -- a legitimate declared module path never traverses upward.""" +def _canonical_rel(rel: str): + """A repo-relative path with ``.`` / ``..`` components collapsed lexically (no filesystem + access), or None when it is empty, the repo root, or escapes the root (a leading ``..`` + after normalization). + + A ``modules.json`` / ``router_config.json`` / load-subdir / ``weight_map`` entry is + repo-controlled. The loader resolves an entry such as ``0/../evil`` to ``evil/`` and + deserializes ``evil/pytorch_model.bin`` -- the offline gate already scopes that SAME + normalized directory (via :func:`_canonical_load_dir`) and the ONLINE scan must agree, or + the raw ``0/../evil`` never equals the canonical ``evil/...`` the Hub reports for the + flagged file and the pickle slips the gate. A legitimate declared path never traverses + upward, so an escaping path is rejected.""" import posixpath norm = posixpath.normpath(_normalize_repo_path(rel).strip("/")) if norm in ("", ".") or norm == ".." or norm.startswith("../"): return None - return base / norm + return norm + + +def _canonical_load_dir(base, rel: str): + """``base`` / ``rel`` with ``.`` / ``..`` collapsed lexically, or None when ``rel`` is + empty, the base itself, or escapes ``base``. String form: :func:`_canonical_rel`.""" + norm = _canonical_rel(rel) + return None if norm is None else base / norm def _file_suffix(path: str) -> str: @@ -122,19 +131,27 @@ def _load_relative_path(norm: str, load_subdirs) -> str: snapshot SUBDIRECTORY (Spark-TTS / BiCodec load ``/LLM``), where a file directly under the subdir is root-level, not nested. Strips the matching load-subdir prefix, or returns ``norm`` unchanged when it is not under one. + + A load-subdir is repo-controlled (the RAG guard unions in each ``modules.json`` module + ``path``), so it is canonicalized (``0/../evil`` -> ``evil``) to match the canonical repo + path the Hub reports for a flagged file -- the raw ``0/../evil`` would never prefix + ``evil/pytorch_model.bin`` and the file would slip through as an unreferenced nested shard, + though the offline gate (which canonicalizes the same path) blocks it. """ for subdir in load_subdirs or (): - prefix = _normalize_repo_path(subdir).strip("/") + prefix = _canonical_rel(subdir) if prefix and norm.startswith(prefix + "/"): return norm[len(prefix) + 1 :] return norm def _index_prefixes(load_subdirs) -> tuple: - """Prefixes to look for weight-index files under: repo root plus each load subdir.""" + """Prefixes to look for weight-index files under: repo root plus each load subdir. Each + subdir is canonicalized (see :func:`_load_relative_path`) so a traversing declared path + resolves to the same directory offline and online.""" prefixes = [""] for subdir in load_subdirs or (): - p = _normalize_repo_path(subdir).strip("/") + p = _canonical_rel(subdir) if p: prefixes.append(p + "/") return tuple(prefixes) @@ -178,7 +195,14 @@ def _indexed_shard_paths( # weight_map paths are relative to the index file's directory. if prefix and not shard_norm.startswith(prefix): shard_norm = prefix + shard_norm - paths.add(shard_norm) + # Collapse . / .. (a repo-controlled weight_map may traverse, e.g. + # "sub/../evil/pytorch_model.bin") so the recorded path is the canonical + # one the Hub reports for the flagged shard -- mirrors the offline gate, + # which resolves the same traversal on disk. An escaping shard (leading + # ..) can never name a repo file, so drop it. + shard_canon = _canonical_rel(shard_norm) + if shard_canon is not None: + paths.add(shard_canon) except Exception: inconclusive = True # Any transient failure -> inconclusive (the shard could be listed only by the index @@ -304,8 +328,14 @@ def _index_weight_map_values(index_path) -> set: # whose every referenced shard is present. A bare adapter (``adapter_model.safetensors``) or # an orphan shard with no index is NOT a loadable base weight -- the loader falls back to and # deserializes the pickle, which therefore stays the live RCE vector. -_SAFETENSORS_BASE_UNSHARDED = ("model.safetensors", "pytorch_model.safetensors") -_SAFETENSORS_BASE_INDEX = ("model.safetensors.index.json", "pytorch_model.safetensors.index.json") +# +# Only the EXACT names the loader resolves count. transformers' SAFE_WEIGHTS_NAME / +# SAFE_WEIGHTS_INDEX_NAME and sentence-transformers' Module.load_torch_weights both look up +# ``model.safetensors`` (then its index); neither ever looks up ``pytorch_model.safetensors``, +# so crediting that name would let a repo shipping an inert ``pytorch_model.safetensors`` decoy +# beside a live ``pytorch_model.bin`` pass unblocked while the loader deserialized the pickle. +_SAFETENSORS_BASE_UNSHARDED = ("model.safetensors",) +_SAFETENSORS_BASE_INDEX = ("model.safetensors.index.json",) def _safetensors_index_complete(index_path) -> bool: @@ -334,14 +364,30 @@ def _safetensors_index_complete(index_path) -> bool: return True +def _exact_named(files: dict, name: str): + """The Path in *files* (lower-name -> Path for one directory) whose REAL basename is + exactly *name* (case-sensitive), or None. A safetensors CREDIT must match the loader's + exact lookup: transformers / sentence-transformers request ``model.safetensors`` verbatim, + so on a case-sensitive filesystem (the Studio default) a mixed-case ``Model.SafeTensors`` + decoy is NOT the file the loader reads -- it falls back to and deserializes the pickle. The + lower-name key would fold the decoy in and fail OPEN, so credit is decided by ``Path.name``. + """ + for path in files.values(): + if path.name == name: + return path + return None + + def _dir_has_loadable_safetensors(files: dict) -> bool: """True when *files* (lower-name -> Path for one directory) hold a safetensors weight a from_pretrained load will read INSTEAD of a pickle sibling: an unsharded base file, or a - complete indexed shard set. A bare adapter or an orphan shard does not qualify.""" - if any(name in files for name in _SAFETENSORS_BASE_UNSHARDED): + complete indexed shard set. A bare adapter or an orphan shard does not qualify. The + safetensors credit is case-SENSITIVE (see :func:`_exact_named`) so a mis-cased decoy the + loader would skip cannot vouch for a live pickle.""" + if any(_exact_named(files, name) is not None for name in _SAFETENSORS_BASE_UNSHARDED): return True for index_name in _SAFETENSORS_BASE_INDEX: - index_path = files.get(index_name) + index_path = _exact_named(files, index_name) if index_path is not None and _safetensors_index_complete(index_path): return True return False @@ -453,8 +499,15 @@ def _cached_pickle_weight_files(snap, load_subdirs = ()) -> list: hits.update(base) # An adapter pickle is deserialized only when from_pretrained auto-detects the adapter # (adapter_config.json present) and there is no adapter_model.safetensors to load instead. + # The safetensors credit is case-SENSITIVE (a mixed-case Adapter_Model.SafeTensors decoy + # is not the file PEFT loads), so it is matched by real basename; the config presence + # stays case-insensitive (over-blocking a mis-cased adapter is the safe direction). adapter = [n for n in names if _ADAPTER_PICKLE_RE.match(n.lower())] - if adapter and "adapter_config.json" in files and "adapter_model.safetensors" not in files: + if ( + adapter + and "adapter_config.json" in files + and _exact_named(files, "adapter_model.safetensors") is None + ): hits.update(adapter) # A load-root pickle index (pytorch_model.bin.index.json) can map shards into SUBDIRECTORIES # that are not themselves load roots; from_pretrained follows the map and deserializes them, From ecbd3256e4e99869b64c704c803aeeaffb306296 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 13:37:10 +0000 Subject: [PATCH 56/77] Treat a conditional __all__ mutation as opaque in the import-hoist linter _collect_dunder_all replayed only top-level module statements, so an __all__ assignment or mutation inside a module-level if / try / for / while / with / match (or a deeper scope) was ignored, leaving the export set understated. A newly hoisted import re-exported only through such a conditional __all__ was then falsely flagged HOISTED-IMPORT-UNUSED, blocking a valid change. A conditional value cannot be replayed statically, so mark the export set opaque when __all__ is bound or mutated anywhere other than a top-level statement. Add a self-test case. --- scripts/verify_import_hoist.py | 49 ++++++++++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index b4d05aec4f6..db394a66e34 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -474,6 +474,41 @@ def _all_mutator_entries(method: str, args: list) -> tuple[set[str], bool]: return entries, False # unknown mutator or dynamic arg -> opaque +def _stmt_binds_all(node: ast.AST) -> bool: + """True when *node* is a statement that BINDS or MUTATES ``__all__`` -- an assignment / + annotated-assign / augmented-assign whose target names ``__all__`` (directly or nested in a + destructuring / item / attr target), or an ``__all__.(...)`` expression call.""" + if isinstance(node, ast.Assign): + targets = node.targets + elif isinstance(node, (ast.AnnAssign, ast.AugAssign)): + targets = [node.target] + elif isinstance(node, ast.Expr): + call = node.value + return ( + isinstance(call, ast.Call) + and isinstance(call.func, ast.Attribute) + and isinstance(call.func.value, ast.Name) + and call.func.value.id == "__all__" + ) + else: + return False + return any( + isinstance(sub, ast.Name) and sub.id == "__all__" for t in targets for sub in ast.walk(t) + ) + + +def _all_bound_outside_module_body(tree: ast.Module) -> bool: + """True when ``__all__`` is bound / mutated anywhere OTHER than a top-level module statement -- + inside a module-level ``if`` / ``try`` / ``for`` / ``while`` / ``with`` / ``match`` (a + conditional whose final value cannot be replayed statically), or inside a nested scope. The + caller then marks ``__all__`` opaque, so a conditionally-added re-export is not mistaken for an + unused hoist. Only ever ADDS opacity, so it cannot create a false blocker.""" + top_level = {id(node) for node in tree.body} + return any( + id(node) not in top_level and _stmt_binds_all(node) for node in ast.walk(tree) + ) + + def _collect_dunder_all(tree: ast.Module) -> tuple[set[str], bool]: """The FINAL module-level ``__all__`` string entries plus an ``opaque`` flag (``= [...]``, ``+= [...]``, or an annotated assign). @@ -502,6 +537,8 @@ def _collect_dunder_all(tree: ast.Module) -> tuple[set[str], bool]: ``__all__`` method call (``insert``/``remove``/...) is opaque. ``__all__`` bound only through a destructuring target (``__all__, meta = [...], v``) or an item/attr target is opaque too -- its value cannot be mapped statically, so the set is not exhaustive. + ``__all__`` bound / mutated in a module-level ``if`` / ``try`` / loop / ``with`` / ``match`` (or + a deeper scope) is opaque -- a conditional value cannot be replayed from ``tree.body`` alone. """ names: set[str] = set() opaque = False @@ -562,6 +599,10 @@ def _collect_dunder_all(tree: ast.Module) -> tuple[set[str], bool]: # extend adds names we cannot see. names |= entries opaque = opaque or not readable + if _all_bound_outside_module_body(tree): + # A conditional / nested `__all__` mutation (module-level if/try/for/while/with/match, or a + # deeper scope) is not in the replayed top-level sequence, so its contribution is unknown. + opaque = True return names, opaque @@ -892,6 +933,14 @@ def used_targets(analysis) -> set[str]: 'from pkg import a\nfrom pkg import b\n__all__, meta = ["a", "b"], 1\n', None, ), + "conditional_all_mutation_is_opaque_keeps_reexport": ( + # __all__ extended inside a module-level `if` is NOT in the replayed top-level sequence, + # so the final value is unknown -> opaque: a re-export added only there must not be flagged + # as an unused hoist (the collector previously scanned only tree.body and missed it) + 'from pkg import a\n__all__ = ["a"]\n', + 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\nif True:\n __all__ += ["b"]\n', + None, + ), "readable_reassign_resets_opacity": ( # an opaque "+=" then a readable replacing "=": the final __all__ is empty, so a # hoisted unused import is a genuine bad hoist -- the reset must un-stick opacity From 3db5c0cdac8cb3f6d43862c3db15a74b01502108 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 13:38:48 +0000 Subject: [PATCH 57/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- scripts/verify_import_hoist.py | 4 +-- studio/backend/tests/test_hf_offline_flags.py | 6 +++- .../tests/test_is_embedding_model_offline.py | 30 ++++++++++++++++--- 3 files changed, 32 insertions(+), 8 deletions(-) diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index db394a66e34..cc6dfe900a7 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -504,9 +504,7 @@ def _all_bound_outside_module_body(tree: ast.Module) -> bool: caller then marks ``__all__`` opaque, so a conditionally-added re-export is not mistaken for an unused hoist. Only ever ADDS opacity, so it cannot create a false blocker.""" top_level = {id(node) for node in tree.body} - return any( - id(node) not in top_level and _stmt_binds_all(node) for node in ast.walk(tree) - ) + return any(id(node) not in top_level and _stmt_binds_all(node) for node in ast.walk(tree)) def _collect_dunder_all(tree: ast.Module) -> tuple[set[str], bool]: diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index 993c01444b2..cfa3b8abbff 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -606,7 +606,11 @@ class _EntryNotFound(Exception): index_file = tmp_path / "pytorch_model.bin.index.json" index_file.write_text('{"weight_map": {"w": "sub/../evil/pytorch_model.bin"}}') - def _download(model_name, filename, token = None): + def _download( + model_name, + filename, + token = None, + ): if filename == "pytorch_model.bin.index.json": return str(index_file) raise _EntryNotFound() diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 6a1cddc7385..1bbf7116e61 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -1564,7 +1564,13 @@ def test_marker_rejects_root_wordembeddings_without_tokenizer(tmp_path, monkeypa _ROOT_STATIC = "org/root-staticembedding" -def _root_staticembedding_repo(tmp_path, monkeypatch, *, include_weights = True, commit = "aaa"): +def _root_staticembedding_repo( + tmp_path, + monkeypatch, + *, + include_weights = True, + commit = "aaa", +): """A StaticEmbedding model whose StaticEmbedding module is saved at the ROOT path "" (save_in_root =True). The root holds config_sentence_transformers.json + tokenizer.json + weights and NO module config. StaticEmbedding already passes the pristine is_root Transformer check because it ships a @@ -1608,7 +1614,13 @@ def test_marker_rejects_root_staticembedding_without_weights(tmp_path, monkeypat _CLIP = "org/clip" -def _clip_repo(tmp_path, monkeypatch, *, complete, commit = "aaa"): +def _clip_repo( + tmp_path, + monkeypatch, + *, + complete, + commit = "aaa", +): """A SentenceTransformer whose modules.json declares a CLIPModel at 0_CLIPModel. CLIPModel(load) -> Transformer.__init__ reads AutoConfig + AutoModel.from_pretrained (weights) + AutoProcessor (processor/tokenizer), so a config-only dir is NOT loadable. complete=False writes only @@ -1670,7 +1682,11 @@ def _sparse_static_repo( mod.mkdir(parents = True) (snap / "modules.json").write_text( _modules_json( - ("0", "0_SparseStaticEmbedding", "sentence_transformers.sparse_encoder.models.SparseStaticEmbedding") + ( + "0", + "0_SparseStaticEmbedding", + "sentence_transformers.sparse_encoder.models.SparseStaticEmbedding", + ) ) ) (snap / "config_sentence_transformers.json").write_text("{}") @@ -1713,7 +1729,13 @@ def test_marker_rejects_sparse_static_embedding_without_tokenizer(tmp_path, monk # ── #7218 P4: a PRESENT but empty / malformed modules.json is not a loadable ST model ── -def _root_weights_snapshot(tmp_path, monkeypatch, modules_body, *, repo_id = "org/present-modules"): +def _root_weights_snapshot( + tmp_path, + monkeypatch, + modules_body, + *, + repo_id = "org/present-modules", +): """A snapshot with a COMPLETE root weight set (config + tokenizer + weights) and a modules.json whose raw body is ``modules_body``. Lets a test assert that a present-but-empty / malformed modules.json is not loadable even though the root weights are complete.""" From fe756ba58feda3fdf3baebcaa4b8e373f3bf8ad6 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 20 Jul 2026 13:47:56 +0000 Subject: [PATCH 58/77] Scope Router child sub-modules as load roots in the online embedding scan The RAG embedding security guard unions the SentenceTransformer module dirs from modules.json into the load roots it scopes for the Hub scan, so a flagged pickle directly under a Transformer module blocks. A Router (legacy Asym) module declares its child sub-modules only in router_config.json, not in modules.json, and Router.load() deserializes each child from its own subdir. The online scan therefore dropped a flagged child pickle (for example query_0_WordEmbeddings/pytorch_model.bin) as an unreferenced nested shard while the loader still deserialized it, the counterpart to the offline gate which already expands router children via _router_child_dirs. _st_module_subdirs now reads router_config.json for any Router-typed module and adds each declared child (joined onto the module path, canonicalized so a traversing entry is dropped) to the load roots. The config is read only for a Router-typed module, so a plain embedder pays no extra fetch, and every failure path still returns () so the guard never bricks the embedder. --- studio/backend/core/rag/embeddings.py | 101 ++++++++++++---- .../test_embedding_model_security_gate.py | 110 ++++++++++++++++++ 2 files changed, 186 insertions(+), 25 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 1fe3987092c..3ddbd896f18 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -81,11 +81,23 @@ def _ambient_hf_token() -> str | None: def _st_module_subdirs(name: str, token: str | None, local_only: bool) -> tuple[str, ...]: - """The module directories a SentenceTransformer load reads weights from, taken from - the repo's ``modules.json`` (each module's non-empty ``path``, e.g. ``0_Transformer``). - ST deserializes ``pytorch_model.bin`` from these dirs, so they are load roots for the - security scan: a flagged pickle directly under one must block. Returns () on any - failure (no modules.json, offline, malformed) so the guard never bricks the embedder. + """The module directories a SentenceTransformer load reads weights from, so the online + security scan scopes them as load roots (a flagged pickle directly under one must block). + Two sources, both repo-controlled: + + * each module's non-empty ``path`` in ``modules.json`` (e.g. ``0_Transformer``), from which + ST deserializes ``pytorch_model.bin``; + * the child sub-modules of any ``Router`` (legacy ``Asym``) module. A Router declares its + children only in ``router_config.json`` (``types`` maps ``{route}_{idx}_{ClassName}`` -> + class), NOT in ``modules.json``, and ``Router.load()`` deserializes each from its own + subdir -- so a flagged pickle in a child (``query_0_WordEmbeddings/pytorch_model.bin``) is + a load root too. This mirrors the offline gate's ``_router_child_dirs`` expansion; without + it the online scan would drop that child pickle as an unreferenced nested shard while the + loader still deserialized it. The config is read only for a Router-typed module, so a plain + embedder pays no extra fetch. + + Returns () on any failure (no modules.json, offline, malformed) so the guard never bricks the + embedder. A traversing declared path (``0/../evil``) is dropped, matching the offline gate. ``local_only`` MUST be the value the caller captured for the load, not a fresh env read (``_hf_offline_if_dns_dead()`` flips the offline vars mid-load): re-reading could force @@ -94,38 +106,77 @@ def _st_module_subdirs(name: str, token: str | None, local_only: bool) -> tuple[ """ try: import json + import posixpath from utils.paths import is_local_path + local_root = None if is_local_path(name): from pathlib import Path from utils.paths import normalize_path - path = Path(normalize_path(name)).expanduser() / "modules.json" - if not path.is_file(): - return () - data = json.loads(path.read_text()) - else: - from huggingface_hub import hf_hub_download - from huggingface_hub.utils import EntryNotFoundError + local_root = Path(normalize_path(name)).expanduser() + def _read_repo_json(rel: str): + """Parse a repo-relative JSON file (local dir or Hub download honoring + ``local_only``), or None when missing / malformed / unreachable.""" try: - # CAPTURED predicate, never a fresh env read (see docstring); hf honors only - # HF_HUB_OFFLINE, so offline would otherwise block on timeouts. - local = hf_hub_download( - name, - "modules.json", - token = token or None, - local_files_only = local_only, - ) - except EntryNotFoundError: - return () - data = json.loads(open(local).read()) + if local_root is not None: + path = local_root.joinpath(*rel.split("/")) + if not path.is_file(): + return None + return json.loads(path.read_text()) + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + try: + # CAPTURED predicate, never a fresh env read (see docstring); hf honors only + # HF_HUB_OFFLINE, so offline would otherwise block on timeouts. + local = hf_hub_download( + name, rel, token = token or None, local_files_only = local_only + ) + except EntryNotFoundError: + return None + return json.loads(open(local).read()) + except Exception: + return None + + def _safe_subdir(rel) -> str | None: + """Canonical repo-relative subdir (``""`` = root), or None if it escapes the repo. + Mirrors the offline gate so a traversing declared path cannot read or scope a + directory outside the repo.""" + norm = posixpath.normpath(str(rel).strip().strip("/")) + if norm in ("", "."): + return "" + if norm == ".." or norm.startswith("../"): + return None + return norm + + data = _read_repo_json("modules.json") + if not isinstance(data, list): + return () subdirs = [] - for module in data or (): - sub = str((module or {}).get("path", "")).strip().strip("/") + for module in data: + if not isinstance(module, dict): + continue + sub = str(module.get("path", "")).strip().strip("/") if sub: subdirs.append(sub) + # Only a Router/Asym module hides load roots in router_config.json; the read is + # scoped to it so a plain embedder incurs no extra fetch. + if str(module.get("type", "")).rsplit(".", 1)[-1].lower() not in ("router", "asym"): + continue + prefix = _safe_subdir(sub) + if prefix is None: + continue + cfg = _read_repo_json(posixpath.join(prefix, "router_config.json")) + types = cfg.get("types") if isinstance(cfg, dict) else None + if not isinstance(types, dict): + continue + for model_id in types: + child = _safe_subdir(model_id) + if child: + subdirs.append(posixpath.join(prefix, child)) return tuple(dict.fromkeys(subdirs)) except Exception: return () diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index 9a40df63d96..520e16867c2 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -462,6 +462,116 @@ def _boom(*a, **k): assert embeddings._st_module_subdirs("acme/no-such-repo-xyz", None, False) == () +def test_st_module_subdirs_expands_root_router_children(tmp_path): + # A Router saved in root (path "") declares its child sub-modules only in + # router_config.json; Router.load() deserializes each from its own subdir, so those subdirs + # are load roots too. The scan must scope them or a flagged child pickle + # (query_0_WordEmbeddings/pytorch_model.bin) slips through as an unreferenced nested shard. + import json + import core.rag.embeddings as embeddings + + (tmp_path / "modules.json").write_text( + json.dumps( + [{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Router.Router"}] + ) + ) + (tmp_path / "router_config.json").write_text( + json.dumps( + { + "types": { + "query_0_WordEmbeddings": "sentence_transformers.models.WordEmbeddings.WordEmbeddings", + "document_0_Transformer": "sentence_transformers.models.Transformer.Transformer", + } + } + ) + ) + subdirs = embeddings._st_module_subdirs(str(tmp_path), None, False) + assert "query_0_WordEmbeddings" in subdirs + assert "document_0_Transformer" in subdirs + + +def test_st_module_subdirs_expands_nested_router_children(tmp_path): + # A Router nested at a module path prefixes its children with that path. + import json + import core.rag.embeddings as embeddings + + (tmp_path / "modules.json").write_text( + json.dumps( + [ + {"idx": 0, "name": "0", "path": "0_Transformer", "type": "..."}, + {"idx": 1, "name": "1", "path": "1_Router", "type": "sentence_transformers.models.Asym.Asym"}, + ] + ) + ) + (tmp_path / "1_Router").mkdir() + (tmp_path / "1_Router" / "router_config.json").write_text( + json.dumps({"types": {"query_0_WordEmbeddings": "..."}}) + ) + subdirs = embeddings._st_module_subdirs(str(tmp_path), None, False) + assert "0_Transformer" in subdirs + assert "1_Router/query_0_WordEmbeddings" in subdirs + + +def test_st_module_subdirs_ignores_router_config_without_router_module(tmp_path): + # A stray router_config.json is read ONLY for a Router-typed module, so a plain embedder + # neither expands children nor pays the extra fetch. + import json + import core.rag.embeddings as embeddings + + (tmp_path / "modules.json").write_text( + json.dumps([{"idx": 0, "name": "0", "path": "0_Transformer", "type": "..."}]) + ) + (tmp_path / "router_config.json").write_text( + json.dumps({"types": {"query_0_WordEmbeddings": "..."}}) + ) + subdirs = embeddings._st_module_subdirs(str(tmp_path), None, False) + assert subdirs == ("0_Transformer",) + + +def test_st_module_subdirs_drops_traversing_router_child(tmp_path): + # A malicious router_config.json child that traverses out of the repo (../evil) is dropped, + # matching the offline gate; a legitimate sibling is still scoped. + import json + import core.rag.embeddings as embeddings + + (tmp_path / "modules.json").write_text( + json.dumps([{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Router.Router"}]) + ) + (tmp_path / "router_config.json").write_text( + json.dumps({"types": {"../evil": "...", "query_0_WordEmbeddings": "..."}}) + ) + subdirs = embeddings._st_module_subdirs(str(tmp_path), None, False) + assert "query_0_WordEmbeddings" in subdirs + assert not any(".." in s for s in subdirs) + + +def test_st_module_subdirs_router_expansion_over_hub(monkeypatch, tmp_path): + # The Hub (non-local) path expands Router children symmetrically: hf_hub_download serves + # modules.json then the Router's router_config.json. + import json + import huggingface_hub + import core.rag.embeddings as embeddings + + (tmp_path / "modules.json").write_text( + json.dumps([{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Router.Router"}]) + ) + (tmp_path / "router_config.json").write_text( + json.dumps({"types": {"query_0_WordEmbeddings": "..."}}) + ) + + def _fake_download(repo_id, filename, **kwargs): + p = tmp_path / filename + if not p.is_file(): + from huggingface_hub.utils import EntryNotFoundError + + raise EntryNotFoundError(f"no {filename}") + return str(p) + + monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fake_download) + subdirs = embeddings._st_module_subdirs("acme/router-embed", None, False) + assert "query_0_WordEmbeddings" in subdirs + + def test_security_block_is_not_swallowed_by_llama_fallback(monkeypatch): # The ST encode fallback must re-raise a security block, not swap to llama-server. import core.rag.embeddings as embeddings From cde4da91a7cc6e7bf739c314760619f99a1621d8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 20 Jul 2026 13:49:30 +0000 Subject: [PATCH 59/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/rag/embeddings.py | 1 - .../test_embedding_model_security_gate.py | 39 ++++++++++++++++--- 2 files changed, 34 insertions(+), 6 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 3ddbd896f18..38f313a70fd 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -114,7 +114,6 @@ def _st_module_subdirs(name: str, token: str | None, local_only: bool) -> tuple[ if is_local_path(name): from pathlib import Path from utils.paths import normalize_path - local_root = Path(normalize_path(name)).expanduser() def _read_repo_json(rel: str): diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index 520e16867c2..de3473c9795 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -472,7 +472,14 @@ def test_st_module_subdirs_expands_root_router_children(tmp_path): (tmp_path / "modules.json").write_text( json.dumps( - [{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Router.Router"}] + [ + { + "idx": 0, + "name": "0", + "path": "", + "type": "sentence_transformers.models.Router.Router", + } + ] ) ) (tmp_path / "router_config.json").write_text( @@ -499,7 +506,12 @@ def test_st_module_subdirs_expands_nested_router_children(tmp_path): json.dumps( [ {"idx": 0, "name": "0", "path": "0_Transformer", "type": "..."}, - {"idx": 1, "name": "1", "path": "1_Router", "type": "sentence_transformers.models.Asym.Asym"}, + { + "idx": 1, + "name": "1", + "path": "1_Router", + "type": "sentence_transformers.models.Asym.Asym", + }, ] ) ) @@ -535,7 +547,16 @@ def test_st_module_subdirs_drops_traversing_router_child(tmp_path): import core.rag.embeddings as embeddings (tmp_path / "modules.json").write_text( - json.dumps([{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Router.Router"}]) + json.dumps( + [ + { + "idx": 0, + "name": "0", + "path": "", + "type": "sentence_transformers.models.Router.Router", + } + ] + ) ) (tmp_path / "router_config.json").write_text( json.dumps({"types": {"../evil": "...", "query_0_WordEmbeddings": "..."}}) @@ -553,7 +574,16 @@ def test_st_module_subdirs_router_expansion_over_hub(monkeypatch, tmp_path): import core.rag.embeddings as embeddings (tmp_path / "modules.json").write_text( - json.dumps([{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Router.Router"}]) + json.dumps( + [ + { + "idx": 0, + "name": "0", + "path": "", + "type": "sentence_transformers.models.Router.Router", + } + ] + ) ) (tmp_path / "router_config.json").write_text( json.dumps({"types": {"query_0_WordEmbeddings": "..."}}) @@ -563,7 +593,6 @@ def _fake_download(repo_id, filename, **kwargs): p = tmp_path / filename if not p.is_file(): from huggingface_hub.utils import EntryNotFoundError - raise EntryNotFoundError(f"no {filename}") return str(p) From 02d843efd7c134f89d2302ec7af04655253cef55 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 21 Jul 2026 02:45:32 +0000 Subject: [PATCH 60/77] Allow a recorded-clean pickle embedder to load offline The offline embedding security gate is fail-closed: with no network to reach Hugging Face's scan, a cached pickle weight cannot be verified, so it is blocked and a model the user already downloaded and used online will not load offline. This adds a persistent cache of clean Hub verdicts so that exact content can load offline, without weakening the gate for an unknown or never-scanned pickle. When an embedding repo is loaded online and HF's scan returns a completed clean verdict, the load roots are hashed and recorded under the scanned commit as an exact map of snapshot-relative pickle name to sha256, in a per-user JSON store at studio_root()/security/embedding_scan_verdicts.json (atomic write, 0600, thread and cross-process locked, 30-day TTL). Offline, a cached pickle model loads only when the active cached commit and every load-root pickle's sha256 match the recorded verdict; a missing record, moved commit, changed or added pickle, expired record, or any error keeps blocking. Online loads always re-query the Hub and an authoritative unsafe verdict deletes any stale record, so a now-flagged commit cannot keep loading on an old clean record. The store binds repo id, full commit, and a per-file sha256 map so a locally swapped pickle at the same commit, a branch advance, or an added load-relevant pickle is detected. A same-user attacker who can rewrite the model cache or the store is outside the enforceable boundary and this is documented; the sha256 is computed just before load, so a narrow verify-to-load window remains, and a Hub scanner false negative is recorded faithfully (safetensors stays the stronger defense). Recording is triggered post-load in the RAG embedder because the settings route only validates and the pre-load guard runs before the constructor downloads; recording is skipped when the loaded commit differs from the scanned commit. The blocked-pickle enumerator now returns snapshot-relative Paths so two module dirs that ship the same pickle basename are hashed and reported distinctly. --- studio/backend/core/rag/embeddings.py | 72 +++-- .../tests/test_embedding_scan_verdicts.py | 150 +++++++++++ .../tests/test_embedding_verdict_cache.py | 233 ++++++++++++++++ studio/backend/tests/test_hf_offline_flags.py | 18 +- studio/backend/tests/test_ssm_runtime.py | 2 +- studio/backend/utils/security/__init__.py | 2 + .../utils/security/embedding_scan_verdicts.py | 248 ++++++++++++++++++ .../backend/utils/security/file_security.py | 159 +++++++++-- 8 files changed, 835 insertions(+), 49 deletions(-) create mode 100644 studio/backend/tests/test_embedding_scan_verdicts.py create mode 100644 studio/backend/tests/test_embedding_verdict_cache.py create mode 100644 studio/backend/utils/security/embedding_scan_verdicts.py diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 38f313a70fd..db3c30c1dab 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -181,11 +181,31 @@ def _safe_subdir(rel) -> str | None: return () -def _guard_model_security(name: str, local_only_load: bool) -> None: +def _security_load_subdirs(name: str, token: str | None, local_only: bool) -> tuple[str, ...]: + """The subdirs a load calls ``from_pretrained`` on, for scoping the security scan: the audio + load roots unioned with the ST module dirs (Transformer module + Router children) from + ``modules.json`` / ``router_config.json``, so a flagged pickle directly under one blocks + instead of passing as an unreferenced nested shard. Shared by the guard and the verdict + recorder so both scope the SAME roots.""" + from utils.security import security_load_subdirs + + return tuple( + dict.fromkeys( + ( + *security_load_subdirs(name, token), + *_st_module_subdirs(name, token, local_only), + ) + ) + ) + + +def _guard_model_security(name: str, local_only_load: bool): """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. + Returns the ``FileSecurityDecision`` (or None on a gate error) so a clean ONLINE load can be + recorded for later offline reuse; raises ``UnsafeEmbeddingModelError`` when blocked. ``local_only_load`` MUST equal the ``local_files_only`` the caller passes to SentenceTransformer -- it is what licenses skipping the Hub scan. ``_hf_offline_if_dns_dead()`` @@ -193,33 +213,38 @@ def _guard_model_security(name: str, local_only_load: bool) -> None: skipped while the constructor fetched the unscanned repo. """ try: - from utils.security import evaluate_file_security, security_load_subdirs + from utils.security import evaluate_file_security 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, local_only_load), - ) - ) - ) - blocked = evaluate_file_security( + decision = evaluate_file_security( name, hf_token = token, - load_subdirs = load_subdirs, + load_subdirs = _security_load_subdirs(name, token, local_only_load), local_only_load = local_only_load, - ).blocked + ) + blocked = decision.blocked # read inside the guard so a gate error never bricks the load except Exception: - return + return None 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." ) + return decision + + +def _record_embedding_verdict_safe(name: str, commit: str | None) -> None: + """After a clean ONLINE load, persist the Hub verdict so a later OFFLINE load of the same + content is not fail-closed. Best-effort: recomputes the scan load roots (files are now cached), + hashes each load-root pickle, and records under the scanned commit. Never disturbs the loaded + model.""" + try: + from utils.security import record_embedding_verdict + + token = _ambient_hf_token() + record_embedding_verdict(name, commit, _security_load_subdirs(name, token, local_only = False)) + except Exception: + pass def _get(model_name: str | None = None): @@ -227,6 +252,7 @@ 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() + record: tuple[str, str | None] | None = None # (load_name, scanned_commit) for a clean online load with _lock: if _model is None or _name != name: _install_torchao_stub_once() @@ -251,7 +277,7 @@ def _get(model_name: str | None = None): # for the constructor could yield False after the guard skipped the Hub scan on # True. local_only = hf_env_offline() - _guard_model_security(load_name, local_only) + decision = _guard_model_security(load_name, local_only) # huggingface_hub honors only HF_HUB_OFFLINE, so a TRANSFORMERS_OFFLINE-only # session would otherwise still fetch missing repo files. _model = SentenceTransformer( @@ -261,7 +287,15 @@ def _get(model_name: str | None = None): local_files_only = local_only, ) _name = name - return _model + # An ONLINE load that HF definitively scanned clean is now fully cached: record the + # verdict (outside the lock, so hashing stays off the hot path) so a later offline + # load of this exact content is not fail-closed. + if not local_only and decision is not None and decision.scanned_clean: + record = (load_name, decision.commit) + model = _model + if record is not None: + _record_embedding_verdict_safe(*record) + return model @lru_cache(maxsize = 1) diff --git a/studio/backend/tests/test_embedding_scan_verdicts.py b/studio/backend/tests/test_embedding_scan_verdicts.py new file mode 100644 index 00000000000..3ea6cdc71bd --- /dev/null +++ b/studio/backend/tests/test_embedding_scan_verdicts.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the persistent embedding "clean verdict" store. + +The store is a fail-safe optimization: a good record lets an offline pickle model load, and ANY +problem (missing / wrong commit / expired / corrupt / disabled) must resolve to "no record" so the +offline gate keeps blocking. These tests pin that contract with an isolated UNSLOTH_STUDIO_HOME. +""" + +import json + +import pytest + + +@pytest.fixture +def verdicts(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path)) + monkeypatch.delenv("UNSLOTH_EMBED_VERDICT_CACHE_DISABLE", raising = False) + import utils.security.embedding_scan_verdicts as v + + v.clear() + return v + + +_SHA = "a" * 64 +_SHA2 = "b" * 64 + + +def test_record_then_lookup_round_trip(verdicts): + verdicts.record_clean("BAAI/bge-m3", "commit1", {"pytorch_model.bin": _SHA}) + assert verdicts.lookup("BAAI/bge-m3", "commit1") == {"pytorch_model.bin": _SHA} + + +def test_lookup_is_case_insensitive_on_repo_id(verdicts): + # Hub repo ids are case-insensitive; a record made under one casing is found under another. + verdicts.record_clean("BAAI/bge-m3", "commit1", {"pytorch_model.bin": _SHA}) + assert verdicts.lookup("baai/bge-m3", "commit1") == {"pytorch_model.bin": _SHA} + + +def test_wrong_commit_misses(verdicts): + verdicts.record_clean("acme/model", "commit1", {"pytorch_model.bin": _SHA}) + assert verdicts.lookup("acme/model", "commit2") is None + + +def test_lookup_without_commit_is_none(verdicts): + verdicts.record_clean("acme/model", "commit1", {"pytorch_model.bin": _SHA}) + assert verdicts.lookup("acme/model", None) is None + + +def test_record_without_commit_is_noop(verdicts): + verdicts.record_clean("acme/model", None, {"pytorch_model.bin": _SHA}) + assert verdicts.lookup("acme/model", "commit1") is None + + +def test_record_with_empty_map_is_noop(verdicts): + verdicts.record_clean("acme/model", "commit1", {}) + assert verdicts.lookup("acme/model", "commit1") is None + + +def test_malformed_hash_is_rejected(verdicts): + # A non-64-hex digest can never be persisted, so it can never seed a spurious match. + verdicts.record_clean("acme/model", "commit1", {"pytorch_model.bin": "not-a-hash"}) + assert verdicts.lookup("acme/model", "commit1") is None + + +def test_ttl_expiry_returns_none(verdicts): + from datetime import datetime, timedelta, timezone + + verdicts.record_clean("acme/model", "commit1", {"pytorch_model.bin": _SHA}) + # Backdate the record past the 30-day TTL by editing the store on disk. + path = verdicts._store_path() + data = json.loads(path.read_text()) + old = (datetime.now(timezone.utc) - timedelta(days = 31)).isoformat() + data["records"]["acme/model"]["recorded_at"] = old + path.write_text(json.dumps(data)) + assert verdicts.lookup("acme/model", "commit1") is None + + +def test_unparseable_timestamp_returns_none(verdicts): + verdicts.record_clean("acme/model", "commit1", {"pytorch_model.bin": _SHA}) + path = verdicts._store_path() + data = json.loads(path.read_text()) + data["records"]["acme/model"]["recorded_at"] = "not-a-timestamp" + path.write_text(json.dumps(data)) + assert verdicts.lookup("acme/model", "commit1") is None + + +def test_forget_drops_record(verdicts): + verdicts.record_clean("acme/model", "commit1", {"pytorch_model.bin": _SHA}) + verdicts.forget("acme/model") + assert verdicts.lookup("acme/model", "commit1") is None + + +def test_corrupt_store_fails_safe(verdicts): + path = verdicts._store_path() + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text("{ this is not json") + assert verdicts.lookup("acme/model", "commit1") is None + + +def test_wrong_schema_version_fails_safe(verdicts): + path = verdicts._store_path() + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(json.dumps({"version": 999, "records": {"acme/model": {"commit": "commit1", + "recorded_at": "2999-01-01T00:00:00+00:00", "pickles": {"pytorch_model.bin": _SHA}}}})) + assert verdicts.lookup("acme/model", "commit1") is None + + +def test_non_dict_records_fails_safe(verdicts): + path = verdicts._store_path() + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(json.dumps({"version": 1, "records": []})) + assert verdicts.lookup("acme/model", "commit1") is None + + +def test_cache_disabled_noops_record_and_lookup(verdicts, monkeypatch): + monkeypatch.setenv("UNSLOTH_EMBED_VERDICT_CACHE_DISABLE", "1") + verdicts.record_clean("acme/model", "commit1", {"pytorch_model.bin": _SHA}) + assert verdicts.lookup("acme/model", "commit1") is None + # And a record made while enabled is not honored once disabled. + monkeypatch.delenv("UNSLOTH_EMBED_VERDICT_CACHE_DISABLE") + verdicts.record_clean("acme/model", "commit1", {"pytorch_model.bin": _SHA}) + monkeypatch.setenv("UNSLOTH_EMBED_VERDICT_CACHE_DISABLE", "true") + assert verdicts.lookup("acme/model", "commit1") is None + + +def test_latest_record_per_repo_wins(verdicts): + verdicts.record_clean("acme/model", "commit1", {"pytorch_model.bin": _SHA}) + verdicts.record_clean("acme/model", "commit2", {"pytorch_model.bin": _SHA2}) + assert verdicts.lookup("acme/model", "commit1") is None + assert verdicts.lookup("acme/model", "commit2") == {"pytorch_model.bin": _SHA2} + + +def test_sha256_file(tmp_path, verdicts): + import hashlib + + f = tmp_path / "w.bin" + f.write_bytes(b"hello world") + assert verdicts.sha256_file(f) == hashlib.sha256(b"hello world").hexdigest() + + +def test_sha256_file_missing_is_none(tmp_path, verdicts): + assert verdicts.sha256_file(tmp_path / "nope.bin") is None + + +def test_record_persists_across_reload(verdicts): + # A fresh _load() (simulating a restart) still sees the record: the write is durable. + verdicts.record_clean("acme/model", "commit1", {"a/pytorch_model.bin": _SHA}) + assert verdicts._load()["records"]["acme/model"]["pickles"] == {"a/pytorch_model.bin": _SHA} diff --git a/studio/backend/tests/test_embedding_verdict_cache.py b/studio/backend/tests/test_embedding_verdict_cache.py new file mode 100644 index 00000000000..cb2fb1ce65e --- /dev/null +++ b/studio/backend/tests/test_embedding_verdict_cache.py @@ -0,0 +1,233 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Integration tests for the offline embedding verdict cache: an ONLINE clean load records a +verdict, and a later OFFLINE load of the SAME content is allowed instead of fail-closed. Every +mismatch (moved commit, changed / added pickle, expired record) must keep blocking. + +Snapshot resolution is stubbed (``_active_snapshot_dir`` / ``_active_commit``) so no HF cache or +network is touched; UNSLOTH_STUDIO_HOME isolates the verdict store per test. +""" + +import json +import sys +import types +from datetime import datetime, timedelta, timezone + +import pytest + +import core.rag.config as config +import core.rag.embeddings as embeddings +import utils.models.model_config as mc +import utils.security.embedding_scan_verdicts as verdicts +import utils.security.file_security as fs + + +@pytest.fixture +def home(tmp_path, monkeypatch): + monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path / "studio")) + monkeypatch.delenv("UNSLOTH_EMBED_VERDICT_CACHE_DISABLE", raising = False) + verdicts.clear() + return tmp_path + + +def _snap(tmp_path, files: dict): + """Build a snapshot dir with *files* (name -> bytes/str; names may be nested like ``a/b.bin``).""" + d = tmp_path / "cache" / "snap" + d.mkdir(parents = True, exist_ok = True) + for name, body in files.items(): + p = d / name + p.parent.mkdir(parents = True, exist_ok = True) + p.write_bytes(body if isinstance(body, bytes) else body.encode()) + return d + + +def _blocked(monkeypatch, snap, commit, load_subdirs = ()): + monkeypatch.setattr(mc, "_active_snapshot_dir", lambda name: snap) + monkeypatch.setattr(mc, "_active_commit", lambda name: commit) + return fs.evaluate_file_security( + "org/model", None, load_subdirs = load_subdirs, local_only_load = True + ).blocked + + +def _sha(path): + return verdicts.sha256_file(path) + + +# ── Offline verify matrix ──────────────────────────────────────────── + +def test_matching_record_allows_offline_pickle(home, tmp_path, monkeypatch): + snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) + verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) + assert _blocked(monkeypatch, snap, "c1") is False + + +def test_no_record_blocks_offline_pickle(home, tmp_path, monkeypatch): + snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) + assert _blocked(monkeypatch, snap, "c1") is True + + +def test_wrong_commit_blocks(home, tmp_path, monkeypatch): + snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) + verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) + assert _blocked(monkeypatch, snap, "c2") is True + + +def test_tampered_pickle_blocks(home, tmp_path, monkeypatch): + snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) + verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) + (snap / "pytorch_model.bin").write_bytes(b"TAMPERED-SAME-COMMIT") + assert _blocked(monkeypatch, snap, "c1") is True + + +def test_partial_record_blocks(home, tmp_path, monkeypatch): + # Two module dirs each ship a pickle; recording only one leaves the set mismatched -> block. + snap = _snap( + tmp_path, + { + "modules.json": json.dumps( + [{"path": "0_A", "type": "..."}, {"path": "0_B", "type": "..."}] + ), + "0_A/pytorch_model.bin": b"aaa", + "0_B/pytorch_model.bin": b"bbb", + }, + ) + verdicts.record_clean("org/model", "c1", {"0_A/pytorch_model.bin": _sha(snap / "0_A/pytorch_model.bin")}) + assert _blocked(monkeypatch, snap, "c1") is True + + +def test_distinct_basenames_in_different_roots_are_keyed_apart(home, tmp_path, monkeypatch): + # Same basename in two module dirs must be recorded/verified as two distinct entries; if the + # enumerator collapsed them by basename, one record could vouch for both. Recording BOTH with + # their real (distinct) hashes allows; the map has two keys. + snap = _snap( + tmp_path, + { + "modules.json": json.dumps( + [{"path": "0_A", "type": "..."}, {"path": "0_B", "type": "..."}] + ), + "0_A/pytorch_model.bin": b"aaa", + "0_B/pytorch_model.bin": b"bbb", + }, + ) + verdicts.record_clean( + "org/model", + "c1", + { + "0_A/pytorch_model.bin": _sha(snap / "0_A/pytorch_model.bin"), + "0_B/pytorch_model.bin": _sha(snap / "0_B/pytorch_model.bin"), + }, + ) + assert _blocked(monkeypatch, snap, "c1") is False + + +def test_expired_record_blocks(home, tmp_path, monkeypatch): + snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) + verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) + path = verdicts._store_path() + data = json.loads(path.read_text()) + data["records"]["org/model"]["recorded_at"] = ( + datetime.now(timezone.utc) - timedelta(days = 31) + ).isoformat() + path.write_text(json.dumps(data)) + assert _blocked(monkeypatch, snap, "c1") is True + + +def test_router_child_pickle_recorded_allows_else_blocks(home, tmp_path, monkeypatch): + # A Router child pickle (declared only in router_config.json) is a load root the offline gate + # scopes. Recording it allows; without a record it blocks. + snap = _snap( + tmp_path, + { + "modules.json": json.dumps([{"path": "", "type": "sentence_transformers.models.Router.Router"}]), + "router_config.json": json.dumps({"types": {"query_0_WordEmbeddings": "..."}}), + "query_0_WordEmbeddings/pytorch_model.bin": b"child-weights", + }, + ) + child = snap / "query_0_WordEmbeddings" / "pytorch_model.bin" + assert _blocked(monkeypatch, snap, "c1") is True + verdicts.record_clean("org/model", "c1", {"query_0_WordEmbeddings/pytorch_model.bin": _sha(child)}) + assert _blocked(monkeypatch, snap, "c1") is False + + +def test_disabled_cache_never_allows(home, tmp_path, monkeypatch): + snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) + verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) + monkeypatch.setenv("UNSLOTH_EMBED_VERDICT_CACHE_DISABLE", "1") + assert _blocked(monkeypatch, snap, "c1") is True + + +# ── record_embedding_verdict (the post-load recorder) ──────────────── + +def _stub_snapshot(monkeypatch, snap, commit): + monkeypatch.setattr(mc, "_active_snapshot_dir", lambda name: snap) + monkeypatch.setattr(mc, "_active_commit", lambda name: commit) + + +def test_record_embedding_verdict_records_and_allows(home, tmp_path, monkeypatch): + snap = _snap(tmp_path, {"pytorch_model.bin": b"w"}) + _stub_snapshot(monkeypatch, snap, "c1") + fs.record_embedding_verdict("org/model", "c1", load_subdirs = ()) + assert verdicts.lookup("org/model", "c1") == {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")} + assert fs.evaluate_file_security("org/model", None, local_only_load = True).blocked is False + + +def test_record_skips_when_commit_moved(home, tmp_path, monkeypatch): + # The active cached commit differs from the scanned commit (branch advanced): record nothing, + # so an unscanned commit is never blessed. + snap = _snap(tmp_path, {"pytorch_model.bin": b"w"}) + _stub_snapshot(monkeypatch, snap, "c2") + fs.record_embedding_verdict("org/model", "c1", load_subdirs = ()) + assert verdicts.lookup("org/model", "c1") is None + assert verdicts.lookup("org/model", "c2") is None + + +def test_record_skips_pickle_free_cache(home, tmp_path, monkeypatch): + snap = _snap(tmp_path, {"model.safetensors": b"\0"}) + _stub_snapshot(monkeypatch, snap, "c1") + fs.record_embedding_verdict("org/model", "c1", load_subdirs = ()) + assert verdicts.lookup("org/model", "c1") is None + + +# ── _get() wiring: record only after a clean ONLINE load ───────────── + +def _drive_get(monkeypatch, decision, *, offline): + fake_st = types.ModuleType("sentence_transformers") + fake_st.SentenceTransformer = lambda name, **k: object() + monkeypatch.setitem(sys.modules, "sentence_transformers", fake_st) + monkeypatch.setattr(config, "effective_embedding_model", lambda: "org/model") + monkeypatch.setattr(embeddings, "_device", lambda: "cpu") + monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) + monkeypatch.setattr("utils.models.resolve_st_cached_repo_id_case", lambda r: r) + monkeypatch.setattr("utils.utils.hf_env_offline", lambda: offline) + monkeypatch.setattr(embeddings, "_guard_model_security", lambda name, lo: decision) + rec = {} + monkeypatch.setattr( + embeddings, "_record_embedding_verdict_safe", lambda n, c: rec.update(name = n, commit = c) + ) + monkeypatch.setattr(embeddings, "_model", None, raising = False) + monkeypatch.setattr(embeddings, "_name", None, raising = False) + embeddings._get() + return rec + + +def test_get_records_after_clean_online_load(home, monkeypatch): + dec = fs.FileSecurityDecision("org/model", False, commit = "c1", scanned_clean = True) + assert _drive_get(monkeypatch, dec, offline = False) == {"name": "org/model", "commit": "c1"} + + +def test_get_does_not_record_when_scan_not_definitive(home, monkeypatch): + # A fail-open ("scan unavailable") online load has scanned_clean False -> record nothing. + dec = fs.FileSecurityDecision("org/model", False, commit = None, scanned_clean = False) + assert _drive_get(monkeypatch, dec, offline = False) == {} + + +def test_get_does_not_record_offline_load(home, monkeypatch): + # Offline: no authoritative scan happened, so nothing is recorded regardless. + dec = fs.FileSecurityDecision("org/model", False, commit = "c1", scanned_clean = True) + assert _drive_get(monkeypatch, dec, offline = True) == {} + + +def test_get_does_not_record_on_gate_error(home, monkeypatch): + # _guard_model_security returns None on a gate error -> record nothing. + assert _drive_get(monkeypatch, None, offline = False) == {} diff --git a/studio/backend/tests/test_hf_offline_flags.py b/studio/backend/tests/test_hf_offline_flags.py index cfa3b8abbff..d3e8f13f723 100644 --- a/studio/backend/tests/test_hf_offline_flags.py +++ b/studio/backend/tests/test_hf_offline_flags.py @@ -104,7 +104,7 @@ def test_shared_gate_still_scans_when_offline_by_default(monkeypatch, var): calls: list = [] _fake_hub(monkeypatch, calls) monkeypatch.setenv(var, "1") - assert fs._fetch_security_status("org/model", None) is None + assert fs._fetch_security_status("org/model", None) == (None, None) assert calls, f"{var} alone must NOT bypass the shared malware gate" @@ -384,7 +384,7 @@ def test_security_scan_runs_when_online(monkeypatch): calls: list = [] _fake_hub(monkeypatch, calls) - assert fs._fetch_security_status("org/model", None) is None + assert fs._fetch_security_status("org/model", None) == (None, None) assert calls, "online must attempt the Hub" @@ -557,9 +557,10 @@ def test_online_scan_blocks_flagged_pickle_under_traversing_module_subdir(monkey monkeypatch.setattr( fs, "_fetch_security_status", - lambda name, token: { - "filesWithIssues": [{"path": "evil/pytorch_model.bin", "level": "unsafe"}] - }, + lambda name, token: ( + {"filesWithIssues": [{"path": "evil/pytorch_model.bin", "level": "unsafe"}]}, + "commitsha", + ), ) # Definitive "no weight index" -> a missed shard is SKIPPED (allowed) on the pristine gate; # the fix instead treats evil/ as a load root so the pickle is root-level and blocks. @@ -577,9 +578,10 @@ def test_online_scan_blocks_flagged_pickle_under_plain_module_subdir(monkeypatch monkeypatch.setattr( fs, "_fetch_security_status", - lambda name, token: { - "filesWithIssues": [{"path": "0_Transformer/pytorch_model.bin", "level": "unsafe"}] - }, + lambda name, token: ( + {"filesWithIssues": [{"path": "0_Transformer/pytorch_model.bin", "level": "unsafe"}]}, + "commitsha", + ), ) monkeypatch.setattr(fs, "_indexed_shard_paths", lambda *a, **k: set()) decision = fs.evaluate_file_security( diff --git a/studio/backend/tests/test_ssm_runtime.py b/studio/backend/tests/test_ssm_runtime.py index bb0caa2887d..f8cad4c2552 100644 --- a/studio/backend/tests/test_ssm_runtime.py +++ b/studio/backend/tests/test_ssm_runtime.py @@ -461,7 +461,7 @@ def _is_gated_module(name: str) -> bool: _sys.modules.pop(m, None) try: - with patch.object(fs, "_fetch_security_status", return_value = None): + with patch.object(fs, "_fetch_security_status", return_value = (None, None)): fs.evaluate_file_security("nvidia/Nemotron-H-8B", load_subdirs = ()) with patch.object( consent, "_load_remote_code_configs", return_value = [{"model_type": "nemotron_h"}] diff --git a/studio/backend/utils/security/__init__.py b/studio/backend/utils/security/__init__.py index 3bcfaebe2f2..bb74b99143f 100644 --- a/studio/backend/utils/security/__init__.py +++ b/studio/backend/utils/security/__init__.py @@ -20,6 +20,7 @@ from utils.security.file_security import ( # noqa: F401 FileSecurityDecision, evaluate_file_security, + record_embedding_verdict, security_load_subdirs, ) from utils.security.remote_code_scan import ( # noqa: F401 @@ -47,6 +48,7 @@ "preflight_remote_code_consent", "preflight_remote_code_consent_for_targets", "evaluate_file_security", + "record_embedding_verdict", "security_load_subdirs", "FileSecurityDecision", "RemoteCodeDecision", diff --git a/studio/backend/utils/security/embedding_scan_verdicts.py b/studio/backend/utils/security/embedding_scan_verdicts.py new file mode 100644 index 00000000000..b880a860118 --- /dev/null +++ b/studio/backend/utils/security/embedding_scan_verdicts.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Persistent, per-user cache of Hugging Face "clean" security verdicts for embedding models. + +The offline embedding-security gate is fail-CLOSED: with no network to reach HF's scan, a cached +pickle weight cannot be verified, so it is blocked. This cache lets a pickle model the user already +loaded ONLINE (and that HF scanned clean) load again OFFLINE, without weakening the gate for an +unknown or never-scanned pickle. A record is honored offline only when the active cached commit and +EVERY load-root pickle's sha256 exactly match what was recorded; anything else keeps blocking. + +Binding: repo id + full commit SHA + an exact map of snapshot-relative pickle name -> sha256. The +commit pins HF's immutable content; the per-file sha256 detects a locally swapped pickle at that +commit; the exact map (not a set) detects an added / renamed / newly load-relevant pickle. + +Threat model. This protects against STALE revisions and ACCIDENTAL / non-concurrent cache tampering. +A same-user attacker with arbitrary filesystem write -- to the HF cache OR to this store under +``studio_root()/security/`` -- is OUTSIDE the enforceable boundary: they can already replace the +weights, imported Python, or loader config, and could forge a record beside a matching poisoned +pickle. A local HMAC/signing key in the same account does not change that. The store is written +0600 in the Studio-private dir and every read failure degrades to "no record" (i.e. block); a +record is never trusted online (the Hub is always re-queried) and is deleted on an authoritative +unsafe verdict. A 30-day TTL bounds offline reuse of a verdict HF may have since revised. The +sha256 is computed locally just before load, so a verify->load TOCTOU window remains (narrow; the +same-user-write threat already dominates), and a Hub scanner false negative is recorded faithfully +-- safetensors remains the stronger, format-level defense and the UI/logs say "previously cleared +by the Hub scan", not "proven safe". +""" + +from __future__ import annotations + +import contextlib +import hashlib +import json +import os +import threading +from datetime import datetime, timezone +from typing import Optional + +from loggers import get_logger +from utils.paths import storage_roots + +logger = get_logger(__name__) + +_SCHEMA_VERSION = 1 +# A recorded verdict is honored for at most this long, bounding offline reuse of a verdict HF may +# have revised since (an offline process cannot learn about a rescan). An online clean load rewrites +# the record and refreshes the window. +_TTL_SECONDS = 30 * 24 * 60 * 60 +_HASH_CHUNK = 1024 * 1024 # 1 MiB streaming reads so a multi-GB weight never loads into memory + +_lock = threading.RLock() + + +def cache_disabled() -> bool: + """True when the verdict cache is turned off by env, forcing the pure fail-closed offline gate.""" + return os.environ.get("UNSLOTH_EMBED_VERDICT_CACHE_DISABLE", "").strip().lower() in ( + "1", + "true", + "yes", + "on", + ) + + +def _store_path(): + return storage_roots.studio_root() / "security" / "embedding_scan_verdicts.json" + + +def _key(repo_id: str) -> str: + """Case-folded store key. Hub repo ids are case-insensitive, so ``BAAI/bge-m3`` and + ``baai/bge-m3`` must resolve to one record.""" + return (repo_id or "").strip().lower() + + +def sha256_file(path) -> Optional[str]: + """Streaming sha256 hex digest of *path*, or None if it cannot be read (an unreadable file must + never verify as a match -> the caller keeps blocking).""" + try: + digest = hashlib.sha256() + with open(path, "rb") as fh: + for chunk in iter(lambda: fh.read(_HASH_CHUNK), b""): + digest.update(chunk) + return digest.hexdigest() + except OSError: + return None + + +def _load() -> dict: + """Parsed store, or an empty skeleton on any error (fail-safe = no records = block offline).""" + try: + with open(_store_path()) as f: + data = json.load(f) + if ( + isinstance(data, dict) + and data.get("version") == _SCHEMA_VERSION + and isinstance(data.get("records"), dict) + ): + return data + except FileNotFoundError: + pass + except Exception as exc: + logger.warning("Could not read embedding scan verdicts (%s); ignoring", exc) + return {"version": _SCHEMA_VERSION, "records": {}} + + +def _save(data: dict) -> None: + """Atomic write (tmp + os.replace), best-effort 0600 in the Studio-private security dir.""" + path = _store_path() + storage_roots.ensure_dir(path.parent) + tmp = path.parent / f".{path.name}.tmp-{os.getpid()}" + try: + with open(tmp, "w") as f: + json.dump(data, f, indent = 2) + try: + os.chmod(tmp, 0o600) + except OSError: + pass + os.replace(tmp, path) + except Exception as exc: + logger.warning("Could not write embedding scan verdicts (%s)", exc) + try: + tmp.unlink(missing_ok = True) + except OSError: + pass + + +@contextlib.contextmanager +def _file_lock(): + """Best-effort cross-process exclusive lock around a read-modify-write. Training / export / + inference record from separate subprocesses, so the in-process RLock alone would let two + processes each read the store and clobber the other's record on ``os.replace``. Degrades to a + no-op when OS locking is unavailable (the only consequence is a lost optimization record, which + just causes a later offline block, never a bypass).""" + path = _store_path() + try: + storage_roots.ensure_dir(path.parent) + fd = os.open(str(path.parent / f"{path.name}.lock"), os.O_CREAT | os.O_RDWR, 0o600) + except Exception: + yield + return + try: + try: + if os.name == "nt": + import msvcrt + msvcrt.locking(fd, msvcrt.LK_LOCK, 1) + else: + import fcntl + fcntl.flock(fd, fcntl.LOCK_EX) + except Exception: + pass # locking unavailable; the thread lock still applies + yield + finally: + try: + if os.name == "nt": + import msvcrt + with contextlib.suppress(Exception): + msvcrt.locking(fd, msvcrt.LK_UNLCK, 1) + else: + import fcntl + fcntl.flock(fd, fcntl.LOCK_UN) + finally: + os.close(fd) + + +def _clean_pickles(pickles) -> Optional[dict]: + """Validate a ``{rel-path: sha256}`` map: a non-empty dict of str->64-hex-char values, or None. + Rejects a hand-edited entry whose shape is wrong so it can never seed a spurious match.""" + if not isinstance(pickles, dict) or not pickles: + return None + out = {} + for rel, digest in pickles.items(): + if not isinstance(rel, str) or not isinstance(digest, str): + return None + d = digest.strip().lower() + if len(d) != 64 or any(c not in "0123456789abcdef" for c in d): + return None + out[rel] = d + return out + + +def record_clean(repo_id: str, commit: Optional[str], pickles: dict) -> None: + """Persist a clean verdict for *repo_id* at *commit*: the exact ``{snapshot-relative pickle + name: sha256}`` map the loader will deserialize. No-op when the cache is disabled, the commit is + unknown, or the map is empty/malformed (nothing to attest).""" + if cache_disabled() or not commit: + return + cleaned = _clean_pickles(pickles) + if cleaned is None: + return + key = _key(repo_id) + if not key: + return + with _lock, _file_lock(): + data = _load() + records = data.setdefault("records", {}) + if not isinstance(records, dict): # tolerate a hand-edited non-dict + records = data["records"] = {} + records[key] = { + "commit": commit, + "recorded_at": datetime.now(timezone.utc).isoformat(), + "pickles": cleaned, + } + _save(data) + + +def lookup(repo_id: str, commit: Optional[str]) -> Optional[dict]: + """The recorded ``{rel: sha256}`` map for *repo_id*, but ONLY when the stored commit equals + *commit* and the record is within the TTL. None otherwise, on a disabled cache, or on any error + -- the offline caller then keeps blocking.""" + if cache_disabled() or not commit: + return None + key = _key(repo_id) + if not key: + return None + with _lock: + entry = _load().get("records", {}).get(key) + if not isinstance(entry, dict) or entry.get("commit") != commit: + return None + recorded_at = entry.get("recorded_at") + try: + age = (datetime.now(timezone.utc) - datetime.fromisoformat(recorded_at)).total_seconds() + except (TypeError, ValueError): + return None # unparseable timestamp -> treat as no record + if age < 0 or age > _TTL_SECONDS: + return None + return _clean_pickles(entry.get("pickles")) + + +def forget(repo_id: str) -> None: + """Drop a repo's recorded verdict (called when HF returns an authoritative unsafe status, so a + now-flagged commit cannot keep loading offline on a stale clean record).""" + key = _key(repo_id) + if not key: + return + with _lock, _file_lock(): + data = _load() + records = data.get("records", {}) + if isinstance(records, dict) and records.pop(key, None) is not None: + _save(data) + + +def clear() -> None: + """Test helper: drop the on-disk store.""" + with _lock: + try: + _store_path().unlink(missing_ok = True) + except OSError: + pass diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 0e3e876cac4..945f04af40b 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -226,6 +226,12 @@ class FileSecurityDecision: blocked: bool unsafe_files: list = field(default_factory = list) # [{"path", "level"}] reason: str = "" + # The commit SHA the Hub scan reported (online only), and whether that scan was a DEFINITIVE + # clean verdict (a completed scan with no load-path issue) -- not merely a non-block. The + # embedding recorder persists a clean verdict only when ``scanned_clean`` and the loaded commit + # equals ``commit``; both stay unset for local-only and fail-open (unavailable) decisions. + commit: Optional[str] = None + scanned_clean: bool = False def response_payload(self) -> dict: """Machine-readable detail merged into the preflight payload the dialog reads.""" @@ -456,17 +462,19 @@ def _st_load_roots(snap, load_subdirs = ()) -> set: return roots -def _cached_pickle_weight_files(snap, load_subdirs = ()) -> list: - """Pickle weight files a from_pretrained load actually deserializes: at a real load root and - with NO loadable safetensors alternative there. A load root is the snapshot root, a directory - ``modules.json`` / ``load_subdirs`` declares, or a plain from_pretrained root (holds - ``config.json``). A stray pickle in a non-load subdir (``archive/``, ``nemo/``) that no load - opens is not a vector, matching the online scan's load-path scoping. Two weight classes are - scoped independently: a BASE pickle (``pytorch_model.bin`` ...) is covered by a loadable base - safetensors the loader picks instead; a PEFT ADAPTER pickle (``adapter_model.bin``), which - from_pretrained auto-loads when ``adapter_config.json`` is present, is covered only by +def _cached_pickle_weight_paths(snap, load_subdirs = ()) -> list: + """The pickle weight FILES (as ``Path`` objects under *snap*) a from_pretrained load actually + deserializes: at a real load root and with NO loadable safetensors alternative there. A load + root is the snapshot root, a directory ``modules.json`` / ``load_subdirs`` declares, or a plain + from_pretrained root (holds ``config.json``). A stray pickle in a non-load subdir (``archive/``, + ``nemo/``) that no load opens is not a vector, matching the online scan's load-path scoping. Two + weight classes are scoped independently: a BASE pickle (``pytorch_model.bin`` ...) is covered by + a loadable base safetensors the loader picks instead; a PEFT ADAPTER pickle (``adapter_model.bin``), + which from_pretrained auto-loads when ``adapter_config.json`` is present, is covered only by ``adapter_model.safetensors`` -- a safetensors base does NOT cover it. A bare adapter or an - orphan shard leaves its pickle live.""" + orphan shard leaves its pickle live. Returns the concrete Paths (not basenames) so the offline + verdict cache can hash exactly the files the loader reads and never conflate two module dirs that + ship the same pickle basename.""" roots = _st_load_roots(snap, load_subdirs) by_dir_pickle: dict = {} by_dir_files: dict = {} # directory -> {lower-name: Path} @@ -496,7 +504,7 @@ def _cached_pickle_weight_files(snap, load_subdirs = ()) -> list: continue base = [n for n in names if _PICKLE_WEIGHT_RE.match(n.lower())] if base and not _dir_has_loadable_safetensors(files): - hits.update(base) + hits.update(files[n.lower()] for n in base) # An adapter pickle is deserialized only when from_pretrained auto-detects the adapter # (adapter_config.json present) and there is no adapter_model.safetensors to load instead. # The safetensors credit is case-SENSITIVE (a mixed-case Adapter_Model.SafeTensors decoy @@ -508,7 +516,7 @@ def _cached_pickle_weight_files(snap, load_subdirs = ()) -> list: and "adapter_config.json" in files and _exact_named(files, "adapter_model.safetensors") is None ): - hits.update(adapter) + hits.update(files[n.lower()] for n in adapter) # A load-root pickle index (pytorch_model.bin.index.json) can map shards into SUBDIRECTORIES # that are not themselves load roots; from_pretrained follows the map and deserializes them, # so include those referenced pickle shards (unless a loadable base safetensors at the index @@ -524,12 +532,89 @@ def _cached_pickle_weight_files(snap, load_subdirs = ()) -> list: shard = root_dir.joinpath(*shard_rel.split("/")) try: if shard.is_file(): - hits.add(shard_rel) + hits.add(shard) except OSError: continue return sorted(hits) +def _snapshot_relative(path, snap) -> str: + """*path* as a snapshot-relative posix string, or its basename if it is somehow not under + *snap* (only for display / the blocked-file list).""" + try: + return path.relative_to(snap).as_posix() + except ValueError: + return path.name + + +def _cached_pickle_weight_files(snap, load_subdirs = ()) -> list: + """Snapshot-relative posix names of the load-root pickle weights (see + :func:`_cached_pickle_weight_paths`). Names are relative to the snapshot so two module dirs that + each ship ``pytorch_model.bin`` are reported (and, in the verdict cache, keyed) distinctly.""" + return [_snapshot_relative(p, snap) for p in _cached_pickle_weight_paths(snap, load_subdirs)] + + +def _pickle_hash_map(snap, paths): + """``{snapshot-relative posix name: sha256}`` for *paths*, or None if any file cannot be hashed + (an unreadable pickle must never verify or record as clean).""" + from utils.security import embedding_scan_verdicts + + out = {} + for p in paths: + digest = embedding_scan_verdicts.sha256_file(p) + if digest is None: + return None + out[_snapshot_relative(p, snap)] = digest + return out + + +def _matches_clean_verdict(model_name: str, snap, paths) -> bool: + """True when a recorded clean Hub verdict covers EXACTLY this load-root pickle set at the active + cached commit: same commit, same relative-name set, same sha256 for every file. A missing + record, moved commit, changed / added / unreadable pickle, or any error -> False (fail-closed).""" + try: + from utils.models.model_config import _active_commit + from utils.security import embedding_scan_verdicts + + commit = _active_commit(model_name) + if not commit: + return False + recorded = embedding_scan_verdicts.lookup(model_name, commit) + if not recorded: + return False + current = _pickle_hash_map(snap, paths) + return current is not None and current == recorded + except Exception: + return False + + +def record_embedding_verdict(model_name: str, scanned_commit, load_subdirs = ()) -> None: + """Record a clean Hub verdict for an embedding repo just loaded ONLINE, so a later offline load + of the same content is not fail-closed. Persists the sha256 of every load-root pickle keyed by + its snapshot-relative name. Records nothing when the snapshot is missing, the active cached + commit differs from the scanned commit (branch moved -> the loaded content was not what HF + scanned), there are no load-root pickles (an inert cache is already allowed), or any file cannot + be hashed. Never raises into the load path.""" + try: + from utils.models.model_config import _active_commit, _active_snapshot_dir + from utils.security import embedding_scan_verdicts + + if not scanned_commit or _active_commit(model_name) != scanned_commit: + return + snap = _active_snapshot_dir(model_name) + if snap is None: + return + paths = _cached_pickle_weight_paths(snap, load_subdirs) + if not paths: + return + pickles = _pickle_hash_map(snap, paths) + if pickles is None: + return + embedding_scan_verdicts.record_clean(model_name, scanned_commit, pickles) + except Exception as exc: + logger.debug("Could not record embedding scan verdict for '%s': %s", model_name, exc) + + def _evaluate_local_only(model_name: str, load_subdirs = ()) -> "FileSecurityDecision": """Fail-CLOSED security decision for an offline (local_files_only) load. @@ -548,12 +633,24 @@ def _evaluate_local_only(model_name: str, load_subdirs = ()) -> "FileSecurityDec if snap is None: return FileSecurityDecision(model_name, False, reason = "offline; nothing cached to scan") - pickles = _cached_pickle_weight_files(snap, load_subdirs) - if not pickles: + paths = _cached_pickle_weight_paths(snap, load_subdirs) + if not paths: return FileSecurityDecision( model_name, False, reason = "offline; cached weights are pickle-free (inert)" ) + # A pickle model the user already loaded ONLINE (and HF scanned clean) may load offline when + # its exact content is unchanged. Allow only when the active cached commit and EVERY load-root + # pickle's sha256 match the recorded clean verdict; a missing record, moved commit, hash + # mismatch, unreadable file, or any error falls through to the fail-closed block below. + if _matches_clean_verdict(model_name, snap, paths): + return FileSecurityDecision( + model_name, + False, + reason = "offline; cached pickle weights match a recorded clean Hub scan", + ) + + pickles = [_snapshot_relative(p, snap) for p in paths] names = ", ".join(pickles) logger.warning( "Blocking offline load of '%s': cached pickle weights cannot be security-scanned " @@ -573,8 +670,10 @@ def _evaluate_local_only(model_name: str, load_subdirs = ()) -> "FileSecurityDec def _fetch_security_status(model_name: str, hf_token: Optional[str]): - """``security_repo_status`` (a dict) or None if unavailable. Hub metadata only; - retries once on a transient error, then returns None so the caller fails open. + """``(security_repo_status, commit_sha)`` from a single Hub metadata call, or ``(None, None)`` + if unavailable. The commit is the SHA the scan applies to (used to bind a recorded clean + verdict to immutable content). Metadata only; retries once on a transient error, then returns + ``(None, None)`` so the caller fails open. """ from huggingface_hub import model_info as hf_model_info @@ -588,7 +687,7 @@ def _fetch_security_status(model_name: str, hf_token: Optional[str]): securityStatus = True, timeout = timeout, ) - return getattr(info, "security_repo_status", None) + return getattr(info, "security_repo_status", None), getattr(info, "sha", None) except Exception as exc: # network/offline/gated/404/unsupported-client last_exc = exc if attempt == 0: @@ -598,7 +697,7 @@ def _fetch_security_status(model_name: str, hf_token: Optional[str]): model_name, type(last_exc).__name__ if last_exc else "unknown", ) - return None + return None, None def evaluate_file_security( @@ -644,11 +743,15 @@ def evaluate_file_security( if local_only_load: return _evaluate_local_only(model_name, load_subdirs) - status = _fetch_security_status(model_name, hf_token) + status, commit = _fetch_security_status(model_name, hf_token) if not isinstance(status, dict): return FileSecurityDecision( model_name, False, reason = "scan unavailable; allowed (fail-open)" ) + # A DEFINITIVE clean verdict (for the durable offline cache) requires a COMPLETED scan, unlike + # the block decision which is not gated on scansDone. An in-progress scan is allowed online but + # must not be persisted as clean. + scans_done = bool(status.get("scansDone")) # Block a non-``safe`` flagged file scoped to the load-path RCE vector (root-level, # code-executing). Not gated on ``scansDone`` (often false even when clean; a flagged @@ -701,7 +804,13 @@ def evaluate_file_security( model_name, ", ".join(f"{s['path']}({s['level']})" for s in skipped), ) - return FileSecurityDecision(model_name, False, reason = "no unsafe files in the load path") + return FileSecurityDecision( + model_name, + False, + reason = "no unsafe files in the load path", + commit = commit, + scanned_clean = scans_done, + ) names = ", ".join(u["path"] for u in unsafe if u["path"]) or "unknown files" logger.warning( @@ -709,9 +818,17 @@ def evaluate_file_security( model_name, names, ) + # An authoritative unsafe verdict revokes any stale clean record for this repo, so a + # now-flagged commit cannot keep loading offline on a previously recorded verdict. + try: + from utils.security import embedding_scan_verdicts + embedding_scan_verdicts.forget(model_name) + except Exception: + pass return FileSecurityDecision( model_name, True, unsafe_files = unsafe, reason = f"Hugging Face security scan flagged unsafe files: {names}", + commit = commit, ) From 831da561f660d581fd8c0f684ecae323f4ed0017 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 21 Jul 2026 03:21:53 +0000 Subject: [PATCH 61/77] Harden the embedding verdict cache against review findings Tighten the offline verdict cache and its enumeration so every uncertain or malformed input fails closed and the recorded hashes always match the files the loader reads: - Hash every case-colliding pickle in a load root, not one representative. On a case-sensitive filesystem pytorch_model.bin and PYTORCH_MODEL.BIN are distinct files; keying by lowered name dropped one and could hash a decoy instead of the loader's target. The enumerator now returns every variant Path. - Only persist a clean verdict for a COMPLETED, entirely-benign scan. Require scansDone to be the boolean True (not a truthy string), filesWithIssues to be a well-formed list, and every flagged file to be a definitively-safe level; a pending, error, unknown, or malformed entry no longer records as clean. The online block decision is unchanged. - Fail closed when the offline cache cannot be inspected: an rglob error now propagates and blocks instead of reading as pickle-free, and a snapshot that errors on resolution (vs a clean not-cached) blocks. The offline guard also raises instead of returning when its own inspection throws, so the constructor never deserializes an unverified cached pickle. - Expand online Router children recursively (bounded BFS with a seen set), mirroring the offline load-root expansion, so a flagged grandchild pickle is scoped online and cannot be recorded clean. - Reject absolute and drive/UNC declared paths in the load-root canonicalizers; the loader would resolve them outside the snapshot, so collapsing them to an in-snapshot relative dir scoped the wrong place. - Pin verdict recording to the scanned commit's snapshot and take the offline verify commit from the snapshot directory name, removing a second refs/main read and the skew it allowed. - Drop the now-unused pickle-name wrapper. --- studio/backend/core/rag/embeddings.py | 54 +++-- .../test_embedding_model_security_gate.py | 37 ++++ .../tests/test_embedding_verdict_cache.py | 193 ++++++++++++++---- .../backend/utils/security/file_security.py | 146 ++++++++----- 4 files changed, 323 insertions(+), 107 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index db3c30c1dab..3af645c24c7 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -107,6 +107,7 @@ def _st_module_subdirs(name: str, token: str | None, local_only: bool) -> tuple[ try: import json import posixpath + import re from utils.paths import is_local_path @@ -141,10 +142,13 @@ def _read_repo_json(rel: str): return None def _safe_subdir(rel) -> str | None: - """Canonical repo-relative subdir (``""`` = root), or None if it escapes the repo. - Mirrors the offline gate so a traversing declared path cannot read or scope a - directory outside the repo.""" - norm = posixpath.normpath(str(rel).strip().strip("/")) + """Canonical repo-relative subdir (``""`` = root), or None if it is absolute or escapes + the repo. Mirrors the offline gate so a traversing (``0/../evil``) or absolute + (``/etc``, ``C:/x``) declared path cannot read or scope a directory outside the repo.""" + raw = str(rel).strip().replace("\\", "/") + if raw.startswith("/") or re.match(r"^[A-Za-z]:/", raw): + return None + norm = posixpath.normpath(raw.strip("/")) if norm in ("", "."): return "" if norm == ".." or norm.startswith("../"): @@ -155,27 +159,37 @@ def _safe_subdir(rel) -> str | None: if not isinstance(data, list): return () subdirs = [] + pending = [] # (canonical-prefix, type-string) nodes that MIGHT be Routers to expand for module in data: if not isinstance(module, dict): continue - sub = str(module.get("path", "")).strip().strip("/") + sub = _safe_subdir(module.get("path", "")) + if sub is None: + continue # absolute / traversing module path -> drop (the offline gate drops it too) if sub: subdirs.append(sub) - # Only a Router/Asym module hides load roots in router_config.json; the read is - # scoped to it so a plain embedder incurs no extra fetch. - if str(module.get("type", "")).rsplit(".", 1)[-1].lower() not in ("router", "asym"): - continue - prefix = _safe_subdir(sub) - if prefix is None: + pending.append((sub, str(module.get("type", "")))) + # A Router child can itself be a Router, so expand recursively (bounded by a seen-set), + # mirroring the offline _st_load_roots BFS -- otherwise a flagged GRANDCHILD pickle + # (child_Router/grand_WordEmbeddings/pytorch_model.bin) is scoped offline but not online and + # could be recorded clean. router_config.json is read only for a Router/Asym-typed node, so + # a plain embedder pays no extra fetch. + seen: set[str] = set() + while pending: + prefix, mtype = pending.pop() + if mtype.rsplit(".", 1)[-1].lower() not in ("router", "asym") or prefix in seen: continue + seen.add(prefix) cfg = _read_repo_json(posixpath.join(prefix, "router_config.json")) types = cfg.get("types") if isinstance(cfg, dict) else None if not isinstance(types, dict): continue - for model_id in types: - child = _safe_subdir(model_id) - if child: - subdirs.append(posixpath.join(prefix, child)) + for model_id, child_type in types.items(): + child = _safe_subdir(posixpath.join(prefix, str(model_id))) + if not child: + continue + subdirs.append(child) + pending.append((child, str(child_type))) return tuple(dict.fromkeys(subdirs)) except Exception: return () @@ -223,7 +237,15 @@ def _guard_model_security(name: str, local_only_load: bool): local_only_load = local_only_load, ) blocked = decision.blocked # read inside the guard so a gate error never bricks the load - except Exception: + except Exception as exc: + # An OFFLINE load cannot be re-scanned, so a gate error must fail CLOSED -- otherwise the + # constructor would deserialize the unscanned cached pickle unguarded. Online, the Hub scan + # is best-effort and a gate error stays fail-open (return None). + if local_only_load: + raise UnsafeEmbeddingModelError( + f"Could not verify cached embedding model {name!r} offline; refusing the load. " + "Reconnect once to scan, or use safetensors weights." + ) from exc return None if blocked: raise UnsafeEmbeddingModelError( diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index de3473c9795..56896949e25 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -566,6 +566,43 @@ def test_st_module_subdirs_drops_traversing_router_child(tmp_path): assert not any(".." in s for s in subdirs) +def test_st_module_subdirs_expands_grandchild_router_children(tmp_path): + # A Router child that is ITSELF a Router must expand recursively (mirroring the offline BFS), so + # a flagged grandchild pickle is scoped online too and cannot be recorded clean. + import json + import core.rag.embeddings as embeddings + + (tmp_path / "modules.json").write_text( + json.dumps([{"path": "", "type": "sentence_transformers.models.Router.Router"}]) + ) + (tmp_path / "router_config.json").write_text( + json.dumps({"types": {"child_Router": "sentence_transformers.models.Router.Router"}}) + ) + (tmp_path / "child_Router").mkdir() + (tmp_path / "child_Router" / "router_config.json").write_text( + json.dumps({"types": {"grand_WordEmbeddings": "..."}}) + ) + subdirs = embeddings._st_module_subdirs(str(tmp_path), None, False) + assert "child_Router" in subdirs + assert "child_Router/grand_WordEmbeddings" in subdirs + + +def test_st_module_subdirs_drops_absolute_module_path(tmp_path): + # An absolute (or drive/UNC) module path is dropped: the loader would resolve it outside the + # snapshot, so canonicalizing it into an in-snapshot relative dir would scope the wrong place. + import json + import core.rag.embeddings as embeddings + + (tmp_path / "modules.json").write_text( + json.dumps( + [{"path": "/etc/evil", "type": "..."}, {"path": "0_Transformer", "type": "..."}] + ) + ) + subdirs = embeddings._st_module_subdirs(str(tmp_path), None, False) + assert subdirs == ("0_Transformer",) + assert not any(s.startswith("/") or "etc" in s for s in subdirs) + + def test_st_module_subdirs_router_expansion_over_hub(monkeypatch, tmp_path): # The Hub (non-local) path expands Router children symmetrically: hf_hub_download serves # modules.json then the Router's router_config.json. diff --git a/studio/backend/tests/test_embedding_verdict_cache.py b/studio/backend/tests/test_embedding_verdict_cache.py index cb2fb1ce65e..dac8dfcb8f1 100644 --- a/studio/backend/tests/test_embedding_verdict_cache.py +++ b/studio/backend/tests/test_embedding_verdict_cache.py @@ -3,10 +3,13 @@ """Integration tests for the offline embedding verdict cache: an ONLINE clean load records a verdict, and a later OFFLINE load of the SAME content is allowed instead of fail-closed. Every -mismatch (moved commit, changed / added pickle, expired record) must keep blocking. +mismatch (moved commit, changed / added pickle, expired record, unreadable/uninspectable cache) +must keep blocking. -Snapshot resolution is stubbed (``_active_snapshot_dir`` / ``_active_commit``) so no HF cache or -network is touched; UNSLOTH_STUDIO_HOME isolates the verdict store per test. +Snapshot resolution is stubbed (``_active_snapshot_dir`` / ``_st_cache_repo_dir``) over a real +``models--org--model/snapshots//`` layout so the commit is taken from the snapshot dir +name, as in production; no HF cache or network is touched, and UNSLOTH_STUDIO_HOME isolates the +verdict store per test. """ import json @@ -22,6 +25,8 @@ import utils.security.embedding_scan_verdicts as verdicts import utils.security.file_security as fs +_REPO = "models--org--model" + @pytest.fixture def home(tmp_path, monkeypatch): @@ -31,9 +36,14 @@ def home(tmp_path, monkeypatch): return tmp_path -def _snap(tmp_path, files: dict): - """Build a snapshot dir with *files* (name -> bytes/str; names may be nested like ``a/b.bin``).""" - d = tmp_path / "cache" / "snap" +def _repo_dir(tmp_path): + return tmp_path / "cache" / _REPO + + +def _snap(tmp_path, files: dict, commit = "c1"): + """Build ``models--org--model/snapshots//`` with *files* (name -> bytes/str; names may + be nested like ``a/b.bin``). Returns the snapshot dir, whose ``.name`` is the commit.""" + d = _repo_dir(tmp_path) / "snapshots" / commit d.mkdir(parents = True, exist_ok = True) for name, body in files.items(): p = d / name @@ -42,9 +52,9 @@ def _snap(tmp_path, files: dict): return d -def _blocked(monkeypatch, snap, commit, load_subdirs = ()): +def _blocked(monkeypatch, snap, load_subdirs = ()): + # The offline verify derives the commit from snap.name; only _active_snapshot_dir is consulted. monkeypatch.setattr(mc, "_active_snapshot_dir", lambda name: snap) - monkeypatch.setattr(mc, "_active_commit", lambda name: commit) return fs.evaluate_file_security( "org/model", None, load_subdirs = load_subdirs, local_only_load = True ).blocked @@ -59,25 +69,26 @@ def _sha(path): def test_matching_record_allows_offline_pickle(home, tmp_path, monkeypatch): snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) - assert _blocked(monkeypatch, snap, "c1") is False + assert _blocked(monkeypatch, snap) is False def test_no_record_blocks_offline_pickle(home, tmp_path, monkeypatch): snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) - assert _blocked(monkeypatch, snap, "c1") is True + assert _blocked(monkeypatch, snap) is True def test_wrong_commit_blocks(home, tmp_path, monkeypatch): - snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) - verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) - assert _blocked(monkeypatch, snap, "c2") is True + snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}, commit = "c1") + # Record under a DIFFERENT commit than the cached snapshot's -> lookup by snap.name misses. + verdicts.record_clean("org/model", "c2", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) + assert _blocked(monkeypatch, snap) is True def test_tampered_pickle_blocks(home, tmp_path, monkeypatch): snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) (snap / "pytorch_model.bin").write_bytes(b"TAMPERED-SAME-COMMIT") - assert _blocked(monkeypatch, snap, "c1") is True + assert _blocked(monkeypatch, snap) is True def test_partial_record_blocks(home, tmp_path, monkeypatch): @@ -93,13 +104,11 @@ def test_partial_record_blocks(home, tmp_path, monkeypatch): }, ) verdicts.record_clean("org/model", "c1", {"0_A/pytorch_model.bin": _sha(snap / "0_A/pytorch_model.bin")}) - assert _blocked(monkeypatch, snap, "c1") is True + assert _blocked(monkeypatch, snap) is True def test_distinct_basenames_in_different_roots_are_keyed_apart(home, tmp_path, monkeypatch): - # Same basename in two module dirs must be recorded/verified as two distinct entries; if the - # enumerator collapsed them by basename, one record could vouch for both. Recording BOTH with - # their real (distinct) hashes allows; the map has two keys. + # Same basename in two module dirs must be recorded/verified as two distinct entries. snap = _snap( tmp_path, { @@ -118,7 +127,25 @@ def test_distinct_basenames_in_different_roots_are_keyed_apart(home, tmp_path, m "0_B/pytorch_model.bin": _sha(snap / "0_B/pytorch_model.bin"), }, ) - assert _blocked(monkeypatch, snap, "c1") is False + assert _blocked(monkeypatch, snap) is False + + +def test_case_variant_pickles_are_both_hashed(home, tmp_path, monkeypatch): + # On a case-sensitive FS pytorch_model.bin and PYTORCH_MODEL.BIN are distinct load-root files; + # both must be hashed. Recording only the exact-case one leaves the collision unmatched -> block; + # recording BOTH (their real distinct hashes) allows. Proves no last-wins collapse. + if (tmp_path / "A").exists(): # skip on a case-insensitive filesystem + pytest.skip("case-insensitive filesystem") + (tmp_path / "casecheck").write_text("x") + if (tmp_path / "CASECHECK").exists(): + pytest.skip("case-insensitive filesystem") + snap = _snap(tmp_path, {"pytorch_model.bin": b"real", "PYTORCH_MODEL.BIN": b"decoy"}) + real = _sha(snap / "pytorch_model.bin") + decoy = _sha(snap / "PYTORCH_MODEL.BIN") + verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": real}) + assert _blocked(monkeypatch, snap) is True # decoy unrecorded -> mismatch + verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": real, "PYTORCH_MODEL.BIN": decoy}) + assert _blocked(monkeypatch, snap) is False def test_expired_record_blocks(home, tmp_path, monkeypatch): @@ -130,12 +157,10 @@ def test_expired_record_blocks(home, tmp_path, monkeypatch): datetime.now(timezone.utc) - timedelta(days = 31) ).isoformat() path.write_text(json.dumps(data)) - assert _blocked(monkeypatch, snap, "c1") is True + assert _blocked(monkeypatch, snap) is True def test_router_child_pickle_recorded_allows_else_blocks(home, tmp_path, monkeypatch): - # A Router child pickle (declared only in router_config.json) is a load root the offline gate - # scopes. Recording it allows; without a record it blocks. snap = _snap( tmp_path, { @@ -145,50 +170,123 @@ def test_router_child_pickle_recorded_allows_else_blocks(home, tmp_path, monkeyp }, ) child = snap / "query_0_WordEmbeddings" / "pytorch_model.bin" - assert _blocked(monkeypatch, snap, "c1") is True + assert _blocked(monkeypatch, snap) is True verdicts.record_clean("org/model", "c1", {"query_0_WordEmbeddings/pytorch_model.bin": _sha(child)}) - assert _blocked(monkeypatch, snap, "c1") is False + assert _blocked(monkeypatch, snap) is False def test_disabled_cache_never_allows(home, tmp_path, monkeypatch): snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) monkeypatch.setenv("UNSLOTH_EMBED_VERDICT_CACHE_DISABLE", "1") - assert _blocked(monkeypatch, snap, "c1") is True + assert _blocked(monkeypatch, snap) is True + + +def test_unreadable_recorded_pickle_blocks(home, tmp_path, monkeypatch): + # A recorded pickle that cannot be re-hashed (sha256_file -> None) must block, not allow. + snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) + verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) + monkeypatch.setattr(verdicts, "sha256_file", lambda p: None) + assert _blocked(monkeypatch, snap) is True + + +def test_uninspectable_cache_blocks(home, tmp_path, monkeypatch): + # If the cache tree cannot be enumerated (rglob OSError), fail CLOSED, not "pickle-free allow". + snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) + + def _boom(*a, **k): + raise OSError("EIO") + + monkeypatch.setattr(fs, "_cached_pickle_weight_paths", _boom) + assert _blocked(monkeypatch, snap) is True + + +def test_unresolvable_snapshot_blocks(home, monkeypatch): + # A snapshot that ERRORS on resolution (not a clean None) fails closed. + def _boom(name): + raise OSError("refs/main unreadable") + + monkeypatch.setattr(mc, "_active_snapshot_dir", _boom) + assert fs.evaluate_file_security("org/model", None, local_only_load = True).blocked is True # ── record_embedding_verdict (the post-load recorder) ──────────────── -def _stub_snapshot(monkeypatch, snap, commit): - monkeypatch.setattr(mc, "_active_snapshot_dir", lambda name: snap) - monkeypatch.setattr(mc, "_active_commit", lambda name: commit) +def _stub_repo(monkeypatch, tmp_path): + monkeypatch.setattr(mc, "_st_cache_repo_dir", lambda name: _repo_dir(tmp_path)) def test_record_embedding_verdict_records_and_allows(home, tmp_path, monkeypatch): - snap = _snap(tmp_path, {"pytorch_model.bin": b"w"}) - _stub_snapshot(monkeypatch, snap, "c1") + snap = _snap(tmp_path, {"pytorch_model.bin": b"w"}, commit = "c1") + _stub_repo(monkeypatch, tmp_path) fs.record_embedding_verdict("org/model", "c1", load_subdirs = ()) assert verdicts.lookup("org/model", "c1") == {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")} - assert fs.evaluate_file_security("org/model", None, local_only_load = True).blocked is False + assert _blocked(monkeypatch, snap) is False -def test_record_skips_when_commit_moved(home, tmp_path, monkeypatch): - # The active cached commit differs from the scanned commit (branch advanced): record nothing, - # so an unscanned commit is never blessed. - snap = _snap(tmp_path, {"pytorch_model.bin": b"w"}) - _stub_snapshot(monkeypatch, snap, "c2") - fs.record_embedding_verdict("org/model", "c1", load_subdirs = ()) +def test_record_skips_when_scanned_commit_not_cached(home, tmp_path, monkeypatch): + # Only snapshots/c1 exists; recording the scanned commit c2 (branch moved) finds no snapshot + # for c2 and records nothing, so an unscanned commit is never blessed. + _snap(tmp_path, {"pytorch_model.bin": b"w"}, commit = "c1") + _stub_repo(monkeypatch, tmp_path) + fs.record_embedding_verdict("org/model", "c2", load_subdirs = ()) assert verdicts.lookup("org/model", "c1") is None assert verdicts.lookup("org/model", "c2") is None def test_record_skips_pickle_free_cache(home, tmp_path, monkeypatch): - snap = _snap(tmp_path, {"model.safetensors": b"\0"}) - _stub_snapshot(monkeypatch, snap, "c1") + _snap(tmp_path, {"model.safetensors": b"\0"}, commit = "c1") + _stub_repo(monkeypatch, tmp_path) fs.record_embedding_verdict("org/model", "c1", load_subdirs = ()) assert verdicts.lookup("org/model", "c1") is None +# ── forget() on an authoritative unsafe verdict (end-to-end) ───────── + +def test_online_unsafe_forgets_recorded_verdict(home, monkeypatch): + # An online scan that flags a load-path pickle must delete any recorded clean verdict. + verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": "a" * 64}) + monkeypatch.setattr( + fs, + "_fetch_security_status", + lambda name, token: ( + {"scansDone": True, "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}]}, + "c1", + ), + ) + decision = fs.evaluate_file_security("org/model", None, local_only_load = False) + assert decision.blocked is True + assert verdicts.lookup("org/model", "c1") is None + + +# ── recordable-clean predicate (only a COMPLETED, entirely-benign scan) ── + +def _scanned_clean(monkeypatch, status): + monkeypatch.setattr(fs, "_fetch_security_status", lambda name, token: (status, "c1")) + return fs.evaluate_file_security("org/model", None, local_only_load = False).scanned_clean + + +def test_recordable_requires_completed_scan(home, monkeypatch): + assert _scanned_clean(monkeypatch, {"scansDone": True, "filesWithIssues": []}) is True + # scansDone must be the boolean True, not a truthy string. + assert _scanned_clean(monkeypatch, {"scansDone": "false", "filesWithIssues": []}) is False + assert _scanned_clean(monkeypatch, {"filesWithIssues": []}) is False + + +def test_recordable_rejects_malformed_or_flagged_manifest(home, monkeypatch): + assert _scanned_clean(monkeypatch, {"scansDone": True}) is False # filesWithIssues not a list + assert _scanned_clean(monkeypatch, {"scansDone": True, "filesWithIssues": "bad"}) is False + assert _scanned_clean(monkeypatch, {"scansDone": True, "filesWithIssues": [None]}) is False + # A pending/error flagged file (even outside the load path) disqualifies a durable clean record. + assert ( + _scanned_clean( + monkeypatch, + {"scansDone": True, "filesWithIssues": [{"path": "archive/x.bin", "level": "error"}]}, + ) + is False + ) + + # ── _get() wiring: record only after a clean ONLINE load ───────────── def _drive_get(monkeypatch, decision, *, offline): @@ -217,17 +315,28 @@ def test_get_records_after_clean_online_load(home, monkeypatch): def test_get_does_not_record_when_scan_not_definitive(home, monkeypatch): - # A fail-open ("scan unavailable") online load has scanned_clean False -> record nothing. dec = fs.FileSecurityDecision("org/model", False, commit = None, scanned_clean = False) assert _drive_get(monkeypatch, dec, offline = False) == {} def test_get_does_not_record_offline_load(home, monkeypatch): - # Offline: no authoritative scan happened, so nothing is recorded regardless. dec = fs.FileSecurityDecision("org/model", False, commit = "c1", scanned_clean = True) assert _drive_get(monkeypatch, dec, offline = True) == {} def test_get_does_not_record_on_gate_error(home, monkeypatch): - # _guard_model_security returns None on a gate error -> record nothing. assert _drive_get(monkeypatch, None, offline = False) == {} + + +# ── guard fails CLOSED offline on a gate error ─────────────────────── + +def test_guard_offline_gate_error_fails_closed(home, monkeypatch): + def _boom(*a, **k): + raise OSError("gate blew up") + + monkeypatch.setattr("utils.security.evaluate_file_security", _boom) + # Offline: a gate error must refuse the load (fail closed), not return None and proceed. + with pytest.raises(embeddings.UnsafeEmbeddingModelError): + embeddings._guard_model_security("org/model", True) + # Online: a gate error stays fail-open (best-effort Hub scan) -> returns None. + assert embeddings._guard_model_security("org/model", False) is None diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 945f04af40b..2c500795ca3 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -45,6 +45,12 @@ {"", "safe", "pending", "scanning", "queued", "unscanned", "error", "unknown", "none"} ) +# Levels that are a DEFINITIVE clean result for a file. A durable clean record (for the offline +# cache) may only be written when every flagged file is at one of these; the other non-blocking +# levels ("pending"/"scanning"/"queued"/"unscanned"/"error"/"unknown") mean the scan did not +# cleanly finish on that file, so they do not block the online load but do disqualify recording. +_DEFINITIVELY_SAFE_LEVELS = frozenset({"", "safe", "none"}) + # Suffixes that cannot execute code on load (tensor-only safetensors, non-pickle gguf, # text/markup/images), so a flag on one is never an RCE vector. _INERT_SUFFIXES = frozenset( @@ -104,10 +110,16 @@ def _canonical_rel(rel: str): normalized directory (via :func:`_canonical_load_dir`) and the ONLINE scan must agree, or the raw ``0/../evil`` never equals the canonical ``evil/...`` the Hub reports for the flagged file and the pickle slips the gate. A legitimate declared path never traverses - upward, so an escaping path is rejected.""" + upward, so an escaping path is rejected. An ABSOLUTE or drive/UNC path is also rejected: the + loader resolves it OUTSIDE the snapshot (``os.path.join`` discards the base on an absolute + join), so collapsing it into an in-snapshot relative directory would scope the wrong dir; a + declared repo path is always relative.""" import posixpath - norm = posixpath.normpath(_normalize_repo_path(rel).strip("/")) + raw = _normalize_repo_path(rel) # backslashes -> "/", "./" prefixes stripped + if raw.startswith("/") or re.match(r"^[A-Za-z]:/", raw): + return None + norm = posixpath.normpath(raw.strip("/")) if norm in ("", ".") or norm == ".." or norm.startswith("../"): return None return norm @@ -474,25 +486,26 @@ def _cached_pickle_weight_paths(snap, load_subdirs = ()) -> list: ``adapter_model.safetensors`` -- a safetensors base does NOT cover it. A bare adapter or an orphan shard leaves its pickle live. Returns the concrete Paths (not basenames) so the offline verdict cache can hash exactly the files the loader reads and never conflate two module dirs that - ship the same pickle basename.""" + ship the same pickle basename. + + Raises ``OSError`` if the snapshot tree cannot be enumerated (``rglob`` failure), so the offline + caller fails CLOSED rather than treat an unreadable cache as pickle-free.""" roots = _st_load_roots(snap, load_subdirs) - by_dir_pickle: dict = {} - by_dir_files: dict = {} # directory -> {lower-name: Path} - try: - for path in snap.rglob("*"): - try: - if not path.is_file(): - continue - low = path.name.lower() - by_dir_files.setdefault(path.parent, {})[low] = path - if _PICKLE_WEIGHT_RE.match(low) or _ADAPTER_PICKLE_RE.match(low): - by_dir_pickle.setdefault(path.parent, []).append(path.name) - except OSError: + by_dir_pickle: dict = {} # directory -> [Path] (EVERY case variant, not last-wins) + by_dir_files: dict = {} # directory -> {lower-name: Path} (safetensors credit; last-wins is safe) + # A whole-tree rglob failure propagates (fail-closed); a single unstattable entry is skipped. + for path in snap.rglob("*"): + try: + if not path.is_file(): continue - except OSError: - return [] + except OSError: + continue + low = path.name.lower() + by_dir_files.setdefault(path.parent, {})[low] = path + if _PICKLE_WEIGHT_RE.match(low) or _ADAPTER_PICKLE_RE.match(low): + by_dir_pickle.setdefault(path.parent, []).append(path) hits: set = set() - for directory, names in by_dir_pickle.items(): + for directory, pickle_paths in by_dir_pickle.items(): files = by_dir_files.get(directory, {}) # A pickle is deserialized only at an actual load root: the snapshot root, a declared # modules.json / load_subdirs dir, or a Router child (all resolved by _st_load_roots). @@ -502,21 +515,25 @@ def _cached_pickle_weight_paths(snap, load_subdirs = ()) -> list: # ignores the same unindexed subdir pickle. if directory not in roots: continue - base = [n for n in names if _PICKLE_WEIGHT_RE.match(n.lower())] + # Hash EVERY case-colliding pickle candidate, not one representative: on a case-sensitive + # filesystem ``pytorch_model.bin`` and ``PYTORCH_MODEL.BIN`` are distinct files; the loader + # reads the exact-case one, and a mis-cased sibling that is never loaded is only over-blocked + # (safe). Keying by lowered name would drop one and could hash a decoy instead of the target. + base = [p for p in pickle_paths if _PICKLE_WEIGHT_RE.match(p.name.lower())] if base and not _dir_has_loadable_safetensors(files): - hits.update(files[n.lower()] for n in base) + hits.update(base) # An adapter pickle is deserialized only when from_pretrained auto-detects the adapter # (adapter_config.json present) and there is no adapter_model.safetensors to load instead. # The safetensors credit is case-SENSITIVE (a mixed-case Adapter_Model.SafeTensors decoy # is not the file PEFT loads), so it is matched by real basename; the config presence # stays case-insensitive (over-blocking a mis-cased adapter is the safe direction). - adapter = [n for n in names if _ADAPTER_PICKLE_RE.match(n.lower())] + adapter = [p for p in pickle_paths if _ADAPTER_PICKLE_RE.match(p.name.lower())] if ( adapter and "adapter_config.json" in files and _exact_named(files, "adapter_model.safetensors") is None ): - hits.update(files[n.lower()] for n in adapter) + hits.update(adapter) # A load-root pickle index (pytorch_model.bin.index.json) can map shards into SUBDIRECTORIES # that are not themselves load roots; from_pretrained follows the map and deserializes them, # so include those referenced pickle shards (unless a loadable base safetensors at the index @@ -547,13 +564,6 @@ def _snapshot_relative(path, snap) -> str: return path.name -def _cached_pickle_weight_files(snap, load_subdirs = ()) -> list: - """Snapshot-relative posix names of the load-root pickle weights (see - :func:`_cached_pickle_weight_paths`). Names are relative to the snapshot so two module dirs that - each ship ``pytorch_model.bin`` are reported (and, in the verdict cache, keyed) distinctly.""" - return [_snapshot_relative(p, snap) for p in _cached_pickle_weight_paths(snap, load_subdirs)] - - def _pickle_hash_map(snap, paths): """``{snapshot-relative posix name: sha256}`` for *paths*, or None if any file cannot be hashed (an unreadable pickle must never verify or record as clean).""" @@ -571,12 +581,15 @@ def _pickle_hash_map(snap, paths): def _matches_clean_verdict(model_name: str, snap, paths) -> bool: """True when a recorded clean Hub verdict covers EXACTLY this load-root pickle set at the active cached commit: same commit, same relative-name set, same sha256 for every file. A missing - record, moved commit, changed / added / unreadable pickle, or any error -> False (fail-closed).""" + record, moved commit, changed / added / unreadable pickle, or any error -> False (fail-closed). + + The commit is taken from the snapshot directory name (``snapshots/``) -- the exact + revision being inspected -- so it needs no second ``refs/main`` read and cannot skew from the + ``snap`` whose files are hashed.""" try: - from utils.models.model_config import _active_commit from utils.security import embedding_scan_verdicts - commit = _active_commit(model_name) + commit = snap.name # snapshots/ if not commit: return False recorded = embedding_scan_verdicts.lookup(model_name, commit) @@ -591,18 +604,22 @@ def _matches_clean_verdict(model_name: str, snap, paths) -> bool: def record_embedding_verdict(model_name: str, scanned_commit, load_subdirs = ()) -> None: """Record a clean Hub verdict for an embedding repo just loaded ONLINE, so a later offline load of the same content is not fail-closed. Persists the sha256 of every load-root pickle keyed by - its snapshot-relative name. Records nothing when the snapshot is missing, the active cached - commit differs from the scanned commit (branch moved -> the loaded content was not what HF - scanned), there are no load-root pickles (an inert cache is already allowed), or any file cannot - be hashed. Never raises into the load path.""" + its snapshot-relative name. Enumeration is PINNED to the scanned commit's snapshot + (``snapshots/``): if that snapshot is not present (the branch moved between the + scan and the load), nothing is recorded, so an unscanned commit is never blessed. Also records + nothing when there are no load-root pickles (an inert cache is already allowed) or any file + cannot be hashed. Never raises into the load path.""" try: - from utils.models.model_config import _active_commit, _active_snapshot_dir + from utils.models.model_config import _st_cache_repo_dir from utils.security import embedding_scan_verdicts - if not scanned_commit or _active_commit(model_name) != scanned_commit: + if not scanned_commit: return - snap = _active_snapshot_dir(model_name) - if snap is None: + repo_dir = _st_cache_repo_dir(model_name) + if repo_dir is None: + return + snap = repo_dir / "snapshots" / str(scanned_commit) + if not snap.is_dir(): return paths = _cached_pickle_weight_paths(snap, load_subdirs) if not paths: @@ -623,17 +640,38 @@ def _evaluate_local_only(model_name: str, load_subdirs = ()) -> "FileSecurityDec deserializes -- and allow only a pickle-free cache (safetensors / gguf are inert). A previously-scanned pickle model must be reloaded online once to pass, or shipped as safetensors. Nothing cached means there is nothing to deserialize, so it is not blocked - (the load fails downstream on its own, which is not a security event). + (the load fails downstream on its own, which is not a security event). But a snapshot that + exists yet cannot be RESOLVED or ENUMERATED (an I/O error, not a clean "not cached") is + blocked: an unreadable cache must not be mistaken for an inert one. """ try: from utils.models.model_config import _active_snapshot_dir snap = _active_snapshot_dir(model_name) except Exception: - snap = None + # Distinct from a clean None (nothing cached): resolution ERRORED, so we cannot vouch for + # the cache. Fail closed rather than allow an unverifiable snapshot. + logger.warning( + "Blocking offline load of '%s': snapshot could not be resolved; cannot verify.", + model_name, + ) + return FileSecurityDecision( + model_name, True, reason = "offline; snapshot could not be resolved, cannot verify" + ) if snap is None: return FileSecurityDecision(model_name, False, reason = "offline; nothing cached to scan") - paths = _cached_pickle_weight_paths(snap, load_subdirs) + try: + paths = _cached_pickle_weight_paths(snap, load_subdirs) + except OSError: + # The cache tree could not be fully enumerated: an undiscovered pickle might exist, so a + # partial "pickle-free" reading is unsafe. Fail closed. + logger.warning( + "Blocking offline load of '%s': cached model could not be inspected; cannot verify.", + model_name, + ) + return FileSecurityDecision( + model_name, True, reason = "offline; cached model could not be inspected, cannot verify" + ) if not paths: return FileSecurityDecision( model_name, False, reason = "offline; cached weights are pickle-free (inert)" @@ -748,10 +786,14 @@ def evaluate_file_security( return FileSecurityDecision( model_name, False, reason = "scan unavailable; allowed (fail-open)" ) - # A DEFINITIVE clean verdict (for the durable offline cache) requires a COMPLETED scan, unlike - # the block decision which is not gated on scansDone. An in-progress scan is allowed online but - # must not be persisted as clean. - scans_done = bool(status.get("scansDone")) + # A DEFINITIVE clean verdict (for the durable offline cache) is STRICTER than the block + # decision: it requires a COMPLETED scan (``scansDone is True`` -- an identity check, so a + # truthy string like ``"false"`` does not qualify), a well-formed ``filesWithIssues`` list, and + # NO concerning flagged file anywhere (any non-``safe`` level, or a malformed entry, + # disqualifies). The block decision below is unchanged (not gated on scansDone); ``recordable`` + # only governs whether the verdict may be persisted for later offline reuse. + issues = status.get("filesWithIssues") + recordable = (status.get("scansDone") is True) and isinstance(issues, list) # Block a non-``safe`` flagged file scoped to the load-path RCE vector (root-level, # code-executing). Not gated on ``scansDone`` (often false even when clean; a flagged @@ -761,10 +803,16 @@ def evaluate_file_security( unsafe = [] skipped = [] # flagged, but not a load-path RCE vector (subdir artifact / inert) maybe_shard = [] # flagged subdir pickle: a load vector ONLY if a root index lists it - for entry in status.get("filesWithIssues") or []: + for entry in issues or []: if not isinstance(entry, dict): + recordable = False # a malformed manifest entry cannot vouch for cleanliness continue level = str(entry.get("level", "")).lower() + if level not in _DEFINITIVELY_SAFE_LEVELS: + # Any non-benign level -- an unsafe finding OR an incomplete/uncertain one + # (pending/error/unknown) -- means this is not a clean COMPLETED scan, so it cannot be + # persisted as a durable clean verdict (the block decision below is unaffected). + recordable = False if level in _NONBLOCKING_LEVELS: continue path = entry.get("path", "") @@ -809,7 +857,7 @@ def evaluate_file_security( False, reason = "no unsafe files in the load path", commit = commit, - scanned_clean = scans_done, + scanned_clean = recordable, ) names = ", ".join(u["path"] for u in unsafe if u["path"]) or "unknown files" From 8c9145831a737899798698643a9488938ca421e1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 03:23:27 +0000 Subject: [PATCH 62/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/core/rag/embeddings.py | 10 +-- .../test_embedding_model_security_gate.py | 4 +- .../tests/test_embedding_scan_verdicts.py | 16 ++++- .../tests/test_embedding_verdict_cache.py | 67 +++++++++++++++---- .../backend/utils/security/file_security.py | 6 +- 5 files changed, 79 insertions(+), 24 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 3af645c24c7..3a8c84accb7 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -202,7 +202,6 @@ def _security_load_subdirs(name: str, token: str | None, local_only: bool) -> tu instead of passing as an unreferenced nested shard. Shared by the guard and the verdict recorder so both scope the SAME roots.""" from utils.security import security_load_subdirs - return tuple( dict.fromkeys( ( @@ -262,9 +261,10 @@ def _record_embedding_verdict_safe(name: str, commit: str | None) -> None: model.""" try: from utils.security import record_embedding_verdict - token = _ambient_hf_token() - record_embedding_verdict(name, commit, _security_load_subdirs(name, token, local_only = False)) + record_embedding_verdict( + name, commit, _security_load_subdirs(name, token, local_only = False) + ) except Exception: pass @@ -274,7 +274,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() - record: tuple[str, str | None] | None = None # (load_name, scanned_commit) for a clean online load + record: tuple[str, str | None] | None = ( + None # (load_name, scanned_commit) for a clean online load + ) with _lock: if _model is None or _name != name: _install_torchao_stub_once() diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index 56896949e25..e0222be88bf 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -594,9 +594,7 @@ def test_st_module_subdirs_drops_absolute_module_path(tmp_path): import core.rag.embeddings as embeddings (tmp_path / "modules.json").write_text( - json.dumps( - [{"path": "/etc/evil", "type": "..."}, {"path": "0_Transformer", "type": "..."}] - ) + json.dumps([{"path": "/etc/evil", "type": "..."}, {"path": "0_Transformer", "type": "..."}]) ) subdirs = embeddings._st_module_subdirs(str(tmp_path), None, False) assert subdirs == ("0_Transformer",) diff --git a/studio/backend/tests/test_embedding_scan_verdicts.py b/studio/backend/tests/test_embedding_scan_verdicts.py index 3ea6cdc71bd..33bc7be5d63 100644 --- a/studio/backend/tests/test_embedding_scan_verdicts.py +++ b/studio/backend/tests/test_embedding_scan_verdicts.py @@ -102,8 +102,20 @@ def test_corrupt_store_fails_safe(verdicts): def test_wrong_schema_version_fails_safe(verdicts): path = verdicts._store_path() path.parent.mkdir(parents = True, exist_ok = True) - path.write_text(json.dumps({"version": 999, "records": {"acme/model": {"commit": "commit1", - "recorded_at": "2999-01-01T00:00:00+00:00", "pickles": {"pytorch_model.bin": _SHA}}}})) + path.write_text( + json.dumps( + { + "version": 999, + "records": { + "acme/model": { + "commit": "commit1", + "recorded_at": "2999-01-01T00:00:00+00:00", + "pickles": {"pytorch_model.bin": _SHA}, + } + }, + } + ) + ) assert verdicts.lookup("acme/model", "commit1") is None diff --git a/studio/backend/tests/test_embedding_verdict_cache.py b/studio/backend/tests/test_embedding_verdict_cache.py index dac8dfcb8f1..b36928598fe 100644 --- a/studio/backend/tests/test_embedding_verdict_cache.py +++ b/studio/backend/tests/test_embedding_verdict_cache.py @@ -40,7 +40,11 @@ def _repo_dir(tmp_path): return tmp_path / "cache" / _REPO -def _snap(tmp_path, files: dict, commit = "c1"): +def _snap( + tmp_path, + files: dict, + commit = "c1", +): """Build ``models--org--model/snapshots//`` with *files* (name -> bytes/str; names may be nested like ``a/b.bin``). Returns the snapshot dir, whose ``.name`` is the commit.""" d = _repo_dir(tmp_path) / "snapshots" / commit @@ -52,7 +56,11 @@ def _snap(tmp_path, files: dict, commit = "c1"): return d -def _blocked(monkeypatch, snap, load_subdirs = ()): +def _blocked( + monkeypatch, + snap, + load_subdirs = (), +): # The offline verify derives the commit from snap.name; only _active_snapshot_dir is consulted. monkeypatch.setattr(mc, "_active_snapshot_dir", lambda name: snap) return fs.evaluate_file_security( @@ -66,9 +74,12 @@ def _sha(path): # ── Offline verify matrix ──────────────────────────────────────────── + def test_matching_record_allows_offline_pickle(home, tmp_path, monkeypatch): snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) - verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) + verdicts.record_clean( + "org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")} + ) assert _blocked(monkeypatch, snap) is False @@ -80,13 +91,17 @@ def test_no_record_blocks_offline_pickle(home, tmp_path, monkeypatch): def test_wrong_commit_blocks(home, tmp_path, monkeypatch): snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}, commit = "c1") # Record under a DIFFERENT commit than the cached snapshot's -> lookup by snap.name misses. - verdicts.record_clean("org/model", "c2", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) + verdicts.record_clean( + "org/model", "c2", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")} + ) assert _blocked(monkeypatch, snap) is True def test_tampered_pickle_blocks(home, tmp_path, monkeypatch): snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) - verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) + verdicts.record_clean( + "org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")} + ) (snap / "pytorch_model.bin").write_bytes(b"TAMPERED-SAME-COMMIT") assert _blocked(monkeypatch, snap) is True @@ -103,7 +118,9 @@ def test_partial_record_blocks(home, tmp_path, monkeypatch): "0_B/pytorch_model.bin": b"bbb", }, ) - verdicts.record_clean("org/model", "c1", {"0_A/pytorch_model.bin": _sha(snap / "0_A/pytorch_model.bin")}) + verdicts.record_clean( + "org/model", "c1", {"0_A/pytorch_model.bin": _sha(snap / "0_A/pytorch_model.bin")} + ) assert _blocked(monkeypatch, snap) is True @@ -144,13 +161,17 @@ def test_case_variant_pickles_are_both_hashed(home, tmp_path, monkeypatch): decoy = _sha(snap / "PYTORCH_MODEL.BIN") verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": real}) assert _blocked(monkeypatch, snap) is True # decoy unrecorded -> mismatch - verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": real, "PYTORCH_MODEL.BIN": decoy}) + verdicts.record_clean( + "org/model", "c1", {"pytorch_model.bin": real, "PYTORCH_MODEL.BIN": decoy} + ) assert _blocked(monkeypatch, snap) is False def test_expired_record_blocks(home, tmp_path, monkeypatch): snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) - verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) + verdicts.record_clean( + "org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")} + ) path = verdicts._store_path() data = json.loads(path.read_text()) data["records"]["org/model"]["recorded_at"] = ( @@ -164,20 +185,26 @@ def test_router_child_pickle_recorded_allows_else_blocks(home, tmp_path, monkeyp snap = _snap( tmp_path, { - "modules.json": json.dumps([{"path": "", "type": "sentence_transformers.models.Router.Router"}]), + "modules.json": json.dumps( + [{"path": "", "type": "sentence_transformers.models.Router.Router"}] + ), "router_config.json": json.dumps({"types": {"query_0_WordEmbeddings": "..."}}), "query_0_WordEmbeddings/pytorch_model.bin": b"child-weights", }, ) child = snap / "query_0_WordEmbeddings" / "pytorch_model.bin" assert _blocked(monkeypatch, snap) is True - verdicts.record_clean("org/model", "c1", {"query_0_WordEmbeddings/pytorch_model.bin": _sha(child)}) + verdicts.record_clean( + "org/model", "c1", {"query_0_WordEmbeddings/pytorch_model.bin": _sha(child)} + ) assert _blocked(monkeypatch, snap) is False def test_disabled_cache_never_allows(home, tmp_path, monkeypatch): snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) - verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) + verdicts.record_clean( + "org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")} + ) monkeypatch.setenv("UNSLOTH_EMBED_VERDICT_CACHE_DISABLE", "1") assert _blocked(monkeypatch, snap) is True @@ -185,7 +212,9 @@ def test_disabled_cache_never_allows(home, tmp_path, monkeypatch): def test_unreadable_recorded_pickle_blocks(home, tmp_path, monkeypatch): # A recorded pickle that cannot be re-hashed (sha256_file -> None) must block, not allow. snap = _snap(tmp_path, {"pytorch_model.bin": b"weights"}) - verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")}) + verdicts.record_clean( + "org/model", "c1", {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")} + ) monkeypatch.setattr(verdicts, "sha256_file", lambda p: None) assert _blocked(monkeypatch, snap) is True @@ -212,6 +241,7 @@ def _boom(name): # ── record_embedding_verdict (the post-load recorder) ──────────────── + def _stub_repo(monkeypatch, tmp_path): monkeypatch.setattr(mc, "_st_cache_repo_dir", lambda name: _repo_dir(tmp_path)) @@ -220,7 +250,9 @@ def test_record_embedding_verdict_records_and_allows(home, tmp_path, monkeypatch snap = _snap(tmp_path, {"pytorch_model.bin": b"w"}, commit = "c1") _stub_repo(monkeypatch, tmp_path) fs.record_embedding_verdict("org/model", "c1", load_subdirs = ()) - assert verdicts.lookup("org/model", "c1") == {"pytorch_model.bin": _sha(snap / "pytorch_model.bin")} + assert verdicts.lookup("org/model", "c1") == { + "pytorch_model.bin": _sha(snap / "pytorch_model.bin") + } assert _blocked(monkeypatch, snap) is False @@ -243,6 +275,7 @@ def test_record_skips_pickle_free_cache(home, tmp_path, monkeypatch): # ── forget() on an authoritative unsafe verdict (end-to-end) ───────── + def test_online_unsafe_forgets_recorded_verdict(home, monkeypatch): # An online scan that flags a load-path pickle must delete any recorded clean verdict. verdicts.record_clean("org/model", "c1", {"pytorch_model.bin": "a" * 64}) @@ -250,7 +283,10 @@ def test_online_unsafe_forgets_recorded_verdict(home, monkeypatch): fs, "_fetch_security_status", lambda name, token: ( - {"scansDone": True, "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}]}, + { + "scansDone": True, + "filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}], + }, "c1", ), ) @@ -261,6 +297,7 @@ def test_online_unsafe_forgets_recorded_verdict(home, monkeypatch): # ── recordable-clean predicate (only a COMPLETED, entirely-benign scan) ── + def _scanned_clean(monkeypatch, status): monkeypatch.setattr(fs, "_fetch_security_status", lambda name, token: (status, "c1")) return fs.evaluate_file_security("org/model", None, local_only_load = False).scanned_clean @@ -289,6 +326,7 @@ def test_recordable_rejects_malformed_or_flagged_manifest(home, monkeypatch): # ── _get() wiring: record only after a clean ONLINE load ───────────── + def _drive_get(monkeypatch, decision, *, offline): fake_st = types.ModuleType("sentence_transformers") fake_st.SentenceTransformer = lambda name, **k: object() @@ -330,6 +368,7 @@ def test_get_does_not_record_on_gate_error(home, monkeypatch): # ── guard fails CLOSED offline on a gate error ─────────────────────── + def test_guard_offline_gate_error_fails_closed(home, monkeypatch): def _boom(*a, **k): raise OSError("gate blew up") diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 2c500795ca3..4d2164b515d 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -601,7 +601,11 @@ def _matches_clean_verdict(model_name: str, snap, paths) -> bool: return False -def record_embedding_verdict(model_name: str, scanned_commit, load_subdirs = ()) -> None: +def record_embedding_verdict( + model_name: str, + scanned_commit, + load_subdirs = (), +) -> None: """Record a clean Hub verdict for an embedding repo just loaded ONLINE, so a later offline load of the same content is not fail-closed. Persists the sha256 of every load-root pickle keyed by its snapshot-relative name. Enumeration is PINNED to the scanned commit's snapshot From 527396a265c9717ca1e3b1a0adee6c23be7a7893 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 21 Jul 2026 04:07:51 +0000 Subject: [PATCH 63/77] Tighten offline embedding classification and the pickle gate Close a set of offline edge cases where validation accepted a cache the local_files_only load then rejects, and one gate bypass: - Credit a sharded model.safetensors.index.json for a pickle sibling only at a from_pretrained root. A non-Transformer SentenceTransformer module (Dense, WordEmbeddings, StaticEmbedding) loads via Module.load_torch_weights, which reads model.safetensors then pytorch_model.bin and never the index, so a sharded safetensors index in such a module dir must not vouch for its pytorch_model.bin. - Stop counting pytorch_model.safetensors as loadable in the offline classifier: the loader probes model.safetensors (then its index) or pytorch_model.bin, never pytorch_model.safetensors, matching the gate that already treats it as a decoy. - Treat a present but unreadable weight index as incomplete: transformers opens and parses any present index, so a malformed one or one without a weight_map fails the load rather than falling back to filename-numbered shards. - Require the CLIP image-processor config (preprocessor_config.json) for a CLIP module: CLIPModel.load builds a CLIPProcessor that needs it, so a tokenizer alone is not enough. - Require a SparseStaticEmbedding config to actually select idf.json (a path ending .json) or ship loadable weights; a bare idf.json the config does not name falls through to load_torch_weights and raises. - Do not use the tag-only recorded-verdict fallback when modules.json is present: with the file present the loader takes the modules.json path, so a present but empty or malformed manifest must not be validated as a plain root Transformer. - Import-hoist linter: only a module-level conditional mutation or a function that declares global __all__ makes the export set opaque; a __all__ bound as a local in a nested function or class no longer masks a genuinely unused hoisted import. --- scripts/verify_import_hoist.py | 68 +++++++++- .../tests/test_embedding_verdict_cache.py | 33 +++++ .../tests/test_is_embedding_model_offline.py | 118 +++++++++++++++--- studio/backend/utils/models/model_config.py | 101 ++++++++++----- .../backend/utils/security/file_security.py | 26 ++-- 5 files changed, 284 insertions(+), 62 deletions(-) diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index cc6dfe900a7..84b40649171 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -497,14 +497,53 @@ def _stmt_binds_all(node: ast.AST) -> bool: ) +def _declares_global_all(func: ast.AST) -> bool: + """True when *func*'s OWN body declares ``global __all__`` (so an ``__all__`` assignment in it + writes the module export set). Does not descend into nested functions/classes -- their + ``global`` statements belong to their own scope.""" + stack = list(ast.iter_child_nodes(func)) + while stack: + node = stack.pop() + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + continue # a nested scope; its `global` decls are its own + if isinstance(node, ast.Global) and "__all__" in node.names: + return True + stack.extend(ast.iter_child_nodes(node)) + return False + + def _all_bound_outside_module_body(tree: ast.Module) -> bool: - """True when ``__all__`` is bound / mutated anywhere OTHER than a top-level module statement -- - inside a module-level ``if`` / ``try`` / ``for`` / ``while`` / ``with`` / ``match`` (a - conditional whose final value cannot be replayed statically), or inside a nested scope. The - caller then marks ``__all__`` opaque, so a conditionally-added re-export is not mistaken for an - unused hoist. Only ever ADDS opacity, so it cannot create a false blocker.""" + """True when the module's runtime ``__all__`` may be mutated somewhere a static read cannot + replay: a MODULE-LEVEL binding nested in a conditional (``if`` / ``try`` / ``for`` / ``while`` + / ``with`` / ``match``, whose final value is not the replayed top-level sequence), or an + ``__all__`` assignment inside a function that declares ``global __all__``. An ``__all__`` bound + as a LOCAL in a nested function or class body (no ``global __all__``) CANNOT change the module + export set, so it is ignored -- otherwise a stray local named ``__all__`` would mask a genuinely + unused hoisted import. The caller marks ``__all__`` opaque on True; this only ever ADDS opacity + at module scope, so it cannot create a false blocker.""" top_level = {id(node) for node in tree.body} - return any(id(node) not in top_level and _stmt_binds_all(node) for node in ast.walk(tree)) + + def _walk(node: ast.AST) -> bool: + # Descends only module-level statements/blocks; nested scopes are handled without recursing. + for child in ast.iter_child_nodes(node): + if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)): + # A __all__ assignment inside the function reaches module scope ONLY via + # `global __all__`; otherwise it is a local and cannot taint the export set. + if _declares_global_all(child) and any( + _stmt_binds_all(n) for n in ast.walk(child) + ): + return True + # else: a purely-local __all__ -> ignore this function entirely. + elif isinstance(child, ast.ClassDef): + continue # a class-body __all__ is a class attribute, not the module's -> ignore + else: + if id(child) not in top_level and _stmt_binds_all(child): + return True # module-level conditional mutation (not the replayed top sequence) + if _walk(child): + return True + return False + + return _walk(tree) def _collect_dunder_all(tree: ast.Module) -> tuple[set[str], bool]: @@ -939,6 +978,23 @@ def used_targets(analysis) -> set[str]: 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\nif True:\n __all__ += ["b"]\n', None, ), + "nested_local_all_does_not_mask_unused_hoist": ( + # a __all__ bound as a LOCAL inside a nested function cannot change the module export set, + # so it must NOT make the module __all__ opaque -- a genuinely unused hoisted import is + # still a bad hoist and must block + 'from pkg import a\n__all__ = ["a"]\n', + 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n' + 'def f():\n __all__ = ["local"]\n return __all__\n', + "BLOCKER", + ), + "global_all_mutation_in_function_is_opaque": ( + # a function that declares `global __all__` and assigns it DOES change the module export + # set at runtime -> opaque, so a re-export supplied that way must not be flagged as unused + 'from pkg import a\n__all__ = ["a"]\n', + 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n' + 'def register():\n global __all__\n __all__ = __all__ + ["b"]\n', + None, + ), "readable_reassign_resets_opacity": ( # an opaque "+=" then a readable replacing "=": the final __all__ is empty, so a # hoisted unused import is a genuine bad hoist -- the reset must un-stick opacity diff --git a/studio/backend/tests/test_embedding_verdict_cache.py b/studio/backend/tests/test_embedding_verdict_cache.py index b36928598fe..69a991a70f1 100644 --- a/studio/backend/tests/test_embedding_verdict_cache.py +++ b/studio/backend/tests/test_embedding_verdict_cache.py @@ -230,6 +230,39 @@ def _boom(*a, **k): assert _blocked(monkeypatch, snap) is True +def test_sharded_safetensors_index_credited_only_at_root(home, tmp_path, monkeypatch): + # A complete model.safetensors.index.json covers a sibling pickle ONLY at a from_pretrained + # root. A non-Transformer ST module (Dense/WordEmbeddings) loads via load_torch_weights, which + # ignores the index and reads pytorch_model.bin -- so a sharded index there must NOT credit the + # pickle, or it deserializes unblocked offline. + index = json.dumps({"weight_map": {"w": "model-00001-of-00001.safetensors"}}) + + # Module dir (not root): index does NOT credit the pickle -> blocked. + mod = _snap( + tmp_path, + { + "modules.json": json.dumps([{"path": "0_Dense", "type": "..."}]), + "0_Dense/pytorch_model.bin": b"pickle", + "0_Dense/model.safetensors.index.json": index, + "0_Dense/model-00001-of-00001.safetensors": b"\0", + }, + commit = "cmod", + ) + assert _blocked(monkeypatch, mod) is True + + # Snapshot root (from_pretrained): the same index DOES credit the pickle -> allowed. + root = _snap( + tmp_path, + { + "pytorch_model.bin": b"pickle", + "model.safetensors.index.json": index, + "model-00001-of-00001.safetensors": b"\0", + }, + commit = "croot", + ) + assert _blocked(monkeypatch, root) is False + + def test_unresolvable_snapshot_blocks(home, monkeypatch): # A snapshot that ERRORS on resolution (not a clean None) fails closed. def _boom(name): diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 1bbf7116e61..f35ba89e156 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -375,10 +375,19 @@ def test_marker_rejects_non_base_weight_bins(tmp_path, monkeypatch): "model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors", ), + ( # #7218: complete numbered shards but a PRESENT-but-unreadable index. transformers opens + # and parses the index to locate the shards, so a stub / truncated one (no weight_map) + # fails the local_files_only load -- it must NOT fall back to the filename-numbering + # heuristic and accept. (_cache_repo_with_files stubs every file as b"\0".) + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors", + "model.safetensors.index.json", + ), ], ) def test_marker_rejects_incomplete_shard_set(tmp_path, monkeypatch, weights): - # An incomplete shard set (or one missing its index map) must not validate. + # An incomplete shard set (one missing its index map, or one whose present index is unreadable) + # must not validate. _cache_repo_with_files(tmp_path, monkeypatch, *weights) assert mc._embedding_marker_in_hf_cache("org/model") is False @@ -386,12 +395,8 @@ def test_marker_rejects_incomplete_shard_set(tmp_path, monkeypatch, weights): @pytest.mark.parametrize( "weights", [ - ("pytorch_model.bin",), # torch .bin - ( # sharded: every index plus the index map - "model-00001-of-00002.safetensors", - "model-00002-of-00002.safetensors", - "model.safetensors.index.json", - ), + ("model.safetensors",), # unsharded safetensors the loader probes + ("pytorch_model.bin",), # torch .bin the loader probes ( # ST module dir that is itself a complete load root "0_Transformer/model.safetensors", "0_Transformer/config.json", @@ -400,11 +405,23 @@ def test_marker_rejects_incomplete_shard_set(tmp_path, monkeypatch, weights): ], ) def test_marker_accepts_recognized_torch_weights(tmp_path, monkeypatch, weights): - # Must not over-reject: single bin, complete sharded set, and a complete module dir. + # Must not over-reject: an unsharded base weight the loader actually probes, and a complete + # module dir. (A complete SHARDED set needs a real weight_map index -- covered separately with + # _sharded_transformer_repo -- because _cache_repo_with_files can only stub the index.) _cache_repo_with_files(tmp_path, monkeypatch, *weights) assert mc._embedding_marker_in_hf_cache("org/model") is True +@pytest.mark.parametrize("decoy", ["pytorch_model.safetensors", "model.bin"]) +def test_marker_rejects_decoy_base_weights(tmp_path, monkeypatch, decoy): + # #7218: from_pretrained / Module.load_torch_weights probe model.safetensors + pytorch_model.bin, + # never pytorch_model.safetensors or model.bin. A cache holding config + tokenizer + only one of + # those decoys validated offline (the pristine loop credited any {model,pytorch_model} x + # {safetensors,bin}) and then 409'd at the local_files_only load, so it must NOT be loadable. + _cache_repo_with_files(tmp_path, monkeypatch, decoy) + assert mc._embedding_marker_in_hf_cache("org/model") is False + + def test_marker_rejects_weights_split_from_their_config(tmp_path, monkeypatch): # modules.json sends SentenceTransformer at 0_Transformer/, which is loaded FROM that # directory. Weights there with the config and tokenizer only at the root is a partial @@ -1322,9 +1339,12 @@ def _tag_only_repo( monkeypatch, commit, repo_id = "org/emb", + modules_json = None, ): """A complete, loadable snapshot with NO modules.json -- the tag-only embedder shape, - which the marker cannot recognize, so detection falls through to the recorded verdict.""" + which the marker cannot recognize, so detection falls through to the recorded verdict. + ``modules_json`` (raw body) writes a modules.json when given, to exercise the case where the + tag-only fallback must NOT run because a manifest is present.""" hf_root = tmp_path / "hf" repo = hf_root / f"models--{repo_id.replace('/', '--')}" snap = repo / "snapshots" / commit @@ -1332,6 +1352,8 @@ def _tag_only_repo( (snap / "config.json").write_text("{}") (snap / "tokenizer.json").write_text("{}") (snap / "model.safetensors").write_bytes(b"\0") + if modules_json is not None: + (snap / "modules.json").write_text(modules_json) (repo / "refs").mkdir(parents = True) (repo / "refs" / "main").write_text(commit) _fake_hf_cache(monkeypatch, hf_root) @@ -1347,6 +1369,21 @@ def test_online_verdict_still_applies_at_the_confirmed_revision(tmp_path, monkey assert mc.is_embedding_model("org/emb") is True +def test_offline_tag_only_fallback_skipped_when_modules_json_present(tmp_path, monkeypatch): + # #7218: the offline tag-only fallback (_known_embedder + _snapshot_has_complete_weights) treats + # a cache like a plain root Transformer, so it must run ONLY when modules.json is ABSENT. With a + # present-but-unloadable modules.json ("[]" builds ZERO modules) the loader takes the modules.json + # path (base/model.py _load_config_modules) and never falls back to the root weights, so even a + # known embedder with a complete root weight set is not loadable offline. Pre-fix the ungated + # fallback fired on the same commit + complete weights and returned True, then the local-only + # load 409'd. Same shape as the control above but with the present manifest. + monkeypatch.setattr(mc, "_load_persisted_embedders", lambda: {"org/emb": "commit_a"}) + monkeypatch.setattr(mc, "_persist_embedder", lambda name, commit: None) + _tag_only_repo(tmp_path, monkeypatch, "commit_a", modules_json = "[]") + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + assert mc.is_embedding_model("org/emb") is False + + def test_online_verdict_stops_applying_when_the_revision_advances(tmp_path, monkeypatch): # A verdict records that the Hub tagged ONE revision an embedder. Once refs/main # advances -- e.g. to a complete but non-embedding Transformer snapshot -- the old @@ -1619,12 +1656,14 @@ def _clip_repo( monkeypatch, *, complete, + processor = True, commit = "aaa", ): """A SentenceTransformer whose modules.json declares a CLIPModel at 0_CLIPModel. CLIPModel(load) -> Transformer.__init__ reads AutoConfig + AutoModel.from_pretrained (weights) + AutoProcessor - (processor/tokenizer), so a config-only dir is NOT loadable. complete=False writes only - config.json (no processor/tokenizer asset, no weights).""" + (image processor + tokenizer), so a config-only dir is NOT loadable. complete=False writes only + config.json (no processor asset, no weights). ``processor=False`` writes the tokenizer + weights + but omits the image-processor config (preprocessor_config.json) AutoProcessor needs for CLIP.""" hf_root = tmp_path / "hf" repo = hf_root / f"models--{_CLIP.replace('/', '--')}" snap = repo / "snapshots" / commit @@ -1635,8 +1674,11 @@ def _clip_repo( (snap / "config_sentence_transformers.json").write_text("{}") (snap / "0_CLIPModel" / "config.json").write_text("{}") if complete: - (snap / "0_CLIPModel" / "tokenizer.json").write_text("{}") # AutoProcessor asset + (snap / "0_CLIPModel" / "tokenizer.json").write_text("{}") # AutoProcessor tokenizer asset (snap / "0_CLIPModel" / "model.safetensors").write_bytes(b"\0") + if processor: + # AutoProcessor's image-processor half reads its own config. + (snap / "0_CLIPModel" / "preprocessor_config.json").write_text("{}") (repo / "refs").mkdir(parents = True) (repo / "refs" / "main").write_text(commit) _fake_hf_cache(monkeypatch, hf_root) @@ -1651,8 +1693,18 @@ def test_marker_rejects_config_only_clip_module(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache(_CLIP) is False +def test_marker_rejects_clip_module_without_processor_config(tmp_path, monkeypatch): + # #7218: CLIP's AutoProcessor is a CLIPProcessor = image processor + tokenizer. The image + # processor reads preprocessor_config.json; a tokenizer alone does not satisfy it. A CLIP dir with + # config + tokenizer + weights but NO preprocessor_config.json validated offline (the generic + # Transformer load-root shape was met) and then failed the local_files_only processor load, so it + # must NOT be loadable. + _clip_repo(tmp_path, monkeypatch, complete = True, processor = False) + assert mc._embedding_marker_in_hf_cache(_CLIP) is False + + def test_marker_accepts_complete_clip_module(tmp_path, monkeypatch): - # Companion: config + processor/tokenizer + weights make the CLIP module a loadable HF load root. + # Companion: config + tokenizer + preprocessor_config + weights make the CLIP module loadable. _clip_repo(tmp_path, monkeypatch, complete = True) assert mc._embedding_marker_in_hf_cache(_CLIP) is True @@ -1672,9 +1724,11 @@ def _sparse_static_repo( commit = "aaa", ): """A SparseEncoder whose modules.json declares a SparseStaticEmbedding at 0_SparseStaticEmbedding. - Its load() reads a tokenizer (AutoTokenizer.from_pretrained) + EITHER an idf.json (from_json - branch, NO weights) OR a complete Torch weight set (load_torch_weights). ``payload`` selects what - backs it: "idf" (idf.json), "weights" (model.safetensors) or "none" (neither).""" + Its load() reads a tokenizer (AutoTokenizer.from_pretrained) + EITHER an idf.json THROUGH a + config.json whose ``path`` names it (from_json branch, NO weights) OR a complete Torch weight set + (load_torch_weights). ``payload`` selects what backs it: "idf" (config path -> idf.json), + "idf_unselected" (idf.json present but config does NOT name a ``.json`` path, so load() falls + through to load_torch_weights and raises), "weights" (model.safetensors) or "none" (neither).""" hf_root = tmp_path / "hf" repo = hf_root / f"models--{_SPARSE_STATIC.replace('/', '--')}" snap = repo / "snapshots" / commit @@ -1690,10 +1744,12 @@ def _sparse_static_repo( ) ) (snap / "config_sentence_transformers.json").write_text("{}") - (mod / "config.json").write_text("{}") + # config.json names the idf json ONLY for the "idf" variant; the loader keys off its ``path``. + config_body = json.dumps({"path": "idf.json"}) if payload == "idf" else "{}" + (mod / "config.json").write_text(config_body) if include_tokenizer: (mod / "tokenizer.json").write_text("{}") - if payload == "idf": + if payload in ("idf", "idf_unselected"): (mod / "idf.json").write_text("{}") elif payload == "weights": (mod / "model.safetensors").write_bytes(b"\0") @@ -1726,6 +1782,16 @@ def test_marker_rejects_sparse_static_embedding_without_tokenizer(tmp_path, monk assert mc._embedding_marker_in_hf_cache(_SPARSE_STATIC) is False +def test_marker_rejects_sparse_static_idf_not_selected_by_config(tmp_path, monkeypatch): + # #7218: SparseStaticEmbedding.load() reads idf.json ONLY when config.json declares a ``path`` + # ending ``.json`` (path = config.pop("path"); if path and path.endswith(".json"): from_json). + # A tokenizer + a bare idf.json whose config does NOT name that path falls through to + # load_torch_weights and RAISES (no Torch weight set), so it is not loadable. Pristine credited + # the idf.json unconditionally and accepted it. + _sparse_static_repo(tmp_path, monkeypatch, payload = "idf_unselected") + assert mc._embedding_marker_in_hf_cache(_SPARSE_STATIC) is False + + # ── #7218 P4: a PRESENT but empty / malformed modules.json is not a loadable ST model ── @@ -1855,6 +1921,22 @@ def test_marker_accepts_sharded_index_with_subdir_shards(tmp_path, monkeypatch): assert mc._embedding_marker_in_hf_cache(_SHARDED) is True +def test_marker_accepts_sharded_index_with_every_mapped_shard(tmp_path, monkeypatch): + # A well-formed index whose weight_map resolves to every present shard is loadable -- the + # positive counterpart to the malformed / missing-shard rejections (a stub index is now rejected, + # so the accept case must carry a real weight_map). + _sharded_transformer_repo( + tmp_path, + monkeypatch, + weight_map = { + "a": "model-00001-of-00002.safetensors", + "b": "model-00002-of-00002.safetensors", + }, + present_shards = ["model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors"], + ) + assert mc._embedding_marker_in_hf_cache(_SHARDED) is True + + def test_marker_accepts_sharded_index_with_all_mapped_shards(tmp_path, monkeypatch): # Companion: a complete uniform shard set the weight_map maps -> loadable. _sharded_transformer_repo( diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 9828ba61eca..d9b336829ef 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2221,14 +2221,14 @@ def _names_have_tokenizer(names: set) -> bool: def _dir_has_complete_torch_weights(names: set) -> bool: - """True when *names* form a COMPLETE Torch weight set: a single ``model.safetensors`` / - ``pytorch_model.bin``, or a full shard set (every index ``1..total`` plus its - ``.index.json`` map). A lone shard or a missing index map is incomplete -- it would - validate then fail under local_files_only.""" - for stem in ("model", "pytorch_model"): - for ext in ("safetensors", "bin"): - if f"{stem}.{ext}" in names: - return True + """True when *names* form a COMPLETE Torch weight set: a single unsharded weight the loader + actually probes (``model.safetensors`` or ``pytorch_model.bin`` -- NOT ``model.bin`` or + ``pytorch_model.safetensors``, which from_pretrained / Module.load_torch_weights never read), or + a full shard set (every index ``1..total`` plus its ``.index.json`` map). A lone shard, a missing + index map or a decoy-named base weight is incomplete -- it would validate then fail under + local_files_only.""" + if "model.safetensors" in names or "pytorch_model.bin" in names: + return True shards: dict = {} for name in names: m = _ST_SHARD_RE.match(name) @@ -2300,23 +2300,27 @@ def _dir_weight_set_is_complete(dir_path: Path, names: set) -> bool: (authoritative, resolved relative to *dir_path*, subdirectory values included); else the filename-numbering shard heuristic (:func:`_dir_has_complete_torch_weights`) as a fallback for an unreadable / stub index. Never raises.""" - # 1. An unsharded base weight is always sufficient. - for stem in ("model", "pytorch_model"): - for ext in ("safetensors", "bin"): - if f"{stem}.{ext}" in names: - return True - # 2. A present, parseable weight-index is authoritative: require every mapped shard present. - saw_parseable_index = False + # 1. An unsharded base weight the loader actually probes: model.safetensors (safetensors) or + # pytorch_model.bin (pickle). from_pretrained / Module.load_torch_weights never read + # ``model.bin`` or ``pytorch_model.safetensors``, so a cache holding only one of those is NOT + # loadable (the security gate treats ``pytorch_model.safetensors`` as a decoy for the same + # reason) -- classifying it as complete accepts a snapshot the local_files_only load then 409s. + if "model.safetensors" in names or "pytorch_model.bin" in names: + return True + # 2. A PRESENT weight-index is authoritative: transformers / sentence-transformers open and parse + # it to locate the sharded weights, so if any index file exists it must fully resolve. One that + # maps a complete shard set validates; one missing a mapped shard OR unreadable / lacking a + # ``weight_map`` (a stub or truncated write) makes the local_files_only load fail -- so do NOT + # fall back to the filename-numbering heuristic when an index file is present. + saw_index = False for index_name in _ST_WEIGHT_INDEX_FILES: if index_name in names: - verdict = _index_weight_set_complete(dir_path / index_name) - if verdict is True: + saw_index = True + if _index_weight_set_complete(dir_path / index_name) is True: return True - if verdict is False: - saw_parseable_index = True - if saw_parseable_index: - return False # a mapped shard is missing -> the local_files_only load would fail - # 3. Fallback: numbering heuristic on basenames (unreadable / stub index, or no index at all). + if saw_index: + return False + # 3. No index file at all: the filename-numbering shard heuristic on basenames. return _dir_has_complete_torch_weights(names) @@ -2504,6 +2508,13 @@ def _names_have_word_embeddings_tokenizer(names: set) -> bool: # the same check without being listed here. _ST_TRANSFORMER_SHAPED_MODULE_NAMES = frozenset({"clipmodel"}) +# ``CLIPModel.load()`` builds its processor via ``AutoProcessor.from_pretrained``, which for CLIP is +# a ``CLIPProcessor`` = image processor + tokenizer. The image processor reads its own config +# (``preprocessor_config.json``); a tokenizer alone does not satisfy it, so a CLIP dir with config + +# weights + only tokenizer files validates offline and then fails the local_files_only load (#7218 +# P3). Required IN ADDITION to the Transformer load-root shape (config + tokenizer + weights). +_ST_CLIP_PROCESSOR_FILE = "preprocessor_config.json" + # ``SparseStaticEmbedding.load()`` (sparse_encoder/modules/sparse_static_embedding.py:157-226) reads # a tokenizer (``AutoTokenizer.from_pretrained``, line 193) and then EITHER, when ``config.json`` # names a ``.json`` ``path``, an IDF json via ``from_json`` -> ``load_file_path(..., "idf.json")`` @@ -2528,6 +2539,21 @@ def _dir_has_module_config(names: set) -> bool: return any(name == "config.json" or name.endswith("_config.json") for name in names) +def _sparse_config_selects_idf_json(dir_path: Path) -> bool: + """True when a ``SparseStaticEmbedding`` dir's ``config.json`` selects the idf-json variant: its + ``path`` value is a string ending ``.json``. ``SparseStaticEmbedding.load()`` uses ``idf.json`` + ONLY through that field (``path = config.pop("path"); if path is not None and + path.endswith(".json"): from_json(...)``, sparse_static_embedding.py:202-205); otherwise it falls + through to ``load_torch_weights`` and RAISES without a Torch weight set. So a tokenizer + bare + ``idf.json`` with a config that does NOT declare the JSON path is not loadable. Never raises.""" + try: + config = json.loads((dir_path / "config.json").read_text(encoding = "utf-8")) + except (OSError, ValueError): + return False + path = config.get("path") if isinstance(config, dict) else None + return isinstance(path, str) and path.endswith(".json") + + def _module_dir_is_loadable(cls: str, is_root: bool, dir_path: Path) -> bool: """True when *dir_path* carries the files the sentence-transformers module class *cls* reads in its own ``load()`` (see sentence_transformers/{base,sentence_transformer,sparse_encoder} @@ -2557,7 +2583,13 @@ def _module_dir_is_loadable(cls: str, is_root: bool, dir_path: Path) -> bool: names = _dir_file_names(dir_path) # Transformer-shaped: a full HF load root (config + tokenizer/processor + weights), root or subdir. if "transformer" in cls or cls in _ST_TRANSFORMER_SHAPED_MODULE_NAMES: - return _dir_is_transformer_load_root(dir_path, names) + if not _dir_is_transformer_load_root(dir_path, names): + return False + # A CLIP-shaped module additionally needs the image-processor config for AutoProcessor; + # a tokenizer alone (checked above) is not enough. + if cls in _ST_TRANSFORMER_SHAPED_MODULE_NAMES: + return _ST_CLIP_PROCESSOR_FILE in names + return True # Router / Asym: loadable when its router_config.json children are (no weights of its own). if cls in _ST_ROUTER_MODULE_NAMES: return _router_dir_is_loadable(dir_path) @@ -2569,11 +2601,15 @@ def _module_dir_is_loadable(cls: str, is_root: bool, dir_path: Path) -> bool: return _ST_STATIC_EMBEDDING_TOKENIZER_FILE in names and _dir_weight_set_is_complete( dir_path, names ) - # SparseStaticEmbedding: a tokenizer + (an idf.json OR a complete Torch weight set). + # SparseStaticEmbedding: a tokenizer + EITHER a complete Torch weight set, OR an idf.json THAT + # config.json actually selects (a "path" ending .json) -- a bare idf.json the config does not + # name falls through to load_torch_weights and raises. if cls == "sparsestaticembedding": if not _names_have_tokenizer(names): return False - return _ST_SPARSE_IDF_DATA_FILE in names or _dir_weight_set_is_complete(dir_path, names) + if _dir_weight_set_is_complete(dir_path, names): + return True + return _ST_SPARSE_IDF_DATA_FILE in names and _sparse_config_selects_idf_json(dir_path) # A weighted non-Transformer module: its module config AND a complete Torch weight set. if cls in _ST_WEIGHTED_MODULE_NAMES: if not _dir_has_module_config(names) or not _dir_weight_set_is_complete(dir_path, names): @@ -2821,13 +2857,18 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: if snapshot is not None and _snapshot_is_loadable_st_model(snapshot): return True active_commit = _active_commit(model_name) - # Trust an online-confirmed positive (memo or persisted allowlist) only when the - # active snapshot has a complete weight set: a recorded True proves the Hub tagged it - # an embedder, not that its files are on disk. The weight gate covers a downloaded - # tag-only embedder (no modules.json) while rejecting a config-only or half-sharded - # snapshot. + # Trust an online-confirmed positive (memo or persisted allowlist) only for a TAG-ONLY + # embedder -- one with NO modules.json -- whose active snapshot has a complete root weight + # set. A recorded True proves the Hub tagged it an embedder, not that its files load. The + # tag-only fallback (_snapshot_has_complete_weights) treats an unreadable manifest like a + # plain root Transformer, so it must NOT run when modules.json is PRESENT: with the file + # present the loader takes the modules.json path (base/model.py _load_config_modules) and a + # present-but-empty / malformed manifest builds zero modules or raises rather than falling + # back to the root weights (#7218). A present, well-formed manifest is already handled + # authoritatively by _snapshot_is_loadable_st_model above. if ( snapshot is not None + and not (snapshot / "modules.json").is_file() and _known_embedder(model_name, cache_key, active_commit) and _snapshot_has_complete_weights(snapshot) ): diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 4d2164b515d..7fb796613d5 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -396,14 +396,24 @@ def _exact_named(files: dict, name: str): return None -def _dir_has_loadable_safetensors(files: dict) -> bool: - """True when *files* (lower-name -> Path for one directory) hold a safetensors weight a - from_pretrained load will read INSTEAD of a pickle sibling: an unsharded base file, or a - complete indexed shard set. A bare adapter or an orphan shard does not qualify. The - safetensors credit is case-SENSITIVE (see :func:`_exact_named`) so a mis-cased decoy the - loader would skip cannot vouch for a live pickle.""" +def _dir_has_loadable_safetensors(files: dict, is_root: bool = False) -> bool: + """True when *files* (lower-name -> Path for one directory) hold a safetensors weight the loader + will read INSTEAD of a pickle sibling: an unsharded ``model.safetensors``, or -- ONLY at a + ``from_pretrained`` root -- a complete ``model.safetensors.index.json`` shard set. A bare + adapter or an orphan shard does not qualify. The safetensors credit is case-SENSITIVE (see + :func:`_exact_named`) so a mis-cased decoy the loader would skip cannot vouch for a live pickle. + + ``is_root`` gates the SHARDED-index credit. A sharded safetensors index is honored only by + ``from_pretrained`` (the snapshot root / a Transformer root loaded via ``AutoModel``). A + non-Transformer SentenceTransformer module (``Dense``, ``WordEmbeddings``, ``StaticEmbedding``) + loads through ``Module.load_torch_weights``, which probes only ``model.safetensors`` then + ``pytorch_model.bin`` and NEVER reads the index -- so crediting a sharded index in such a module + dir would let its ``pytorch_model.bin`` deserialize unblocked. The unsharded ``model.safetensors`` + credit is honored by both loaders and applies everywhere.""" if any(_exact_named(files, name) is not None for name in _SAFETENSORS_BASE_UNSHARDED): return True + if not is_root: + return False for index_name in _SAFETENSORS_BASE_INDEX: index_path = _exact_named(files, index_name) if index_path is not None and _safetensors_index_complete(index_path): @@ -520,7 +530,7 @@ def _cached_pickle_weight_paths(snap, load_subdirs = ()) -> list: # reads the exact-case one, and a mis-cased sibling that is never loaded is only over-blocked # (safe). Keying by lowered name would drop one and could hash a decoy instead of the target. base = [p for p in pickle_paths if _PICKLE_WEIGHT_RE.match(p.name.lower())] - if base and not _dir_has_loadable_safetensors(files): + if base and not _dir_has_loadable_safetensors(files, is_root = directory == snap): hits.update(base) # An adapter pickle is deserialized only when from_pretrained auto-detects the adapter # (adapter_config.json present) and there is no adapter_model.safetensors to load instead. @@ -541,7 +551,7 @@ def _cached_pickle_weight_paths(snap, load_subdirs = ()) -> list: for root_dir in roots: files = by_dir_files.get(root_dir, {}) index_path = files.get(_PICKLE_INDEX_FILE) - if index_path is None or _dir_has_loadable_safetensors(files): + if index_path is None or _dir_has_loadable_safetensors(files, is_root = root_dir == snap): continue for shard_rel in _index_weight_map_values(index_path): if _file_suffix(shard_rel) not in _PICKLE_SUFFIXES: From 0704c6ea70a90ffe9db0cdf16a4e47740f96836b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 04:09:04 +0000 Subject: [PATCH 64/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- scripts/verify_import_hoist.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index 84b40649171..8d1940340ce 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -529,9 +529,7 @@ def _walk(node: ast.AST) -> bool: if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)): # A __all__ assignment inside the function reaches module scope ONLY via # `global __all__`; otherwise it is a local and cannot taint the export set. - if _declares_global_all(child) and any( - _stmt_binds_all(n) for n in ast.walk(child) - ): + if _declares_global_all(child) and any(_stmt_binds_all(n) for n in ast.walk(child)): return True # else: a purely-local __all__ -> ignore this function entirely. elif isinstance(child, ast.ClassDef): From 529876b2439093fa94d0f782e54254d068726148 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 21 Jul 2026 04:25:40 +0000 Subject: [PATCH 65/77] Scope Router-child pickles to their deepest load root and gate the ST offline kwarg The online scan stripped the first matching load-subdir prefix from a flagged file, so a nested Router child pickle (0_Router/query_0_WordEmbeddings/pytorch_model.bin) matched the parent 0_Router root, looked like an unreferenced nested shard, and slipped the gate even though Router.load() deserializes that child directly. Match the deepest (longest) load subdir instead, so the child becomes root-level under its own load root and blocks. pyproject sets no lower bound on sentence-transformers and the local_files_only constructor arg is absent on older releases, so always forwarding it broke every embedder warm on those installs. Pass it only for an offline load; an online warm never forwards it and works as before, while the offline capability still requires a version that supports it. --- studio/backend/core/rag/embeddings.py | 19 +++++--- .../tests/test_embedding_verdict_cache.py | 48 +++++++++++++++++++ studio/backend/tests/test_file_security.py | 23 +++++++++ .../backend/utils/security/file_security.py | 15 ++++-- 4 files changed, 96 insertions(+), 9 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 3a8c84accb7..e768f3267ed 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -304,12 +304,19 @@ def _get(model_name: str | None = None): decision = _guard_model_security(load_name, local_only) # huggingface_hub honors only HF_HUB_OFFLINE, so a TRANSFORMERS_OFFLINE-only # session would otherwise still fetch missing repo files. - _model = SentenceTransformer( - load_name, - device = device, - model_kwargs = dtype_kwargs("float16"), - local_files_only = local_only, - ) + # + # pyproject sets no lower bound on sentence-transformers and the + # ``local_files_only`` constructor arg is absent on older releases, so pass it + # ONLY when we actually need an offline load. Online (local_only False) then never + # forwards it -- an old install warms the Hub-backed embedder exactly as before, + # while offline (the new capability) still requires a version that supports it. + st_kwargs = { + "device": device, + "model_kwargs": dtype_kwargs("float16"), + } + if local_only: + st_kwargs["local_files_only"] = True + _model = SentenceTransformer(load_name, **st_kwargs) _name = name # An ONLINE load that HF definitively scanned clean is now fully cached: record the # verdict (outside the lock, so hashing stays off the hot path) so a later offline diff --git a/studio/backend/tests/test_embedding_verdict_cache.py b/studio/backend/tests/test_embedding_verdict_cache.py index 69a991a70f1..4712c55698a 100644 --- a/studio/backend/tests/test_embedding_verdict_cache.py +++ b/studio/backend/tests/test_embedding_verdict_cache.py @@ -399,6 +399,54 @@ def test_get_does_not_record_on_gate_error(home, monkeypatch): assert _drive_get(monkeypatch, None, offline = False) == {} +def _drive_get_capturing_st(monkeypatch, st_factory, *, offline): + """Drive ``_get`` with a caller-supplied SentenceTransformer factory (to probe the exact + kwargs the constructor receives). Returns nothing; the factory records what it needs.""" + fake_st = types.ModuleType("sentence_transformers") + fake_st.SentenceTransformer = st_factory + monkeypatch.setitem(sys.modules, "sentence_transformers", fake_st) + monkeypatch.setattr(config, "effective_embedding_model", lambda: "org/model") + monkeypatch.setattr(embeddings, "_device", lambda: "cpu") + monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) + monkeypatch.setattr("utils.models.resolve_st_cached_repo_id_case", lambda r: r) + monkeypatch.setattr("utils.utils.hf_env_offline", lambda: offline) + monkeypatch.setattr( + embeddings, "_guard_model_security", lambda name, lo: fs.FileSecurityDecision("org/model", False) + ) + monkeypatch.setattr(embeddings, "_record_embedding_verdict_safe", lambda n, c: None) + monkeypatch.setattr(embeddings, "_model", None, raising = False) + monkeypatch.setattr(embeddings, "_name", None, raising = False) + embeddings._get() + + +def test_get_online_omits_local_files_only_for_old_st(home, monkeypatch): + # pyproject pins no minimum sentence-transformers, and older releases lack the + # local_files_only constructor arg. An ONLINE warm must not forward it, or such an install + # raises TypeError on every embedder load. The strict-signature factory (no **kwargs) raises + # exactly that if the kwarg is passed. + seen = {} + + def _old_st(name, *, device, model_kwargs): + seen["called"] = True + return object() + + _drive_get_capturing_st(monkeypatch, _old_st, offline = False) + assert seen.get("called") is True + + +def test_get_offline_still_passes_local_files_only(home, monkeypatch): + # Offline (the new capability) still forwards local_files_only=True, so a modern install + # loads purely from cache. + seen = {} + + def _st(name, **k): + seen.update(k) + return object() + + _drive_get_capturing_st(monkeypatch, _st, offline = True) + assert seen.get("local_files_only") is True + + # ── guard fails CLOSED offline on a gate error ─────────────────────── diff --git a/studio/backend/tests/test_file_security.py b/studio/backend/tests/test_file_security.py index b4c8f5d2429..5d1f7cdcecc 100644 --- a/studio/backend/tests/test_file_security.py +++ b/studio/backend/tests/test_file_security.py @@ -438,6 +438,29 @@ def test_indexed_shard_under_load_subdir_blocks(): assert d.blocked is True +def test_flagged_router_child_pickle_blocks_via_deepest_load_root(): + # A nested Router emits BOTH its own dir and each child as load subdirs, and Router.load() + # deserializes the child pickle directly. The flagged child must scope against the DEEPEST + # matching load root (its own dir) so it lands root-level there and blocks -- stripping only + # the parent prefix would leave "query_0_WordEmbeddings/pytorch_model.bin" looking like an + # unreferenced nested shard and, with no index, slip the online gate. + status = { + "scansDone": False, + "filesWithIssues": [ + {"path": "0_Router/query_0_WordEmbeddings/pytorch_model.bin", "level": "unsafe"} + ], + } + with _patch_status(status), _patch_no_index(): + d = evaluate_file_security( + "org/router-embedder", + load_subdirs = ("0_Router", "0_Router/query_0_WordEmbeddings"), + ) + assert d.blocked is True + assert d.unsafe_files == [ + {"path": "0_Router/query_0_WordEmbeddings/pytorch_model.bin", "level": "unsafe"} + ] + + # -- Source files are the consent gate's domain, not a deserialization vector -- diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 7fb796613d5..fc7dc854772 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -149,12 +149,21 @@ def _load_relative_path(norm: str, load_subdirs) -> str: path the Hub reports for a flagged file -- the raw ``0/../evil`` would never prefix ``evil/pytorch_model.bin`` and the file would slip through as an unreferenced nested shard, though the offline gate (which canonicalizes the same path) blocks it. + + The DEEPEST matching load-subdir wins: a nested Router emits both its own dir and its child + (``0_Router`` and ``0_Router/query_0_WordEmbeddings``), and ``Router.load()`` deserializes the + child pickle directly. Stripping only the parent prefix would leave ``query_0_WordEmbeddings/...`` + looking like an unreferenced nested shard (allowed); stripping the longest matching prefix makes + the child pickle root-level under its own load root and blocks. """ + best = norm + best_len = -1 for subdir in load_subdirs or (): prefix = _canonical_rel(subdir) - if prefix and norm.startswith(prefix + "/"): - return norm[len(prefix) + 1 :] - return norm + if prefix and norm.startswith(prefix + "/") and len(prefix) > best_len: + best = norm[len(prefix) + 1 :] + best_len = len(prefix) + return best def _index_prefixes(load_subdirs) -> tuple: From 6c5854002c3c5f89e6b5f2014596f3f35ed630e0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 04:26:32 +0000 Subject: [PATCH 66/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/tests/test_embedding_verdict_cache.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/studio/backend/tests/test_embedding_verdict_cache.py b/studio/backend/tests/test_embedding_verdict_cache.py index 4712c55698a..84cbc6bf36a 100644 --- a/studio/backend/tests/test_embedding_verdict_cache.py +++ b/studio/backend/tests/test_embedding_verdict_cache.py @@ -411,7 +411,9 @@ def _drive_get_capturing_st(monkeypatch, st_factory, *, offline): monkeypatch.setattr("utils.models.resolve_st_cached_repo_id_case", lambda r: r) monkeypatch.setattr("utils.utils.hf_env_offline", lambda: offline) monkeypatch.setattr( - embeddings, "_guard_model_security", lambda name, lo: fs.FileSecurityDecision("org/model", False) + embeddings, + "_guard_model_security", + lambda name, lo: fs.FileSecurityDecision("org/model", False), ) monkeypatch.setattr(embeddings, "_record_embedding_verdict_safe", lambda n, c: None) monkeypatch.setattr(embeddings, "_model", None, raising = False) From e55a1c632d29d31e5842679fc26a1cf65f87e497 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 21 Jul 2026 04:48:56 +0000 Subject: [PATCH 67/77] Reject snapshot-escaping shard paths and credit Transformer submodule safetensors The offline pickle enumerator joined a weight-index weight_map value straight to the load root and followed it, so a repo-controlled index mapping "../.." into a sibling snapshot made an offline from_pretrained deserialize an out-of-snapshot pickle, and an online load would then hash and record that external file as the scanned commit's clean content. Reject any shard path that escapes the snapshot root and fail closed, mirroring the canonical-root check the online shard scan already applies. A complete model.safetensors.index.json was credited over a sibling pickle only at the snapshot root, but a Transformer module subdirectory (0_Transformer/) is loaded via AutoModel.from_pretrained, which honors that shard set and never reads the pickle. Credit the sharded index for Transformer-typed modules declared in modules.json so a cached model that ships both a sharded safetensors checkpoint and an unused PyTorch checkpoint is no longer falsely blocked offline. Non-Transformer modules (Dense, WordEmbeddings, StaticEmbedding) read a flat weight with no index and keep their pickle blocked. Limit the import-hoist verifier's global __all__ scan to the declaring function's own scope so a nested inner-scope local __all__ no longer marks the module export set opaque and mask an unused hoisted import. --- scripts/verify_import_hoist.py | 31 ++++- .../tests/test_embedding_verdict_cache.py | 68 +++++++++++ .../backend/utils/security/file_security.py | 107 +++++++++++++++--- 3 files changed, 189 insertions(+), 17 deletions(-) diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index 8d1940340ce..c0e83b0e9de 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -512,6 +512,23 @@ def _declares_global_all(func: ast.AST) -> bool: return False +def _binds_all_in_own_scope(func: ast.AST) -> bool: + """True when *func*'s OWN body (not a nested function/class) binds ``__all__``. Paired with + :func:`_declares_global_all`, this identifies a function that writes the MODULE export set. A + nested scope's ``__all__`` is a local of THAT scope and cannot reach module scope, so it is not + counted -- otherwise ``def f(): global __all__`` with an inner ``def g(): __all__ = [...]`` would + wrongly mark the module ``__all__`` opaque and mask a genuinely unused hoisted import.""" + stack = list(ast.iter_child_nodes(func)) + while stack: + node = stack.pop() + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): + continue # a nested scope; its `__all__` binds are its own + if _stmt_binds_all(node): + return True + stack.extend(ast.iter_child_nodes(node)) + return False + + def _all_bound_outside_module_body(tree: ast.Module) -> bool: """True when the module's runtime ``__all__`` may be mutated somewhere a static read cannot replay: a MODULE-LEVEL binding nested in a conditional (``if`` / ``try`` / ``for`` / ``while`` @@ -528,8 +545,9 @@ def _walk(node: ast.AST) -> bool: for child in ast.iter_child_nodes(node): if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)): # A __all__ assignment inside the function reaches module scope ONLY via - # `global __all__`; otherwise it is a local and cannot taint the export set. - if _declares_global_all(child) and any(_stmt_binds_all(n) for n in ast.walk(child)): + # `global __all__` AND a bind in the function's OWN scope; a bind nested in a + # further inner scope is that scope's local and cannot taint the export set. + if _declares_global_all(child) and _binds_all_in_own_scope(child): return True # else: a purely-local __all__ -> ignore this function entirely. elif isinstance(child, ast.ClassDef): @@ -993,6 +1011,15 @@ def used_targets(analysis) -> set[str]: 'def register():\n global __all__\n __all__ = __all__ + ["b"]\n', None, ), + "nested_global_all_bind_does_not_mask_unused_hoist": ( + # `global __all__` in the OUTER function but the only `__all__` bind is in a further inner + # scope: that bind is the inner scope's local and never reaches module scope, so it must NOT + # mark the module __all__ opaque -- the unused hoisted import stays a genuine bad hoist + 'from pkg import a\n__all__ = ["a"]\n', + 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n' + 'def f():\n global __all__\n def g():\n __all__ = ["x"]\n return __all__\n return g\n', + "BLOCKER", + ), "readable_reassign_resets_opacity": ( # an opaque "+=" then a readable replacing "=": the final __all__ is empty, so a # hoisted unused import is a genuine bad hoist -- the reset must un-stick opacity diff --git a/studio/backend/tests/test_embedding_verdict_cache.py b/studio/backend/tests/test_embedding_verdict_cache.py index 84cbc6bf36a..0bc28e39676 100644 --- a/studio/backend/tests/test_embedding_verdict_cache.py +++ b/studio/backend/tests/test_embedding_verdict_cache.py @@ -263,6 +263,74 @@ def test_sharded_safetensors_index_credited_only_at_root(home, tmp_path, monkeyp assert _blocked(monkeypatch, root) is False +def test_sharded_index_credited_in_transformer_submodule(home, tmp_path, monkeypatch): + # A 0_Transformer/ module is loaded via AutoModel.from_pretrained, which honors a complete + # model.safetensors.index.json and never reads the sibling pytorch_model.bin -- so that pickle + # must NOT block offline even though the index is in a SUBDIR, not the snapshot root. + index = json.dumps({"weight_map": {"w": "model-00001-of-00001.safetensors"}}) + snap = _snap( + tmp_path, + { + "modules.json": json.dumps( + [ + { + "path": "0_Transformer", + "type": "sentence_transformers.models.Transformer.Transformer", + } + ] + ), + "0_Transformer/pytorch_model.bin": b"pickle", + "0_Transformer/model.safetensors.index.json": index, + "0_Transformer/model-00001-of-00001.safetensors": b"\0", + }, + commit = "ctf", + ) + assert _blocked(monkeypatch, snap) is False + + +def test_sharded_index_not_credited_in_non_transformer_submodule(home, tmp_path, monkeypatch): + # The same layout under a WordEmbeddings module (load_torch_weights, no index) still blocks: the + # credit is extended ONLY to Transformer-typed dirs, never to a module that reads the pickle. + index = json.dumps({"weight_map": {"w": "model-00001-of-00001.safetensors"}}) + snap = _snap( + tmp_path, + { + "modules.json": json.dumps( + [ + { + "path": "0_WordEmbeddings", + "type": "sentence_transformers.models.WordEmbeddings.WordEmbeddings", + } + ] + ), + "0_WordEmbeddings/pytorch_model.bin": b"pickle", + "0_WordEmbeddings/model.safetensors.index.json": index, + "0_WordEmbeddings/model-00001-of-00001.safetensors": b"\0", + }, + commit = "cwe", + ) + assert _blocked(monkeypatch, snap) is True + + +def test_index_shard_escaping_snapshot_blocks(home, tmp_path, monkeypatch): + # A weight index whose weight_map points OUT of the snapshot ("../..") must fail closed, not be + # followed: offline from_pretrained would resolve it on disk and deserialize an out-of-snapshot + # pickle, and an online load would hash+record that external file as this commit's clean content. + snap = _snap( + tmp_path, + { + "pytorch_model.bin.index.json": json.dumps( + {"weight_map": {"w": "../../evil/pytorch_model.bin"}} + ), + }, + commit = "cesc", + ) + monkeypatch.setattr(mc, "_active_snapshot_dir", lambda name: snap) + decision = fs.evaluate_file_security("org/model", None, local_only_load = True) + assert decision.blocked is True + assert "outside the snapshot" in decision.reason + + def test_unresolvable_snapshot_blocks(home, monkeypatch): # A snapshot that ERRORS on resolution (not a clean None) fails closed. def _boom(name): diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index fc7dc854772..49ddf0e5bb0 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -405,23 +405,23 @@ def _exact_named(files: dict, name: str): return None -def _dir_has_loadable_safetensors(files: dict, is_root: bool = False) -> bool: +def _dir_has_loadable_safetensors(files: dict, honors_sharded_index: bool = False) -> bool: """True when *files* (lower-name -> Path for one directory) hold a safetensors weight the loader - will read INSTEAD of a pickle sibling: an unsharded ``model.safetensors``, or -- ONLY at a - ``from_pretrained`` root -- a complete ``model.safetensors.index.json`` shard set. A bare - adapter or an orphan shard does not qualify. The safetensors credit is case-SENSITIVE (see + will read INSTEAD of a pickle sibling: an unsharded ``model.safetensors``, or -- ONLY where the + loader honors it -- a complete ``model.safetensors.index.json`` shard set. A bare adapter or an + orphan shard does not qualify. The safetensors credit is case-SENSITIVE (see :func:`_exact_named`) so a mis-cased decoy the loader would skip cannot vouch for a live pickle. - ``is_root`` gates the SHARDED-index credit. A sharded safetensors index is honored only by - ``from_pretrained`` (the snapshot root / a Transformer root loaded via ``AutoModel``). A - non-Transformer SentenceTransformer module (``Dense``, ``WordEmbeddings``, ``StaticEmbedding``) - loads through ``Module.load_torch_weights``, which probes only ``model.safetensors`` then - ``pytorch_model.bin`` and NEVER reads the index -- so crediting a sharded index in such a module - dir would let its ``pytorch_model.bin`` deserialize unblocked. The unsharded ``model.safetensors`` - credit is honored by both loaders and applies everywhere.""" + ``honors_sharded_index`` gates the SHARDED-index credit. A sharded safetensors index is honored + only by ``AutoModel.from_pretrained`` -- the snapshot root, or a Transformer-typed module + subdirectory (``0_Transformer/``). A non-Transformer SentenceTransformer module (``Dense``, + ``WordEmbeddings``, ``StaticEmbedding``) loads through ``Module.load_torch_weights``, which probes + only ``model.safetensors`` then ``pytorch_model.bin`` and NEVER reads the index -- so crediting a + sharded index in such a module dir would let its ``pytorch_model.bin`` deserialize unblocked. The + unsharded ``model.safetensors`` credit is honored by both loaders and applies everywhere.""" if any(_exact_named(files, name) is not None for name in _SAFETENSORS_BASE_UNSHARDED): return True - if not is_root: + if not honors_sharded_index: return False for index_name in _SAFETENSORS_BASE_INDEX: index_path = _exact_named(files, index_name) @@ -493,6 +493,48 @@ def _st_load_roots(snap, load_subdirs = ()) -> set: return roots +def _is_transformer_module_type(type_str) -> bool: + """True for the SentenceTransformer ``Transformer`` module class, whose loader is + ``AutoModel.from_pretrained`` and therefore honors a ``model.safetensors.index.json`` shard set. + ``Dense`` / ``WordEmbeddings`` / ``StaticEmbedding`` read a single flat weight file through + ``Module.load_torch_weights`` and never consult an index, so their sharded-index credit stays + gated off.""" + return str(type_str).rsplit(".", 1)[-1] == "Transformer" + + +def _transformer_load_roots(snap) -> set: + """The load roots whose loader honors a sharded ``model.safetensors.index.json``: the + Transformer-typed modules ``modules.json`` declares (loaded via ``AutoModel.from_pretrained``). + A complete safetensors shard set in such a SUBMODULE (e.g. ``0_Transformer/``) means the loader + reads safetensors and never the sibling ``pytorch_model.bin``, so that pickle must not block. + The declared type is authoritative -- the loader instantiates the module by the SAME type, so a + mis-declared dir loads through its declared loader too. A missing / malformed ``modules.json`` + yields an empty set, leaving the shard credit unextended (the conservative over-block + direction). The snapshot root is credited separately by the caller.""" + import json + + roots: set = set() + try: + modules = json.loads((snap / "modules.json").read_text(encoding = "utf-8")) + except (OSError, ValueError): + return roots + if isinstance(modules, list): + for module in modules: + if isinstance(module, dict) and _is_transformer_module_type(module.get("type")): + root = _canonical_load_dir(snap, str(module.get("path") or "")) + if root is not None: + roots.add(root) + return roots + + +class _EscapingShardError(OSError): + """A weight-index ``weight_map`` value resolves OUTSIDE the snapshot (a repo-controlled ``../`` + traversal). Subclasses ``OSError`` so the offline caller's existing fail-closed enumeration + handler blocks the load; raised instead of following the path, because offline the loader would + resolve it on the local filesystem and an online load would hash+record an out-of-snapshot file + as this commit's clean content.""" + + def _cached_pickle_weight_paths(snap, load_subdirs = ()) -> list: """The pickle weight FILES (as ``Path`` objects under *snap*) a from_pretrained load actually deserializes: at a real load root and with NO loadable safetensors alternative there. A load @@ -510,6 +552,10 @@ def _cached_pickle_weight_paths(snap, load_subdirs = ()) -> list: Raises ``OSError`` if the snapshot tree cannot be enumerated (``rglob`` failure), so the offline caller fails CLOSED rather than treat an unreadable cache as pickle-free.""" roots = _st_load_roots(snap, load_subdirs) + # Dirs whose loader (AutoModel.from_pretrained) honors a sharded safetensors index: the snapshot + # root plus each Transformer-typed module subdir. Elsewhere a sharded index is NOT a substitute + # for a live pickle (see _dir_has_loadable_safetensors). + index_honoring = _transformer_load_roots(snap) | {snap} by_dir_pickle: dict = {} # directory -> [Path] (EVERY case variant, not last-wins) by_dir_files: dict = {} # directory -> {lower-name: Path} (safetensors credit; last-wins is safe) # A whole-tree rglob failure propagates (fail-closed); a single unstattable entry is skipped. @@ -539,7 +585,9 @@ def _cached_pickle_weight_paths(snap, load_subdirs = ()) -> list: # reads the exact-case one, and a mis-cased sibling that is never loaded is only over-blocked # (safe). Keying by lowered name would drop one and could hash a decoy instead of the target. base = [p for p in pickle_paths if _PICKLE_WEIGHT_RE.match(p.name.lower())] - if base and not _dir_has_loadable_safetensors(files, is_root = directory == snap): + if base and not _dir_has_loadable_safetensors( + files, honors_sharded_index = directory in index_honoring + ): hits.update(base) # An adapter pickle is deserialized only when from_pretrained auto-detects the adapter # (adapter_config.json present) and there is no adapter_model.safetensors to load instead. @@ -560,12 +608,30 @@ def _cached_pickle_weight_paths(snap, load_subdirs = ()) -> list: for root_dir in roots: files = by_dir_files.get(root_dir, {}) index_path = files.get(_PICKLE_INDEX_FILE) - if index_path is None or _dir_has_loadable_safetensors(files, is_root = root_dir == snap): + if index_path is None or _dir_has_loadable_safetensors( + files, honors_sharded_index = root_dir in index_honoring + ): continue + try: + root_rel = root_dir.relative_to(snap).as_posix() + except ValueError: + continue # a load root not under the snapshot (should not happen) -> skip conservatively for shard_rel in _index_weight_map_values(index_path): if _file_suffix(shard_rel) not in _PICKLE_SUFFIXES: continue - shard = root_dir.joinpath(*shard_rel.split("/")) + # A weight_map value is relative to the index dir (root_dir). Canonicalize it against the + # snapshot root and REFUSE any that escape it: a repo-controlled index can map "../.." into + # a sibling snapshot, and unlike the Hub (which cannot serve "../"), an offline + # from_pretrained resolves it on the local filesystem and would deserialize an + # out-of-snapshot pickle -- an online load would then hash+record that external file as + # this commit's clean content. Fail closed rather than follow it. + combined = shard_rel if root_rel in ("", ".") else root_rel + "/" + shard_rel + shard_canon = _canonical_rel(combined) + if shard_canon is None: + raise _EscapingShardError( + f"weight index maps a shard outside the snapshot: {shard_rel!r}" + ) + shard = snap.joinpath(*shard_canon.split("/")) try: if shard.is_file(): hits.add(shard) @@ -685,6 +751,17 @@ def _evaluate_local_only(model_name: str, load_subdirs = ()) -> "FileSecurityDec try: paths = _cached_pickle_weight_paths(snap, load_subdirs) + except _EscapingShardError as exc: + # A weight index maps a shard outside the snapshot (a "../" traversal). Following it would + # deserialize an out-of-snapshot pickle, so refuse the load outright. + logger.warning( + "Blocking offline load of '%s': %s; cannot verify.", model_name, exc + ) + return FileSecurityDecision( + model_name, + True, + reason = "offline; a weight index maps a shard outside the snapshot, cannot verify", + ) except OSError: # The cache tree could not be fully enumerated: an undiscovered pickle might exist, so a # partial "pickle-free" reading is unsafe. Fail closed. From e6674191585a50c505d41e13e735c885f58a9aca Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 04:50:15 +0000 Subject: [PATCH 68/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/utils/security/file_security.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 49ddf0e5bb0..1b67e59f0ce 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -754,9 +754,7 @@ def _evaluate_local_only(model_name: str, load_subdirs = ()) -> "FileSecurityDec except _EscapingShardError as exc: # A weight index maps a shard outside the snapshot (a "../" traversal). Following it would # deserialize an out-of-snapshot pickle, so refuse the load outright. - logger.warning( - "Blocking offline load of '%s': %s; cannot verify.", model_name, exc - ) + logger.warning("Blocking offline load of '%s': %s; cannot verify.", model_name, exc) return FileSecurityDecision( model_name, True, From 845ba4bc04967d9ffd58cc1328ecfd8a9e553eb0 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 21 Jul 2026 05:10:24 +0000 Subject: [PATCH 69/77] Scope Router children against the snapshot and mirror the ST alias rewrite Router.load resolves each child at Path(subfolder, model_id) relative to the Router dir, so a nested 1_Router with a "../evil" child points at evil/ inside the snapshot and the loader deserializes evil/pytorch_model.bin. The offline enumerator canonicalized the child against the Router dir alone and dropped anything with "..", so that pickle was never scanned and the gate reported the cache pickle-free. Canonicalize router children against the snapshot, retaining in-snapshot siblings as load roots and failing closed on a child that escapes the snapshot itself, matching the online scan which already joins the prefix before normalizing. The security gate resolved a slashless model id by probing the bare cache dir first, but the SentenceTransformer constructor rewrites a non-basic slashless name to sentence-transformers/ and loads THAT snapshot (only the basic ORIGINAL_TRANSFORMER_MODELS load bare). With both models-- and models--sentence-transformers-- cached, the gate inspected the bare dir while the loader read the namespaced one, so a pickle there bypassed the local-only gate. Mirror the constructor: try the namespaced candidate first for non-basic slashless names. Add the same not (snapshot / modules.json).is_file() guard to the transient-Hub-failure tag-only fallback that the offline branch already carries, so a cache whose present manifest is empty or malformed is no longer reported as a loadable embedder. --- .../tests/test_embedding_verdict_cache.py | 40 ++++++++++++ .../tests/test_is_embedding_model_offline.py | 37 +++++++++++ studio/backend/utils/models/model_config.py | 41 ++++++++++-- .../backend/utils/security/file_security.py | 65 ++++++++++++------- 4 files changed, 153 insertions(+), 30 deletions(-) diff --git a/studio/backend/tests/test_embedding_verdict_cache.py b/studio/backend/tests/test_embedding_verdict_cache.py index 0bc28e39676..b21e6d3941c 100644 --- a/studio/backend/tests/test_embedding_verdict_cache.py +++ b/studio/backend/tests/test_embedding_verdict_cache.py @@ -312,6 +312,46 @@ def test_sharded_index_not_credited_in_non_transformer_submodule(home, tmp_path, assert _blocked(monkeypatch, snap) is True +def test_router_child_escaping_router_dir_but_in_snapshot_blocks(home, tmp_path, monkeypatch): + # Router.load resolves each child at Path(subfolder, model_id): a nested 1_Router with a + # "../evil" child points at evil/ INSIDE the snapshot, which the loader deserializes. The gate + # must scan evil/ as a load root, not drop the child for escaping the router dir. + snap = _snap( + tmp_path, + { + "modules.json": json.dumps( + [{"path": "1_Router", "type": "sentence_transformers.models.Router.Router"}] + ), + "1_Router/router_config.json": json.dumps( + {"types": {"../evil": "sentence_transformers.models.WordEmbeddings.WordEmbeddings"}} + ), + "evil/pytorch_model.bin": b"pickle", + }, + commit = "crte", + ) + assert _blocked(monkeypatch, snap) is True + + +def test_router_child_escaping_snapshot_blocks(home, tmp_path, monkeypatch): + # A ROOT Router (subfolder "") with a "../evil" child resolves OUTSIDE the snapshot; the local + # loader would follow it out of the cache, so the gate fails closed even though the snapshot's + # own weights are inert safetensors. + snap = _snap( + tmp_path, + { + "modules.json": json.dumps( + [{"path": "", "type": "sentence_transformers.models.Router.Router"}] + ), + "router_config.json": json.dumps( + {"types": {"../evil": "sentence_transformers.models.WordEmbeddings.WordEmbeddings"}} + ), + "model.safetensors": b"\0", + }, + commit = "crse", + ) + assert _blocked(monkeypatch, snap) is True + + def test_index_shard_escaping_snapshot_blocks(home, tmp_path, monkeypatch): # A weight index whose weight_map points OUT of the snapshot ("../..") must fail closed, not be # followed: offline from_pretrained would resolve it on disk and deserialize an out-of-snapshot diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index f35ba89e156..4be2fbf3c0c 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -1384,6 +1384,43 @@ def test_offline_tag_only_fallback_skipped_when_modules_json_present(tmp_path, m assert mc.is_embedding_model("org/emb") is False +def test_transient_tag_only_fallback_skipped_when_modules_json_present(tmp_path, monkeypatch): + # #7218: the TRANSIENT-Hub-failure tag-only fallback must mirror the offline branch and run ONLY + # when modules.json is ABSENT. With a present-but-unloadable modules.json ("[]") the loader takes + # the modules.json path and never falls back to the root weights, so a known embedder with a + # complete root weight set is still not loadable. Pre-fix the ungated transient fallback fired on + # the same commit + complete weights and returned True, then the local-only load 409'd. + monkeypatch.setattr(mc, "_load_persisted_embedders", lambda: {"org/emb": "commit_a"}) + monkeypatch.setattr(mc, "_persist_embedder", lambda name, commit: None) + _tag_only_repo(tmp_path, monkeypatch, "commit_a", modules_json = "[]") + + def _transient(model_name, token = None, **kwargs): + raise OSError("Temporary failure in name resolution") # non-permanent -> transient branch + + _fake_hf_model_info(monkeypatch, _transient) + mc._embedding_detection_cache.clear() + assert mc.is_embedding_model("org/emb") is False + + +def test_st_repo_id_candidates_prefers_namespaced_for_non_basic(): + # The constructor rewrites a non-basic slashless name to sentence-transformers/ and loads + # THAT snapshot, so the gate must probe it first (probing the bare name first lets a pickle in + # the namespaced dir slip when a bare dir also exists). A slashed id is used verbatim. + assert ( + mc._st_repo_id_candidates("all-MiniLM-L6-v2")[0] + == "sentence-transformers/all-MiniLM-L6-v2" + ) + assert mc._st_repo_id_candidates("BAAI/bge-m3") == ["BAAI/bge-m3"] + + +def test_st_repo_id_candidates_keeps_basic_model_bare(monkeypatch): + # A basic transformer model (ORIGINAL_TRANSFORMER_MODELS) is loaded bare by the constructor, so + # the bare id is tried first. + monkeypatch.setattr(mc, "_original_transformer_models", lambda: frozenset({"bert-base-uncased"})) + assert mc._st_repo_id_candidates("bert-base-uncased")[0] == "bert-base-uncased" + assert mc._st_repo_id_candidates("some-st-model")[0] == "sentence-transformers/some-st-model" + + def test_online_verdict_stops_applying_when_the_revision_advances(tmp_path, monkeypatch): # A verdict records that the Hub tagged ONE revision an embedder. Once refs/main # advances -- e.g. to a complete but non-embedding Transformer snapshot -- the old diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index d9b336829ef..4212ef2d370 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1704,17 +1704,45 @@ def resolve_st_cached_repo_id_case(repo_id: str) -> str: return name[len(prefix) :].replace("--", "/") +_ORIGINAL_TRANSFORMER_MODELS_CACHE: Optional[frozenset] = None + + +def _original_transformer_models() -> frozenset: + """The lowercased ``ORIGINAL_TRANSFORMER_MODELS`` list SentenceTransformer treats as plain HF + models: these slashless names load BARE, while every OTHER slashless name is rewritten to the + ``sentence-transformers/`` org. Imported once and memoized; an import failure yields an empty + set, so every slashless name then resolves to the namespaced org -- the constructor's dominant + behavior.""" + global _ORIGINAL_TRANSFORMER_MODELS_CACHE + if _ORIGINAL_TRANSFORMER_MODELS_CACHE is None: + try: + from sentence_transformers.util import ORIGINAL_TRANSFORMER_MODELS + + _ORIGINAL_TRANSFORMER_MODELS_CACHE = frozenset( + str(m).lower() for m in ORIGINAL_TRANSFORMER_MODELS + ) + except Exception: + _ORIGINAL_TRANSFORMER_MODELS_CACHE = frozenset() + return _ORIGINAL_TRANSFORMER_MODELS_CACHE + + def _st_repo_id_candidates(repo_id: str) -> list[str]: """Repo ids the ST loader could resolve *repo_id* to, in the order it tries them. - A slashless short name like ``all-MiniLM-L6-v2`` is a supported Sentence Transformers - alias: the loader falls back to the ``sentence-transformers/`` organization, so the - snapshot is cached under ``models--sentence-transformers--all-MiniLM-L6-v2``. Probing - only the bare name would report a cache miss -- a 409 offline -- for a model that is - cached and loadable.""" + Mirrors the SentenceTransformer constructor: a slashless short name like ``all-MiniLM-L6-v2`` is + rewritten to ``sentence-transformers/`` and loaded from THERE (never the bare id), UNLESS + it is one of the basic transformer models (``ORIGINAL_TRANSFORMER_MODELS``, e.g. + ``bert-base-uncased``), which load bare. The local-only security gate must inspect the SAME + snapshot the constructor loads, so the constructor's target is tried FIRST; the other spelling + stays only as a defensive fallback for an unusually laid-out cache. Probing the bare name first + would let a pickle in ``models--sentence-transformers--`` -- the dir the constructor + actually loads -- slip the gate when a bare ``models--`` also exists.""" if "/" in repo_id: return [repo_id] - return [repo_id, f"sentence-transformers/{repo_id}"] + namespaced = f"sentence-transformers/{repo_id}" + if repo_id.lower() in _original_transformer_models(): + return [repo_id, namespaced] # a basic transformer model loads bare + return [namespaced, repo_id] # the constructor rewrites a non-basic slashless name to the org def _st_cache_repo_dir(repo_id: str) -> Optional[Path]: @@ -2936,6 +2964,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: snapshot = _active_snapshot_dir(model_name) if ( snapshot is not None + and not (snapshot / "modules.json").is_file() and _known_embedder(model_name, cache_key, _active_commit(model_name)) and _snapshot_has_complete_weights(snapshot) ): diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 1b67e59f0ce..032653e2e3a 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -430,14 +430,21 @@ def _dir_has_loadable_safetensors(files: dict, honors_sharded_index: bool = Fals return False -def _router_child_dirs(root) -> set: +def _router_child_dirs(root, snap) -> set: """Child sub-module directories a SentenceTransformer ``Router`` (legacy ``Asym``) at *root* deserializes. A Router declares its children only in ``router_config.json`` (``types`` maps ``{route}_{idx}_{ClassName}`` -> module class), NOT in the top-level ``modules.json``, and - ``Router.load()`` calls ``module_class.load(subfolder=model_id)`` on each. A child such as - ``query_0_WordEmbeddings/`` holds ``wordembedding_config.json`` + ``pytorch_model.bin`` and no - ``config.json``, so its pickle is still deserialized and its dir must be a load root. Returns - an empty set when there is no readable router config.""" + ``Router.load()`` loads each from ``Path(subfolder, model_id)`` -- the child ``model_id`` + resolved RELATIVE TO the Router dir *root*. A child such as ``query_0_WordEmbeddings/`` holds + ``wordembedding_config.json`` + ``pytorch_model.bin`` and no ``config.json``, so its pickle is + still deserialized and its dir must be a load root. + + A ``model_id`` may use ``..`` to point at a SIBLING dir that is still inside the snapshot + (``1_Router`` + ``../evil`` -> ``evil/``): the local loader follows it and deserializes + ``evil/pytorch_model.bin``, so it is canonicalized against the SNAPSHOT (not the Router dir) and + retained as a load root. A child that escapes the snapshot ITSELF (a root Router with + ``../evil``) is a traversal the local loader would follow out of the cache, so it fails closed + (:class:`_SnapshotEscapeError`). Returns an empty set when there is no readable router config.""" import json try: @@ -447,12 +454,19 @@ def _router_child_dirs(root) -> set: if not isinstance(config, dict): return set() types = config.get("types") + if not isinstance(types, dict): + return set() + try: + root_rel = root.relative_to(snap).as_posix() + except ValueError: + return set() # a Router dir not under the snapshot (should not happen) -> scan nothing children: set = set() - if isinstance(types, dict): - for model_id in types: - child = _canonical_load_dir(root, str(model_id)) - if child is not None: - children.add(child) + for model_id in types: + combined = str(model_id) if root_rel in ("", ".") else root_rel + "/" + str(model_id) + child_rel = _canonical_rel(combined) + if child_rel is None: + raise _SnapshotEscapeError(f"router child escapes the snapshot: {str(model_id)!r}") + children.add(snap.joinpath(*child_rel.split("/"))) return children @@ -481,12 +495,13 @@ def _st_load_roots(snap, load_subdirs = ()) -> set: # A Router/Asym module declares its child sub-modules in router_config.json, not modules.json, # and Router.load() deserializes each child's weights from its own subdir. Treat those child # dirs as load roots too so a pickle in a config.json-less child (e.g. - # query_0_WordEmbeddings/pytorch_model.bin) is scanned. Bounded BFS: every child is a strict - # subpath and a visited set stops any cycle; scanning extra dirs only tightens the gate. + # query_0_WordEmbeddings/pytorch_model.bin) is scanned. Bounded BFS: each child canonicalizes to + # an in-snapshot dir and a visited set (roots) stops any cycle; the finite set of in-snapshot + # dirs bounds it. Scanning extra dirs only tightens the gate. pending = list(roots) while pending: current = pending.pop() - for child in _router_child_dirs(current): + for child in _router_child_dirs(current, snap): if child not in roots: roots.add(child) pending.append(child) @@ -527,12 +542,13 @@ def _transformer_load_roots(snap) -> set: return roots -class _EscapingShardError(OSError): - """A weight-index ``weight_map`` value resolves OUTSIDE the snapshot (a repo-controlled ``../`` - traversal). Subclasses ``OSError`` so the offline caller's existing fail-closed enumeration - handler blocks the load; raised instead of following the path, because offline the loader would - resolve it on the local filesystem and an online load would hash+record an out-of-snapshot file - as this commit's clean content.""" +class _SnapshotEscapeError(OSError): + """A repo-controlled load path -- a weight-index ``weight_map`` value or a ``Router`` + ``router_config.json`` child -- resolves OUTSIDE the snapshot via a ``../`` traversal. + Subclasses ``OSError`` so the offline caller's fail-closed enumeration handler blocks the load; + raised instead of following the path, because offline the loader resolves it on the local + filesystem (reading an out-of-snapshot pickle) and an online load would hash+record that + external file as this commit's clean content.""" def _cached_pickle_weight_paths(snap, load_subdirs = ()) -> list: @@ -628,7 +644,7 @@ def _cached_pickle_weight_paths(snap, load_subdirs = ()) -> list: combined = shard_rel if root_rel in ("", ".") else root_rel + "/" + shard_rel shard_canon = _canonical_rel(combined) if shard_canon is None: - raise _EscapingShardError( + raise _SnapshotEscapeError( f"weight index maps a shard outside the snapshot: {shard_rel!r}" ) shard = snap.joinpath(*shard_canon.split("/")) @@ -751,14 +767,15 @@ def _evaluate_local_only(model_name: str, load_subdirs = ()) -> "FileSecurityDec try: paths = _cached_pickle_weight_paths(snap, load_subdirs) - except _EscapingShardError as exc: - # A weight index maps a shard outside the snapshot (a "../" traversal). Following it would - # deserialize an out-of-snapshot pickle, so refuse the load outright. + except _SnapshotEscapeError as exc: + # A repo-controlled load path (weight-index shard or Router child) resolves outside the + # snapshot via "../". Following it would deserialize an out-of-snapshot pickle, so refuse + # the load outright. logger.warning("Blocking offline load of '%s': %s; cannot verify.", model_name, exc) return FileSecurityDecision( model_name, True, - reason = "offline; a weight index maps a shard outside the snapshot, cannot verify", + reason = "offline; a cached load path resolves outside the snapshot, cannot verify", ) except OSError: # The cache tree could not be fully enumerated: an undiscovered pickle might exist, so a From f08dc34bc6783d87d73682c6153dd5bb3e3cb0cc Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 05:13:16 +0000 Subject: [PATCH 70/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../tests/test_is_embedding_model_offline.py | 13 +++++++++---- studio/backend/utils/models/model_config.py | 1 - 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 4be2fbf3c0c..d1b6e029a44 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -1394,7 +1394,11 @@ def test_transient_tag_only_fallback_skipped_when_modules_json_present(tmp_path, monkeypatch.setattr(mc, "_persist_embedder", lambda name, commit: None) _tag_only_repo(tmp_path, monkeypatch, "commit_a", modules_json = "[]") - def _transient(model_name, token = None, **kwargs): + def _transient( + model_name, + token = None, + **kwargs, + ): raise OSError("Temporary failure in name resolution") # non-permanent -> transient branch _fake_hf_model_info(monkeypatch, _transient) @@ -1407,8 +1411,7 @@ def test_st_repo_id_candidates_prefers_namespaced_for_non_basic(): # THAT snapshot, so the gate must probe it first (probing the bare name first lets a pickle in # the namespaced dir slip when a bare dir also exists). A slashed id is used verbatim. assert ( - mc._st_repo_id_candidates("all-MiniLM-L6-v2")[0] - == "sentence-transformers/all-MiniLM-L6-v2" + mc._st_repo_id_candidates("all-MiniLM-L6-v2")[0] == "sentence-transformers/all-MiniLM-L6-v2" ) assert mc._st_repo_id_candidates("BAAI/bge-m3") == ["BAAI/bge-m3"] @@ -1416,7 +1419,9 @@ def test_st_repo_id_candidates_prefers_namespaced_for_non_basic(): def test_st_repo_id_candidates_keeps_basic_model_bare(monkeypatch): # A basic transformer model (ORIGINAL_TRANSFORMER_MODELS) is loaded bare by the constructor, so # the bare id is tried first. - monkeypatch.setattr(mc, "_original_transformer_models", lambda: frozenset({"bert-base-uncased"})) + monkeypatch.setattr( + mc, "_original_transformer_models", lambda: frozenset({"bert-base-uncased"}) + ) assert mc._st_repo_id_candidates("bert-base-uncased")[0] == "bert-base-uncased" assert mc._st_repo_id_candidates("some-st-model")[0] == "sentence-transformers/some-st-model" diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 4212ef2d370..06a417eea09 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1717,7 +1717,6 @@ def _original_transformer_models() -> frozenset: if _ORIGINAL_TRANSFORMER_MODELS_CACHE is None: try: from sentence_transformers.util import ORIGINAL_TRANSFORMER_MODELS - _ORIGINAL_TRANSFORMER_MODELS_CACHE = frozenset( str(m).lower() for m in ORIGINAL_TRANSFORMER_MODELS ) From 1c7170d58febd651db98d8f80055e1d9e76526cd Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 21 Jul 2026 05:58:14 +0000 Subject: [PATCH 71/77] Tighten root shard credit, module-path escapes, and weight-set probe order Credit the sharded safetensors index at the snapshot ROOT only when the root is actually loaded through an AutoModel/from_pretrained path. A modules.json root module of a non-Transformer type (StaticEmbedding / WordEmbeddings / Dense) loads via load_torch_weights, which reads pytorch_model.bin and ignores the index, so crediting a root shard index there suppressed a live root pickle and let the offline gate report the cache pickle-free. Recognize the Transformer subclasses CLIPModel and MLMTransformer as index-honoring load roots (they load via from_pretrained), so a sharded-safetensors CLIP/MLM submodule with a legacy pytorch_model.bin sibling is no longer falsely blocked offline. Mirrors the classifier dispatch. Fail closed on an absolute or snapshot-escaping modules.json module path (or load_subdirs entry) instead of silently dropping it: SentenceTransformer resolves such a path outside the snapshot and would deserialize an external pytorch_model.bin the gate cannot scan. On the classifier side, walk the weight set in the exact from_pretrained probe order (model.safetensors, its index, pytorch_model.bin, its index) so a pickle behind a malformed safetensors index is no longer accepted as complete, and restrict shard names to the loader-probed stem/ext pairs so a decoy model-*.bin / pytorch_model-*.safetensors set is not treated as loadable. --- .../tests/test_embedding_verdict_cache.py | 58 ++++++++++++++ .../tests/test_is_embedding_model_offline.py | 41 ++++++++++ studio/backend/utils/models/model_config.py | 60 +++++++-------- .../backend/utils/security/file_security.py | 77 ++++++++++++++++--- 4 files changed, 191 insertions(+), 45 deletions(-) diff --git a/studio/backend/tests/test_embedding_verdict_cache.py b/studio/backend/tests/test_embedding_verdict_cache.py index b21e6d3941c..b2441d970e4 100644 --- a/studio/backend/tests/test_embedding_verdict_cache.py +++ b/studio/backend/tests/test_embedding_verdict_cache.py @@ -288,6 +288,64 @@ def test_sharded_index_credited_in_transformer_submodule(home, tmp_path, monkeyp assert _blocked(monkeypatch, snap) is False +def test_sharded_index_credited_in_clip_and_mlm_submodules(home, tmp_path, monkeypatch): + # CLIPModel and MLMTransformer are Transformer subclasses loaded via AutoModel.from_pretrained, + # which honors model.safetensors.index.json -- so a sharded-safetensors CLIP/MLM submodule with a + # legacy pytorch_model.bin sibling must NOT block offline. + index = json.dumps({"weight_map": {"w": "model-00001-of-00001.safetensors"}}) + for path, mtype, commit in ( + ("0_CLIPModel", "sentence_transformers.models.CLIPModel.CLIPModel", "cclip"), + ("0_MLM", "sentence_transformers.sparse_encoder.models.MLMTransformer.MLMTransformer", "cmlm"), + ): + snap = _snap( + tmp_path, + { + "modules.json": json.dumps([{"path": path, "type": mtype}]), + f"{path}/pytorch_model.bin": b"pickle", + f"{path}/model.safetensors.index.json": index, + f"{path}/model-00001-of-00001.safetensors": b"\0", + }, + commit = commit, + ) + assert _blocked(monkeypatch, snap) is False, mtype + + +def test_root_non_transformer_module_pickle_blocks_despite_shard_index(home, tmp_path, monkeypatch): + # A modules.json ROOT StaticEmbedding module loads via load_torch_weights, which ignores the + # sharded safetensors index and reads pytorch_model.bin -- so a complete ROOT shard index must + # NOT suppress the live root pickle (crediting snap unconditionally would). + index = json.dumps({"weight_map": {"w": "model-00001-of-00001.safetensors"}}) + snap = _snap( + tmp_path, + { + "modules.json": json.dumps( + [{"path": "", "type": "sentence_transformers.models.StaticEmbedding.StaticEmbedding"}] + ), + "pytorch_model.bin": b"pickle", + "model.safetensors.index.json": index, + "model-00001-of-00001.safetensors": b"\0", + }, + commit = "crootse", + ) + assert _blocked(monkeypatch, snap) is True + + +def test_escaping_module_path_blocks(home, tmp_path, monkeypatch): + # A modules.json module path that escapes the snapshot ("../evil") is resolved OUTSIDE the + # snapshot by SentenceTransformer, so the gate fails closed rather than silently dropping it. + snap = _snap( + tmp_path, + { + "modules.json": json.dumps( + [{"path": "../evil", "type": "sentence_transformers.models.Transformer.Transformer"}] + ), + "model.safetensors": b"\0", + }, + commit = "cescmod", + ) + assert _blocked(monkeypatch, snap) is True + + def test_sharded_index_not_credited_in_non_transformer_submodule(home, tmp_path, monkeypatch): # The same layout under a WordEmbeddings module (load_torch_weights, no index) still blocks: the # credit is extended ONLY to Transformer-typed dirs, never to a module that reads the pickle. diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index d1b6e029a44..4f93d03be75 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -1426,6 +1426,47 @@ def test_st_repo_id_candidates_keeps_basic_model_bare(monkeypatch): assert mc._st_repo_id_candidates("some-st-model")[0] == "sentence-transformers/some-st-model" +def test_dir_weight_set_pickle_behind_bad_safetensors_index_is_incomplete(tmp_path): + # #7218: transformers probes the safetensors index BEFORE the pickle bin, so pytorch_model.bin + # behind an incomplete model.safetensors.index.json is NOT complete -- the loader follows the bad + # index and fails instead of falling back to the bin. + (tmp_path / "pytorch_model.bin").write_bytes(b"\0") + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"w": "model-00001-of-00002.safetensors"}}) # a mapped shard absent + ) + names = {p.name for p in tmp_path.iterdir()} + assert mc._dir_weight_set_is_complete(tmp_path, names) is False + + +def test_dir_weight_set_pickle_alone_is_complete(tmp_path): + # Control: a plain pytorch_model.bin with no safetensors index is complete. + (tmp_path / "pytorch_model.bin").write_bytes(b"\0") + names = {p.name for p in tmp_path.iterdir()} + assert mc._dir_weight_set_is_complete(tmp_path, names) is True + + +def test_dir_weight_set_decoy_shard_set_is_incomplete(tmp_path): + # #7218: model-*.bin + model.bin.index.json is a decoy shard set no loader probes (real + # safetensors shards are model-*.safetensors; real pickle shards pytorch_model-*.bin), so it is + # NOT complete. + (tmp_path / "model-00001-of-00001.bin").write_bytes(b"\0") + (tmp_path / "model.bin.index.json").write_text( + json.dumps({"weight_map": {"w": "model-00001-of-00001.bin"}}) + ) + names = {p.name for p in tmp_path.iterdir()} + assert mc._dir_weight_set_is_complete(tmp_path, names) is False + + +def test_dir_weight_set_real_sharded_safetensors_is_complete(tmp_path): + # Control: a real sharded safetensors set still validates after the decoy tightening. + (tmp_path / "model-00001-of-00001.safetensors").write_bytes(b"\0") + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": {"w": "model-00001-of-00001.safetensors"}}) + ) + names = {p.name for p in tmp_path.iterdir()} + assert mc._dir_weight_set_is_complete(tmp_path, names) is True + + def test_online_verdict_stops_applying_when_the_revision_advances(tmp_path, monkeypatch): # A verdict records that the Hub tagged ONE revision an embedder. Once refs/main # advances -- e.g. to a complete but non-embedding Transformer snapshot -- the old diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 06a417eea09..fb7395634c5 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2218,6 +2218,12 @@ def _known_embedder(model_name: str, cache_key: tuple, active_commit: Optional[s # training_args.bin / adapter-only artifact can't pass offline validation then fail on load. _ST_WEIGHT_FILE_RE = re.compile(r"^(model|pytorch_model)(-\d+-of-\d+)?\.(safetensors|bin)$") _ST_SHARD_RE = re.compile(r"^(model|pytorch_model)-(\d+)-of-(\d+)\.(safetensors|bin)$") +# Only these (stem, ext) shard sets exist in transformers / sentence-transformers: safetensors shards +# are named ``model-*.safetensors`` (index ``model.safetensors.index.json``), pickle shards +# ``pytorch_model-*.bin`` (index ``pytorch_model.bin.index.json``). The crossed combinations +# (``model-*.bin`` / ``pytorch_model-*.safetensors`` and their index maps) are decoys no loader ever +# probes, so a cache carrying only one of them is NOT loadable. +_ST_SHARD_EXT_FOR_STEM = {"model": "safetensors", "pytorch_model": "bin"} # A Transformer module also loads an AutoTokenizer, so weights without any tokenizer asset # fail the local_files_only load. Any ONE of these suffices -- a permissive union, so only @@ -2259,7 +2265,8 @@ def _dir_has_complete_torch_weights(names: set) -> bool: shards: dict = {} for name in names: m = _ST_SHARD_RE.match(name) - if m: + # Only a valid stem<->ext shard set counts; a decoy such as model-*.bin is never probed. + if m and _ST_SHARD_EXT_FOR_STEM.get(m.group(1)) == m.group(4): shards.setdefault((m.group(1), m.group(4), int(m.group(3))), set()).add(int(m.group(2))) for (stem, ext, total), indices in shards.items(): index_map = f"{stem}.{ext}.index.json" @@ -2268,19 +2275,6 @@ def _dir_has_complete_torch_weights(names: set) -> bool: return False -# Weight-index files a ``from_pretrained`` / sentence-transformers load follows to locate sharded -# weights. Its ``weight_map`` is AUTHORITATIVE about which shards the load reads, so every mapped -# shard must be present -- the ``model-*-of-*`` FILENAME numbering above is only a proxy for it. -# All (stem, ext) combinations of the base weight set are listed so a ``pytorch_model.bin`` shard -# set is validated the same way as a ``model.safetensors`` one. -_ST_WEIGHT_INDEX_FILES = ( - "model.safetensors.index.json", - "model.bin.index.json", - "pytorch_model.safetensors.index.json", - "pytorch_model.bin.index.json", -) - - def _index_weight_set_complete(index_path: Path) -> Optional[bool]: """Parse a weight-index and confirm every shard its ``weight_map`` references is present, resolved RELATIVE TO the index's own directory (a ``weight_map`` value may name a @@ -2327,27 +2321,25 @@ def _dir_weight_set_is_complete(dir_path: Path, names: set) -> bool: (authoritative, resolved relative to *dir_path*, subdirectory values included); else the filename-numbering shard heuristic (:func:`_dir_has_complete_torch_weights`) as a fallback for an unreadable / stub index. Never raises.""" - # 1. An unsharded base weight the loader actually probes: model.safetensors (safetensors) or - # pytorch_model.bin (pickle). from_pretrained / Module.load_torch_weights never read - # ``model.bin`` or ``pytorch_model.safetensors``, so a cache holding only one of those is NOT - # loadable (the security gate treats ``pytorch_model.safetensors`` as a decoy for the same - # reason) -- classifying it as complete accepts a snapshot the local_files_only load then 409s. - if "model.safetensors" in names or "pytorch_model.bin" in names: + # Mirror the from_pretrained / sentence-transformers probe ORDER exactly: the FIRST present weight + # source wins, and a present safetensors weight/index makes the loader IGNORE a sibling pickle. + # 1. unsharded model.safetensors + # 2. sharded model.safetensors.index.json (present index is authoritative -> must fully resolve) + # 3. unsharded pytorch_model.bin + # 4. sharded pytorch_model.bin.index.json (present index is authoritative -> must fully resolve) + # from_pretrained / Module.load_torch_weights never read ``model.bin`` or + # ``pytorch_model.safetensors`` (nor their decoy index maps), so those do not count. A pickle + # sibling behind a malformed/incomplete safetensors index does NOT rescue the load -- the loader + # follows the bad index and fails rather than falling back to the bin. + if "model.safetensors" in names: return True - # 2. A PRESENT weight-index is authoritative: transformers / sentence-transformers open and parse - # it to locate the sharded weights, so if any index file exists it must fully resolve. One that - # maps a complete shard set validates; one missing a mapped shard OR unreadable / lacking a - # ``weight_map`` (a stub or truncated write) makes the local_files_only load fail -- so do NOT - # fall back to the filename-numbering heuristic when an index file is present. - saw_index = False - for index_name in _ST_WEIGHT_INDEX_FILES: - if index_name in names: - saw_index = True - if _index_weight_set_complete(dir_path / index_name) is True: - return True - if saw_index: - return False - # 3. No index file at all: the filename-numbering shard heuristic on basenames. + if "model.safetensors.index.json" in names: + return _index_weight_set_complete(dir_path / "model.safetensors.index.json") is True + if "pytorch_model.bin" in names: + return True + if "pytorch_model.bin.index.json" in names: + return _index_weight_set_complete(dir_path / "pytorch_model.bin.index.json") is True + # No loader-probed base weight or index by exact name: the filename-numbering shard heuristic. return _dir_has_complete_torch_weights(names) diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 032653e2e3a..360d3d792cd 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -470,15 +470,36 @@ def _router_child_dirs(root, snap) -> set: return children +def _load_root_or_raise(snap, raw_rel: str): + """Resolve a repo-declared load path (a ``modules.json`` module ``path`` or a ``load_subdirs`` + entry) to an in-snapshot directory, or None for a ROOT-like path (empty / ``.`` -- loads from + *snap*, already a root). Raises :class:`_SnapshotEscapeError` for an ABSOLUTE or + snapshot-ESCAPING path: ``SentenceTransformer`` resolves such a declared path OUTSIDE the + snapshot and would deserialize an external ``pytorch_model.bin`` the gate cannot scan, so the + cache is unverifiable and must fail closed rather than silently drop the module.""" + root = _canonical_load_dir(snap, raw_rel) + if root is not None: + return root + import posixpath + + norm = posixpath.normpath(_normalize_repo_path(raw_rel).strip("/")) + if norm in ("", "."): + return None # a root module loads from snap, already a root + raise _SnapshotEscapeError(f"declared load path escapes the snapshot: {raw_rel!r}") + + def _st_load_roots(snap, load_subdirs = ()) -> set: """Directories a load opens ``from_pretrained`` on: the snapshot root, every module path ``modules.json`` declares, and each passed-in ``load_subdirs`` entry. A SentenceTransformer module can load from a directory without ``config.json`` -- e.g. a ``0_WordEmbeddings/`` module with ``wordembedding_config.json`` + ``pytorch_model.bin`` -- so a pickle there is - still deserialized and must be treated as a load root.""" + still deserialized and must be treated as a load root. + + An absolute or snapshot-escaping declared path fails closed (:class:`_SnapshotEscapeError`): the + loader resolves it outside the snapshot, so the cache cannot be verified.""" roots = {snap} for subdir in load_subdirs or (): - root = _canonical_load_dir(snap, str(subdir)) + root = _load_root_or_raise(snap, str(subdir)) if root is not None: roots.add(root) try: @@ -489,7 +510,7 @@ def _st_load_roots(snap, load_subdirs = ()) -> set: if isinstance(modules, list): for module in modules: if isinstance(module, dict): - root = _canonical_load_dir(snap, str(module.get("path") or "")) + root = _load_root_or_raise(snap, str(module.get("path") or "")) if root is not None: roots.add(root) # A Router/Asym module declares its child sub-modules in router_config.json, not modules.json, @@ -509,12 +530,15 @@ def _st_load_roots(snap, load_subdirs = ()) -> set: def _is_transformer_module_type(type_str) -> bool: - """True for the SentenceTransformer ``Transformer`` module class, whose loader is - ``AutoModel.from_pretrained`` and therefore honors a ``model.safetensors.index.json`` shard set. - ``Dense`` / ``WordEmbeddings`` / ``StaticEmbedding`` read a single flat weight file through + """True for a Transformer-SHAPED SentenceTransformer module -- ``Transformer`` and its subclasses + ``MLMTransformer`` / ``CLIPModel``, all of which load via ``AutoModel.from_pretrained`` and so + honor a ``model.safetensors.index.json`` shard set. Mirrors the classifier's + ``_ST_TRANSFORMER_SHAPED_MODULE_NAMES`` + ``"transformer" in cls`` dispatch. ``Dense`` / + ``WordEmbeddings`` / ``StaticEmbedding`` read a single flat weight file through ``Module.load_torch_weights`` and never consult an index, so their sharded-index credit stays gated off.""" - return str(type_str).rsplit(".", 1)[-1] == "Transformer" + cls = str(type_str).rsplit(".", 1)[-1].strip().lower() + return "transformer" in cls or cls == "clipmodel" def _transformer_load_roots(snap) -> set: @@ -542,6 +566,32 @@ def _transformer_load_roots(snap) -> set: return roots +def _root_honors_sharded_index(snap) -> bool: + """Whether the snapshot ROOT is loaded through an index-honoring path + (``AutoModel.from_pretrained``). True with no / unreadable ``modules.json`` (a plain + ``from_pretrained`` root, or one that loads no root weights at all), or when a ROOT module is + Transformer-shaped. False when ``modules.json`` declares a ROOT module of a NON-Transformer type + (``StaticEmbedding`` / ``WordEmbeddings`` / ``Dense`` ...), whose ``load()`` reads + ``pytorch_model.bin`` and ignores the sharded index -- crediting a root shard index would then + suppress a LIVE root pickle.""" + import json + import posixpath + + try: + modules = json.loads((snap / "modules.json").read_text(encoding = "utf-8")) + except (OSError, ValueError): + return True + if not isinstance(modules, list): + return True + for module in modules: + if not isinstance(module, dict): + continue + norm = posixpath.normpath(_normalize_repo_path(str(module.get("path") or "")).strip("/")) + if norm in ("", ".") and not _is_transformer_module_type(module.get("type")): + return False # a non-Transformer ROOT module reads the root pickle, ignoring the index + return True + + class _SnapshotEscapeError(OSError): """A repo-controlled load path -- a weight-index ``weight_map`` value or a ``Router`` ``router_config.json`` child -- resolves OUTSIDE the snapshot via a ``../`` traversal. @@ -568,10 +618,15 @@ def _cached_pickle_weight_paths(snap, load_subdirs = ()) -> list: Raises ``OSError`` if the snapshot tree cannot be enumerated (``rglob`` failure), so the offline caller fails CLOSED rather than treat an unreadable cache as pickle-free.""" roots = _st_load_roots(snap, load_subdirs) - # Dirs whose loader (AutoModel.from_pretrained) honors a sharded safetensors index: the snapshot - # root plus each Transformer-typed module subdir. Elsewhere a sharded index is NOT a substitute - # for a live pickle (see _dir_has_loadable_safetensors). - index_honoring = _transformer_load_roots(snap) | {snap} + # Dirs whose loader (AutoModel.from_pretrained) honors a sharded safetensors index: each + # Transformer-typed module subdir, plus the snapshot root -- but ONLY when the root is actually + # loaded through a Transformer/from_pretrained path. A modules.json root module of a + # non-Transformer type reads pytorch_model.bin and ignores the index, so crediting the root shard + # index there would suppress a live root pickle. Elsewhere a sharded index is NOT a substitute for + # a live pickle (see _dir_has_loadable_safetensors). + index_honoring = _transformer_load_roots(snap) + if _root_honors_sharded_index(snap): + index_honoring = index_honoring | {snap} by_dir_pickle: dict = {} # directory -> [Path] (EVERY case variant, not last-wins) by_dir_files: dict = {} # directory -> {lower-name: Path} (safetensors credit; last-wins is safe) # A whole-tree rglob failure propagates (fail-closed); a single unstattable entry is skipped. From 2571e1b7054505ef8dee8e9d2669bf33ff7b6931 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 05:59:47 +0000 Subject: [PATCH 72/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../tests/test_embedding_verdict_cache.py | 20 ++++++++++++++++--- .../tests/test_is_embedding_model_offline.py | 4 +++- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/studio/backend/tests/test_embedding_verdict_cache.py b/studio/backend/tests/test_embedding_verdict_cache.py index b2441d970e4..7195b144ee4 100644 --- a/studio/backend/tests/test_embedding_verdict_cache.py +++ b/studio/backend/tests/test_embedding_verdict_cache.py @@ -295,7 +295,11 @@ def test_sharded_index_credited_in_clip_and_mlm_submodules(home, tmp_path, monke index = json.dumps({"weight_map": {"w": "model-00001-of-00001.safetensors"}}) for path, mtype, commit in ( ("0_CLIPModel", "sentence_transformers.models.CLIPModel.CLIPModel", "cclip"), - ("0_MLM", "sentence_transformers.sparse_encoder.models.MLMTransformer.MLMTransformer", "cmlm"), + ( + "0_MLM", + "sentence_transformers.sparse_encoder.models.MLMTransformer.MLMTransformer", + "cmlm", + ), ): snap = _snap( tmp_path, @@ -319,7 +323,12 @@ def test_root_non_transformer_module_pickle_blocks_despite_shard_index(home, tmp tmp_path, { "modules.json": json.dumps( - [{"path": "", "type": "sentence_transformers.models.StaticEmbedding.StaticEmbedding"}] + [ + { + "path": "", + "type": "sentence_transformers.models.StaticEmbedding.StaticEmbedding", + } + ] ), "pytorch_model.bin": b"pickle", "model.safetensors.index.json": index, @@ -337,7 +346,12 @@ def test_escaping_module_path_blocks(home, tmp_path, monkeypatch): tmp_path, { "modules.json": json.dumps( - [{"path": "../evil", "type": "sentence_transformers.models.Transformer.Transformer"}] + [ + { + "path": "../evil", + "type": "sentence_transformers.models.Transformer.Transformer", + } + ] ), "model.safetensors": b"\0", }, diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 4f93d03be75..74f4b640ec6 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -1432,7 +1432,9 @@ def test_dir_weight_set_pickle_behind_bad_safetensors_index_is_incomplete(tmp_pa # index and fails instead of falling back to the bin. (tmp_path / "pytorch_model.bin").write_bytes(b"\0") (tmp_path / "model.safetensors.index.json").write_text( - json.dumps({"weight_map": {"w": "model-00001-of-00002.safetensors"}}) # a mapped shard absent + json.dumps( + {"weight_map": {"w": "model-00001-of-00002.safetensors"}} + ) # a mapped shard absent ) names = {p.name for p in tmp_path.iterdir()} assert mc._dir_weight_set_is_complete(tmp_path, names) is False From b76eda18c20d2f428b44fadc0391eeea3e84c52e Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 21 Jul 2026 10:03:13 +0000 Subject: [PATCH 73/77] Restore scripts/verify_import_hoist.py to main The offline embedding cache fix does not depend on the __all__ scope handling that had accumulated in this linter, so revert the file to its main version and keep the PR focused on the feature. The feature modules still pass the existing import hoist check unchanged. --- scripts/verify_import_hoist.py | 346 --------------------------------- 1 file changed, 346 deletions(-) diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index c0e83b0e9de..22a21a2ebce 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -447,218 +447,6 @@ def _legb_chain(scope: Scope) -> list[Scope]: # ---------------------------------------------------------------- analysis -def _all_mutator_entries(method: str, args: list) -> tuple[set[str], bool]: - """String entries a ``__all__.(args)`` call adds, and whether it was fully - readable. ``append(x)`` adds one string constant; ``extend([...])`` adds a list literal's - string constants. A dynamic arg, or any other mutator (``insert``/``remove``/...), is not - statically readable -> caller marks the export set opaque.""" - entries: set[str] = set() - if method == "append" and len(args) == 1: - arg = args[0] - if isinstance(arg, ast.Constant) and isinstance(arg.value, str): - entries.add(arg.value) - return entries, True - return entries, False - if ( - method == "extend" - and len(args) == 1 - and isinstance(args[0], (ast.List, ast.Tuple, ast.Set)) - ): - readable = True - for elt in args[0].elts: - if isinstance(elt, ast.Constant) and isinstance(elt.value, str): - entries.add(elt.value) - else: - readable = False - return entries, readable - return entries, False # unknown mutator or dynamic arg -> opaque - - -def _stmt_binds_all(node: ast.AST) -> bool: - """True when *node* is a statement that BINDS or MUTATES ``__all__`` -- an assignment / - annotated-assign / augmented-assign whose target names ``__all__`` (directly or nested in a - destructuring / item / attr target), or an ``__all__.(...)`` expression call.""" - if isinstance(node, ast.Assign): - targets = node.targets - elif isinstance(node, (ast.AnnAssign, ast.AugAssign)): - targets = [node.target] - elif isinstance(node, ast.Expr): - call = node.value - return ( - isinstance(call, ast.Call) - and isinstance(call.func, ast.Attribute) - and isinstance(call.func.value, ast.Name) - and call.func.value.id == "__all__" - ) - else: - return False - return any( - isinstance(sub, ast.Name) and sub.id == "__all__" for t in targets for sub in ast.walk(t) - ) - - -def _declares_global_all(func: ast.AST) -> bool: - """True when *func*'s OWN body declares ``global __all__`` (so an ``__all__`` assignment in it - writes the module export set). Does not descend into nested functions/classes -- their - ``global`` statements belong to their own scope.""" - stack = list(ast.iter_child_nodes(func)) - while stack: - node = stack.pop() - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): - continue # a nested scope; its `global` decls are its own - if isinstance(node, ast.Global) and "__all__" in node.names: - return True - stack.extend(ast.iter_child_nodes(node)) - return False - - -def _binds_all_in_own_scope(func: ast.AST) -> bool: - """True when *func*'s OWN body (not a nested function/class) binds ``__all__``. Paired with - :func:`_declares_global_all`, this identifies a function that writes the MODULE export set. A - nested scope's ``__all__`` is a local of THAT scope and cannot reach module scope, so it is not - counted -- otherwise ``def f(): global __all__`` with an inner ``def g(): __all__ = [...]`` would - wrongly mark the module ``__all__`` opaque and mask a genuinely unused hoisted import.""" - stack = list(ast.iter_child_nodes(func)) - while stack: - node = stack.pop() - if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)): - continue # a nested scope; its `__all__` binds are its own - if _stmt_binds_all(node): - return True - stack.extend(ast.iter_child_nodes(node)) - return False - - -def _all_bound_outside_module_body(tree: ast.Module) -> bool: - """True when the module's runtime ``__all__`` may be mutated somewhere a static read cannot - replay: a MODULE-LEVEL binding nested in a conditional (``if`` / ``try`` / ``for`` / ``while`` - / ``with`` / ``match``, whose final value is not the replayed top-level sequence), or an - ``__all__`` assignment inside a function that declares ``global __all__``. An ``__all__`` bound - as a LOCAL in a nested function or class body (no ``global __all__``) CANNOT change the module - export set, so it is ignored -- otherwise a stray local named ``__all__`` would mask a genuinely - unused hoisted import. The caller marks ``__all__`` opaque on True; this only ever ADDS opacity - at module scope, so it cannot create a false blocker.""" - top_level = {id(node) for node in tree.body} - - def _walk(node: ast.AST) -> bool: - # Descends only module-level statements/blocks; nested scopes are handled without recursing. - for child in ast.iter_child_nodes(node): - if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)): - # A __all__ assignment inside the function reaches module scope ONLY via - # `global __all__` AND a bind in the function's OWN scope; a bind nested in a - # further inner scope is that scope's local and cannot taint the export set. - if _declares_global_all(child) and _binds_all_in_own_scope(child): - return True - # else: a purely-local __all__ -> ignore this function entirely. - elif isinstance(child, ast.ClassDef): - continue # a class-body __all__ is a class attribute, not the module's -> ignore - else: - if id(child) not in top_level and _stmt_binds_all(child): - return True # module-level conditional mutation (not the replayed top sequence) - if _walk(child): - return True - return False - - return _walk(tree) - - -def _collect_dunder_all(tree: ast.Module) -> tuple[set[str], bool]: - """The FINAL module-level ``__all__`` string entries plus an ``opaque`` flag - (``= [...]``, ``+= [...]``, or an annotated assign). - - A name listed in ``__all__`` is a public re-export, which is a real use of the - import that binds it. ``__all__`` entries are string constants, not ``Name`` - loads, so the load-based use scan never sees them -- without this a package - ``__init__`` that adds ``from .x import y`` purely to re-export ``y`` looks - like an unused hoist and trips ``HOISTED-IMPORT-UNUSED``. - - Only the final value counts, so assignments are replayed in order: a plain ``=`` - REPLACES the list and ``+=`` extends it. Unioning every assignment instead would - let ``__all__ = ["y"]`` followed by ``__all__ = []`` still mark ``y`` used, so a - genuinely unused hoist would slip through. - - ``opaque`` is True when the FINAL ``__all__`` value is not statically exhaustive: - it is rebound or extended by a value we cannot read (a call, name, comprehension, - spread). The caller must then NOT treat a name's absence from the static set as - proof the import is unexported -- a re-export supplied dynamically - (``__all__ += _exports()``) would otherwise trip ``HOISTED-IMPORT-UNUSED``. A - replacing ``=`` resets opacity (it discards the prior list AND any prior dynamic - part), so a later readable ``__all__ = [...]`` clears an earlier opaque ``+=``; an - ``+=`` extends, so prior opacity persists. A bare ``__all__: list[str]`` annotation - has no runtime value and is skipped entirely. ``__all__.append("X")`` / - ``.extend([...])`` are runtime extenders read the same way as ``+=``; any other - ``__all__`` method call (``insert``/``remove``/...) is opaque. ``__all__`` bound only - through a destructuring target (``__all__, meta = [...], v``) or an item/attr target - is opaque too -- its value cannot be mapped statically, so the set is not exhaustive. - ``__all__`` bound / mutated in a module-level ``if`` / ``try`` / loop / ``with`` / ``match`` (or - a deeper scope) is opaque -- a conditional value cannot be replayed from ``tree.body`` alone. - """ - names: set[str] = set() - opaque = False - for node in tree.body: - if isinstance(node, ast.Expr): - # __all__.append("X") / .extend([...]) / other mutators: a runtime EXTEND of the - # export list, never a replace. Read append/extend string args; anything else opaque. - call = node.value - if ( - isinstance(call, ast.Call) - and isinstance(call.func, ast.Attribute) - and isinstance(call.func.value, ast.Name) - and call.func.value.id == "__all__" - ): - entries, readable = _all_mutator_entries(call.func.attr, call.args) - names |= entries - opaque = opaque or not readable - continue - if isinstance(node, ast.Assign): - targets, replaces = node.targets, True - elif isinstance(node, ast.AnnAssign): - if node.value is None: - continue # bare `__all__: list[str]` annotation: no runtime value - targets, replaces = [node.target], True - elif isinstance(node, ast.AugAssign): - targets, replaces = [node.target], False - else: - continue - # A direct `__all__ = [...]` / `+= [...]` target lets us read the value below. `__all__` - # reached only through a destructuring target (`__all__, meta = [...], v` -> an ast.Tuple) - # or an item/attr target cannot be mapped to its value statically, so mark the export set - # opaque instead of silently missing the binding and false-positiving HOISTED-IMPORT-UNUSED. - if not any(isinstance(t, ast.Name) and t.id == "__all__" for t in targets): - if any( - isinstance(sub, ast.Name) and sub.id == "__all__" - for t in targets - for sub in ast.walk(t) - ): - opaque = True - continue - value = node.value - entries: set[str] = set() - readable = True # every element is a statically-known string constant - if isinstance(value, (ast.List, ast.Tuple, ast.Set)): - for elt in value.elts: - if isinstance(elt, ast.Constant) and isinstance(elt.value, str): - entries.add(elt.value) - else: - readable = False # a spread / computed element: unknown extra exports - else: - readable = False # a call, a name, a comprehension: contents unknown - if replaces: - # A replace discards the prior list and its opacity: an unreadable rebind - # is opaque, a readable one is fully known again. - names, opaque = entries, not readable - else: - # An extend keeps the prior list, so prior opacity persists; an unreadable - # extend adds names we cannot see. - names |= entries - opaque = opaque or not readable - if _all_bound_outside_module_body(tree): - # A conditional / nested `__all__` mutation (module-level if/try/for/while/with/match, or a - # deeper scope) is not in the replayed top-level sequence, so its contribution is unknown. - opaque = True - return names, opaque - - def _analyze(src: str): tree = ast.parse(src) b = _Builder() @@ -688,23 +476,6 @@ def _analyze(src: str): for n, bs in module.bindings.items() if any(x.kind in ("import", "importfrom") for x in bs) } - # Re-exports count as uses: a name listed in module-level __all__ that is bound - # by a module import is deliberately exported, not a dangling hoist. Fold its - # targets into the used set so HOISTED-IMPORT-UNUSED does not fire on a - # legitimately-added `from .x import y` in a package __init__. When __all__ is - # opaque (a dynamically-computed part), the static list is not known to be - # exhaustive, so credit EVERY module import rather than risk flagging a - # dynamically-exported hoist as unused. - _all_names, _all_opaque = _collect_dunder_all(tree) - _reexported = ( - module_imports.values() - if _all_opaque - else (_bs for _n in _all_names if (_bs := module_imports.get(_n))) - ) - for _bs in _reexported: - targets_by_scope.setdefault(module.qualname, set()).update( - x.target for x in _bs if x.target - ) module_dup = { n for n, bs in module.bindings.items() @@ -935,123 +706,6 @@ def used_targets(analysis) -> set[str]: "import os\nimport sys\ndef f(x):\n return x._b + sys.argv[0]\n", None, ), - "reexport_in_all_is_used": ( - # a new re-export added to a package __init__ (name in __all__, no load) is - # a deliberate export, NOT a botched hoist -> must not block - 'from pkg import a\n__all__ = ["a"]\n', - 'from pkg import a\nfrom pkg import b\n__all__ = ["a", "b"]\n', - None, - ), - "unused_import_not_in_all_still_blocks": ( - # the fix is precise: a newly-added module import that is neither loaded nor - # listed in __all__ is still a dangling/unused hoist - 'from pkg import a\n__all__ = ["a"]\n', - 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n', - "BLOCKER", - ), - "reassigned_all_drops_the_reexport": ( - # only the FINAL __all__ exports anything: a later plain "=" REPLACES the list, - # so b is not re-exported and its import is a genuine unused hoist. Unioning - # every assignment would have marked it used and let this through. - 'from pkg import a\n__all__ = ["a"]\n', - 'from pkg import a\nfrom pkg import b\n__all__ = ["a", "b"]\n__all__ = ["a"]\n', - "BLOCKER", - ), - "augmented_all_extends_the_reexport": ( - # "+=" extends rather than replaces, so b IS re-exported -> must not block - 'from pkg import a\n__all__ = ["a"]\n', - 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n__all__ += ["b"]\n', - None, - ), - "unreadable_all_keeps_earlier_reexports": ( - # rebound to something not statically readable: contents unknown, so keep what - # we had rather than flag a real re-export as unused - 'from pkg import a\n__all__ = ["a"]\n', - 'from pkg import a\nfrom pkg import b\n__all__ = ["a", "b"]\n__all__ = sorted(__all__)\n', - None, - ), - "unreadable_augmented_all_keeps_the_reexport": ( - # "+=" extended by a value we cannot read statically: the augmented contents are - # unknown, so a dynamically-supplied re-export must not be flagged as unused - # (mirrors the unreadable "=" case, which also preserves uncertainty) - 'from pkg import a\n__all__ = ["a"]\n', - 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n__all__ += sorted(["b"])\n', - None, - ), - "destructured_all_is_opaque_keeps_reexport": ( - # __all__ bound only through a tuple-unpacking target cannot be mapped to its value - # statically, so it is opaque: a new import re-exported only through it must not be - # flagged as an unused hoist (mirrors the unreadable "+=" case) - 'from pkg import a\n__all__, meta = ["a"], 1\n', - 'from pkg import a\nfrom pkg import b\n__all__, meta = ["a", "b"], 1\n', - None, - ), - "conditional_all_mutation_is_opaque_keeps_reexport": ( - # __all__ extended inside a module-level `if` is NOT in the replayed top-level sequence, - # so the final value is unknown -> opaque: a re-export added only there must not be flagged - # as an unused hoist (the collector previously scanned only tree.body and missed it) - 'from pkg import a\n__all__ = ["a"]\n', - 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\nif True:\n __all__ += ["b"]\n', - None, - ), - "nested_local_all_does_not_mask_unused_hoist": ( - # a __all__ bound as a LOCAL inside a nested function cannot change the module export set, - # so it must NOT make the module __all__ opaque -- a genuinely unused hoisted import is - # still a bad hoist and must block - 'from pkg import a\n__all__ = ["a"]\n', - 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n' - 'def f():\n __all__ = ["local"]\n return __all__\n', - "BLOCKER", - ), - "global_all_mutation_in_function_is_opaque": ( - # a function that declares `global __all__` and assigns it DOES change the module export - # set at runtime -> opaque, so a re-export supplied that way must not be flagged as unused - 'from pkg import a\n__all__ = ["a"]\n', - 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n' - 'def register():\n global __all__\n __all__ = __all__ + ["b"]\n', - None, - ), - "nested_global_all_bind_does_not_mask_unused_hoist": ( - # `global __all__` in the OUTER function but the only `__all__` bind is in a further inner - # scope: that bind is the inner scope's local and never reaches module scope, so it must NOT - # mark the module __all__ opaque -- the unused hoisted import stays a genuine bad hoist - 'from pkg import a\n__all__ = ["a"]\n', - 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n' - 'def f():\n global __all__\n def g():\n __all__ = ["x"]\n return __all__\n return g\n', - "BLOCKER", - ), - "readable_reassign_resets_opacity": ( - # an opaque "+=" then a readable replacing "=": the final __all__ is empty, so a - # hoisted unused import is a genuine bad hoist -- the reset must un-stick opacity - 'from pkg import a\n__all__ = ["a"]\n', - 'from pkg import a\nfrom pkg import b\n__all__ += sorted(["x"])\n__all__ = []\n', - "BLOCKER", - ), - "annotation_only_all_is_not_opaque": ( - # a bare `__all__: list[str]` annotation has no runtime value: it must be skipped, - # not treated as an unreadable assignment, so a new unused hoist is still caught - "from pkg import a\n__all__: list[str]\n", - "from pkg import a\nfrom pkg import b\n__all__: list[str]\n", - "BLOCKER", - ), - "all_append_reexport_is_used": ( - # __all__.append("b") re-exports b -> must not flag the hoisted import as unused - 'from pkg import a\n__all__ = ["a"]\n', - 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n__all__.append("b")\n', - None, - ), - "all_extend_reexport_is_used": ( - # __all__.extend(["b"]) re-exports b as well - 'from pkg import a\n__all__ = ["a"]\n', - 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n__all__.extend(["b"])\n', - None, - ), - "all_append_unrelated_still_blocks": ( - # a readable append of a DIFFERENT name does not credit an unrelated unused hoist - 'from pkg import a\n__all__ = ["a"]\n', - 'from pkg import a\nfrom pkg import b\n__all__ = ["a"]\n__all__.append("a")\n', - "BLOCKER", - ), } From cadd262048393896cc71799f98d3ffe9099e3850 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Tue, 21 Jul 2026 10:03:13 +0000 Subject: [PATCH 74/77] Reuse a shared HF cache skeleton in the offline classification tests Extract _mk_repo and _activate helpers for the repeated snapshot cache setup that every per-type builder duplicated, and fold the two StaticEmbedding missing-asset cases into one parametrized test. Same 125 collected items, all still passing. --- .../tests/test_is_embedding_model_offline.py | 212 +++++------------- 1 file changed, 61 insertions(+), 151 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 74f4b640ec6..089e64631a6 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -194,6 +194,23 @@ def _fake_hf_cache(monkeypatch, root): monkeypatch.setitem(sys.modules, "huggingface_hub.constants", fake.constants) +def _mk_repo(tmp_path, repo_id = "org/model", commit = "aaa"): + """(repo_dir, snapshot_dir) for a fresh HF cache repo under ``tmp_path/hf``; the snapshot dir + is created. Shared skeleton for the per-type repo builders below.""" + repo = tmp_path / "hf" / f"models--{repo_id.replace('/', '--')}" + snap = repo / "snapshots" / commit + snap.mkdir(parents = True) + return repo, snap + + +def _activate(repo, monkeypatch, commit = "aaa"): + """Write refs/main and point the HF/ST cache at *repo*'s root, ignoring SENTENCE_TRANSFORMERS_HOME.""" + (repo / "refs").mkdir(parents = True, exist_ok = True) + (repo / "refs" / "main").write_text(commit) + _fake_hf_cache(monkeypatch, repo.parent) + monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + + def _st_snapshot( root, repo_dir, @@ -263,19 +280,13 @@ def test_marker_only_snapshot_is_not_loadable(tmp_path, monkeypatch): def test_marker_onnx_only_snapshot_is_not_loadable(tmp_path, monkeypatch): # Marker + config but only an ONNX export: the Torch backend needs safetensors/bin. - hf_root = tmp_path / "hf" - repo = hf_root / "models--org--model" - snap = repo / "snapshots" / "aaa" - snap.mkdir(parents = True) + repo, snap = _mk_repo(tmp_path, "org/model", "aaa") (snap / "modules.json").write_text("[]") (snap / "config.json").write_text("{}") (snap / "tokenizer.json").write_text("{}") (snap / "model.onnx").write_bytes(b"\0") # refs/main so the probe reaches the weight check. - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text("aaa") - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch) assert mc._embedding_marker_in_hf_cache("org/model") is False @@ -320,10 +331,7 @@ def _cache_repo_with_files( unrealistic snapshot (#7218). Root-level files declare a single Transformer at path ``""`` -- the save_in_root layout a plain load reads from the snapshot root -- rather than an empty ``[]`` list, which builds ZERO modules and is not a loadable ST model (#7218 P4).""" - hf_root = tmp_path / "hf" - repo = hf_root / "models--org--model" - snap = repo / "snapshots" / commit - snap.mkdir(parents = True) + repo, snap = _mk_repo(tmp_path, commit = commit) (snap / "config.json").write_text("{}") if tokenizer: (snap / "tokenizer.json").write_text("{}") @@ -347,10 +355,7 @@ def _cache_repo_with_files( ) else: (snap / "modules.json").write_text(_ROOT_TRANSFORMER_MODULES) - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text(commit) - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch, commit) def test_marker_rejects_non_base_weight_bins(tmp_path, monkeypatch): @@ -436,9 +441,7 @@ def test_marker_rejects_complete_dir_at_an_undeclared_path(tmp_path, monkeypatch # UNDECLARED sibling dir (stray_complete/) the loader never opens. Judging any complete # directory would accept this snapshot and then 409 at the first local_files_only load, so # it must be restricted to the declared load roots and read as NOT loadable (#7218 P2). - hf_root = tmp_path / "hf" - repo = hf_root / "models--org--model" - snap = repo / "snapshots" / "aaa" + repo, snap = _mk_repo(tmp_path, "org/model", "aaa") (snap / "0_Transformer").mkdir(parents = True) (snap / "modules.json").write_text( _modules_json(("0", "0_Transformer", "sentence_transformers.models.Transformer")) @@ -449,10 +452,7 @@ def test_marker_rejects_complete_dir_at_an_undeclared_path(tmp_path, monkeypatch (stray / "config.json").write_text("{}") (stray / "tokenizer.json").write_text("{}") (stray / "model.safetensors").write_bytes(b"\0") - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text("aaa") - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch) assert mc._embedding_marker_in_hf_cache("org/model") is False @@ -460,9 +460,7 @@ def test_marker_accepts_complete_dir_at_the_declared_module_path(tmp_path, monke # Companion to the rejection above: when the complete config+tokenizer+weights sit at the # DECLARED module path (0_Transformer/), the loader opens exactly that directory, so a # normal Transformer model must still be recognized (#7218 P2). - hf_root = tmp_path / "hf" - repo = hf_root / "models--org--model" - snap = repo / "snapshots" / "aaa" + repo, snap = _mk_repo(tmp_path, "org/model", "aaa") (snap / "0_Transformer").mkdir(parents = True) (snap / "modules.json").write_text( _modules_json(("0", "0_Transformer", "sentence_transformers.models.Transformer")) @@ -470,10 +468,7 @@ def test_marker_accepts_complete_dir_at_the_declared_module_path(tmp_path, monke (snap / "0_Transformer" / "config.json").write_text("{}") (snap / "0_Transformer" / "tokenizer.json").write_text("{}") (snap / "0_Transformer" / "model.safetensors").write_bytes(b"\0") - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text("aaa") - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch) assert mc._embedding_marker_in_hf_cache("org/model") is True @@ -523,9 +518,7 @@ def _wordembeddings_repo( ``1_Pooling`` module dir holding ``config.json``. It carries NO HF ``config.json`` / tokenizer, so the Transformer-shaped weight check alone misclassifies it as non-embedding and the settings route 409s it (#7218).""" - hf_root = tmp_path / "hf" - repo = hf_root / f"models--{repo_id.replace('/', '--')}" - snap = repo / "snapshots" / commit + repo, snap = _mk_repo(tmp_path, repo_id, commit) (snap / "0_WordEmbeddings").mkdir(parents = True) (snap / "modules.json").write_text( _modules_json( @@ -543,10 +536,7 @@ def _wordembeddings_repo( (snap / "0_WordEmbeddings" / weight_file).write_bytes(b"\0") (snap / "1_Pooling").mkdir(parents = True) (snap / "1_Pooling" / "config.json").write_text("{}") - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text(commit) - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch, commit) return snap @@ -628,10 +618,7 @@ def _router_repo( weights of its own -- ``Router.load()`` reads router_config.json and loads each child from its subdir -- so the Transformer-shaped root check alone misclassifies it (#7218). ``complete=False`` strips the Transformer child's weights so the load would fail.""" - hf_root = tmp_path / "hf" - repo = hf_root / f"models--{repo_id.replace('/', '--')}" - snap = repo / "snapshots" / commit - snap.mkdir(parents = True) + repo, snap = _mk_repo(tmp_path, repo_id, commit) router_dir = snap if root_path in ("", ".") else snap / root_path router_dir.mkdir(parents = True, exist_ok = True) (snap / "modules.json").write_text( @@ -658,10 +645,7 @@ def _router_repo( pooling = router_dir / "query_1_Pooling" pooling.mkdir(parents = True) (pooling / "config.json").write_text("{}") - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text(commit) - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch, commit) return snap @@ -699,20 +683,14 @@ def test_marker_rejects_self_referential_router_without_recursing(tmp_path, monk # must be rejected, not recursed into forever: the child path resolves back to the same dir, so # without the self-reference guard _router_dir_is_loadable would re-enter until RecursionError # and turn is_embedding_model into a 500. It must instead return False gracefully (never raise). - hf_root = tmp_path / "hf" - repo = hf_root / "models--org--self-router" - snap = repo / "snapshots" / "aaa" - snap.mkdir(parents = True) + repo, snap = _mk_repo(tmp_path, "org/self-router", "aaa") (snap / "modules.json").write_text( _modules_json(("0", "", "sentence_transformers.models.Router")) ) (snap / "router_config.json").write_text( json.dumps({"types": {".": "sentence_transformers.models.Router"}}) ) - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text("aaa") - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch) assert mc._embedding_marker_in_hf_cache("org/self-router") is False @@ -738,9 +716,7 @@ def _staticembedding_repo( ``tokenizer.json`` + ``model.safetensors`` -- exactly what ``StaticEmbedding.load()`` reads, which writes NO config. It carries NO HF/module ``config.json``, so the config-gated non-Transformer path alone misclassifies it and the settings route 409s it offline (#7218).""" - hf_root = tmp_path / "hf" - repo = hf_root / f"models--{repo_id.replace('/', '--')}" - snap = repo / "snapshots" / commit + repo, snap = _mk_repo(tmp_path, repo_id, commit) (snap / "0_StaticEmbedding").mkdir(parents = True) (snap / "modules.json").write_text( _modules_json( @@ -752,10 +728,7 @@ def _staticembedding_repo( (snap / "0_StaticEmbedding" / "tokenizer.json").write_text("{}") if include_weights: (snap / "0_StaticEmbedding" / weight_file).write_bytes(b"\0") - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text(commit) - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch, commit) return snap @@ -768,17 +741,12 @@ def test_marker_accepts_complete_staticembedding_model(tmp_path, monkeypatch, we assert mc._embedding_marker_in_hf_cache(_STATIC) is True -def test_marker_rejects_staticembedding_without_weights(tmp_path, monkeypatch): - # StaticEmbedding.load() hard-reads model.safetensors / pytorch_model.bin, so a tokenizer-only - # module dir (weights pruned) would fail the offline load and must NOT validate. - _staticembedding_repo(tmp_path, monkeypatch, include_weights = False) - assert mc._embedding_marker_in_hf_cache(_STATIC) is False - - -def test_marker_rejects_staticembedding_without_tokenizer(tmp_path, monkeypatch): - # StaticEmbedding.load() hard-reads tokenizer.json (Tokenizer.from_file), so a weights-only - # module dir would fail the offline load and must NOT validate. - _staticembedding_repo(tmp_path, monkeypatch, include_tokenizer = False) +@pytest.mark.parametrize("kwargs", [{"include_weights": False}, {"include_tokenizer": False}]) +def test_marker_rejects_staticembedding_missing_asset(tmp_path, monkeypatch, kwargs): + # StaticEmbedding.load() hard-reads BOTH tokenizer.json (Tokenizer.from_file) and + # model.safetensors / pytorch_model.bin, so a module dir missing either would fail the offline + # load and must NOT validate. + _staticembedding_repo(tmp_path, monkeypatch, **kwargs) assert mc._embedding_marker_in_hf_cache(_STATIC) is False @@ -797,9 +765,7 @@ def _we_dense_repo( commit = "aaa", ): # WordEmbeddings + Pooling + a Dense projection module, with the Dense weights ABSENT. - hf_root = tmp_path / "hf" - repo = hf_root / "models--org--we-dense" - snap = repo / "snapshots" / commit + repo, snap = _mk_repo(tmp_path, "org/we-dense", commit) (snap / "0_WordEmbeddings").mkdir(parents = True) (snap / "modules.json").write_text( _modules_json( @@ -818,10 +784,7 @@ def _we_dense_repo( (snap / "1_Pooling" / "config.json").write_text("{}") (snap / "2_Dense").mkdir(parents = True) (snap / "2_Dense" / "config.json").write_text("{}") # config present, weights absent - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text(commit) - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch, commit) return snap @@ -858,9 +821,7 @@ def _transformer_plus_module_repo( declared module -- each of these modules' ``load()`` ends in ``load_torch_weights``, which RAISES without a weight file -- so a weightless sibling must fail offline validation even though the Transformer alone is complete.""" - hf_root = tmp_path / "hf" - repo = hf_root / f"models--{repo_id.replace('/', '--')}" - snap = repo / "snapshots" / commit + repo, snap = _mk_repo(tmp_path, repo_id, commit) transformer = snap / "0_Transformer" transformer.mkdir(parents = True) (transformer / "config.json").write_text("{}") @@ -877,10 +838,7 @@ def _transformer_plus_module_repo( ("1", module_dir, module_type), ) ) - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text(commit) - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch, commit) return snap @@ -932,9 +890,7 @@ def test_marker_accepts_complete_transformer_with_complete_sibling( def test_marker_accepts_complete_bow_model(tmp_path, monkeypatch): # A BoW module keeps its vocab in config.json and writes NO weight file; a complete cache is # still loadable via BoW.load(config.json), so it must validate. - hf_root = tmp_path / "hf" - repo = hf_root / "models--org--bow" - snap = repo / "snapshots" / "aaa" + repo, snap = _mk_repo(tmp_path, "org/bow", "aaa") (snap / "0_BoW").mkdir(parents = True) (snap / "modules.json").write_text( _modules_json( @@ -946,19 +902,14 @@ def test_marker_accepts_complete_bow_model(tmp_path, monkeypatch): (snap / "0_BoW" / "config.json").write_text('{"vocab": ["a", "b"]}') (snap / "1_Pooling").mkdir(parents = True) (snap / "1_Pooling" / "config.json").write_text("{}") - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text("aaa") - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch) assert mc._embedding_marker_in_hf_cache("org/bow") is True def test_marker_rejects_structural_only_module_list(tmp_path, monkeypatch): # A degenerate modules.json with only structural modules (Pooling / Normalize) has no source # of embeddings and must NOT validate on the non-Transformer path. - hf_root = tmp_path / "hf" - repo = hf_root / "models--org--degenerate" - snap = repo / "snapshots" / "aaa" + repo, snap = _mk_repo(tmp_path, "org/degenerate", "aaa") (snap / "1_Pooling").mkdir(parents = True) (snap / "modules.json").write_text( _modules_json( @@ -967,10 +918,7 @@ def test_marker_rejects_structural_only_module_list(tmp_path, monkeypatch): ) ) (snap / "1_Pooling" / "config.json").write_text("{}") - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text("aaa") - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch) assert mc._embedding_marker_in_hf_cache("org/degenerate") is False @@ -1345,19 +1293,13 @@ def _tag_only_repo( which the marker cannot recognize, so detection falls through to the recorded verdict. ``modules_json`` (raw body) writes a modules.json when given, to exercise the case where the tag-only fallback must NOT run because a manifest is present.""" - hf_root = tmp_path / "hf" - repo = hf_root / f"models--{repo_id.replace('/', '--')}" - snap = repo / "snapshots" / commit - snap.mkdir(parents = True) + repo, snap = _mk_repo(tmp_path, repo_id, commit) (snap / "config.json").write_text("{}") (snap / "tokenizer.json").write_text("{}") (snap / "model.safetensors").write_bytes(b"\0") if modules_json is not None: (snap / "modules.json").write_text(modules_json) - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text(commit) - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch, commit) def test_online_verdict_still_applies_at_the_confirmed_revision(tmp_path, monkeypatch): @@ -1643,9 +1585,7 @@ def _root_wordembeddings_repo( WordEmbeddings.load() reads exactly those files, but the Transformer-shaped root check fails on the tokenizer, so the pristine `if is_root: _dir_is_transformer_load_root(...)` returned False and the offline settings route 409'd it (#7218 P1). include_tokenizer=False prunes the tokenizer.""" - hf_root = tmp_path / "hf" - repo = hf_root / f"models--{repo_id.replace('/', '--')}" - snap = repo / "snapshots" / commit + repo, snap = _mk_repo(tmp_path, repo_id, commit) (snap / "1_Pooling").mkdir(parents = True) (snap / "modules.json").write_text( _modules_json( @@ -1659,10 +1599,7 @@ def _root_wordembeddings_repo( (snap / "whitespacetokenizer_config.json").write_text("{}") (snap / "model.safetensors").write_bytes(b"\0") (snap / "1_Pooling" / "config.json").write_text("{}") - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text(commit) - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch, commit) return snap @@ -1698,10 +1635,7 @@ def _root_staticembedding_repo( config. StaticEmbedding already passes the pristine is_root Transformer check because it ships a tokenizer.json + config_sentence_transformers.json + weights; dispatching StaticEmbedding through its own branch BEFORE the root fallback must keep it loadable (#7218 P1).""" - hf_root = tmp_path / "hf" - repo = hf_root / f"models--{_ROOT_STATIC.replace('/', '--')}" - snap = repo / "snapshots" / commit - snap.mkdir(parents = True) + repo, snap = _mk_repo(tmp_path, _ROOT_STATIC, commit) (snap / "modules.json").write_text( _modules_json(("0", "", "sentence_transformers.models.StaticEmbedding")) ) @@ -1709,10 +1643,7 @@ def _root_staticembedding_repo( (snap / "tokenizer.json").write_text("{}") if include_weights: (snap / "model.safetensors").write_bytes(b"\0") - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text(commit) - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch, commit) return snap @@ -1749,9 +1680,7 @@ def _clip_repo( (image processor + tokenizer), so a config-only dir is NOT loadable. complete=False writes only config.json (no processor asset, no weights). ``processor=False`` writes the tokenizer + weights but omits the image-processor config (preprocessor_config.json) AutoProcessor needs for CLIP.""" - hf_root = tmp_path / "hf" - repo = hf_root / f"models--{_CLIP.replace('/', '--')}" - snap = repo / "snapshots" / commit + repo, snap = _mk_repo(tmp_path, _CLIP, commit) (snap / "0_CLIPModel").mkdir(parents = True) (snap / "modules.json").write_text( _modules_json(("0", "0_CLIPModel", "sentence_transformers.models.CLIPModel")) @@ -1764,10 +1693,7 @@ def _clip_repo( if processor: # AutoProcessor's image-processor half reads its own config. (snap / "0_CLIPModel" / "preprocessor_config.json").write_text("{}") - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text(commit) - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch, commit) def test_marker_rejects_config_only_clip_module(tmp_path, monkeypatch): @@ -1814,9 +1740,7 @@ def _sparse_static_repo( (load_torch_weights). ``payload`` selects what backs it: "idf" (config path -> idf.json), "idf_unselected" (idf.json present but config does NOT name a ``.json`` path, so load() falls through to load_torch_weights and raises), "weights" (model.safetensors) or "none" (neither).""" - hf_root = tmp_path / "hf" - repo = hf_root / f"models--{_SPARSE_STATIC.replace('/', '--')}" - snap = repo / "snapshots" / commit + repo, snap = _mk_repo(tmp_path, _SPARSE_STATIC, commit) mod = snap / "0_SparseStaticEmbedding" mod.mkdir(parents = True) (snap / "modules.json").write_text( @@ -1838,10 +1762,7 @@ def _sparse_static_repo( (mod / "idf.json").write_text("{}") elif payload == "weights": (mod / "model.safetensors").write_bytes(b"\0") - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text(commit) - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch, commit) @pytest.mark.parametrize("payload", ["idf", "weights"]) @@ -1890,18 +1811,12 @@ def _root_weights_snapshot( """A snapshot with a COMPLETE root weight set (config + tokenizer + weights) and a modules.json whose raw body is ``modules_body``. Lets a test assert that a present-but-empty / malformed modules.json is not loadable even though the root weights are complete.""" - hf_root = tmp_path / "hf" - repo = hf_root / f"models--{repo_id.replace('/', '--')}" - snap = repo / "snapshots" / "aaa" - snap.mkdir(parents = True) + repo, snap = _mk_repo(tmp_path, repo_id, "aaa") (snap / "modules.json").write_text(modules_body) (snap / "config.json").write_text("{}") (snap / "tokenizer.json").write_text("{}") (snap / "model.safetensors").write_bytes(b"\0") - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text("aaa") - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch) return repo_id @@ -1947,9 +1862,7 @@ def _sharded_transformer_repo( tokenizer + a model.safetensors.index.json whose ``weight_map`` is the given dict, plus ``present_shards`` written to disk (paths relative to the module dir). Lets a test reference a shard the index maps but leave it absent, or place shards in a subdirectory.""" - hf_root = tmp_path / "hf" - repo = hf_root / f"models--{_SHARDED.replace('/', '--')}" - snap = repo / "snapshots" / commit + repo, snap = _mk_repo(tmp_path, _SHARDED, commit) mod = snap / module_path mod.mkdir(parents = True) (snap / "modules.json").write_text( @@ -1963,10 +1876,7 @@ def _sharded_transformer_repo( target = mod / Path(shard_rel) target.parent.mkdir(parents = True, exist_ok = True) target.write_bytes(b"\0") - (repo / "refs").mkdir(parents = True) - (repo / "refs" / "main").write_text(commit) - _fake_hf_cache(monkeypatch, hf_root) - monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False) + _activate(repo, monkeypatch, commit) def test_marker_rejects_sharded_index_with_a_missing_mapped_shard(tmp_path, monkeypatch): From fd7c2de841e1e5fc425f8cef8cb3607300b48efe Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 10:04:10 +0000 Subject: [PATCH 75/77] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../backend/tests/test_is_embedding_model_offline.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/studio/backend/tests/test_is_embedding_model_offline.py b/studio/backend/tests/test_is_embedding_model_offline.py index 089e64631a6..8a5e09cf8bb 100644 --- a/studio/backend/tests/test_is_embedding_model_offline.py +++ b/studio/backend/tests/test_is_embedding_model_offline.py @@ -194,7 +194,11 @@ def _fake_hf_cache(monkeypatch, root): monkeypatch.setitem(sys.modules, "huggingface_hub.constants", fake.constants) -def _mk_repo(tmp_path, repo_id = "org/model", commit = "aaa"): +def _mk_repo( + tmp_path, + repo_id = "org/model", + commit = "aaa", +): """(repo_dir, snapshot_dir) for a fresh HF cache repo under ``tmp_path/hf``; the snapshot dir is created. Shared skeleton for the per-type repo builders below.""" repo = tmp_path / "hf" / f"models--{repo_id.replace('/', '--')}" @@ -203,7 +207,11 @@ def _mk_repo(tmp_path, repo_id = "org/model", commit = "aaa"): return repo, snap -def _activate(repo, monkeypatch, commit = "aaa"): +def _activate( + repo, + monkeypatch, + commit = "aaa", +): """Write refs/main and point the HF/ST cache at *repo*'s root, ignoring SENTENCE_TRANSFORMERS_HOME.""" (repo / "refs").mkdir(parents = True, exist_ok = True) (repo / "refs" / "main").write_text(commit) From ef49462d884db2ef855183eb22d9a050fcc26592 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 22 Jul 2026 06:40:12 +0000 Subject: [PATCH 76/77] Reclassify embedding models from the cache on every offline call is_embedding_model consulted its process memo before the offline branch, so an online lookup that memoized True from tags (without caching any weights) was returned unchanged once the session went offline -- the studio flips HF_HUB_OFFLINE in-process on a dead DNS, and the ungated check-embedding route can populate the memo. Settings would then accept a repo the offline loader cannot open. Run the offline cache-marker reclassification ahead of the memo and never record it, so an offline verdict always reflects the local cache and a later cache materialization is not masked by a stale negative. Add regression tests. --- .../tests/test_offline_embedding_minimal.py | 29 +++++++++++++++++++ studio/backend/utils/models/model_config.py | 23 ++++++++------- 2 files changed, 42 insertions(+), 10 deletions(-) diff --git a/studio/backend/tests/test_offline_embedding_minimal.py b/studio/backend/tests/test_offline_embedding_minimal.py index eec8d4578fd..43f0751d8b1 100644 --- a/studio/backend/tests/test_offline_embedding_minimal.py +++ b/studio/backend/tests/test_offline_embedding_minimal.py @@ -274,6 +274,35 @@ def test_offline_slashless_resolves_via_alias(hf_cache, monkeypatch): 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 + # the session goes offline (the studio flips HF_HUB_OFFLINE in-process on a dead DNS), + # 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): + # Offline, an uncached repo is False; because the offline branch never records a memo, once + # its snapshot materializes (another process populates the shared cache) the next call + # re-reports True instead of returning a stale negative. + 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) ────────────── diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 720cec7d063..03d02491724 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2111,6 +2111,19 @@ 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 HF cache on EVERY call, before and without + # consulting the process memo. Never call the Hub -- a DNS-dead session hangs on model_info + # retries -- and never trust or record a memo here: an online lookup can memoize True from + # tags without any weights cached, so returning that once the session goes offline (the studio + # flips HF_HUB_OFFLINE in-process on a dead DNS) would accept a repo _get() cannot load; a + # cached negative could likewise be invalidated by a later cache materialization. A cached + # modules.json marks a Sentence-Transformers repo, mirroring the local-path check below. The + # probe is local-filesystem only, so recomputing per call is 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] @@ -2122,16 +2135,6 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: _embedding_detection_cache[cache_key] = is_emb return is_emb - # Offline: never call the Hub -- a DNS-dead session hangs on model_info retries. - # Classify from the local cache the load will use: a cached modules.json marks a - # Sentence-Transformers repo, mirroring the local-path check above. - from utils.utils import hf_env_offline - - if hf_env_offline(): - is_emb = _embedding_marker_in_hf_cache(model_name) - _embedding_detection_cache[cache_key] = is_emb - return is_emb - try: from huggingface_hub import model_info as hf_model_info From 48309dc2953672b4d08ed55787294d155e59d153 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 22 Jul 2026 08:37:30 +0000 Subject: [PATCH 77/77] Tighten comments on the offline embedding path Condense the offline-embedding helper docstrings and inline comments added in this PR to fewer, clearer lines, keeping the non-obvious security and offline rationale. Comments and docstrings only; no code change. --- studio/backend/core/rag/embeddings.py | 25 ++++--- studio/backend/routes/settings.py | 29 ++++----- .../test_embedding_model_security_gate.py | 4 +- .../tests/test_offline_embedding_minimal.py | 65 +++++++++---------- studio/backend/utils/models/model_config.py | 24 +++---- .../backend/utils/security/file_security.py | 58 +++++++---------- studio/backend/utils/utils.py | 45 +++++-------- 7 files changed, 103 insertions(+), 147 deletions(-) diff --git a/studio/backend/core/rag/embeddings.py b/studio/backend/core/rag/embeddings.py index 7d2ef7a470e..0c743e4ea49 100644 --- a/studio/backend/core/rag/embeddings.py +++ b/studio/backend/core/rag/embeddings.py @@ -126,9 +126,8 @@ def _guard_model_security(name: str, local_only: bool = False) -> None: /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 and fails closed on an unscanned pickle - weight; the subdir probes are skipped because they would hit the network and hang, and - the offline gate walks the whole snapshot anyway. + ``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 @@ -137,9 +136,8 @@ def _guard_model_security(name: str, local_only: bool = False) -> None: if local_only: load_subdirs = () else: - # 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. + # 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)) @@ -164,8 +162,8 @@ def _guard_model_security(name: str, local_only: bool = False) -> None: def _st_accepts_local_files_only(st_cls) -> bool: - """Whether this SentenceTransformer version accepts ``local_files_only`` (added in newer - releases). Passing it to an older constructor raises, so gate on the signature.""" + """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 @@ -178,8 +176,8 @@ 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 the offline state once so the security gate and the load agree (no window - # where the gate is skipped as offline but the constructor then reaches the network). + # 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: @@ -195,10 +193,9 @@ def _get(model_name: str | None = None): from utils.utils import hf_cache_snapshot_dir snapshot = hf_cache_snapshot_dir(name) if snapshot is not None: - # Load from the resolved 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). If nothing is cached we fall through to a - # cache-only repo load, which fails fast offline instead of hanging. + # 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 diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 5bea5c864c4..f36c8870e31 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -418,8 +418,8 @@ def update_embedding_model( hf_token = (payload.hf_token or "").strip() or None from utils.utils import hf_env_offline - # Capture the offline state once: offline, both the Hub malware scan and the is-embedding - # metadata check are unreachable, so both degrade to the local cache below. + # 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 @@ -444,14 +444,13 @@ 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() - # Offline the Hub scan is unreachable and the subdir probes below would hit the - # network and hang on a dead DNS; the offline gate walks the whole cached snapshot, - # so no load-subdir hints are needed. + # 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 the ST module dirs (0_Transformer/) so a flagged pickle directly under - # one blocks instead of passing as an unreferenced nested shard. + # 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( ( @@ -489,15 +488,14 @@ 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): - # Offline, is_embedding_model can only confirm the Sentence-Transformers layout - # (modules.json); a transformers-native embedder (e.g. a feature-extraction model - # like gte-modernbert) is unverifiable without the Hub metadata. If the repo is - # already cached and loadable, accept it rather than raising a 409 that online - # would not -- SentenceTransformer can load any cached encoder. Uncached -> 409. + # 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. + # 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( @@ -508,8 +506,7 @@ def update_embedding_model( "you may be offline)." ), ) - # The GGUF availability probe below calls the Hub (list_repo_files); skip it offline - # so a dead-DNS session cannot hang. A local GGUF check stays (no network). + # 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) diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index 61ba01bcfbe..a6c18bd8dee 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -108,7 +108,7 @@ def test_hard_block_uses_non_forceable_status(client, monkeypatch): 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; since SentenceTransformer can load any cached encoder, accept it (no 409). + # 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") @@ -123,7 +123,7 @@ def test_offline_cached_non_st_model_is_accepted(client, monkeypatch): def test_offline_partial_or_uncached_model_still_409(client, monkeypatch): - # Offline but NOT loadable (uncached or a metadata-only partial cache): keep the forceable + # 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)) diff --git a/studio/backend/tests/test_offline_embedding_minimal.py b/studio/backend/tests/test_offline_embedding_minimal.py index 43f0751d8b1..8862e231e50 100644 --- a/studio/backend/tests/test_offline_embedding_minimal.py +++ b/studio/backend/tests/test_offline_embedding_minimal.py @@ -3,15 +3,11 @@ """Offline RAG embedding-model handling (issue #6817). -Offline, the studio must never call the Hub: a DNS-dead session hangs on model_info / -download retries. These tests build a fake HF cache under a temp ``HF_HUB_CACHE`` and -assert that, offline: - * ``is_embedding_model`` classifies from the cached ``modules.json`` and never calls - the Hub (``model_info`` patched to raise if reached); - * the file-security gate fails CLOSED on an unscanned pickle weight with no safetensors - alternative, and allows an inert (safetensors/gguf) cache; - * the embedder threads ``local_files_only`` into the SentenceTransformer load. -Online behavior is unchanged: a bounded ``model_info`` timeout with a local-cache fallback. +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 @@ -30,14 +26,14 @@ st_repo_id_candidates, ) -# A minimal sentence-transformers modules.json (the marker the local-path check keys on). +# 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): - """A modules.json listing one Transformer module per ``path`` (a load root).""" + """modules.json listing one Transformer module per path (a load root).""" import json return json.dumps( [ @@ -61,8 +57,8 @@ def _make_cache( files, commit = _COMMIT, ): - """Build a canonical HF-cache snapshot (``refs/main`` + ``snapshots//``) for - ``repo_id`` under ``root`` with ``{relpath: contents}``. Returns the snapshot dir.""" + """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") @@ -78,7 +74,7 @@ def _make_cache( def _no_network(): - """Patch the Hub metadata call to fail loudly if any offline path reaches it.""" + """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")) @@ -89,7 +85,7 @@ def _is_embedding_model(*args, **kwargs): @pytest.fixture def hf_cache(tmp_path, monkeypatch): - """Point the HF cache at a fresh temp dir the resolver reads at call time.""" + """Point the HF cache at a fresh temp dir.""" root = tmp_path / "hub" root.mkdir() monkeypatch.setenv("HF_HOME", str(tmp_path)) @@ -99,7 +95,7 @@ def hf_cache(tmp_path, monkeypatch): @pytest.fixture(autouse = True) def _clean_env(monkeypatch): - """Each test starts online with an empty detection cache; offline tests opt in.""" + """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 @@ -180,7 +176,7 @@ def test_snapshot_dir_none_when_snapshot_missing(hf_cache): def test_snapshot_dir_expands_env_vars_in_cache_path(tmp_path, monkeypatch): - # HF_HUB_CACHE with an unexpanded $VAR must resolve to the same place the loader uses. + # 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)) @@ -192,8 +188,7 @@ def test_snapshot_dir_expands_env_vars_in_cache_path(tmp_path, monkeypatch): def test_snapshot_dir_uses_sentence_transformers_home(tmp_path, monkeypatch): - # SentenceTransformer uses SENTENCE_TRANSFORMERS_HOME as its cache_folder, so the gate must - # inspect it too. + # 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)) @@ -204,8 +199,8 @@ def test_snapshot_dir_uses_sentence_transformers_home(tmp_path, monkeypatch): def test_snapshot_dir_st_home_is_exclusive(tmp_path, monkeypatch): - # With SENTENCE_TRANSFORMERS_HOME set, ST loads only from it; a model that lives only under - # HF_HUB_CACHE is not what the loader would use, so the resolver must not report it. + # 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" @@ -223,7 +218,7 @@ def test_snapshot_is_loadable_with_config_and_weights(hf_cache): def test_snapshot_is_not_loadable_when_metadata_only(hf_cache): - # A partial cache (refs/main resolves, but no weight files) is not loadable. + # 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 @@ -233,7 +228,7 @@ def test_snapshot_is_not_loadable_when_uncached(hf_cache): def test_gate_blocks_pickle_in_sentence_transformers_home(tmp_path, monkeypatch): - # A pickle cached under SENTENCE_TRANSFORMERS_HOME must still fail closed offline. + # 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)) @@ -276,9 +271,8 @@ def test_offline_slashless_resolves_via_alias(hf_cache, monkeypatch): 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 - # the session goes offline (the studio flips HF_HUB_OFFLINE in-process on a dead DNS), - # 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. + # 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( @@ -293,9 +287,8 @@ def test_offline_ignores_stale_online_memo(hf_cache, monkeypatch): def test_offline_recomputes_after_cache_materializes(hf_cache, monkeypatch): - # Offline, an uncached repo is False; because the offline branch never records a memo, once - # its snapshot materializes (another process populates the shared cache) the next call - # re-reports True instead of returning a stale negative. + # 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 @@ -411,8 +404,8 @@ def test_gate_allows_pickle_in_subdir_with_safetensors(hf_cache): 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 by - # SentenceTransformer, so it must not block the offline load (matches the online gate). + # 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", @@ -441,8 +434,8 @@ def test_gate_allows_adapter_pickle_with_adapter_safetensors(hf_cache): 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. + # 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 @@ -535,8 +528,8 @@ def test_get_offline_loads_from_local_snapshot(hf_cache, monkeypatch): snapshot = _make_cache( hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"} ) - # TRANSFORMERS_OFFLINE only: a cached model is loaded from its local snapshot dir, a local - # path that never touches the Hub -- offline-safe on ANY sentence-transformers version. + # 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) @@ -564,7 +557,7 @@ def test_get_offline_uncached_uses_local_files_only(tmp_path, monkeypatch): monkeypatch.setattr(embeddings, "_name", None, raising = False) monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None) monkeypatch.setattr(embeddings, "_device", lambda: "cpu") - # No cache -> fall back to a repo-id load forced cache-only (fails fast offline, not a hang). + # 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) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 03d02491724..50a997218f2 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -2076,16 +2076,13 @@ def download_gguf_file( _embedding_detection_cache: Dict[tuple, bool] = {} -# Bound the Hub metadata lookup so a DNS-dead session fails fast (and falls back to the -# local cache) instead of hanging on huggingface_hub's default retry loop. +# 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 active cached snapshot carries a ``modules.json`` -- the - Sentence-Transformers marker, mirroring the local-path check. Cache-only, no network; - used offline and as a fallback when a bounded Hub lookup times out. - """ + """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) @@ -2113,14 +2110,10 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: """ from utils.utils import hf_env_offline - # Offline (remote repo): reclassify from the local HF cache on EVERY call, before and without - # consulting the process memo. Never call the Hub -- a DNS-dead session hangs on model_info - # retries -- and never trust or record a memo here: an online lookup can memoize True from - # tags without any weights cached, so returning that once the session goes offline (the studio - # flips HF_HUB_OFFLINE in-process on a dead DNS) would accept a repo _get() cannot load; a - # cached negative could likewise be invalidated by a later cache materialization. A cached - # modules.json marks a Sentence-Transformers repo, mirroring the local-path check below. The - # probe is local-filesystem only, so recomputing per call is cheap. + # 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) @@ -2159,8 +2152,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: return is_emb except Exception as e: - # A bounded timeout or transient network error must not hang or hard-fail the UI. - # Best-effort fall back to the local cache marker before giving up. + # 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}") is_emb = _embedding_marker_in_hf_cache(model_name) _embedding_detection_cache[cache_key] = is_emb diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 2047afac2b7..0490d38d7cb 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -38,23 +38,21 @@ logger = get_logger(__name__) -# Pickle-format weight files (plain or sharded) that execute code when a load deserializes -# them; safetensors/gguf are inert (tensor-only). Grouped by weight family so an inert -# safetensors only suppresses the pickle it actually replaces (below): the base loader will -# not use an adapter's safetensors in place of pytorch_model.bin, and vice versa. +# 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 weight set. HF saves the torch pickle as ``pytorch_model.bin`` but the -# safetensors as ``model.safetensors`` (the stems differ), so a base pickle is inert-replaced by -# ``model.safetensors``, its shards, or the sharded index -- not by an adapter's safetensors. +# 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 weight set: ``adapter_model.safetensors``, its shards, or index. +# 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, @@ -290,20 +288,18 @@ def _fetch_security_status(model_name: str, hf_token: Optional[str]): def _st_load_roots(snapshot: Path) -> list: - """Directories a SentenceTransformer load actually deserializes weights from: the snapshot - root plus each module ``path`` in ``modules.json`` (e.g. ``0_Transformer``), read locally, - no network. Mirrors the online gate, which ignores unreferenced nested pickles (an example - artifact or a ``nemo/`` pickle ST never loads), so the offline gate does not over-block. - """ + """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 -> the snapshot root is the only load root + 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("/") - # A relative module path only; ignore a crafted "../" escape. + # Relative module path only; ignore a crafted "../" escape. if path and ".." not in path.split("/"): candidate = snapshot / path if candidate not in roots: @@ -312,14 +308,11 @@ def _st_load_roots(snapshot: Path) -> list: def _cached_pickle_weight_files(snapshot: Path) -> list: - """Pickle-format weight files in ``snapshot``'s ST load roots a load would deserialize, - EXCLUDING those whose own weight family also ships an inert ``safetensors`` in the same - directory (the load prefers the safetensors). A base pickle (``pytorch_model.bin`` / - ``model.bin`` / ``consolidated``) is suppressed only by a base ``model.safetensors`` weight; - an ``adapter_model`` pickle only by an ``adapter_model.safetensors`` -- an unrelated - safetensors is not a substitute the loader would pick. Scans only the load roots (not every - nested file). Raises ``OSError`` if the snapshot root cannot be read (caller blocks). - """ + """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: @@ -327,7 +320,7 @@ def _cached_pickle_weight_files(snapshot: Path) -> list: except OSError: if root == snapshot: raise # top-level unreadable -> fail closed - continue # an unreadable module subdir: nothing loadable to attest here + 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: @@ -341,11 +334,9 @@ def _cached_pickle_weight_files(snapshot: Path) -> list: def _evaluate_local_only(model_name: str) -> FileSecurityDecision: - """Offline security gate. The Hub malware scan is unreachable, so inspect the local cache - and fail CLOSED on an unscanned pickle weight (code-executing on load) that has no inert - ``safetensors`` alternative, instead of failing open or hanging on the network. A - safetensors/gguf-only cache loads normally; nothing cached -> nothing to load -> allowed. - """ + """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: @@ -372,8 +363,7 @@ def _evaluate_local_only(model_name: str) -> FileSecurityDecision: model_name, False, reason = "offline; cached weights are inert (safetensors/gguf)" ) - # Report snapshot-relative posix paths (matching the online gate, and disambiguating a - # same-named pickle in different module dirs). + # 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( @@ -407,8 +397,7 @@ def evaluate_file_security( 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 (``HF_HUB_OFFLINE`` / ``TRANSFORMERS_OFFLINE``): - the Hub scan is unreachable, so instead of hanging or failing open, inspect the local + ``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 @@ -425,8 +414,7 @@ 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: the Hub scan is unreachable; inspect the local cache and fail closed on an - # unscanned pickle weight instead of hanging on model_info or failing open. + # 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) diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index 75414164a24..21e11c67064 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -17,17 +17,17 @@ # ── Offline / HF-cache helpers ────────────────────────────────── -# An offline model load must never touch the network: a DNS-dead session hangs on -# huggingface_hub download retries. These read the local HF cache the load itself uses. +# 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. + """True when HF_HUB_OFFLINE or TRANSFORMERS_OFFLINE requests offline mode. - Broader than huggingface_hub, which honors only ``HF_HUB_OFFLINE``; the studio also - honors ``TRANSFORMERS_OFFLINE`` because users set it to keep transformers loads local. + 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: @@ -36,10 +36,8 @@ def hf_env_offline() -> bool: def st_repo_id_candidates(model_name: str) -> list: - """Repo ids a Sentence-Transformers load may resolve ``model_name`` to. A slashless - name (``all-MiniLM-L6-v2``) is loaded as ``sentence-transformers/all-MiniLM-L6-v2``, - so both are candidate cache repos. - """ + """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 [] @@ -50,18 +48,14 @@ def st_repo_id_candidates(model_name: str) -> list: def _expand_path(raw: str) -> Path: - """Expand ``~`` and ``$VARS`` the way huggingface_hub does for its cache paths, so the gate - resolves the same directory the loader will.""" + """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 single local cache root a Sentence-Transformers / huggingface_hub load resolves to, - by the loader's own precedence (it selects ONE ``cache_folder`` and does not fall through): - an explicit ``SENTENCE_TRANSFORMERS_HOME`` (ST's cache_folder), else ``HF_HUB_CACHE``, else - ``HF_HOME/hub``, else ``~/.cache/huggingface/hub`` (mirroring huggingface_hub). Expanded - (``~`` / ``$VARS``), read from the env at call time. Returned as a one-element 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)] @@ -75,11 +69,8 @@ def _hf_cache_roots() -> list: def hf_cache_snapshot_dir(model_name: str) -> Optional[Path]: - """Active local snapshot dir for ``model_name``'s ``main`` revision, or None when the repo - is not cached in any known cache root (including ``SENTENCE_TRANSFORMERS_HOME``). Reads - ``refs/main`` then ``snapshots/``; never touches the network. Tries the - ``sentence-transformers/`` alias for a slashless name. - """ + """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: @@ -106,16 +97,14 @@ def hf_cache_snapshot_dir(model_name: str) -> Optional[Path]: return None -# Weight file suffixes a load can consume. Presence of one (plus a config) distinguishes a -# real cached model from a metadata-only partial cache that would fail at load time. +# 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 active snapshot is cached AND actually loadable -- it has a - config (``config.json`` or ``modules.json``) and at least one weight file -- rather than a - metadata-only partial cache that resolves a ``refs/main`` but would fail at load. No network. - """ + """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