From c2114d64dd962f94dd54ccf27376e37cb112ec39 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 23 Jul 2026 20:06:30 -0700 Subject: [PATCH] Studio: fail closed on index-referenced nested pickle shards in the offline embedding gate (#7366) * Studio: fail closed on index-referenced nested pickle shards in the offline embedding gate The offline embedding security gate (HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE) only scanned the direct files of each SentenceTransformer load root and never parsed local weight indexes, so a cached snapshot whose pytorch_model.bin.index.json maps a weight to a nested shard (e.g. shards/pytorch_model-00001-of-00001.bin) was treated as inert and allowed. The loader then follows the index into the subdir and unpickles the shard. The online gate already blocks index-referenced subdir pickles, so the offline path was strictly weaker. Parse each local weight index in a load root and follow weight_map into nested dirs, flagging any referenced pickle-extension shard. Paths resolve lexically (normpath), never Path.resolve(), since HF cache snapshot files symlink into blobs/ and resolving would leave the snapshot dir and false-block every sharded model offline. An absolute path, a .. traversal that escapes the snapshot, or an unreadable/invalid index fails closed. The existing safetensors-sibling suppression is kept. * Studio: classify offline indexed shards by torch.load path, not pickle extension load_state_dict picks safetensors vs torch.load per shard by the shard's own suffix, so two offline-gate gaps remained: - A model.safetensors.index.json whose weight_map points at a .bin shard was suppressed by has_base_safetensors (the index file itself matches the base safetensors regex), yet Transformers still torch.loads that shard. Only the pytorch index is superseded by a base safetensors now; a safetensors index is the chosen archive, so its non-safetensors targets are always flagged. - A pytorch index can map weights to arbitrary names (shards/payload, weights.data); the loader torch.loads any target not ending in .safetensors. Flag indexed shards by that rule instead of a pickle-extension allowlist. Restrict the scan to the two torch-family indexes (tf/flax load via non-pickle loaders). Add regression tests for both cases. * Studio: match offline weight-index filenames case-insensitively The index-name check compared the on-disk filename exactly, while the surrounding weight and safetensors matches use case-insensitive rules. On a case-insensitive volume (Windows or macOS) from_pretrained opens an oddly-cased cache file such as PYTORCH_MODEL.BIN.INDEX.JSON when it requests the canonical lowercase name, so the exact-case check skipped it and a nested pickle shard it referenced was allowed through. Lower-case the index name before matching, as the rest of the gate does, and add a regression test. * Studio: match load_state_dict format/selection exactly in the offline index scan Two edge cases in the offline weight-index scan: - load_state_dict decides safetensors vs torch.load with a case-sensitive endswith(".safetensors"), so a shard named payload.SAFETENSORS still deserializes via torch.load. Classify indexed shard suffixes case-sensitively to match, instead of lower-casing (which treated such a shard as inert). - A complete direct model.safetensors is selected before either sharded index, so a stale model.safetensors.index.json referencing a .bin shard never loads. Skip both indexes when a direct model.safetensors is present, so an otherwise loadable model is not over-blocked. Add regression tests for both. * Studio: read the offline weight index as UTF-8 Path.read_text() uses the locale default, which is cp1252 on Windows, so a UTF-8 weight index with non-ASCII bytes raised UnicodeDecodeError and the gate blocked an otherwise loadable model. JSON is UTF-8 by spec (and how the loader reads it), so pin the encoding. * Studio: resolve safetensors alternatives via the loader's own filename lookup The offline gate decided a safetensors alternative existed by case-folding the directory listing. On a case-sensitive filesystem that let an uppercase decoy such as MODEL.SAFETENSORS suppress the pickle scan, yet from_pretrained asks for the canonical lowercase model.safetensors, does not find the decoy, and selects the pickle (a direct pytorch_model.bin or the pytorch index) and deserializes it. Probe each alternative with (root / name).is_file() instead, mirroring the loader: is_file() honors the platform's case rules, so a decoy suppresses only where the loader would truly open it. Suppression must never fail open; detection stays case-insensitive (fail closed). Add regression tests for the direct and indexed pickle decoys (skipped on case-insensitive volumes, where no bypass exists). * Studio: resolve indexes and shards exactly as from_pretrained does Two more loader-fidelity gaps in the offline index scan: - Shard lookup normalized backslashes to forward slashes. On POSIX a backslash is a literal filename character, so an index naming dir\payload.bin matches a real pickle of that exact name that Transformers joins and deserializes, while the normalized dir/payload.bin missed it. Join the raw weight_map value with os.path.join so the probe mirrors the loader on each platform. - Index detection case-folded the directory listing, so on a case-sensitive filesystem an uppercase PYTORCH_MODEL.BIN.INDEX.JSON artifact the loader never opens was treated as live and its shard blocked. Probe the canonical name with the loader's own is_file lookup instead, so an index counts only where from_pretrained would actually load it. Update the uppercase-index tests to assert the correct per-filesystem behavior and add a POSIX backslash-shard regression test. --------- Co-authored-by: danielhanchen --- .../tests/test_offline_embedding_minimal.py | 344 ++++++++++++++++++ .../backend/utils/security/file_security.py | 113 +++++- 2 files changed, 438 insertions(+), 19 deletions(-) diff --git a/studio/backend/tests/test_offline_embedding_minimal.py b/studio/backend/tests/test_offline_embedding_minimal.py index 7ff580c36d..ccc6b5f76a 100644 --- a/studio/backend/tests/test_offline_embedding_minimal.py +++ b/studio/backend/tests/test_offline_embedding_minimal.py @@ -51,6 +51,27 @@ def _modules_json(*paths): _COMMIT = "0123456789abcdef0123456789abcdef01234567" +def _fs_case_sensitive(root): + """Whether root's filesystem is case-sensitive (Linux yes; macOS/Windows usually no). The gate + mirrors the loader, whose file lookups follow the same rule, so some cases only exist on one.""" + probe = Path(root) / "_case_probe" + probe.write_text("x") + try: + return not (Path(root) / "_CASE_PROBE").exists() + finally: + probe.unlink() + + +def _requires_case_sensitive_fs(root): + if not _fs_case_sensitive(root): + pytest.skip("requires a case-sensitive filesystem") + + +def _requires_case_insensitive_fs(root): + if _fs_case_sensitive(root): + pytest.skip("requires a case-insensitive filesystem") + + def _make_cache( root, repo_id, @@ -382,6 +403,329 @@ def test_gate_blocks_sharded_pickle(hf_cache): assert _offline_decision("org/shard").blocked is True +def test_gate_blocks_indexed_pickle_shard_in_subdirectory(hf_cache): + # from_pretrained follows weight_map paths relative to the root index, so these nested shards + # are deserialized even though they are not direct children of the load root (iterdir misses + # them). The online gate blocks index-referenced subdir pickles; the offline gate must too. + _make_cache( + hf_cache, + "org/indexed-shard", + { + "pytorch_model.bin.index.json": ( + '{"weight_map": {"layer.weight": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-shard") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_indexed_pickle_shard_with_nonstandard_stem(hf_cache): + # The index tells the loader to deserialize this file, so a pickle EXTENSION is enough -- the + # shard's stem need not match the on-disk weight-name heuristic (which only guesses bare files). + _make_cache( + hf_cache, + "org/indexed-odd", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/evil-00001-of-00001.bin"}}', + "shards/evil-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-odd") + assert decision.blocked is True + assert any(u["path"] == "shards/evil-00001-of-00001.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_safetensors_index_pointing_to_pickle_shard(hf_cache): + # load_state_dict picks safetensors vs torch.load by each shard's own suffix, so a + # model.safetensors.index.json that maps a weight to a .bin shard still deserializes it. The + # index's own existence must not suppress the shard it names. + _make_cache( + hf_cache, + "org/st-index-pickle", + { + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/st-index-pickle") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_indexed_shard_with_no_pickle_extension(hf_cache): + # Transformers torch.loads any indexed shard not ending in .safetensors, so an unconventional + # extensionless name is still a deserialization target. + _make_cache( + hf_cache, + "org/indexed-noext", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/payload"}}', + "shards/payload": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/indexed-noext") + assert decision.blocked is True + assert any(u["path"] == "shards/payload" for u in decision.unsafe_files) + + +_UPPER_INDEX_FILES = { + "PYTORCH_MODEL.BIN.INDEX.JSON": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", +} + + +def test_gate_blocks_uppercase_index_on_case_insensitive_fs(hf_cache): + # On a case-insensitive volume (Windows/macOS) from_pretrained opens an oddly-cased index when it + # requests the canonical lowercase name, so the loader-mirror lookup resolves it and blocks. + _requires_case_insensitive_fs(hf_cache) + _make_cache(hf_cache, "org/upper-index", _UPPER_INDEX_FILES) + with _no_network(): + decision = _offline_decision("org/upper-index") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_allows_uppercase_index_on_case_sensitive_fs(hf_cache): + # On a case-sensitive FS from_pretrained's os.path.isfile of the canonical lowercase name misses + # the uppercase artifact and never loads its shard, so the gate must not over-block it. + _requires_case_sensitive_fs(hf_cache) + _make_cache(hf_cache, "org/upper-index", _UPPER_INDEX_FILES) + with _no_network(): + assert _offline_decision("org/upper-index").blocked is False + + +def test_gate_blocks_indexed_shard_named_with_backslash(hf_cache): + # On POSIX a backslash is a literal filename char, so from_pretrained joins the raw weight_map + # value and deserializes a file actually named "dir\payload.bin"; the gate must probe it verbatim. + import os + + if os.sep != "/": + pytest.skip("backslash is a path separator off POSIX") + _make_cache( + hf_cache, + "org/backslash", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "dir\\\\payload.bin"}}', + "dir\\payload.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/backslash") + assert decision.blocked is True + assert any(u["path"] == "dir\\payload.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_indexed_shard_with_uppercase_safetensors_suffix(hf_cache): + # load_state_dict's endswith(".safetensors") is case-sensitive, so a shard named payload.SAFETENSORS + # falls to torch.load. The gate must classify shard suffixes case-sensitively to match it. + _make_cache( + hf_cache, + "org/upper-suffix", + { + "pytorch_model.bin.index.json": '{"weight_map": {"w": "shards/payload.SAFETENSORS"}}', + "shards/payload.SAFETENSORS": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/upper-suffix") + assert decision.blocked is True + assert any(u["path"] == "shards/payload.SAFETENSORS" for u in decision.unsafe_files) + + +def test_gate_allows_stale_safetensors_index_beside_direct_safetensors(hf_cache): + # A complete direct model.safetensors is selected before either index, so a stale + # model.safetensors.index.json referencing a .bin shard never deserializes -> must not block. + _make_cache( + hf_cache, + "org/direct-plus-stale-index", + { + "model.safetensors": "tensors", + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + assert _offline_decision("org/direct-plus-stale-index").blocked is False + + +def test_gate_blocks_pytorch_index_with_uppercase_safetensors_decoy(hf_cache): + # On a case-sensitive FS, from_pretrained asks for the canonical lowercase model.safetensors, does + # not find an uppercase decoy, and selects the pytorch index instead. The decoy must not suppress. + _requires_case_sensitive_fs(hf_cache) + _make_cache( + hf_cache, + "org/upper-decoy", + { + "MODEL.SAFETENSORS": "decoy", + "pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/upper-decoy") + assert decision.blocked is True + assert any( + u["path"] == "shards/pytorch_model-00001-of-00001.bin" for u in decision.unsafe_files + ) + + +def test_gate_blocks_direct_pickle_with_uppercase_safetensors_decoy(hf_cache): + # Same decoy against a direct pytorch_model.bin: the loader selects the pickle, so the uppercase + # safetensors must not suppress it on a case-sensitive FS. + _requires_case_sensitive_fs(hf_cache) + _make_cache( + hf_cache, + "org/upper-decoy-direct", + {"MODEL.SAFETENSORS": "decoy", "pytorch_model.bin": "pickle"}, + ) + with _no_network(): + decision = _offline_decision("org/upper-decoy-direct") + assert decision.blocked is True + assert any(u["path"] == "pytorch_model.bin" for u in decision.unsafe_files) + + +def test_gate_blocks_indexed_pickle_shard_in_module_subdir(hf_cache): + # A weight index inside a sentence-transformers module load root points at a nested pickle shard. + _make_cache( + hf_cache, + "org/mod-indexed", + { + "modules.json": _modules_json("0_Transformer"), + "0_Transformer/pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "0_Transformer/shards/pytorch_model-00001-of-00001.bin": "pickle", + }, + ) + with _no_network(): + decision = _offline_decision("org/mod-indexed") + assert decision.blocked is True + assert any( + u["path"] == "0_Transformer/shards/pytorch_model-00001-of-00001.bin" + for u in decision.unsafe_files + ) + + +def test_gate_allows_indexed_pickle_shard_with_safetensors_sibling(hf_cache): + # A base model.safetensors makes the loader ignore the pickle index entirely, so it must not + # block (mirrors the direct-file safetensors-sibling suppression). + _make_cache( + hf_cache, + "org/indexed-both", + { + "pytorch_model.bin.index.json": ( + '{"weight_map": {"w": "shards/pytorch_model-00001-of-00001.bin"}}' + ), + "shards/pytorch_model-00001-of-00001.bin": "pickle", + "model.safetensors": "y", + }, + ) + with _no_network(): + assert _offline_decision("org/indexed-both").blocked is False + + +def test_gate_allows_indexed_safetensors_shard_in_subdirectory(hf_cache): + # A safetensors index lists inert shards -- following it must never block (guards against a + # scanner that flags every indexed shard regardless of format). + _make_cache( + hf_cache, + "org/st-indexed", + { + "model.safetensors.index.json": ( + '{"weight_map": {"w": "shards/model-00001-of-00001.safetensors"}}' + ), + "shards/model-00001-of-00001.safetensors": "tensors", + }, + ) + with _no_network(): + assert _offline_decision("org/st-indexed").blocked is False + + +def test_gate_blocks_on_index_path_traversal(hf_cache): + # A weight_map entry escaping the snapshot via ".." is abnormal/hostile -> fail closed. + _make_cache( + hf_cache, + "org/escape", + {"pytorch_model.bin.index.json": '{"weight_map": {"w": "../../../../etc/evil.bin"}}'}, + ) + with _no_network(): + assert _offline_decision("org/escape").blocked is True + + +def test_gate_allows_symlinked_sharded_safetensors(tmp_path, monkeypatch): + # Real HF caches store snapshot files as symlinks into blobs/. A resolve()-based containment + # check would escape the snapshot and false-block every sharded model; the lexical gate must not. + import hashlib + import os + + from huggingface_hub.file_download import repo_folder_name + + root = tmp_path / "hub" + root.mkdir() + monkeypatch.setenv("HF_HOME", str(tmp_path)) + monkeypatch.setenv("HF_HUB_CACHE", str(root)) + monkeypatch.setattr( + "utils.hf_cache_settings.get_hf_cache_paths", + lambda: SimpleNamespace(hub_cache = root), + ) + repo_dir = root / repo_folder_name(repo_id = "org/sym", repo_type = "model") + (repo_dir / "refs").mkdir(parents = True) + (repo_dir / "refs" / "main").write_text(_COMMIT) + blobs = repo_dir / "blobs" + blobs.mkdir() + snapshot = repo_dir / "snapshots" / _COMMIT + (snapshot / "shards").mkdir(parents = True) + + def _blobbed(rel, content): + digest = hashlib.sha256(content.encode()).hexdigest() + (blobs / digest).write_text(content) + target = snapshot / rel + target.parent.mkdir(parents = True, exist_ok = True) + target.symlink_to(os.path.relpath(blobs / digest, target.parent)) + + _blobbed("config.json", "{}") + _blobbed( + "model.safetensors.index.json", + '{"weight_map": {"w": "shards/model-00001-of-00001.safetensors"}}', + ) + _blobbed("shards/model-00001-of-00001.safetensors", "tensors") + with _no_network(): + assert _offline_decision("org/sym").blocked is False + + +def test_gate_allows_index_without_weight_map(hf_cache): + # An index whose top-level JSON has no dict weight_map lets the loader resolve no shards, so it + # must not crash or block on its own (only inert safetensors are cached here). + _make_cache( + hf_cache, + "org/no-wm", + {"model.safetensors.index.json": "[]", "model.safetensors": "x"}, + ) + with _no_network(): + assert _offline_decision("org/no-wm").blocked is False + + def test_gate_allows_nothing_cached(hf_cache): with _no_network(): assert _offline_decision("org/missing").blocked is False diff --git a/studio/backend/utils/security/file_security.py b/studio/backend/utils/security/file_security.py index 91d7ad8f0e..892f7862a9 100644 --- a/studio/backend/utils/security/file_security.py +++ b/studio/backend/utils/security/file_security.py @@ -46,17 +46,6 @@ _PICKLE_WEIGHT_RE = re.compile( 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. @@ -94,6 +83,13 @@ _INERT_SUFFIXES = frozenset( _SOURCE_SUFFIXES = frozenset({".py", ".pyc", ".pyx", ".pyi"}) +# Torch-family weight indexes: from_pretrained feeds each shard they name to load_state_dict, which +# torch.load()s (pickle) any shard whose name does not end in .safetensors, whatever its stem. A +# pytorch index is superseded when a base safetensors is present (the loader prefers it); a +# safetensors index IS the chosen archive, so a non-safetensors target it names still loads. tf/flax +# indexes load via non-pickle loaders, so they are not a torch.load vector here. +_TORCH_INDEX_FILES = ("pytorch_model.bin.index.json", "model.safetensors.index.json") + # Root weight-index files. from_pretrained reads these to find sharded weights, so a # flagged subdir pickle is a load vector iff a root index references it. _TRANSFORMERS_INDEX_FILES = ( @@ -313,13 +309,72 @@ def _st_load_roots(snapshot: Path) -> list: return roots +def _indexed_pickle_shards(index_path: Path, root: Path, snapshot: Path) -> list: + """Shards a torch weight index points a ``from_pretrained`` load at that load_state_dict would + torch.load (pickle): every ``weight_map`` target NOT ending in ``.safetensors``, whatever its + stem (an arbitrary name like ``shards/payload`` still deserializes). Resolved relative to the + index dir (``root``) like the loader, so a shard in a nested dir is followed (iterdir misses it). + Lexical only, never ``Path.resolve()`` (HF snapshot files symlink into ``blobs/``, so resolving + escapes the snapshot and false-blocks every shard). Raises OSError -> caller fails CLOSED on an + unreadable/invalid index or a target escaping the snapshot.""" + import json + import os + + try: + # JSON is UTF-8 by spec; pin it so a non-ASCII index is not misdecoded (and needlessly + # blocked) under Windows' cp1252 default. + parsed = json.loads(index_path.read_text(encoding = "utf-8")) + except (OSError, ValueError) as exc: + raise OSError(f"unreadable weight index: {index_path}") from exc + weight_map = parsed.get("weight_map") if isinstance(parsed, dict) else None + if not isinstance(weight_map, dict): + return [] # no dict weight_map -> the loader resolves no shards from this index + snapshot_norm = os.path.normpath(str(snapshot)) + shards = [] + for shard in weight_map.values(): + raw = str(shard) + if not raw: + continue + # Join the RAW weight_map value like from_pretrained's os.path.join: on POSIX a backslash is a + # literal filename char (not a separator), so normalizing it would probe a different path than + # the loader opens. normpath + containment stay platform-aware (os.sep) to block "..". + joined = os.path.normpath(os.path.join(str(root), raw)) + if joined != snapshot_norm and not joined.startswith(snapshot_norm + os.sep): + raise OSError(f"weight index escapes the snapshot: {index_path}") + shard_path = Path(joined) + # Case-SENSITIVE, mirroring load_state_dict's own endswith(".safetensors"): a shard named + # payload.SAFETENSORS is not treated as safetensors by the loader and falls to torch.load. + if not shard_path.name.endswith(".safetensors") and shard_path.is_file(): + shards.append(shard_path) + return shards + + +def _loader_resolves(root: Path, name: str) -> bool: + """True iff from_pretrained would open ``name`` under ``root``. ``is_file()`` honors the platform + (case-sensitive on Linux, case-insensitive on Windows/macOS), so it mirrors the loader's own + lookup: an oddly-cased decoy counts as an alternative only where the loader would truly open it. + A name-fold instead would let an uppercase MODEL.SAFETENSORS suppress the scan on Linux while the + loader, asking for the canonical lowercase name, silently falls through to a pickle index.""" + return (root / name).is_file() + + 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).""" + """Pickle weight files a SentenceTransformer/Transformers load deserializes from 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. Covers + both direct-child pickles AND pickle shards referenced by a local weight index (which the loader + follows into nested dirs, matching the online gate). Raises OSError -- caller fails CLOSED -- if + the snapshot root or a weight index is unreadable, or an index reference escapes the snapshot.""" blocked = [] + seen = set() + + def _add(path: Path): + key = str(path) + if key not in seen: + seen.add(key) + blocked.append(path) + for root in _st_load_roots(snapshot): try: entries = [p for p in root.iterdir() if p.is_file()] @@ -327,15 +382,35 @@ def _cached_pickle_weight_files(snapshot: Path) -> list: 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) + # Safetensors alternatives the loader would actually resolve (never a bare name-fold, which + # fails OPEN: see _loader_resolves). A base pickle is replaced only by a base safetensors, an + # adapter pickle only by an adapter one. A single model.safetensors also outranks BOTH indexes. + has_direct_base_safetensors = _loader_resolves(root, "model.safetensors") + has_base_safetensors = has_direct_base_safetensors or _loader_resolves( + root, "model.safetensors.index.json" + ) + has_adapter_safetensors = _loader_resolves(root, "adapter_model.safetensors") 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) + _add(path) + # A torch weight index makes from_pretrained load nested shards iterdir never sees; the loader + # torch.loads any not ending in .safetensors. Probe the canonical index name with the loader's + # own lookup (_loader_resolves), so an oddly-cased artifact it would never open does not block. + # A direct model.safetensors wins over BOTH indexes; failing that a base safetensors still + # outranks the pytorch index, while a safetensors index is itself the chosen archive. + for index_name in _TORCH_INDEX_FILES: + if not _loader_resolves(root, index_name): + continue + if has_direct_base_safetensors: + continue + if index_name == "pytorch_model.bin.index.json" and has_base_safetensors: + continue + for shard_path in _indexed_pickle_shards(root / index_name, root, snapshot): + _add(shard_path) return blocked