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 <unslothai@gmail.com>
This commit is contained in:
Daniel Han
GitHub
parent
0807d03ed0
commit
c2114d64dd
2 changed files with 438 additions and 19 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue