studio: classify embedding models from the HF cache and honor offline mode (#7218)
* 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. * 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * 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. * 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: 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: 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. * 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: 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: 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: 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: 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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: 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: 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. * 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. * 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: 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * 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: 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: 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. * 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. * 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: 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * 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. * 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: 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. * 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: 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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.) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * 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. * 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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). * 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). * 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). * 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). * 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. * 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. * 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. * 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. * 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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/ <name> and loads THAT snapshot (only the basic ORIGINAL_TRANSFORMER_MODELS load bare). With both models--<name> and models--sentence-transformers--<name> 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 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. * 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. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <unslothai@gmail.com>
This commit is contained in:
Hakan Baysal
GitHub
parent
55433bd7b8
commit
aa49c0710e
7 changed files with 1001 additions and 37 deletions
|
|
@ -22,6 +22,7 @@ from typing import Callable
|
|||
|
||||
from utils.hardware.hardware import DeviceType, get_device
|
||||
from utils.transformers_dtype import dtype_kwargs
|
||||
from utils.utils import hf_env_offline
|
||||
|
||||
from . import config
|
||||
|
||||
|
|
@ -119,30 +120,55 @@ def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]:
|
|||
return ()
|
||||
|
||||
|
||||
def _guard_model_security(name: str) -> None:
|
||||
def _guard_model_security(name: str, local_only: bool = False) -> None:
|
||||
"""Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside
|
||||
SentenceTransformer regardless of trust_remote_code. Defense in depth behind the
|
||||
/settings gate (a name can also arrive via env/default); local paths and unreachable
|
||||
scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error.
|
||||
|
||||
``local_only`` (offline) inspects the local cache; subdir probes are skipped (they'd hit the
|
||||
network and hang, and the offline gate walks the whole snapshot anyway).
|
||||
"""
|
||||
try:
|
||||
from utils.security import evaluate_file_security, security_load_subdirs
|
||||
|
||||
token = _ambient_hf_token()
|
||||
# Union the audio-model load roots with the ST module dirs so a flagged pickle
|
||||
# directly under a Transformer module dir (0_Transformer/) blocks instead of
|
||||
# passing as an unreferenced nested shard.
|
||||
load_subdirs = tuple(
|
||||
dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token)))
|
||||
)
|
||||
blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked
|
||||
if local_only:
|
||||
load_subdirs = ()
|
||||
else:
|
||||
# Union audio-model load roots with ST module dirs so a flagged pickle under a
|
||||
# Transformer module dir blocks instead of passing as an unreferenced nested shard.
|
||||
load_subdirs = tuple(
|
||||
dict.fromkeys(
|
||||
(*security_load_subdirs(name, token), *_st_module_subdirs(name, token))
|
||||
)
|
||||
)
|
||||
blocked = evaluate_file_security(
|
||||
name, hf_token = token, load_subdirs = load_subdirs, local_only_load = local_only
|
||||
).blocked
|
||||
except Exception:
|
||||
return
|
||||
if blocked:
|
||||
raise UnsafeEmbeddingModelError(
|
||||
f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security "
|
||||
"scan; refusing to load. Set a different RAG embedding model."
|
||||
reason = (
|
||||
"has cached pickle weights that cannot be security-scanned offline and no "
|
||||
"safetensors alternative"
|
||||
if local_only
|
||||
else "is flagged as unsafe by Hugging Face's security scan"
|
||||
)
|
||||
raise UnsafeEmbeddingModelError(
|
||||
f"Embedding model {name!r} {reason}; refusing to load. "
|
||||
"Set a different RAG embedding model."
|
||||
)
|
||||
|
||||
|
||||
def _st_accepts_local_files_only(st_cls) -> bool:
|
||||
"""Whether this SentenceTransformer version accepts local_files_only; passing it to an
|
||||
older constructor raises, so gate on the signature."""
|
||||
try:
|
||||
import inspect
|
||||
return "local_files_only" in inspect.signature(st_cls.__init__).parameters
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def _get(model_name: str | None = None):
|
||||
|
|
@ -150,6 +176,9 @@ def _get(model_name: str | None = None):
|
|||
for a ~1.5x speedup at negligible accuracy loss."""
|
||||
global _model, _name
|
||||
name = model_name or config.effective_embedding_model()
|
||||
# Capture offline state once so the gate and the load agree (no window where the gate is
|
||||
# skipped as offline but the constructor then reaches the network).
|
||||
local_only = hf_env_offline()
|
||||
with _lock:
|
||||
if _model is None or _name != name:
|
||||
_install_torchao_stub_once()
|
||||
|
|
@ -157,8 +186,20 @@ def _get(model_name: str | None = None):
|
|||
|
||||
device = _device()
|
||||
logger.info("loading embedding model %s on %s", name, device)
|
||||
_guard_model_security(name)
|
||||
_model = SentenceTransformer(name, device = device, model_kwargs = dtype_kwargs("float16"))
|
||||
_guard_model_security(name, local_only)
|
||||
st_kwargs = dict(device = device, model_kwargs = dtype_kwargs("float16"))
|
||||
load_target = name
|
||||
if local_only:
|
||||
from utils.utils import hf_cache_snapshot_dir
|
||||
snapshot = hf_cache_snapshot_dir(name)
|
||||
if snapshot is not None:
|
||||
# Load from the local snapshot dir: a local path never touches the Hub, so
|
||||
# this is offline-safe on ANY sentence-transformers version (even ones
|
||||
# predating local_files_only).
|
||||
load_target = str(snapshot)
|
||||
elif _st_accepts_local_files_only(SentenceTransformer):
|
||||
st_kwargs["local_files_only"] = True
|
||||
_model = SentenceTransformer(load_target, **st_kwargs)
|
||||
_name = name
|
||||
return _model
|
||||
|
||||
|
|
|
|||
|
|
@ -416,6 +416,11 @@ def update_embedding_model(
|
|||
log = logger,
|
||||
) from exc
|
||||
hf_token = (payload.hf_token or "").strip() or None
|
||||
from utils.utils import hf_env_offline
|
||||
|
||||
# Offline, both the Hub malware scan and the is-embedding check are unreachable and degrade
|
||||
# to the local cache below; capture the state once.
|
||||
local_only_load = hf_env_offline()
|
||||
# The env/default model needs no verification; saving it is a no-op override.
|
||||
# A local GGUF on the llama-server backend is accepted as-is: it is exactly
|
||||
# what the backend loads, and HF metadata cannot verify a local path.
|
||||
|
|
@ -439,26 +444,41 @@ def update_embedding_model(
|
|||
# Fall back to the loader's own token so a gated/private repo is actually scanned
|
||||
# (a token-less scan fails open for exactly the repo that would still load).
|
||||
scan_token = hf_token or _ambient_hf_token()
|
||||
# Include the ST module dirs (0_Transformer/) so a flagged pickle directly under
|
||||
# one blocks instead of passing as an unreferenced nested shard.
|
||||
load_subdirs = tuple(
|
||||
dict.fromkeys(
|
||||
(
|
||||
*security_load_subdirs(model, scan_token),
|
||||
*_st_module_subdirs(model, scan_token),
|
||||
# Offline: subdir probes would hit the network and hang; the offline gate walks the
|
||||
# whole cached snapshot, so no load-subdir hints are needed.
|
||||
if local_only_load:
|
||||
load_subdirs = ()
|
||||
else:
|
||||
# Include ST module dirs (0_Transformer/) so a flagged pickle directly under one
|
||||
# blocks instead of passing as an unreferenced nested shard.
|
||||
load_subdirs = tuple(
|
||||
dict.fromkeys(
|
||||
(
|
||||
*security_load_subdirs(model, scan_token),
|
||||
*_st_module_subdirs(model, scan_token),
|
||||
)
|
||||
)
|
||||
)
|
||||
)
|
||||
if evaluate_file_security(model, hf_token = scan_token, load_subdirs = load_subdirs).blocked:
|
||||
if evaluate_file_security(
|
||||
model,
|
||||
hf_token = scan_token,
|
||||
load_subdirs = load_subdirs,
|
||||
local_only_load = local_only_load,
|
||||
).blocked:
|
||||
# 403, not 409: the client routes every 409 into the forceable "save anyway"
|
||||
# flow, but this block is a hard, non-forceable security refusal.
|
||||
raise HTTPException(
|
||||
status_code = 403,
|
||||
if local_only_load:
|
||||
detail = (
|
||||
f"{model!r} has cached pickle weights that cannot be security-scanned "
|
||||
"offline and no safetensors alternative, so it cannot be used as the "
|
||||
"embedding model. Re-download it with safetensors weights while online."
|
||||
)
|
||||
else:
|
||||
detail = (
|
||||
f"{model!r} is flagged as unsafe by Hugging Face's security scan and "
|
||||
"cannot be used as the embedding model."
|
||||
),
|
||||
)
|
||||
)
|
||||
raise HTTPException(status_code = 403, detail = detail)
|
||||
if model != default_embedding_model() and not payload.force and not is_local_gguf:
|
||||
from core.rag import config as rag_config
|
||||
|
||||
|
|
@ -468,15 +488,28 @@ def update_embedding_model(
|
|||
# which would wrongly 409 a valid online GGUF embedder.
|
||||
gguf_named = _llama_backend_active() and rag_config._names_gguf(model)
|
||||
if not gguf_named and not is_embedding_model(model, hf_token = hf_token):
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = (
|
||||
f"Could not verify {model!r} as an embedding model on "
|
||||
"Hugging Face (it may be the wrong model type, gated, or "
|
||||
"you may be offline)."
|
||||
),
|
||||
)
|
||||
gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error(model, hf_token)
|
||||
# Offline, is_embedding_model can only confirm the ST layout (modules.json); a
|
||||
# transformers-native embedder (e.g. gte-modernbert) is unverifiable without Hub
|
||||
# metadata. If already cached and loadable, accept it rather than raising a 409 that
|
||||
# online would not (ST can load any cached encoder). Uncached -> 409.
|
||||
from utils.utils import hf_cache_snapshot_is_loadable
|
||||
|
||||
# Require a genuinely loadable cache (config + weights), not just a resolved refs/main,
|
||||
# so a metadata-only partial cache still gets the forceable 409.
|
||||
offline_cached = local_only_load and hf_cache_snapshot_is_loadable(model)
|
||||
if not offline_cached:
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = (
|
||||
f"Could not verify {model!r} as an embedding model on "
|
||||
"Hugging Face (it may be the wrong model type, gated, or "
|
||||
"you may be offline)."
|
||||
),
|
||||
)
|
||||
# The Hub GGUF probe (list_repo_files) can hang offline; skip it. Local check stays.
|
||||
gguf_error = _local_gguf_backend_error(model)
|
||||
if gguf_error is None and not local_only_load:
|
||||
gguf_error = _hf_gguf_backend_error(model, hf_token)
|
||||
if gguf_error:
|
||||
raise HTTPException(status_code = 409, detail = gguf_error)
|
||||
set_rag_embedding_model(model)
|
||||
|
|
|
|||
|
|
@ -106,6 +106,56 @@ def test_hard_block_uses_non_forceable_status(client, monkeypatch):
|
|||
assert unverified.status_code == 409
|
||||
|
||||
|
||||
def test_offline_cached_non_st_model_is_accepted(client, monkeypatch):
|
||||
# Offline, a cached transformers-native embedder (no modules.json) is unverifiable via HF
|
||||
# metadata, but ST can load any cached encoder, so accept it (no 409).
|
||||
c, saved = client
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
||||
import utils.models as _models
|
||||
import utils.utils as _uu
|
||||
|
||||
monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False)
|
||||
monkeypatch.setattr(_uu, "hf_cache_snapshot_is_loadable", lambda name: True)
|
||||
r = c.put("/embedding-model", json = {"embedding_model": "acme/gte-modernbert"})
|
||||
assert r.status_code == 200
|
||||
assert saved.get("model") == "acme/gte-modernbert"
|
||||
|
||||
|
||||
def test_offline_partial_or_uncached_model_still_409(client, monkeypatch):
|
||||
# Offline but not loadable (uncached or metadata-only partial cache): keep the forceable
|
||||
# 409, since the cache-only load would fail anyway.
|
||||
c, _saved = client
|
||||
monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False))
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
||||
import utils.models as _models
|
||||
import utils.utils as _uu
|
||||
|
||||
monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: False)
|
||||
monkeypatch.setattr(_uu, "hf_cache_snapshot_is_loadable", lambda name: False)
|
||||
r = c.put("/embedding-model", json = {"embedding_model": "acme/uncached-embedder"})
|
||||
assert r.status_code == 409
|
||||
|
||||
|
||||
def test_offline_skips_remote_gguf_probe(client, monkeypatch):
|
||||
# Offline + llama backend: the remote GGUF probe (list_repo_files) must be skipped so a
|
||||
# dead-DNS session cannot hang.
|
||||
c, _saved = client
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
||||
monkeypatch.setattr(settings, "_llama_backend_active", lambda: True)
|
||||
monkeypatch.setattr(settings, "_local_gguf_backend_error", lambda model: None)
|
||||
|
||||
def _boom(*a, **k):
|
||||
raise AssertionError("hit the network for the GGUF probe")
|
||||
|
||||
monkeypatch.setattr(settings, "_hf_gguf_backend_error", _boom)
|
||||
import utils.models as _models
|
||||
|
||||
monkeypatch.setattr(_models, "is_embedding_model", lambda *a, **k: True)
|
||||
r = c.put("/embedding-model", json = {"embedding_model": "acme/embedder"})
|
||||
assert r.status_code == 200
|
||||
|
||||
|
||||
def test_llama_backend_skips_the_st_pickle_scan(monkeypatch):
|
||||
# On the llama-server backend the embedder loads GGUF (inert), not the ST repo's
|
||||
# pickle, so a flagged ST repo with a clean GGUF companion must not be rejected here.
|
||||
|
|
|
|||
583
studio/backend/tests/test_offline_embedding_minimal.py
Normal file
583
studio/backend/tests/test_offline_embedding_minimal.py
Normal file
|
|
@ -0,0 +1,583 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Offline RAG embedding-model handling (issue #6817).
|
||||
|
||||
Offline the studio must never call the Hub (a DNS-dead session hangs on retries). Using a fake
|
||||
HF cache under a temp HF_HUB_CACHE, assert that offline: is_embedding_model classifies from the
|
||||
cached modules.json without the Hub; the file-security gate fails CLOSED on an unscanned pickle
|
||||
weight with no safetensors alternative and allows an inert cache; the embedder threads
|
||||
local_files_only into the load. Online behavior is unchanged (bounded timeout + cache fallback).
|
||||
"""
|
||||
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from utils.security import evaluate_file_security
|
||||
from utils.utils import (
|
||||
hf_cache_snapshot_dir,
|
||||
hf_cache_snapshot_is_loadable,
|
||||
hf_env_offline,
|
||||
st_repo_id_candidates,
|
||||
)
|
||||
|
||||
# Minimal sentence-transformers modules.json (the marker the gate keys on).
|
||||
MODULES_JSON = (
|
||||
'[{"idx": 0, "name": "0", "path": "", "type": "sentence_transformers.models.Transformer"}]'
|
||||
)
|
||||
|
||||
|
||||
def _modules_json(*paths):
|
||||
"""modules.json listing one Transformer module per path (a load root)."""
|
||||
import json
|
||||
return json.dumps(
|
||||
[
|
||||
{
|
||||
"idx": i,
|
||||
"name": str(i),
|
||||
"path": p,
|
||||
"type": "sentence_transformers.models.Transformer",
|
||||
}
|
||||
for i, p in enumerate(paths)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
_COMMIT = "0123456789abcdef0123456789abcdef01234567"
|
||||
|
||||
|
||||
def _make_cache(
|
||||
root,
|
||||
repo_id,
|
||||
files,
|
||||
commit = _COMMIT,
|
||||
):
|
||||
"""Build a canonical HF-cache snapshot (refs/main + snapshots/<commit>/) for repo_id under
|
||||
root from {relpath: contents}; returns the snapshot dir."""
|
||||
from huggingface_hub.file_download import repo_folder_name
|
||||
|
||||
repo_dir = Path(root) / repo_folder_name(repo_id = repo_id, repo_type = "model")
|
||||
(repo_dir / "refs").mkdir(parents = True, exist_ok = True)
|
||||
(repo_dir / "refs" / "main").write_text(commit)
|
||||
snapshot = repo_dir / "snapshots" / commit
|
||||
snapshot.mkdir(parents = True, exist_ok = True)
|
||||
for rel, contents in files.items():
|
||||
path = snapshot / rel
|
||||
path.parent.mkdir(parents = True, exist_ok = True)
|
||||
path.write_text(contents)
|
||||
return snapshot
|
||||
|
||||
|
||||
def _no_network():
|
||||
"""Patch model_info to fail loudly if any offline path reaches the network."""
|
||||
return patch("huggingface_hub.model_info", side_effect = AssertionError("hit the network"))
|
||||
|
||||
|
||||
def _is_embedding_model(*args, **kwargs):
|
||||
from utils.models.model_config import is_embedding_model
|
||||
return is_embedding_model(*args, **kwargs)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def hf_cache(tmp_path, monkeypatch):
|
||||
"""Point the HF cache at a fresh temp dir."""
|
||||
root = tmp_path / "hub"
|
||||
root.mkdir()
|
||||
monkeypatch.setenv("HF_HOME", str(tmp_path))
|
||||
monkeypatch.setenv("HF_HUB_CACHE", str(root))
|
||||
return root
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def _clean_env(monkeypatch):
|
||||
"""Start each test online with an empty detection cache; offline tests opt in."""
|
||||
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
|
||||
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
|
||||
from utils.models import model_config as mc
|
||||
|
||||
mc._embedding_detection_cache.clear()
|
||||
yield
|
||||
mc._embedding_detection_cache.clear()
|
||||
|
||||
|
||||
# ── hf_env_offline ───────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on", " On "])
|
||||
def test_hf_env_offline_true(monkeypatch, value):
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", value)
|
||||
assert hf_env_offline() is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["0", "false", "no", "off", ""])
|
||||
def test_hf_env_offline_false(monkeypatch, value):
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", value)
|
||||
assert hf_env_offline() is False
|
||||
|
||||
|
||||
def test_hf_env_offline_honors_transformers_flag(monkeypatch):
|
||||
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
|
||||
assert hf_env_offline() is True
|
||||
|
||||
|
||||
def test_hf_env_offline_default_false():
|
||||
assert hf_env_offline() is False
|
||||
|
||||
|
||||
# ── st_repo_id_candidates ────────────────────────────────────────
|
||||
|
||||
|
||||
def test_candidates_slashless_adds_st_alias():
|
||||
assert st_repo_id_candidates("all-MiniLM-L6-v2") == [
|
||||
"all-MiniLM-L6-v2",
|
||||
"sentence-transformers/all-MiniLM-L6-v2",
|
||||
]
|
||||
|
||||
|
||||
def test_candidates_with_org_is_verbatim():
|
||||
assert st_repo_id_candidates("org/model") == ["org/model"]
|
||||
|
||||
|
||||
def test_candidates_empty_name():
|
||||
assert st_repo_id_candidates(" ") == []
|
||||
|
||||
|
||||
# ── hf_cache_snapshot_dir ────────────────────────────────────────
|
||||
|
||||
|
||||
def test_snapshot_dir_resolves_active_commit(hf_cache):
|
||||
snapshot = _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON})
|
||||
assert hf_cache_snapshot_dir("org/emb") == snapshot
|
||||
|
||||
|
||||
def test_snapshot_dir_none_when_uncached(hf_cache):
|
||||
assert hf_cache_snapshot_dir("org/missing") is None
|
||||
|
||||
|
||||
def test_snapshot_dir_uses_st_alias_for_slashless(hf_cache):
|
||||
snapshot = _make_cache(
|
||||
hf_cache, "sentence-transformers/all-MiniLM-L6-v2", {"modules.json": MODULES_JSON}
|
||||
)
|
||||
assert hf_cache_snapshot_dir("all-MiniLM-L6-v2") == snapshot
|
||||
|
||||
|
||||
def test_snapshot_dir_none_when_snapshot_missing(hf_cache):
|
||||
from huggingface_hub.file_download import repo_folder_name
|
||||
|
||||
repo_dir = hf_cache / repo_folder_name(repo_id = "org/broken", repo_type = "model")
|
||||
(repo_dir / "refs").mkdir(parents = True)
|
||||
(repo_dir / "refs" / "main").write_text("deadbeef") # no snapshots/deadbeef dir
|
||||
assert hf_cache_snapshot_dir("org/broken") is None
|
||||
|
||||
|
||||
def test_snapshot_dir_expands_env_vars_in_cache_path(tmp_path, monkeypatch):
|
||||
# An unexpanded $VAR in HF_HUB_CACHE must resolve where the loader looks.
|
||||
real = tmp_path / "hub"
|
||||
real.mkdir()
|
||||
monkeypatch.setenv("MY_HF_CACHE", str(real))
|
||||
monkeypatch.setenv("HF_HUB_CACHE", "$MY_HF_CACHE")
|
||||
monkeypatch.delenv("HF_HOME", raising = False)
|
||||
monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False)
|
||||
snapshot = _make_cache(real, "org/emb", {"modules.json": MODULES_JSON})
|
||||
assert hf_cache_snapshot_dir("org/emb") == snapshot
|
||||
|
||||
|
||||
def test_snapshot_dir_uses_sentence_transformers_home(tmp_path, monkeypatch):
|
||||
# ST uses SENTENCE_TRANSFORMERS_HOME as its cache_folder, so the gate must inspect it too.
|
||||
st_home = tmp_path / "st_home"
|
||||
st_home.mkdir()
|
||||
monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home))
|
||||
monkeypatch.delenv("HF_HUB_CACHE", raising = False)
|
||||
monkeypatch.delenv("HF_HOME", raising = False)
|
||||
snapshot = _make_cache(st_home, "org/emb", {"modules.json": MODULES_JSON})
|
||||
assert hf_cache_snapshot_dir("org/emb") == snapshot
|
||||
|
||||
|
||||
def test_snapshot_dir_st_home_is_exclusive(tmp_path, monkeypatch):
|
||||
# With SENTENCE_TRANSFORMERS_HOME set, ST loads only from it, so a model living only under
|
||||
# HF_HUB_CACHE must not be reported.
|
||||
st_home = tmp_path / "st_home"
|
||||
st_home.mkdir()
|
||||
hub = tmp_path / "hub"
|
||||
hub.mkdir()
|
||||
monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home))
|
||||
monkeypatch.setenv("HF_HUB_CACHE", str(hub))
|
||||
monkeypatch.delenv("HF_HOME", raising = False)
|
||||
_make_cache(hub, "org/emb", {"modules.json": MODULES_JSON}) # only in the HF hub cache
|
||||
assert hf_cache_snapshot_dir("org/emb") is None
|
||||
|
||||
|
||||
def test_snapshot_is_loadable_with_config_and_weights(hf_cache):
|
||||
_make_cache(hf_cache, "org/emb", {"config.json": "{}", "model.safetensors": "x"})
|
||||
assert hf_cache_snapshot_is_loadable("org/emb") is True
|
||||
|
||||
|
||||
def test_snapshot_is_not_loadable_when_metadata_only(hf_cache):
|
||||
# A partial cache (refs/main resolves but no weights) is not loadable.
|
||||
_make_cache(hf_cache, "org/partial", {"config.json": "{}", "modules.json": MODULES_JSON})
|
||||
assert hf_cache_snapshot_is_loadable("org/partial") is False
|
||||
|
||||
|
||||
def test_snapshot_is_not_loadable_when_uncached(hf_cache):
|
||||
assert hf_cache_snapshot_is_loadable("org/missing") is False
|
||||
|
||||
|
||||
def test_gate_blocks_pickle_in_sentence_transformers_home(tmp_path, monkeypatch):
|
||||
# A pickle under SENTENCE_TRANSFORMERS_HOME must still fail closed offline.
|
||||
st_home = tmp_path / "st_home"
|
||||
st_home.mkdir()
|
||||
monkeypatch.setenv("SENTENCE_TRANSFORMERS_HOME", str(st_home))
|
||||
monkeypatch.delenv("HF_HUB_CACHE", raising = False)
|
||||
monkeypatch.delenv("HF_HOME", raising = False)
|
||||
_make_cache(st_home, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"})
|
||||
with _no_network():
|
||||
assert evaluate_file_security("org/pk", local_only_load = True).blocked is True
|
||||
|
||||
|
||||
# ── is_embedding_model: offline (no network) ─────────────────────
|
||||
|
||||
|
||||
def test_offline_true_for_cached_st_model(hf_cache, monkeypatch):
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
||||
_make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON, "config.json": "{}"})
|
||||
with _no_network():
|
||||
assert _is_embedding_model("org/emb") is True
|
||||
|
||||
|
||||
def test_offline_false_for_cached_non_st_model(hf_cache, monkeypatch):
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
||||
_make_cache(hf_cache, "org/plain", {"config.json": "{}", "model.safetensors": "x"})
|
||||
with _no_network():
|
||||
assert _is_embedding_model("org/plain") is False
|
||||
|
||||
|
||||
def test_offline_false_when_uncached(hf_cache, monkeypatch):
|
||||
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
|
||||
with _no_network():
|
||||
assert _is_embedding_model("org/missing") is False
|
||||
|
||||
|
||||
def test_offline_slashless_resolves_via_alias(hf_cache, monkeypatch):
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
||||
_make_cache(hf_cache, "sentence-transformers/all-MiniLM-L6-v2", {"modules.json": MODULES_JSON})
|
||||
with _no_network():
|
||||
assert _is_embedding_model("all-MiniLM-L6-v2") is True
|
||||
|
||||
|
||||
def test_offline_ignores_stale_online_memo(hf_cache, monkeypatch):
|
||||
# An online lookup memoizes True for an UNCACHED repo (tags say embedding, no weights). Once
|
||||
# offline, is_embedding_model must reclassify from the empty cache and return False, not the
|
||||
# stale online True that would make settings accept a repo _get() cannot load.
|
||||
with patch(
|
||||
"huggingface_hub.model_info",
|
||||
side_effect = lambda *a, **k: SimpleNamespace(
|
||||
tags = ["sentence-transformers"], pipeline_tag = None
|
||||
),
|
||||
):
|
||||
assert _is_embedding_model("org/uncached-emb") is True # memoized True online
|
||||
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
||||
with _no_network():
|
||||
assert _is_embedding_model("org/uncached-emb") is False # recomputed from empty cache
|
||||
|
||||
|
||||
def test_offline_recomputes_after_cache_materializes(hf_cache, monkeypatch):
|
||||
# Because the offline branch never records a memo, once an uncached repo's snapshot
|
||||
# materializes (another process populates the cache) the next call re-reports True.
|
||||
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
|
||||
with _no_network():
|
||||
assert _is_embedding_model("org/later") is False # uncached
|
||||
_make_cache(hf_cache, "org/later", {"modules.json": MODULES_JSON})
|
||||
assert _is_embedding_model("org/later") is True # cache now present, no stale negative
|
||||
|
||||
|
||||
# ── is_embedding_model: online (bounded + fallback) ──────────────
|
||||
|
||||
|
||||
def test_online_passes_bounded_timeout(hf_cache):
|
||||
seen = {}
|
||||
|
||||
def _mi(
|
||||
name,
|
||||
token = None,
|
||||
timeout = None,
|
||||
**kw,
|
||||
):
|
||||
seen["timeout"] = timeout
|
||||
return SimpleNamespace(tags = ["sentence-transformers"], pipeline_tag = None)
|
||||
|
||||
with patch("huggingface_hub.model_info", side_effect = _mi):
|
||||
assert _is_embedding_model("org/emb") is True
|
||||
assert seen["timeout"] == 15.0
|
||||
|
||||
|
||||
def test_online_error_falls_back_to_cache_marker(hf_cache):
|
||||
_make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON})
|
||||
with patch("huggingface_hub.model_info", side_effect = RuntimeError("dns dead")):
|
||||
assert _is_embedding_model("org/emb") is True
|
||||
|
||||
|
||||
def test_online_error_without_cache_returns_false(hf_cache):
|
||||
with patch("huggingface_hub.model_info", side_effect = RuntimeError("dns dead")):
|
||||
assert _is_embedding_model("org/missing") is False
|
||||
|
||||
|
||||
# ── evaluate_file_security: offline fail-closed gate ─────────────
|
||||
|
||||
|
||||
def _offline_decision(name):
|
||||
return evaluate_file_security(name, local_only_load = True)
|
||||
|
||||
|
||||
def test_gate_allows_safetensors_only(hf_cache):
|
||||
_make_cache(hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"})
|
||||
with _no_network():
|
||||
assert _offline_decision("org/st").blocked is False
|
||||
|
||||
|
||||
def test_gate_blocks_pickle_without_safetensors(hf_cache):
|
||||
_make_cache(hf_cache, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"})
|
||||
with _no_network():
|
||||
decision = _offline_decision("org/pk")
|
||||
assert decision.blocked is True
|
||||
assert any(u["path"] == "pytorch_model.bin" for u in decision.unsafe_files)
|
||||
|
||||
|
||||
def test_gate_allows_pickle_with_safetensors_sibling(hf_cache):
|
||||
_make_cache(hf_cache, "org/both", {"pytorch_model.bin": "x", "model.safetensors": "y"})
|
||||
with _no_network():
|
||||
assert _offline_decision("org/both").blocked is False
|
||||
|
||||
|
||||
def test_gate_blocks_sharded_pickle(hf_cache):
|
||||
_make_cache(
|
||||
hf_cache,
|
||||
"org/shard",
|
||||
{
|
||||
"pytorch_model-00001-of-00002.bin": "a",
|
||||
"pytorch_model-00002-of-00002.bin": "b",
|
||||
},
|
||||
)
|
||||
with _no_network():
|
||||
assert _offline_decision("org/shard").blocked is True
|
||||
|
||||
|
||||
def test_gate_allows_nothing_cached(hf_cache):
|
||||
with _no_network():
|
||||
assert _offline_decision("org/missing").blocked is False
|
||||
|
||||
|
||||
def test_gate_allows_gguf_only(hf_cache):
|
||||
_make_cache(hf_cache, "org/gg", {"model.gguf": "x"})
|
||||
with _no_network():
|
||||
assert _offline_decision("org/gg").blocked is False
|
||||
|
||||
|
||||
def test_gate_blocks_pickle_in_module_subdir(hf_cache):
|
||||
# 0_Transformer is a module load root (listed in modules.json), so its pickle blocks.
|
||||
_make_cache(
|
||||
hf_cache,
|
||||
"org/mod",
|
||||
{"modules.json": _modules_json("0_Transformer"), "0_Transformer/pytorch_model.bin": "x"},
|
||||
)
|
||||
with _no_network():
|
||||
assert _offline_decision("org/mod").blocked is True
|
||||
|
||||
|
||||
def test_gate_allows_pickle_in_subdir_with_safetensors(hf_cache):
|
||||
_make_cache(
|
||||
hf_cache,
|
||||
"org/mod2",
|
||||
{
|
||||
"modules.json": _modules_json("0_Transformer"),
|
||||
"0_Transformer/pytorch_model.bin": "x",
|
||||
"0_Transformer/model.safetensors": "y",
|
||||
},
|
||||
)
|
||||
with _no_network():
|
||||
assert _offline_decision("org/mod2").blocked is False
|
||||
|
||||
|
||||
def test_gate_allows_unreferenced_nested_pickle(hf_cache):
|
||||
# A pickle in a dir NOT referenced by modules.json (e.g. nemo/) is never deserialized, so it
|
||||
# must not block the offline load (matches the online gate).
|
||||
_make_cache(
|
||||
hf_cache,
|
||||
"org/aux",
|
||||
{
|
||||
"modules.json": MODULES_JSON, # Transformer at the root only
|
||||
"model.safetensors": "w",
|
||||
"nemo/pytorch_model.bin": "x",
|
||||
},
|
||||
)
|
||||
with _no_network():
|
||||
assert _offline_decision("org/aux").blocked is False
|
||||
|
||||
|
||||
def test_gate_blocks_adapter_pickle_without_safetensors(hf_cache):
|
||||
_make_cache(hf_cache, "org/ad", {"config.json": "{}", "adapter_model.bin": "x"})
|
||||
with _no_network():
|
||||
decision = _offline_decision("org/ad")
|
||||
assert decision.blocked is True
|
||||
assert any(u["path"] == "adapter_model.bin" for u in decision.unsafe_files)
|
||||
|
||||
|
||||
def test_gate_allows_adapter_pickle_with_adapter_safetensors(hf_cache):
|
||||
_make_cache(hf_cache, "org/ad2", {"adapter_model.bin": "x", "adapter_model.safetensors": "y"})
|
||||
with _no_network():
|
||||
assert _offline_decision("org/ad2").blocked is False
|
||||
|
||||
|
||||
def test_gate_blocks_base_pickle_with_only_adapter_safetensors_decoy(hf_cache):
|
||||
# A decoy adapter_model.safetensors must NOT suppress a base pytorch_model.bin (the base
|
||||
# loader would still deserialize the unscanned pickle).
|
||||
_make_cache(hf_cache, "org/decoy", {"pytorch_model.bin": "x", "adapter_model.safetensors": "y"})
|
||||
with _no_network():
|
||||
assert _offline_decision("org/decoy").blocked is True
|
||||
|
||||
|
||||
def test_gate_blocks_adapter_pickle_with_only_base_safetensors_decoy(hf_cache):
|
||||
# Symmetric: a base model.safetensors must NOT suppress an adapter_model.bin.
|
||||
_make_cache(hf_cache, "org/decoy2", {"adapter_model.bin": "x", "model.safetensors": "y"})
|
||||
with _no_network():
|
||||
assert _offline_decision("org/decoy2").blocked is True
|
||||
|
||||
|
||||
def test_gate_reports_snapshot_relative_path(hf_cache):
|
||||
_make_cache(
|
||||
hf_cache,
|
||||
"org/mod3",
|
||||
{"modules.json": _modules_json("0_Transformer"), "0_Transformer/pytorch_model.bin": "x"},
|
||||
)
|
||||
with _no_network():
|
||||
decision = _offline_decision("org/mod3")
|
||||
assert decision.blocked is True
|
||||
assert any(u["path"] == "0_Transformer/pytorch_model.bin" for u in decision.unsafe_files)
|
||||
|
||||
|
||||
# ── evaluate_file_security: online path unchanged ────────────────
|
||||
|
||||
|
||||
def test_online_default_blocks_unsafe():
|
||||
status = {
|
||||
"scansDone": True,
|
||||
"filesWithIssues": [{"path": "pytorch_model.bin", "level": "unsafe"}],
|
||||
}
|
||||
with patch(
|
||||
"huggingface_hub.model_info",
|
||||
side_effect = lambda *a, **k: SimpleNamespace(security_repo_status = status),
|
||||
):
|
||||
assert evaluate_file_security("org/x").blocked is True
|
||||
|
||||
|
||||
def test_online_default_allows_clean():
|
||||
status = {"scansDone": True, "filesWithIssues": []}
|
||||
with patch(
|
||||
"huggingface_hub.model_info",
|
||||
side_effect = lambda *a, **k: SimpleNamespace(security_repo_status = status),
|
||||
):
|
||||
assert evaluate_file_security("org/x").blocked is False
|
||||
|
||||
|
||||
# ── embeddings guard + loader ────────────────────────────────────
|
||||
|
||||
|
||||
def test_guard_offline_blocks_pickle_only(hf_cache):
|
||||
from core.rag.embeddings import UnsafeEmbeddingModelError, _guard_model_security
|
||||
_make_cache(hf_cache, "org/pk", {"config.json": "{}", "pytorch_model.bin": "x"})
|
||||
with _no_network():
|
||||
with pytest.raises(UnsafeEmbeddingModelError):
|
||||
_guard_model_security("org/pk", local_only = True)
|
||||
|
||||
|
||||
def test_guard_offline_allows_safetensors(hf_cache):
|
||||
from core.rag.embeddings import _guard_model_security
|
||||
_make_cache(hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"})
|
||||
with _no_network():
|
||||
_guard_model_security("org/st", local_only = True) # must not raise
|
||||
|
||||
|
||||
def _install_fake_sentence_transformers(monkeypatch, captured):
|
||||
class FakeSentenceTransformer:
|
||||
def __init__(
|
||||
self,
|
||||
name,
|
||||
*,
|
||||
device = None,
|
||||
model_kwargs = None,
|
||||
local_files_only = False,
|
||||
**kw,
|
||||
):
|
||||
captured["name"] = name
|
||||
captured["device"] = device
|
||||
captured["local_files_only"] = local_files_only
|
||||
|
||||
module = types.ModuleType("sentence_transformers")
|
||||
module.SentenceTransformer = FakeSentenceTransformer
|
||||
monkeypatch.setitem(sys.modules, "sentence_transformers", module)
|
||||
|
||||
|
||||
def test_get_offline_loads_from_local_snapshot(hf_cache, monkeypatch):
|
||||
from core.rag import embeddings
|
||||
|
||||
snapshot = _make_cache(
|
||||
hf_cache, "org/st", {"modules.json": MODULES_JSON, "model.safetensors": "x"}
|
||||
)
|
||||
# TRANSFORMERS_OFFLINE only: a cached model loads from its local snapshot dir (a local path,
|
||||
# never the Hub), offline-safe on ANY sentence-transformers version.
|
||||
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
|
||||
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
|
||||
monkeypatch.setattr(embeddings, "_model", None, raising = False)
|
||||
monkeypatch.setattr(embeddings, "_name", None, raising = False)
|
||||
monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None)
|
||||
monkeypatch.setattr(embeddings, "_device", lambda: "cpu")
|
||||
captured = {}
|
||||
_install_fake_sentence_transformers(monkeypatch, captured)
|
||||
with _no_network():
|
||||
embeddings._get("org/st")
|
||||
assert captured["name"] == str(snapshot)
|
||||
|
||||
|
||||
def test_get_offline_uncached_uses_local_files_only(tmp_path, monkeypatch):
|
||||
from core.rag import embeddings
|
||||
|
||||
empty = tmp_path / "hub"
|
||||
empty.mkdir()
|
||||
monkeypatch.setenv("HF_HUB_CACHE", str(empty))
|
||||
monkeypatch.delenv("HF_HOME", raising = False)
|
||||
monkeypatch.delenv("SENTENCE_TRANSFORMERS_HOME", raising = False)
|
||||
monkeypatch.setenv("TRANSFORMERS_OFFLINE", "1")
|
||||
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
|
||||
monkeypatch.setattr(embeddings, "_model", None, raising = False)
|
||||
monkeypatch.setattr(embeddings, "_name", None, raising = False)
|
||||
monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None)
|
||||
monkeypatch.setattr(embeddings, "_device", lambda: "cpu")
|
||||
# No cache -> repo-id load forced cache-only (fails fast offline, not a hang).
|
||||
monkeypatch.setattr(embeddings, "_guard_model_security", lambda name, local_only = False: None)
|
||||
captured = {}
|
||||
_install_fake_sentence_transformers(monkeypatch, captured)
|
||||
embeddings._get("org/uncached-xyz")
|
||||
assert captured["name"] == "org/uncached-xyz"
|
||||
assert captured["local_files_only"] is True
|
||||
|
||||
|
||||
def test_get_online_omits_local_files_only(monkeypatch):
|
||||
from core.rag import embeddings
|
||||
|
||||
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
|
||||
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
|
||||
monkeypatch.setattr(embeddings, "_model", None, raising = False)
|
||||
monkeypatch.setattr(embeddings, "_name", None, raising = False)
|
||||
monkeypatch.setattr(embeddings, "_install_torchao_stub_once", lambda: None)
|
||||
monkeypatch.setattr(embeddings, "_device", lambda: "cpu")
|
||||
# Isolate the loader wiring from the online guard's network calls.
|
||||
monkeypatch.setattr(embeddings, "_guard_model_security", lambda name, local_only = False: None)
|
||||
captured = {}
|
||||
_install_fake_sentence_transformers(monkeypatch, captured)
|
||||
embeddings._get("org/online")
|
||||
assert captured["local_files_only"] is False
|
||||
|
|
@ -2076,6 +2076,24 @@ def download_gguf_file(
|
|||
_embedding_detection_cache: Dict[tuple, bool] = {}
|
||||
|
||||
|
||||
# Bound the Hub lookup so a DNS-dead session fails fast to the cache instead of hanging on retries.
|
||||
_HUB_MODEL_INFO_TIMEOUT = 15.0
|
||||
|
||||
|
||||
def _embedding_marker_in_hf_cache(model_name: str) -> bool:
|
||||
"""True when model_name's cached snapshot carries a modules.json (the ST marker).
|
||||
Cache-only, no network; used offline and as a fallback when the Hub lookup times out."""
|
||||
from utils.utils import hf_cache_snapshot_dir
|
||||
|
||||
snapshot = hf_cache_snapshot_dir(model_name)
|
||||
if snapshot is None:
|
||||
return False
|
||||
try:
|
||||
return (snapshot / "modules.json").is_file()
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
|
||||
"""Detect embedding/sentence-transformer models via HF metadata.
|
||||
|
||||
|
|
@ -2090,6 +2108,15 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
|
|||
Returns:
|
||||
True if embedding model, else False (default for local paths or errors).
|
||||
"""
|
||||
from utils.utils import hf_env_offline
|
||||
|
||||
# Offline (remote repo): reclassify from the local cache on every call, before/without the
|
||||
# memo. An online lookup can memoize True from tags with no weights cached, so trusting it once
|
||||
# the session goes offline would accept a repo _get() cannot load; a cached negative can also be
|
||||
# invalidated by later cache materialization. The cache probe is local-only, so it's cheap.
|
||||
if not is_local_path(model_name) and hf_env_offline():
|
||||
return _embedding_marker_in_hf_cache(model_name)
|
||||
|
||||
cache_key = (model_name, hf_token)
|
||||
if cache_key in _embedding_detection_cache:
|
||||
return _embedding_detection_cache[cache_key]
|
||||
|
|
@ -2104,7 +2131,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
|
|||
try:
|
||||
from huggingface_hub import model_info as hf_model_info
|
||||
|
||||
info = hf_model_info(model_name, token = hf_token)
|
||||
info = hf_model_info(model_name, token = hf_token, timeout = _HUB_MODEL_INFO_TIMEOUT)
|
||||
tags = set(info.tags or [])
|
||||
pipeline_tag = info.pipeline_tag or ""
|
||||
|
||||
|
|
@ -2125,9 +2152,11 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
|
|||
return is_emb
|
||||
|
||||
except Exception as e:
|
||||
# Timeout or transient network error: fall back to the local cache marker, don't hard-fail.
|
||||
logger.warning(f"Could not determine if {model_name} is embedding model: {e}")
|
||||
_embedding_detection_cache[cache_key] = False
|
||||
return False
|
||||
is_emb = _embedding_marker_in_hf_cache(model_name)
|
||||
_embedding_detection_cache[cache_key] = is_emb
|
||||
return is_emb
|
||||
|
||||
|
||||
def _has_model_weight_files(model_dir: Path) -> bool:
|
||||
|
|
|
|||
|
|
@ -29,13 +29,35 @@ Policy:
|
|||
scanned so a repo cannot dodge the gate by suffixing its name.
|
||||
"""
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
from loggers import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Pickle-format weight files (plain or sharded) that execute code on load; safetensors/gguf
|
||||
# are inert. Grouped by weight family so an inert safetensors only suppresses the pickle it
|
||||
# actually replaces: the loader won't use an adapter's safetensors for pytorch_model.bin.
|
||||
_PICKLE_WEIGHT_RE = re.compile(
|
||||
r"^(model|pytorch_model|adapter_model|consolidated)(-\d+-of-\d+)?"
|
||||
r"\.(bin|pt|pth|ckpt|pkl|pickle)$",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
# Base-model safetensors set: HF names the base pickle pytorch_model.bin but the safetensors
|
||||
# model.safetensors (stems differ), so a base pickle is replaced only by these, not an adapter's.
|
||||
_BASE_SAFETENSORS_RE = re.compile(
|
||||
r"^(model(-\d+-of-\d+)?\.safetensors|model\.safetensors\.index\.json)$",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
# Adapter (PEFT) safetensors set: adapter_model.safetensors, its shards, or index.
|
||||
_ADAPTER_SAFETENSORS_RE = re.compile(
|
||||
r"^(adapter_model(-\d+-of-\d+)?\.safetensors|adapter_model\.safetensors\.index\.json)$",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
# Non-blocking levels: clean or not-yet-finished. Anything else (unsafe/suspicious/
|
||||
# malicious or a future label) blocks, so Hub schema drift fails CLOSED.
|
||||
_NONBLOCKING_LEVELS = frozenset(
|
||||
|
|
@ -265,11 +287,105 @@ def _fetch_security_status(model_name: str, hf_token: Optional[str]):
|
|||
return None
|
||||
|
||||
|
||||
def _st_load_roots(snapshot: Path) -> list:
|
||||
"""Directories a SentenceTransformer load deserializes weights from: the snapshot root plus
|
||||
each module path in modules.json. Local, no network. Mirrors the online gate (which ignores
|
||||
unreferenced nested pickles ST never loads) so the offline gate doesn't over-block."""
|
||||
roots = [snapshot]
|
||||
try:
|
||||
import json
|
||||
modules = json.loads((snapshot / "modules.json").read_text())
|
||||
except (OSError, ValueError):
|
||||
return roots # no / invalid modules.json -> snapshot root is the only load root
|
||||
for module in modules or ():
|
||||
path = str((module or {}).get("path", "")).strip().strip("/")
|
||||
# Relative module path only; ignore a crafted "../" escape.
|
||||
if path and ".." not in path.split("/"):
|
||||
candidate = snapshot / path
|
||||
if candidate not in roots:
|
||||
roots.append(candidate)
|
||||
return roots
|
||||
|
||||
|
||||
def _cached_pickle_weight_files(snapshot: Path) -> list:
|
||||
"""Pickle weight files in snapshot's ST load roots, EXCLUDING those whose weight family also
|
||||
ships an inert safetensors in the same dir (the loader prefers it): a base pickle is suppressed
|
||||
only by a base model.safetensors, an adapter pickle only by adapter_model.safetensors -- an
|
||||
unrelated safetensors is no substitute. Load roots only. Raises OSError if the snapshot root is
|
||||
unreadable (caller blocks)."""
|
||||
blocked = []
|
||||
for root in _st_load_roots(snapshot):
|
||||
try:
|
||||
entries = [p for p in root.iterdir() if p.is_file()]
|
||||
except OSError:
|
||||
if root == snapshot:
|
||||
raise # top-level unreadable -> fail closed
|
||||
continue # unreadable module subdir: nothing loadable to attest here
|
||||
has_base_safetensors = any(_BASE_SAFETENSORS_RE.match(p.name) for p in entries)
|
||||
has_adapter_safetensors = any(_ADAPTER_SAFETENSORS_RE.match(p.name) for p in entries)
|
||||
for path in entries:
|
||||
if not _PICKLE_WEIGHT_RE.match(path.name):
|
||||
continue
|
||||
is_adapter = path.name.lower().startswith("adapter_model")
|
||||
has_alternative = has_adapter_safetensors if is_adapter else has_base_safetensors
|
||||
if not has_alternative:
|
||||
blocked.append(path)
|
||||
return blocked
|
||||
|
||||
|
||||
def _evaluate_local_only(model_name: str) -> FileSecurityDecision:
|
||||
"""Offline security gate. The Hub scan is unreachable, so inspect the local cache and fail
|
||||
CLOSED on an unscanned pickle weight with no inert safetensors alternative, rather than
|
||||
failing open or hanging. Safetensors/gguf-only cache loads; nothing cached -> allowed."""
|
||||
from utils.utils import hf_cache_snapshot_dir
|
||||
|
||||
try:
|
||||
snapshot = hf_cache_snapshot_dir(model_name)
|
||||
except Exception:
|
||||
logger.warning("Offline gate: could not resolve the cache for '%s'; blocking.", model_name)
|
||||
return FileSecurityDecision(
|
||||
model_name, True, reason = "offline; could not inspect the local cache"
|
||||
)
|
||||
|
||||
if snapshot is None:
|
||||
return FileSecurityDecision(model_name, False, reason = "offline; nothing cached to load")
|
||||
|
||||
try:
|
||||
pickles = _cached_pickle_weight_files(snapshot)
|
||||
except OSError:
|
||||
logger.warning("Offline gate: could not read the cache for '%s'; blocking.", model_name)
|
||||
return FileSecurityDecision(
|
||||
model_name, True, reason = "offline; could not read the local cache"
|
||||
)
|
||||
|
||||
if not pickles:
|
||||
return FileSecurityDecision(
|
||||
model_name, False, reason = "offline; cached weights are inert (safetensors/gguf)"
|
||||
)
|
||||
|
||||
# Snapshot-relative posix paths (match the online gate; disambiguate same-named pickles).
|
||||
rel_paths = sorted(p.relative_to(snapshot).as_posix() for p in pickles)
|
||||
names = ", ".join(rel_paths)
|
||||
logger.warning(
|
||||
"Blocking offline load of '%s': cached pickle weight(s) cannot be malware-scanned "
|
||||
"offline and have no safetensors alternative (%s).",
|
||||
model_name,
|
||||
names,
|
||||
)
|
||||
return FileSecurityDecision(
|
||||
model_name,
|
||||
True,
|
||||
unsafe_files = [{"path": rel, "level": "unscanned"} for rel in rel_paths],
|
||||
reason = f"offline; unscanned pickle weights with no safetensors alternative: {names}",
|
||||
)
|
||||
|
||||
|
||||
def evaluate_file_security(
|
||||
model_name: str,
|
||||
hf_token: Optional[str] = None,
|
||||
*,
|
||||
load_subdirs = (),
|
||||
local_only_load: bool = False,
|
||||
) -> FileSecurityDecision:
|
||||
"""Block a load when HF's security scan flags unsafe serialized files.
|
||||
|
||||
|
|
@ -280,6 +396,9 @@ def evaluate_file_security(
|
|||
``load_subdirs`` names subdirs the load calls ``from_pretrained`` on (e.g. ``("LLM",)``
|
||||
for Spark-TTS / BiCodec, loading ``<snapshot>/LLM``): a flagged file directly under one
|
||||
is root-level there and blocks, and an index inside it is honored when scoping shards.
|
||||
|
||||
``local_only_load`` marks an offline load: with the Hub scan unreachable, inspect the local
|
||||
cache and fail CLOSED on an unscanned pickle weight with no safetensors alternative.
|
||||
"""
|
||||
# Scan the repo the load actually fetches, not the literal alias (which 404s and
|
||||
# fails open): the Spark-TTS "<parent>/LLM" alias is really unsloth/<parent> from LLM/.
|
||||
|
|
@ -295,6 +414,10 @@ def evaluate_file_security(
|
|||
# Cannot classify the path -> do not block on that account.
|
||||
return FileSecurityDecision(model_name, False, reason = "path check failed; not blocked")
|
||||
|
||||
# Offline: inspect the local cache and fail closed rather than hang on model_info or fail open.
|
||||
if local_only_load:
|
||||
return _evaluate_local_only(model_name)
|
||||
|
||||
status = _fetch_security_status(model_name, hf_token)
|
||||
if not isinstance(status, dict):
|
||||
return FileSecurityDecision(
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import structlog
|
|||
from loggers import get_logger
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
|
|
@ -15,6 +16,110 @@ import tempfile
|
|||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
# ── Offline / HF-cache helpers ──────────────────────────────────
|
||||
# An offline load must never touch the network (a DNS-dead session hangs on hub retries);
|
||||
# these read the local HF cache the load itself uses.
|
||||
|
||||
_HF_OFFLINE_TRUE_VALUES = frozenset({"1", "true", "yes", "on"})
|
||||
|
||||
|
||||
def hf_env_offline() -> bool:
|
||||
"""True when HF_HUB_OFFLINE or TRANSFORMERS_OFFLINE requests offline mode.
|
||||
|
||||
Also honors TRANSFORMERS_OFFLINE (hub honors only HF_HUB_OFFLINE) since users set it
|
||||
to keep transformers loads local.
|
||||
"""
|
||||
for var in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"):
|
||||
if os.environ.get(var, "").strip().lower() in _HF_OFFLINE_TRUE_VALUES:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def st_repo_id_candidates(model_name: str) -> list:
|
||||
"""Repo ids a Sentence-Transformers load may resolve model_name to; a slashless name
|
||||
also resolves under the sentence-transformers/ namespace, so both are candidates."""
|
||||
name = (model_name or "").strip().strip("/")
|
||||
if not name:
|
||||
return []
|
||||
candidates = [name]
|
||||
if "/" not in name:
|
||||
candidates.append(f"sentence-transformers/{name}")
|
||||
return candidates
|
||||
|
||||
|
||||
def _expand_path(raw: str) -> Path:
|
||||
"""Expand ~ and $VARS as huggingface_hub does, so the gate resolves the loader's dir."""
|
||||
return Path(os.path.expandvars(os.path.expanduser(raw)))
|
||||
|
||||
|
||||
def _hf_cache_roots() -> list:
|
||||
"""The one cache root the loader resolves to, by its own precedence (it picks ONE
|
||||
cache_folder, no fall-through): SENTENCE_TRANSFORMERS_HOME, else HF_HUB_CACHE, else
|
||||
HF_HOME/hub, else ~/.cache/huggingface/hub. Expanded, read from env, one-element list."""
|
||||
st_home = os.environ.get("SENTENCE_TRANSFORMERS_HOME")
|
||||
if st_home:
|
||||
return [_expand_path(st_home)]
|
||||
hub = os.environ.get("HF_HUB_CACHE") or os.environ.get("HUGGINGFACE_HUB_CACHE")
|
||||
if hub:
|
||||
return [_expand_path(hub)]
|
||||
hf_home = os.environ.get("HF_HOME")
|
||||
if hf_home:
|
||||
return [_expand_path(hf_home) / "hub"]
|
||||
return [Path.home() / ".cache" / "huggingface" / "hub"]
|
||||
|
||||
|
||||
def hf_cache_snapshot_dir(model_name: str) -> Optional[Path]:
|
||||
"""Active local snapshot dir for model_name's main revision, or None if not cached.
|
||||
Reads refs/main then snapshots/<commit>; no network. Tries the ST alias for slashless names."""
|
||||
try:
|
||||
from huggingface_hub.file_download import repo_folder_name
|
||||
except Exception:
|
||||
repo_folder_name = None
|
||||
for cache_root in _hf_cache_roots():
|
||||
for repo_id in st_repo_id_candidates(model_name):
|
||||
try:
|
||||
if repo_folder_name is not None:
|
||||
folder = repo_folder_name(repo_id = repo_id, repo_type = "model")
|
||||
else:
|
||||
folder = "models--" + repo_id.replace("/", "--")
|
||||
repo_dir = cache_root / folder
|
||||
ref = repo_dir / "refs" / "main"
|
||||
if not ref.is_file():
|
||||
continue
|
||||
commit = ref.read_text().strip()
|
||||
if not commit:
|
||||
continue
|
||||
snapshot = repo_dir / "snapshots" / commit
|
||||
if snapshot.is_dir():
|
||||
return snapshot
|
||||
except OSError:
|
||||
continue
|
||||
return None
|
||||
|
||||
|
||||
# A weight file plus a config distinguishes a real cached model from a metadata-only
|
||||
# partial cache that resolves refs/main but would fail at load time.
|
||||
_LOADABLE_WEIGHT_SUFFIXES = frozenset({".safetensors", ".bin", ".gguf", ".pt", ".pth", ".ckpt"})
|
||||
|
||||
|
||||
def hf_cache_snapshot_is_loadable(model_name: str) -> bool:
|
||||
"""True when model_name's snapshot is cached and loadable: a config (config.json or
|
||||
modules.json) plus at least one weight file, not a metadata-only partial cache. No network."""
|
||||
snapshot = hf_cache_snapshot_dir(model_name)
|
||||
if snapshot is None:
|
||||
return False
|
||||
try:
|
||||
has_config = (snapshot / "config.json").is_file() or (snapshot / "modules.json").is_file()
|
||||
if not has_config:
|
||||
return False
|
||||
for path in snapshot.rglob("*"):
|
||||
if path.suffix.lower() in _LOADABLE_WEIGHT_SUFFIXES and path.is_file():
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
return False
|
||||
|
||||
|
||||
# ── Client-safe error helpers ───────────────────────────────────
|
||||
# Never return raw exception text to clients; log server-side, return generic.
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue