fix: sniff .bin magic before treating weights as pickle

src/quantpreflight/loader.py

@@ -34,6 +34,68 @@ 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")
+# 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:
+    lower = filename.lower()
+    return lower.endswith(_LEGACY_WEIGHT_SUFFIXES) and any(
+        hint in lower for hint in _LEGACY_WEIGHT_NAME_HINTS
+    )
+
+
+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 non-pickle 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(8)
+        except OSError:
+            # Unreadable files still warn — a scanner should err on the side of caution.
+            found.append(fn)
+            continue
+        if not _header_is_known_non_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]:
+    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, 7), token=token, timeout=timeout)
+        except Exception:
+            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):
     """Preserve `Range` across cross-origin CDN redirects while stripping `Authorization`."""
 
@@ -290,11 +352,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 +776,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")

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
@@ -285,6 +286,12 @@ 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"
+_PICKLE_PROTO0 = pickle.dumps({"data": [1, 2]}, protocol=0)
+
+
 @pytest.mark.parametrize(
     ("name", "legacy"),
     [
@@ -301,7 +308,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 +323,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 +336,27 @@ 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_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": "<x>"}}})
     report = analyze_checkpoint(load_checkpoint_metadata(str(ckpt)))

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
@@ -177,12 +178,32 @@ 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"
+_PICKLE_PROTO0 = pickle.dumps({"data": [1, 2]}, protocol=0)
+_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, 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_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 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(
@@ -207,6 +228,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 +238,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"):

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"