Skills › Scientific Agent Skills › geniml › model_artifact_inspector.py
model_artifact_inspector.py
377 lines · skills/geniml/scripts/model_artifact_inspector.py at 5f2bc03
1#!/usr/bin/env python3
2"""Inspect local model artifacts and checksums without deserialization."""
4from __future__ import annotations
6import argparse
7import json
8import math
9import os
10import re
11import sys
12from collections import Counter
13from pathlib import Path
15from _common import (
16 HARD_MAX_BYTES,
17 HARD_MAX_FILES,
18 SafetyError,
19 add_path_mode_argument,
20 display_path,
21 fail_json,
22 int_type,
23 iter_text_lines,
24 local_path,
25 print_json,
26 sha256_file,
27 simple_yaml_mapping,
28)
31TOOL = "model-artifact-inspector"
32_HASH_LINE = re.compile(r"^([0-9a-fA-F]{64})[ \t]+[* ]?(.+?)$")
33_DESERIALIZATION_RISK = {
34 ".pt",
35 ".pth",
36 ".ckpt",
37 ".bin",
38 ".pkl",
39 ".pickle",
40 ".joblib",
41 ".model",
42}
43_NATIVE_RISK = {".so", ".dylib", ".dll", ".exe"}
44_ARCHIVE_RISK = {".tar", ".zip", ".tgz", ".gz", ".bz2", ".xz"}
45_METADATA_KEYS = {
46 "architectures",
47 "embedding_dim",
48 "embedding_size",
49 "model_type",
50 "pooling_method",
51 "vocab_size",
52}
55def build_parser() -> argparse.ArgumentParser:
56 parser = argparse.ArgumentParser(
57 description=(
58 "Inventory and hash a local Geniml model bundle without importing "
59 "ML libraries, extracting archives, or deserializing model files."
60 )
61 )
62 parser.add_argument("--model-dir", required=True)
63 parser.add_argument(
64 "--verify-manifest",
65 help="Optional local SHA256SUMS-style manifest.",
66 )
67 parser.add_argument(
68 "--max-files",
69 type=int_type(minimum=1, maximum=HARD_MAX_FILES, label="max-files"),
70 default=1_000,
71 )
72 parser.add_argument(
73 "--max-file-bytes",
74 type=int_type(minimum=1, maximum=HARD_MAX_BYTES, label="max-file-bytes"),
75 default=2 * 1024**3,
76 )
77 parser.add_argument(
78 "--max-total-bytes",
79 type=int_type(minimum=1, maximum=HARD_MAX_BYTES, label="max-total-bytes"),
80 default=4 * 1024**3,
81 )
82 parser.add_argument(
83 "--max-output-files",
84 type=int_type(minimum=0, maximum=5_000, label="max-output-files"),
85 default=200,
86 )
87 add_path_mode_argument(parser)
88 return parser
91def _walk_regular_files(root: Path, max_files: int) -> list[Path]:
92 files: list[Path] = []
93 pending = [root]
94 while pending:
95 directory = pending.pop()
96 with os.scandir(directory) as entries:
97 ordered = sorted(entries, key=lambda entry: entry.name)
98 for entry in ordered:
99 path = Path(entry.path)
100 if entry.is_symlink():
101 raise SafetyError("symlink is not allowed in model bundle")
102 if entry.is_dir(follow_symlinks=False):
103 pending.append(path)
104 elif entry.is_file(follow_symlinks=False):
105 files.append(path)
106 if len(files) > max_files:
107 raise SafetyError("model bundle exceeds max-files")
108 else:
109 raise SafetyError("special filesystem entry is not allowed")
110 return sorted(files)
113def _risk_for(path: Path) -> str:
114 suffix = path.suffix.lower()
115 if suffix in _DESERIALIZATION_RISK:
116 return "deserialization"
117 if suffix in _NATIVE_RISK:
118 return "native-executable"
119 if suffix in _ARCHIVE_RISK:
120 return "archive-expansion"
121 if suffix == ".safetensors":
122 return "bounded-format-still-untrusted"
123 return "data-or-metadata"
126def _safe_json(path: Path, max_bytes: int):
127 lines = [
128 line
129 for _, line in iter_text_lines(
130 path,
131 max_bytes=max_bytes,
132 max_records=100_000,
133 )
134 ]
135 def finite_float(value):
136 number = float(value)
137 if not math.isfinite(number):
138 raise SafetyError("JSON numeric metadata must be finite")
139 return number
141 def reject_constant(value):
142 raise SafetyError("JSON numeric metadata must be finite")
144 def unique_keys(pairs):
145 mapping = {}
146 for key, value in pairs:
147 if key in mapping:
148 raise SafetyError("duplicate JSON metadata key")
149 mapping[key] = value
150 return mapping
152 try:
153 return json.loads("\n".join(lines), parse_float=finite_float,
154 parse_constant=reject_constant, object_pairs_hook=unique_keys)
155 except (ValueError, RecursionError) as exc:
156 raise SafetyError("invalid JSON metadata") from exc
159def _read_checksum_manifest(
160 path: Path,
161 *,
162 max_bytes: int,
163 max_files: int,
164) -> dict[str, str]:
165 expected: dict[str, str] = {}
166 for line_number, line in iter_text_lines(
167 path,
168 max_bytes=max_bytes,
169 max_records=max_files + 100,
170 ):
171 if not line or line.startswith("#"):
172 continue
173 match = _HASH_LINE.fullmatch(line)
174 if not match:
175 raise SafetyError(f"invalid checksum manifest line {line_number}")
176 digest, name = match.groups()
177 relative = Path(name)
178 if (
179 relative.is_absolute()
180 or ".." in relative.parts
181 or not relative.parts
182 or "://" in name
183 or name.startswith("~")
184 ):
185 raise SafetyError(f"unsafe checksum path at line {line_number}")
186 normalized = relative.as_posix()
187 if normalized in expected:
188 raise SafetyError(f"duplicate checksum entry at line {line_number}")
189 expected[normalized] = digest.lower()
190 if len(expected) > max_files:
191 raise SafetyError("checksum manifest exceeds max-files")
192 return expected
195def _redacted_metadata(value) -> tuple[dict[str, object], int]:
196 if not isinstance(value, dict):
197 raise SafetyError("config metadata must be a mapping")
198 selected: dict[str, object] = {}
199 for key in _METADATA_KEYS:
200 if key not in value:
201 continue
202 item = value[key]
203 if isinstance(item, (str, int, float, bool)) or item is None:
204 if isinstance(item, str) and len(item) > 200:
205 selected[key] = "<redacted-long-string>"
206 else:
207 selected[key] = item
208 elif key == "architectures" and isinstance(item, list):
209 selected[key] = [
210 entry if isinstance(entry, str) and len(entry) <= 100 else "<redacted>"
211 for entry in item[:20]
212 ]
213 return selected, max(0, len(value) - len(selected))
216def inspect(args: argparse.Namespace) -> tuple[dict, int]:
217 model_dir = local_path(args.model_dir, kind="dir")
218 files = _walk_regular_files(model_dir, args.max_files)
219 if not files:
220 raise SafetyError("model directory is empty")
222 errors: Counter[str] = Counter()
223 warnings: Counter[str] = Counter()
224 risk_counts: Counter[str] = Counter()
225 computed: dict[str, str] = {}
226 summaries: list[dict] = []
227 metadata: dict[str, object] = {}
228 total_bytes = 0
230 for index, path in enumerate(files, start=1):
231 relative = path.relative_to(model_dir).as_posix()
232 size = path.lstat().st_size
233 if size > args.max_file_bytes:
234 errors["file_exceeds_byte_limit"] += 1
235 continue
236 total_bytes += size
237 if total_bytes > args.max_total_bytes:
238 raise SafetyError("model bundle exceeds max-total-bytes")
239 digest, hashed_size = sha256_file(path, max_bytes=args.max_file_bytes)
240 computed[relative] = digest
241 risk = _risk_for(path)
242 risk_counts[risk] += 1
243 if risk in {"deserialization", "native-executable", "archive-expansion"}:
244 warnings[f"risky_artifact:{risk}"] += 1
246 if path.name in {"config.yaml", "config.yml"}:
247 parsed = simple_yaml_mapping(
248 path,
249 max_bytes=min(args.max_file_bytes, 16 * 1024**2),
250 )
251 selected, omitted = _redacted_metadata(parsed)
252 metadata[path.name] = selected
253 if omitted:
254 warnings["config_metadata_fields_redacted"] += omitted
255 elif path.name in {"config.json", "tokenizer_config.json"}:
256 parsed = _safe_json(
257 path,
258 max_bytes=min(args.max_file_bytes, 16 * 1024**2),
259 )
260 selected, omitted = _redacted_metadata(parsed)
261 metadata[path.name] = selected
262 if omitted:
263 warnings["config_metadata_fields_redacted"] += omitted
265 if len(summaries) < args.max_output_files:
266 if args.path_mode == "redacted":
267 shown_path = f"artifact_{index:04d}"
268 elif args.path_mode == "basename":
269 shown_path = path.name
270 else:
271 shown_path = str(path)
272 summaries.append(
273 {
274 "path": shown_path,
275 "sha256": digest,
276 "size_bytes": hashed_size,
277 "risk_class": risk,
278 "deserialized": False,
279 }
280 )
282 standard = {
283 name: name in computed for name in ("checkpoint.pt", "config.yaml", "universe.bed")
284 }
285 for name, present in standard.items():
286 if not present:
287 errors[f"missing_standard_artifact:{name}"] += 1
289 verification = None
290 if args.verify_manifest:
291 manifest = local_path(args.verify_manifest, kind="file")
292 expected = _read_checksum_manifest(
293 manifest,
294 max_bytes=min(args.max_file_bytes, 64 * 1024**2),
295 max_files=args.max_files,
296 )
297 missing = sorted(set(expected) - set(computed))
298 unlisted = sorted(set(computed) - set(expected))
299 mismatched = sorted(
300 name
301 for name in set(expected) & set(computed)
302 if expected[name] != computed[name]
303 )
304 if missing:
305 errors["checksum_manifest_missing_files"] = len(missing)
306 if mismatched:
307 errors["checksum_mismatches"] = len(mismatched)
308 if unlisted:
309 warnings["files_not_listed_in_checksum_manifest"] = len(unlisted)
310 verification = {
311 "manifest": display_path(manifest, 2, args.path_mode),
312 "expected_entries": len(expected),
313 "missing_count": len(missing),
314 "mismatch_count": len(mismatched),
315 "unlisted_count": len(unlisted),
316 "verified": not missing and not mismatched,
317 }
319 config = metadata.get("config.yaml")
320 if isinstance(config, dict):
321 vocab_size = config.get("vocab_size")
322 embedding_dim = config.get("embedding_dim", config.get("embedding_size"))
323 if type(vocab_size) is not int or vocab_size <= 0:
324 errors["config:invalid_vocab_size"] += 1
325 if type(embedding_dim) is not int or embedding_dim <= 0:
326 errors["config:invalid_embedding_dim"] += 1
327 if "embedding_size" in config and "embedding_dim" not in config:
328 warnings["config:deprecated_embedding_size_key"] += 1
330 report = {
331 "ok": not errors,
332 "tool": TOOL,
333 "contract": {
334 "network_used": False,
335 "model_deserialized": False,
336 "archives_extracted": False,
337 "dynamic_imports_used": False,
338 "symlinks_allowed": False,
339 },
340 "model_dir": display_path(model_dir, 1, args.path_mode),
341 "summary": {
342 "file_count": len(files),
343 "hashed_file_count": len(computed),
344 "total_bytes": total_bytes,
345 "output_file_summaries": len(summaries),
346 "omitted_file_summaries": max(0, len(files) - len(summaries)),
347 "risk_classes": dict(sorted(risk_counts.items())),
348 "standard_artifacts": standard,
349 },
350 "metadata": metadata,
351 "checksum_verification": verification,
352 "errors": dict(sorted(errors.items())),
353 "warnings": dict(sorted(warnings.items())),
354 "files": summaries,
355 "loading_policy": [
356 "verify provenance and checksums before loading",
357 "never inspect .pt/.model/pickle/joblib by deserializing it",
358 "load only in an isolated environment with CPU, RAM, disk, and time bounds",
359 "compare universe bytes/order, tokenizer special IDs, vocab size, and assembly",
360 "treat native binaries and archives as independently untrusted",
361 ],
362 }
363 return report, 0 if not errors else 2
366def main() -> int:
367 args = build_parser().parse_args()
368 try:
369 report, exit_code = inspect(args)
370 print_json(report)
371 return exit_code
372 except (OSError, SafetyError, UnicodeError) as exc:
373 return fail_json(TOOL, exc)
376if __name__ == "__main__":
377 sys.exit(main())