diff --git a/src/quantpreflight/loader.py b/src/quantpreflight/loader.py index ada3e27..5039cca 100644 --- a/src/quantpreflight/loader.py +++ b/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") diff --git a/tests/test_loader_local.py b/tests/test_loader_local.py index 5b6c87d..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 @@ -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": ""}}}) report = analyze_checkpoint(load_checkpoint_metadata(str(ckpt))) diff --git a/tests/test_loader_remote.py b/tests/test_loader_remote.py index a522cba..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 @@ -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"): 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"