studio: classify embedding models from the HF cache and honor offline mode - #7218
Conversation
… 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 (unslothai#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.
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 033d60720b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…p 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.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 763ef72ceb
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…che 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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: abf5874990
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 21bef982a5
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…e 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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e9c7da63fc
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…fault 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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bfc1cff1de
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3468fa974f
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…asing 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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 202b23f5de
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cb1f5381d7
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: df32b909a2
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 38b78c8d72
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
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.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f94f3ecb86
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…e 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.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f08dc34bc6
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| # Dirs whose loader (AutoModel.from_pretrained) honors a sharded safetensors index: the snapshot | ||
| # root plus each Transformer-typed module subdir. Elsewhere a sharded index is NOT a substitute | ||
| # for a live pickle (see _dir_has_loadable_safetensors). | ||
| index_honoring = _transformer_load_roots(snap) | {snap} |
There was a problem hiding this comment.
Fresh evidence after the earlier shard-index fix: this still unconditionally treats the snapshot root as an AutoModel.from_pretrained root. When modules.json declares a root StaticEmbedding/WordEmbeddings/other non-Transformer module, SentenceTransformer calls that module's load() on snap, and those loaders ignore model.safetensors.index.json and fall back to pytorch_model.bin; a complete root safetensors index therefore suppresses the live root pickle here, so the offline gate reports the cache pickle-free and allows an unscanned pickle to deserialize. Only add snap to index_honoring when the root is actually loaded through a Transformer/from_pretrained path.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 1c7170d. _cached_pickle_weight_paths now credits the sharded index at the snapshot root only when _root_honors_sharded_index(snap) is true: no/unreadable modules.json (a plain from_pretrained root) or a Transformer-shaped root module. A modules.json root module of a non-Transformer type (StaticEmbedding/WordEmbeddings/Dense) no longer credits the root index, so a live root pickle behind a complete shard index still blocks. Added test_root_non_transformer_module_pickle_blocks_despite_shard_index.
| # model/pytorch_model .safetensors/.bin, plain or sharded. Matched by NAME not suffix so a | ||
| # training_args.bin / adapter-only artifact can't pass offline validation then fail on load. | ||
| _ST_WEIGHT_FILE_RE = re.compile(r"^(model|pytorch_model)(-\d+-of-\d+)?\.(safetensors|bin)$") | ||
| _ST_SHARD_RE = re.compile(r"^(model|pytorch_model)-(\d+)-of-(\d+)\.(safetensors|bin)$") |
There was a problem hiding this comment.
Fresh evidence after the unsharded decoy fix: this shard regex still accepts invalid load names such as model-00001-of-00001.bin and pytorch_model-00001-of-00001.safetensors (with matching .index.json files below). Transformers/SentenceTransformers only probe sharded model.safetensors.index.json or pytorch_model.bin.index.json, so an offline cache with config/tokenizer plus one of these decoy shard sets is classified as complete, then the first local_files_only load fails because no real weight set exists.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 1c7170d. Shard matching is now restricted to the loader-probed stem/ext pairs (model-.safetensors + model.safetensors.index.json, pytorch_model-.bin + pytorch_model.bin.index.json); the decoy sets model-.bin and pytorch_model-.safetensors and their index maps are no longer treated as a complete weight set. Added test_dir_weight_set_decoy_shard_set_is_incomplete plus a real-sharded-safetensors control.
| root = _canonical_load_dir(snap, str(module.get("path") or "")) | ||
| if root is not None: | ||
| roots.add(root) |
There was a problem hiding this comment.
Fresh evidence after the traversal fix: _canonical_load_dir() returns None for a modules.json path that is absolute or escapes the snapshot, but this branch treats that like a harmless root path and silently omits it from roots. SentenceTransformer resolves such declared paths outside the snapshot, so with a forced/offline custom embedder that points a module at an existing external directory containing pytorch_model.bin, _evaluate_local_only() can report no cached pickle while the loader deserializes that external pickle. Treat invalid or escaping module paths as an unverifiable cache and block instead of dropping them.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 1c7170d. _st_load_roots now routes each declared module path and load_subdirs entry through _load_root_or_raise, which fails closed (_SnapshotEscapeError) on an absolute or snapshot-escaping path instead of silently dropping it, while still treating an empty/'.' path as the root. Added test_escaping_module_path_blocks.
| ``Dense`` / ``WordEmbeddings`` / ``StaticEmbedding`` read a single flat weight file through | ||
| ``Module.load_torch_weights`` and never consult an index, so their sharded-index credit stays | ||
| gated off.""" | ||
| return str(type_str).rsplit(".", 1)[-1] == "Transformer" |
There was a problem hiding this comment.
Fresh evidence after the exact Transformer-submodule fix: _transformer_load_roots() only credits sharded safetensors when the declared class name is exactly Transformer. CLIPModel and MLMTransformer modules are also Transformer-shaped load roots that use the from_pretrained path and honor model.safetensors.index.json, so a complete sharded-safetensors CLIP/MLM module with a legacy pytorch_model.bin sibling is still treated as an active pickle and blocked offline even though the loader chooses the inert shard set.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 1c7170d. _is_transformer_module_type now recognizes the Transformer subclasses (any class name containing 'transformer', plus 'clipmodel'), mirroring the classifier's _ST_TRANSFORMER_SHAPED_MODULE_NAMES + 'transformer' in cls dispatch, so a sharded-safetensors CLIP/MLM submodule with a legacy pytorch_model.bin sibling is credited and no longer blocked. Added test_sharded_index_credited_in_clip_and_mlm_submodules.
| if "model.safetensors" in names or "pytorch_model.bin" in names: | ||
| return True |
There was a problem hiding this comment.
Fresh evidence after the unreadable-index fix: this returns True for pytorch_model.bin before considering a present model.safetensors.index.json. Transformers probes the safetensors index before the PyTorch bin, so a cache with pytorch_model.bin plus a malformed or incomplete safetensors index is accepted by offline validation and persisted, then the first local_files_only load follows the bad index and fails instead of falling back to the bin.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 1c7170d. _dir_weight_set_is_complete now walks the weight set in the exact from_pretrained probe order (model.safetensors, model.safetensors.index.json, pytorch_model.bin, pytorch_model.bin.index.json), so a pytorch_model.bin behind a malformed/incomplete safetensors index is no longer accepted as complete (the loader follows the bad index and fails). Added test_dir_weight_set_pickle_behind_bad_safetensors_index_is_incomplete plus a pickle-alone control.
| rel = str(model_id or "").strip().strip("/") | ||
| if not child_cls or not rel or ".." in Path(rel).parts: | ||
| return False # malformed / traversing child reference |
There was a problem hiding this comment.
Fresh evidence after the security-side Router traversal fix: a Router saved in 1_Router/ can declare a child like ../evil, which resolves to snap/evil and is loaded by Router.load(); the offline pickle gate now handles that in-snapshot sibling, but this classifier still rejects any .. before resolving it. A complete cached Router with such a sibling child is therefore reported as not loadable offline even though SentenceTransformer(..., local_files_only=True) can load it.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The observation is accurate that _router_dir_is_loadable rejects a '..' child before resolving it, but this is below the bar for this PR. A Router with a '../evil' sibling child is not producible by any real SentenceTransformer save (children are always nested subdir names), so it requires a hand-crafted modules layout, and the consequence is benign: the model is classified 'not an embedding model' (an over-reject), not mis-loaded. The security-side gate independently handles the actual in-snapshot '../' router child. Relaxing the '..' guard here is also risky: it doubles as the router recursion-termination guard ('children resolve deeper'), so allowing '..' without threading the snapshot root through the mutually-recursive loader-shape validators for containment, plus a visited-set/depth cap, could regress into unbounded recursion or over-acceptance. That risk is disproportionate to a hand-crafted-only, benign over-reject, so leaving as-is. Happy to revisit if a real saved model exhibits this layout.
…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.
for more information, see https://pre-commit.ci
|
@codex review |
|
Codex Review: Didn't find any major issues. Keep them coming! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
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.
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.
for more information, see https://pre-commit.ci
|
@codex review |
|
Codex Review: Didn't find any major issues. What shall we delve into next? Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
…lothai#6817 Update the branch onto current main and reduce the change to the minimal fix for the offline hang reported in issue unslothai#6817: - is_embedding_model: classify from the local HF cache when offline (no model_info call), and bound the online model_info with a timeout so an unflagged DNS-dead session fails fast and falls back to the cache. - evaluate_file_security: add a fail-closed offline path that inspects the cached snapshot instead of running the Hub malware scan. - RAG warm and the settings route: read the offline flag once and load the Sentence-Transformers model from the local snapshot with no network calls. The net change against main is 7 files (+1013 / -37).
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 92cc0e7f35
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| except Exception: | ||
| return |
There was a problem hiding this comment.
When offline inspection itself raises, this blanket handler returns and _get() immediately constructs SentenceTransformer from the cached snapshot, so an unverified pickle can still be deserialized. This is especially problematic for malformed/unreadable cache metadata that causes evaluate_file_security(..., local_only_load=True) to raise rather than return a blocking decision; offline mode has no Hub scan to compensate. Re-raise an UnsafeEmbeddingModelError (or otherwise block) when local_only is true.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The offline gate already fails closed on every reachable case: _evaluate_local_only returns a blocking decision when the cache cannot be resolved or read, rather than raising. The only way evaluate_file_security(local_only_load=True) raises is a valid-JSON-but-wrong-shape modules.json, and that same file makes SentenceTransformer's own _load_sbert_model raise while reading modules.json, before any module weight is deserialized, so there is no reachable path where the swallow lets an unscanned pickle load. I have left the guard unchanged; glad to add the offline fail-closed invariant as defense in depth if preferred, but it is not fixing a reachable bug.
| 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 |
There was a problem hiding this comment.
For a transformers-native cached embedder with config.json and weights but no tokenizer files, this returns True, so the offline settings fallback accepts and persists the model. _get() then loads the snapshot path through SentenceTransformer, whose Transformer module initializes an AutoTokenizer from that same local directory and fails. Validate the tokenizer assets alongside config and weights before using this fallback as evidence that the cache is loadable.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
hf_cache_snapshot_is_loadable only gates whether a non-force save shows a forceable 409 (settings.py:483 gates the check on not force), and the security gate at settings.py:436 already 403-blocks a pickle-weight partial cache first, so this needs inert safetensors weights. There the SentenceTransformer load fails identically at first encode whether or not the 409 was shown, and force persists either way, so it is a bypassable verify-time warning rather than a reachable crash or wrong result. Leaving the loadability check unchanged.
| if hf_env_offline(): | ||
| is_emb = _embedding_marker_in_hf_cache(model_name) | ||
| _embedding_detection_cache[cache_key] = is_emb |
There was a problem hiding this comment.
After an online metadata check has cached True for a repo, this early return bypasses the new offline cache-marker check entirely. In particular, the model-picker routes can populate this memo without downloading weights; if the process is then put offline, settings accepts the uncached repo because is_embedding_model() remains true, while _get() has no snapshot and the subsequent local-only SentenceTransformer load fails. Check cache materialization before trusting an online memo in offline mode (and avoid retaining stale negative results as well).
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in ef49462. is_embedding_model consulted the process memo before the offline branch, so an online lookup that memoized True from tags (no weights cached) was returned once the session went offline, and the ungated check-embedding route can populate that memo. The offline cache-marker reclassification now runs ahead of the memo and is never recorded, so an offline verdict always reflects the local cache and a later cache materialization is not masked by a stale negative. Regression tests added.
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.
|
@codex review |
|
Codex Review: Didn't find any major issues. Nice work! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
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.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 48309dc295
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # 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: |
There was a problem hiding this comment.
When offline, a custom cached SentenceTransformer repo containing only modules.json (and possibly config metadata) makes is_embedding_model() return true, so this entire block is skipped and the setting is saved without force. The loadability check therefore only protects transformers-native models; _get() subsequently selects that partial snapshot and fails when SentenceTransformer needs the missing weights. Apply the same hf_cache_snapshot_is_loadable() requirement to the offline positive-ST path so incomplete cached downloads return the intended forceable 409.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Offline, hub is forced into local-files-only mode (the constant above), so a partial ST snapshot (modules.json, no weights) does not hang: _get() loads the local snapshot path and SentenceTransformer raises an immediate local error on the missing weights, no network. That is a clean failure on an incomplete cache (not well-formed input), not a reachable hang or wrong result, so it is below the bar for a code change here.
| 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) |
There was a problem hiding this comment.
With the llama-server backend and offline mode enabled, this skips the only remote-GGUF validation and accepts an uncached GGUF-named model without force. The later LlamaServerBackend._resolve_model_path() still calls list_repo_files() before hf_hub_download(), so the first ingestion/retrieval either hangs under TRANSFORMERS_OFFLINE (which huggingface_hub does not honor) or fails under HF_HUB_OFFLINE. Check the selected GGUF candidate in the local cache before saving, or reject it with the existing forceable 409 when no cached GGUF is available.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
With the offline flag set (always at launch for the studio process; nothing sets it at runtime here), huggingface_hub honors it, so list_repo_files()/hf_hub_download() in _resolve_model_path() fail fast offline (LocalEntryNotFoundError, ~0ms) rather than hang. The residual is a clean error on an uncached GGUF, and the llama-server GGUF preflight is outside this PR's sentence-transformers embedding scope (a separate #6817 follow-up).
| # 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): |
There was a problem hiding this comment.
When only TRANSFORMERS_OFFLINE=1 is set and the installed SentenceTransformer predates local_files_only, this branch leaves load_target as the remote repo ID and passes no offline flag. SentenceTransformer then delegates to huggingface_hub, which does not honor TRANSFORMERS_OFFLINE, so an uncached model performs the network request this change is intended to avoid and can hang on DNS retries. Enforce Hub offline mode around this fallback or fail locally before calling the older constructor.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The premise is not accurate for huggingface_hub 0.36.2: constants.HF_HUB_OFFLINE = _is_true(HF_HUB_OFFLINE or TRANSFORMERS_OFFLINE) is evaluated at import, so hub DOES honor TRANSFORMERS_OFFLINE. Verified: with TRANSFORMERS_OFFLINE=1 set at launch, hf_hub_download() on an uncached repo with NO local_files_only kwarg raises LocalEntryNotFoundError in 0.00s (no network, no hang) - the download forces offline from the constant. The studio sets the offline env at launch (nothing sets it at runtime for the RAG embedder; _hf_offline_if_dns_dead wraps only llama inference), so an old-ST uncached load fails fast rather than reaching the network.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 48309dc295
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| elif _st_accepts_local_files_only(SentenceTransformer): | ||
| st_kwargs["local_files_only"] = True |
There was a problem hiding this comment.
When only TRANSFORMERS_OFFLINE is set and hf_cache_snapshot_dir() cannot resolve a snapshot (for example, an uncached model), SentenceTransformers versions that predate local_files_only take neither branch and receive the remote repo id unchanged. TRANSFORMERS_OFFLINE is not honored by huggingface_hub, so those versions can still attempt network requests and hang instead of failing locally, despite this code classifying the session as offline. Set HF_HUB_OFFLINE for the load or raise a cache-only error when the constructor lacks local_files_only.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Same as the P1 above: hub 0.36.2 folds TRANSFORMERS_OFFLINE into HF_HUB_OFFLINE at import and the underlying download forces offline from that constant, so an uncached load on an ST predating local_files_only raises LocalEntryNotFoundError immediately (0.00s) instead of reaching the network. Not reachable for the launch-set offline session the studio actually runs.
| 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 |
There was a problem hiding this comment.
When model_info times out or otherwise fails without offline mode being set, this fallback returns True for any snapshot containing modules.json, including a metadata-only partial cache with no weights. The settings route then skips its hf_cache_snapshot_is_loadable() check because that check only runs after a False classification, so a normal non-force save succeeds even though the embedding load still cannot run while the Hub is unavailable. Require a loadable snapshot for this fallback or preserve the forceable 409 response.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
is_embedding_model is a classifier, not a loadability gate - conflating them would mis-report a legitimately-cached embedder mid-download. The over-accept is the same partial-cache case: offline the load fails fast locally (no hang), and the online model_info-timeout path is the separate unflagged-network follow-up (out of this PR's scope). A clean error on an incomplete cache is below the bar.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 48309dc295
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| elif _st_accepts_local_files_only(SentenceTransformer): | ||
| st_kwargs["local_files_only"] = True |
There was a problem hiding this comment.
When only TRANSFORMERS_OFFLINE=1 is set and the embedding model is not already resolved to a snapshot, legacy SentenceTransformers versions without an explicit local_files_only constructor parameter receive neither a local path nor a cache-only flag. Those versions proceed through their Hugging Face download path, which honors HF_HUB_OFFLINE rather than TRANSFORMERS_OFFLINE, so this path can still attempt DNS/network retries—the exact offline hang this change is intended to avoid. This affects installs using an older supported SentenceTransformers release with an uncached default or force-saved embedding model.
Useful? React with 👍 / 👎.
3 participants
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.
Fixes #6817