From 09ed797356e3672eff493503f981293668f36555 Mon Sep 17 00:00:00 2001 From: ShahadatTest Date: Tue, 3 Nov 2026 21:04:45 +0000 Subject: [PATCH 1/2] fix: sniff .bin magic before treating weights as pickle Only count name-matched .bin/.pt/.pth files as legacy torch pickle when the first bytes are a zip (PK) or pickle protocol (0x80). Metal gpt-oss weights and other non-pickle formats no longer trip QPF014 (#6). --- src/quantpreflight/loader.py | 75 +++++++++++++++++++++++++++++++----- tests/test_loader_local.py | 21 ++++++++-- tests/test_loader_remote.py | 28 ++++++++++++++ tests/test_rules_stage3.py | 2 +- 4 files changed, 112 insertions(+), 14 deletions(-) diff --git a/src/quantpreflight/loader.py b/src/quantpreflight/loader.py index ada3e27..d8cc6eb 100644 --- a/src/quantpreflight/loader.py +++ b/src/quantpreflight/loader.py @@ -34,6 +34,67 @@ class CheckpointLoadError(RuntimeError): """Raised when checkpoint headers or metadata cannot be located or fetched.""" +_LEGACY_WEIGHT_SUFFIXES = (".bin", ".pt", ".pth") +_LEGACY_WEIGHT_NAME_HINTS = ("model", "pytorch", "weight", "consolidated") +# torch.save since 1.6 writes a zip; classic pickle protocols start with opcode 0x80. +_TORCH_ZIP_MAGIC = b"PK\x03\x04" + + +def _looks_like_legacy_weight_name(filename: str) -> bool: + lower = filename.lower() + return lower.endswith(_LEGACY_WEIGHT_SUFFIXES) and any( + hint in lower for hint in _LEGACY_WEIGHT_NAME_HINTS + ) + + +def _header_looks_like_torch_pickle(header: bytes) -> bool: + """True when the first bytes match a torch zip or a pickle protocol frame.""" + if not header: + return False + return header.startswith(_TORCH_ZIP_MAGIC) or header[0] == 0x80 + + +def _local_legacy_weight_files(directory: Path, filenames: list[str]) -> list[str]: + """Name-filter then sniff magic bytes so Metal/other .bin formats are not flagged.""" + found: list[str] = [] + for fn in filenames: + if not _looks_like_legacy_weight_name(fn): + continue + path = directory / fn + try: + with path.open("rb") as handle: + header = handle.read(4) + except OSError: + # Unreadable files still warn — a scanner should err on the side of caution. + found.append(fn) + continue + if _header_looks_like_torch_pickle(header): + found.append(fn) + return found + + +def _remote_legacy_weight_files( + filenames: list[str], + *, + base_url: str, + token: str | None, + timeout: float, +) -> list[str]: + found: list[str] = [] + for fn in filenames: + if not _looks_like_legacy_weight_name(fn): + continue + url = f"{base_url}/{fn}" + try: + header = _http_fetch_bytes(url, byte_range=(0, 3), token=token, timeout=timeout) + except Exception: + found.append(fn) + continue + if _header_looks_like_torch_pickle(header): + found.append(fn) + return found + + class _RangePreservingRedirectHandler(urllib.request.HTTPRedirectHandler): """Preserve `Range` across cross-origin CDN redirects while stripping `Authorization`.""" @@ -290,11 +351,7 @@ def _read_opt_json(filename: str) -> dict[str, Any]: tok_json = _read_opt_json("tokenizer.json") if tok_json_path.is_file() else None all_dir_files = [p.name for p in target_path.iterdir() if p.is_file()] - legacy_files = [ - fn for fn in all_dir_files - if fn.lower().endswith((".bin", ".pt", ".pth")) - and any(k in fn.lower() for k in ("model", "pytorch", "weight", "consolidated")) - ] + legacy_files = _local_legacy_weight_files(target_path, all_dir_files) custom_code_files = [ fn for fn in all_dir_files if fn.lower().endswith(".py") @@ -718,11 +775,9 @@ def _srange(start: int, end: int) -> bytes: duplicate_tensors.append(tname) raw_entries[tname] = (entry, sname) - remote_legacy = [ - fn for fn in remote_siblings - if fn.lower().endswith((".bin", ".pt", ".pth")) - and any(k in fn.lower() for k in ("model", "pytorch", "weight", "consolidated")) - ] + remote_legacy = _remote_legacy_weight_files( + remote_siblings, base_url=base_url, token=resolved_token, timeout=timeout + ) remote_code = [ fn for fn in remote_siblings if fn.lower().endswith(".py") diff --git a/tests/test_loader_local.py b/tests/test_loader_local.py index 5b6c87d..cd3ef46 100644 --- a/tests/test_loader_local.py +++ b/tests/test_loader_local.py @@ -285,6 +285,11 @@ def test_invalid_json_names_the_file(tmp_path: Path) -> None: load_checkpoint_metadata(str(ckpt)) +# torch.save zip / classic pickle protocol magic (issue #6) +_PICKLE_ZIP = b"PK\x03\x04rest" +_PICKLE_PROTO = b"\x80\x02abc" + + @pytest.mark.parametrize( ("name", "legacy"), [ @@ -301,7 +306,7 @@ def test_invalid_json_names_the_file(tmp_path: Path) -> None: ) def test_legacy_weight_names(tmp_path: Path, name: str, legacy: bool) -> None: ckpt = _ckpt(tmp_path) - (ckpt / name).write_bytes(b"x") + (ckpt / name).write_bytes(_PICKLE_PROTO) assert (name in load_checkpoint_metadata(str(ckpt)).legacy_weight_files) is legacy @@ -316,8 +321,8 @@ def test_pickle_only_folder(tmp_path: Path) -> None: ckpt = tmp_path / "ckpt" ckpt.mkdir() (ckpt / "config.json").write_text('{"vocab_size": 8, "quantization_config": {"quant_method": "gptq", "bits": 4}}', encoding="utf-8") - (ckpt / "b_model.bin").write_bytes(b"x") - (ckpt / "a_model.bin").write_bytes(b"x") + (ckpt / "b_model.bin").write_bytes(_PICKLE_ZIP) + (ckpt / "a_model.bin").write_bytes(_PICKLE_PROTO) (ckpt / "run.py").write_text("x", encoding="utf-8") meta = load_checkpoint_metadata(str(ckpt)) assert meta.container_format == "pytorch_bin" @@ -329,6 +334,16 @@ def test_pickle_only_folder(tmp_path: Path) -> None: assert {f.rule_id for f in analyze_checkpoint(meta).findings} >= {"QPF014"} +def test_metal_model_bin_is_not_treated_as_pickle(tmp_path: Path) -> None: + """gpt-oss Metal weights start with 'GPT-OSS …', not torch pickle magic (#6).""" + ckpt = _ckpt(tmp_path) + (ckpt / "metal_model.bin").write_bytes(b"GPT-OSS v1.0\x00weights") + (ckpt / "pytorch_model.bin").write_bytes(_PICKLE_ZIP) + meta = load_checkpoint_metadata(str(ckpt)) + assert meta.legacy_weight_files == ("pytorch_model.bin",) + assert "metal_model.bin" not in meta.legacy_weight_files + + def test_added_tokens_decoder_reaches_the_vocab_rule(tmp_path: Path) -> None: ckpt = _ckpt(tmp_path, tokenizer_config={"added_tokens_decoder": {"8": {"content": ""}}}) report = analyze_checkpoint(load_checkpoint_metadata(str(ckpt))) diff --git a/tests/test_loader_remote.py b/tests/test_loader_remote.py index a522cba..4521127 100644 --- a/tests/test_loader_remote.py +++ b/tests/test_loader_remote.py @@ -177,12 +177,29 @@ def test_any_safetensors_files_are_used_as_a_fallback(hub: FakeHub) -> None: assert load_checkpoint_metadata("org/model").shard_files == ("one.safetensors", "two.SAFETENSORS") +_PICKLE_ZIP = b"PK\x03\x04rest" +_METAL_BIN = b"GPT-OSS v1.0\x00weights" + + def test_remote_legacy_and_code_files_are_recorded(hub: FakeHub) -> None: hub.add("model.safetensors", _shard("a.weight")) + hub.add("pytorch_model.bin", _PICKLE_ZIP) hub.siblings("model.safetensors", "pytorch_model.bin", "training_args.bin", "modeling_x.py", "b.pt") meta = load_checkpoint_metadata("org/model") assert meta.legacy_weight_files == ("pytorch_model.bin",) assert meta.custom_code_files == ("modeling_x.py",) + assert (BASE + "/pytorch_model.bin") in hub.requests + assert (0, 3) in hub.ranges + + +def test_remote_metal_model_bin_is_not_treated_as_pickle(hub: FakeHub) -> None: + hub.add("model.safetensors", _shard("a.weight")) + hub.add("metal_model.bin", _METAL_BIN) + hub.add("pytorch_model.bin", _PICKLE_ZIP) + hub.siblings("model.safetensors", "metal_model.bin", "pytorch_model.bin") + meta = load_checkpoint_metadata("org/model") + assert meta.legacy_weight_files == ("pytorch_model.bin",) + assert "metal_model.bin" not in meta.legacy_weight_files @pytest.mark.parametrize( @@ -207,6 +224,7 @@ def test_bare_gguf_repository_picks_a_common_variant(hub: FakeHub, names: list[s def test_pickle_only_repository(hub: FakeHub) -> None: hub.add("config.json", {"vocab_size": 8, "quantization_config": {"quant_method": "gptq"}}) + hub.add("pytorch_model.bin", _PICKLE_ZIP) hub.siblings("pytorch_model.bin", "run.py") meta = load_checkpoint_metadata("org/model") assert meta.container_format == "pytorch_bin" @@ -216,6 +234,16 @@ def test_pickle_only_repository(hub: FakeHub) -> None: assert meta.tokenizer_meta.max_token_id == 7 +def test_unreadable_remote_legacy_candidate_is_still_counted(hub: FakeHub) -> None: + """When the first bytes cannot be fetched, keep warning (issue #6).""" + hub.add("config.json", {"vocab_size": 8}) + hub.siblings("pytorch_model.bin", "run.py") + # No body for pytorch_model.bin → FakeHub 404 on the magic sniff. + meta = load_checkpoint_metadata("org/model") + assert meta.container_format == "pytorch_bin" + assert meta.legacy_weight_files == ("pytorch_model.bin",) + + def test_repository_with_nothing_to_read(hub: FakeHub) -> None: hub.siblings("README.md") with pytest.raises(CheckpointLoadError, match="No model.safetensors"): diff --git a/tests/test_rules_stage3.py b/tests/test_rules_stage3.py index bf1f75c..48bc6e0 100644 --- a/tests/test_rules_stage3.py +++ b/tests/test_rules_stage3.py @@ -118,7 +118,7 @@ def test_qpf014_safetensors_with_lingering_pickle_file(tmp_path: Path) -> None: mixed_dir.mkdir() (mixed_dir / "config.json").write_text('{"architectures": ["LlamaForCausalLM"]}', encoding="utf-8") (mixed_dir / "model.safetensors").write_bytes(build_safetensors_header_file({"embed": ("F16", [10, 10])})) - (mixed_dir / "pytorch_model.bin").write_bytes(b"legacy_pickle") + (mixed_dir / "pytorch_model.bin").write_bytes(b"PK\x03\x04mock_zip_pickle") meta = load_checkpoint_metadata(str(mixed_dir)) assert meta.container_format == "safetensors" From 5ce85cecd792b62232f0195537e7756d5013a33d Mon Sep 17 00:00:00 2001 From: ShahadatTest Date: Sat, 7 Nov 2026 01:42:10 +0600 Subject: [PATCH 2/2] fix: exclude known non-pickle formats and fetch remote magic in parallel Count candidate legacy weight files as pickle unless their first 8 bytes match known non-pickle prefixes (GPT-OSS or GGUF). This avoids missing legacy pickle protocols 0 and 1 while preserving zero-warning behavior for Metal weights. In addition, fetch remote legacy weight headers in parallel with ThreadPoolExecutor (#6). --- src/quantpreflight/loader.py | 41 ++++++++++++++++++------------------ tests/test_loader_local.py | 13 ++++++++++++ tests/test_loader_remote.py | 12 +++++++---- 3 files changed, 42 insertions(+), 24 deletions(-) diff --git a/src/quantpreflight/loader.py b/src/quantpreflight/loader.py index d8cc6eb..5039cca 100644 --- a/src/quantpreflight/loader.py +++ b/src/quantpreflight/loader.py @@ -36,8 +36,8 @@ class CheckpointLoadError(RuntimeError): _LEGACY_WEIGHT_SUFFIXES = (".bin", ".pt", ".pth") _LEGACY_WEIGHT_NAME_HINTS = ("model", "pytorch", "weight", "consolidated") -# torch.save since 1.6 writes a zip; classic pickle protocols start with opcode 0x80. -_TORCH_ZIP_MAGIC = b"PK\x03\x04" +# Formats known not to be torch pickle (issue #6) +_KNOWN_NON_PICKLE_PREFIXES: tuple[bytes, ...] = (b"GPT-OSS ", b"GGUF") def _looks_like_legacy_weight_name(filename: str) -> bool: @@ -47,15 +47,13 @@ def _looks_like_legacy_weight_name(filename: str) -> bool: ) -def _header_looks_like_torch_pickle(header: bytes) -> bool: - """True when the first bytes match a torch zip or a pickle protocol frame.""" - if not header: - return False - return header.startswith(_TORCH_ZIP_MAGIC) or header[0] == 0x80 +def _header_is_known_non_pickle(header: bytes) -> bool: + """True when the first bytes match a format known not to be torch pickle.""" + return any(header.startswith(prefix) for prefix in _KNOWN_NON_PICKLE_PREFIXES) def _local_legacy_weight_files(directory: Path, filenames: list[str]) -> list[str]: - """Name-filter then sniff magic bytes so Metal/other .bin formats are not flagged.""" + """Name-filter then sniff magic bytes so Metal/other non-pickle formats are not flagged.""" found: list[str] = [] for fn in filenames: if not _looks_like_legacy_weight_name(fn): @@ -63,12 +61,12 @@ def _local_legacy_weight_files(directory: Path, filenames: list[str]) -> list[st path = directory / fn try: with path.open("rb") as handle: - header = handle.read(4) + header = handle.read(8) except OSError: # Unreadable files still warn — a scanner should err on the side of caution. found.append(fn) continue - if _header_looks_like_torch_pickle(header): + if not _header_is_known_non_pickle(header): found.append(fn) return found @@ -80,19 +78,22 @@ def _remote_legacy_weight_files( token: str | None, timeout: float, ) -> list[str]: - found: list[str] = [] - for fn in filenames: - if not _looks_like_legacy_weight_name(fn): - continue + candidates = [fn for fn in filenames if _looks_like_legacy_weight_name(fn)] + if not candidates: + return [] + + def _probe_candidate(fn: str) -> bool: url = f"{base_url}/{fn}" try: - header = _http_fetch_bytes(url, byte_range=(0, 3), token=token, timeout=timeout) + header = _http_fetch_bytes(url, byte_range=(0, 7), token=token, timeout=timeout) except Exception: - found.append(fn) - continue - if _header_looks_like_torch_pickle(header): - found.append(fn) - return found + return True + return not _header_is_known_non_pickle(header) + + with concurrent.futures.ThreadPoolExecutor(max_workers=min(8, len(candidates))) as pool: + matches = list(pool.map(_probe_candidate, candidates)) + + return [fn for fn, is_pickle in zip(candidates, matches) if is_pickle] class _RangePreservingRedirectHandler(urllib.request.HTTPRedirectHandler): diff --git a/tests/test_loader_local.py b/tests/test_loader_local.py index cd3ef46..1a84f09 100644 --- a/tests/test_loader_local.py +++ b/tests/test_loader_local.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import pickle import os from pathlib import Path from typing import Any @@ -288,6 +289,7 @@ def test_invalid_json_names_the_file(tmp_path: Path) -> None: # torch.save zip / classic pickle protocol magic (issue #6) _PICKLE_ZIP = b"PK\x03\x04rest" _PICKLE_PROTO = b"\x80\x02abc" +_PICKLE_PROTO0 = pickle.dumps({"data": [1, 2]}, protocol=0) @pytest.mark.parametrize( @@ -344,6 +346,17 @@ def test_metal_model_bin_is_not_treated_as_pickle(tmp_path: Path) -> None: assert "metal_model.bin" not in meta.legacy_weight_files +def test_protocol_0_pickle_is_treated_as_pickle(tmp_path: Path) -> None: + """Pickle protocol 0 starts with ASCII opcode bytes (issue #6), must still be flagged.""" + ckpt = _ckpt(tmp_path) + (ckpt / "metal_model.bin").write_bytes(b"GPT-OSS v1.0\x00weights") + (ckpt / "zip_model.bin").write_bytes(_PICKLE_ZIP) + (ckpt / "pytorch_model.bin").write_bytes(_PICKLE_PROTO0) + meta = load_checkpoint_metadata(str(ckpt)) + assert set(meta.legacy_weight_files) == {"pytorch_model.bin", "zip_model.bin"} + assert "metal_model.bin" not in meta.legacy_weight_files + + def test_added_tokens_decoder_reaches_the_vocab_rule(tmp_path: Path) -> None: ckpt = _ckpt(tmp_path, tokenizer_config={"added_tokens_decoder": {"8": {"content": ""}}}) report = analyze_checkpoint(load_checkpoint_metadata(str(ckpt))) diff --git a/tests/test_loader_remote.py b/tests/test_loader_remote.py index 4521127..42d24b9 100644 --- a/tests/test_loader_remote.py +++ b/tests/test_loader_remote.py @@ -3,6 +3,7 @@ from __future__ import annotations import io +import pickle import json import urllib.error import urllib.request @@ -178,6 +179,7 @@ def test_any_safetensors_files_are_used_as_a_fallback(hub: FakeHub) -> None: _PICKLE_ZIP = b"PK\x03\x04rest" +_PICKLE_PROTO0 = pickle.dumps({"data": [1, 2]}, protocol=0) _METAL_BIN = b"GPT-OSS v1.0\x00weights" @@ -189,17 +191,19 @@ def test_remote_legacy_and_code_files_are_recorded(hub: FakeHub) -> None: assert meta.legacy_weight_files == ("pytorch_model.bin",) assert meta.custom_code_files == ("modeling_x.py",) assert (BASE + "/pytorch_model.bin") in hub.requests - assert (0, 3) in hub.ranges + assert (0, 7) in hub.ranges def test_remote_metal_model_bin_is_not_treated_as_pickle(hub: FakeHub) -> None: hub.add("model.safetensors", _shard("a.weight")) hub.add("metal_model.bin", _METAL_BIN) - hub.add("pytorch_model.bin", _PICKLE_ZIP) - hub.siblings("model.safetensors", "metal_model.bin", "pytorch_model.bin") + hub.add("pytorch_model.bin", _PICKLE_PROTO0) + hub.add("zip_model.bin", _PICKLE_ZIP) + hub.siblings("model.safetensors", "metal_model.bin", "pytorch_model.bin", "zip_model.bin") meta = load_checkpoint_metadata("org/model") - assert meta.legacy_weight_files == ("pytorch_model.bin",) + assert set(meta.legacy_weight_files) == {"pytorch_model.bin", "zip_model.bin"} assert "metal_model.bin" not in meta.legacy_weight_files + assert (0, 7) in hub.ranges @pytest.mark.parametrize(