Skills › Scientific Agent Skills › pufferlib › inspect_checkpoint.py
inspect_checkpoint.py
229 lines · skills/pufferlib/scripts/inspect_checkpoint.py at 5f2bc03
1#!/usr/bin/env python3
2"""Inspect checkpoint file metadata without deserializing checkpoint contents."""
4from __future__ import annotations
6import argparse
7import hashlib
8import os
9import stat
10from pathlib import Path
11from typing import Any, BinaryIO
13try:
14 from ._common import (
15 UserInputError,
16 bounded_int,
17 emit_json,
18 load_json_object,
19 resolve_local_path,
20 secret_key_paths,
21 validate_sha256,
22 )
23except ImportError: # Direct script execution.
24 from _common import (
25 UserInputError,
26 bounded_int,
27 emit_json,
28 load_json_object,
29 resolve_local_path,
30 secret_key_paths,
31 validate_sha256,
32 )
34_SIDECAR_FIELDS = {
35 "created_at",
36 "environment",
37 "format",
38 "framework",
39 "framework_version",
40 "license",
41 "notes",
42 "parent_sha256",
43 "policy",
44 "schema_version",
45 "seed",
46 "sha256",
47 "source_commit",
48 "source_url",
49 "training_steps",
50}
53def _open_regular_no_follow(path: Path) -> tuple[BinaryIO, os.stat_result]:
54 flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
55 try:
56 descriptor = os.open(path, flags)
57 except OSError as exc:
58 raise UserInputError(f"cannot open checkpoint safely: {path}") from exc
59 try:
60 file_stat = os.fstat(descriptor)
61 if not stat.S_ISREG(file_stat.st_mode):
62 raise UserInputError("checkpoint must be a regular file")
63 return os.fdopen(descriptor, "rb"), file_stat
64 except Exception:
65 os.close(descriptor)
66 raise
69def _detect_format(prefix: bytes, suffix: str) -> dict[str, Any]:
70 if prefix.startswith(b"PK\x03\x04"):
71 return {
72 "family": "zip-container",
73 "risk": "may contain a torch.save pickle payload; not opened",
74 }
75 if prefix.startswith(b"\x80"):
76 protocol = prefix[1] if len(prefix) > 1 else None
77 return {
78 "family": "pickle-like",
79 "pickle_protocol_byte": protocol,
80 "risk": "unsafe to deserialize unless provenance is trusted",
81 }
82 if suffix.lower() == ".bin":
83 return {
84 "family": "opaque-bin",
85 "risk": "could be PufferLib native weights or another binary format",
86 }
87 return {
88 "family": "opaque",
89 "risk": "format not identified; no deserialization attempted",
90 }
93def _hash_and_prefix(handle: BinaryIO, *, chunk_bytes: int = 1_048_576) -> tuple[str, bytes]:
94 digest = hashlib.sha256()
95 prefix = b""
96 while True:
97 chunk = handle.read(chunk_bytes)
98 if not chunk:
99 break
100 if not prefix:
101 prefix = chunk[:16]
102 digest.update(chunk)
103 return digest.hexdigest(), prefix
106def _safe_sidecar(
107 metadata_path: str | None, *, root: str | Path
108) -> tuple[dict[str, Any] | None, list[str]]:
109 if metadata_path is None:
110 return None, []
111 raw = load_json_object(metadata_path, root=root, max_bytes=262_144)
112 secret_paths = secret_key_paths(raw)
113 if secret_paths:
114 raise UserInputError(
115 "sidecar contains credential-bearing keys: " + ", ".join(secret_paths)
116 )
117 safe = {key: raw[key] for key in sorted(raw) if key in _SIDECAR_FIELDS}
118 unknown = sorted(set(raw) - _SIDECAR_FIELDS)
119 return safe, unknown
122def inspect_checkpoint(
123 checkpoint_path: str,
124 *,
125 root: str | Path,
126 metadata_path: str | None,
127 expected_sha256: str | None,
128 max_bytes: int,
129) -> dict[str, Any]:
130 """Hash and classify one local regular file without importing torch or pickle."""
131 resolved = resolve_local_path(
132 checkpoint_path,
133 root=root,
134 must_exist=True,
135 reject_symlink=True,
136 )
137 handle, file_stat = _open_regular_no_follow(resolved)
138 with handle:
139 if file_stat.st_size > max_bytes:
140 raise UserInputError(
141 f"checkpoint is {file_stat.st_size} bytes; cap is {max_bytes}"
142 )
143 digest, prefix = _hash_and_prefix(handle)
145 expected = validate_sha256(expected_sha256) if expected_sha256 else None
146 sidecar, unknown_fields = _safe_sidecar(metadata_path, root=root)
147 sidecar_digest = sidecar.get("sha256") if sidecar else None
148 if sidecar_digest is not None:
149 validate_sha256(sidecar_digest, name="sidecar sha256")
151 return {
152 "checkpoint": {
153 "format_detection": _detect_format(prefix, resolved.suffix),
154 "name": resolved.name,
155 "sha256": digest,
156 "size_bytes": file_stat.st_size,
157 },
158 "deserialized": False,
159 "expected_sha256_matches": None if expected is None else digest == expected,
160 "metadata": sidecar,
161 "metadata_sha256_matches": (
162 None if sidecar_digest is None else digest == sidecar_digest
163 ),
164 "network_used": False,
165 "sidecar_unknown_fields_omitted": unknown_fields,
166 "warning": (
167 "Inspection does not establish trust. Verify source, license, signature or "
168 "attestation, and checksum before sandboxed loading."
169 ),
170 }
173def build_parser() -> argparse.ArgumentParser:
174 parser = argparse.ArgumentParser(
175 description=(
176 "Hash and classify one local checkpoint without torch.load, pickle, "
177 "archive extraction, dynamic imports, or network access."
178 )
179 )
180 parser.add_argument("checkpoint", help="Checkpoint path beneath --root")
181 parser.add_argument("--root", default=".", help="Allowed local path root")
182 parser.add_argument("--metadata", help="Explicit strict-JSON sidecar beneath --root")
183 parser.add_argument("--expected-sha256")
184 parser.add_argument(
185 "--max-bytes",
186 type=int,
187 default=2_147_483_648,
188 help="1..68719476736",
189 )
190 parser.add_argument("--compact", action="store_true")
191 return parser
194def main(argv: list[str] | None = None) -> int:
195 parser = build_parser()
196 args = parser.parse_args(argv)
197 try:
198 max_bytes = bounded_int(
199 args.max_bytes,
200 name="max_bytes",
201 minimum=1,
202 maximum=68_719_476_736,
203 )
204 report = inspect_checkpoint(
205 args.checkpoint,
206 root=args.root,
207 metadata_path=args.metadata,
208 expected_sha256=args.expected_sha256,
209 max_bytes=max_bytes,
210 )
211 except (UserInputError, OSError, ValueError) as exc:
212 report = {
213 "deserialized": False,
214 "errors": [str(exc)],
215 "network_used": False,
216 "status": "invalid",
217 }
218 emit_json(report, pretty=not args.compact)
219 return 1
220 emit_json(report, pretty=not args.compact)
221 # A requested integrity check must fail the CLI when bytes do not match.
222 return 1 if (
223 report["expected_sha256_matches"] is False
224 or report["metadata_sha256_matches"] is False
225 ) else 0
228if __name__ == "__main__":
229 raise SystemExit(main())