| |
| """Check the outer distribution file set and hashes; never alter model files.""" |
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| NAME = "DISTRIBUTION_MANIFEST.json" |
| MANIFEST_SHA256 = "eef92679b27ed4fac41873d7aafb4f67371de4f875b74c023db9fa50d5975b3a" |
|
|
|
|
| def digest(path): |
| h = hashlib.sha256() |
| with path.open("rb") as stream: |
| for block in iter(lambda: stream.read(8 * 1024 * 1024), b""): |
| h.update(block) |
| return h.hexdigest() |
|
|
|
|
| def check(root, expected_outer_sha256): |
| root = root.resolve(strict=True) |
| outer = root / NAME |
| if outer.is_symlink() or not outer.is_file(): |
| raise ValueError("missing or linked outer manifest") |
| if digest(outer) != expected_outer_sha256: |
| raise ValueError("outer manifest SHA-256 differs from the separately trusted value") |
| spec = json.loads(outer.read_text(encoding="utf-8")) |
| if spec.get("schema") != 1 or spec.get("original_checkpoint_manifest_sha256") != MANIFEST_SHA256: |
| raise ValueError("wrong package schema or original checkpoint identity") |
| files = spec.get("files") |
| if not isinstance(files, dict) or not files or NAME in files: |
| raise ValueError("invalid file map or self-entry") |
| if not any(key.startswith("checkpoint/") for key in files): |
| raise ValueError("checkpoint file set is absent") |
| for name, item in files.items(): |
| p = Path(name) |
| if (not isinstance(name, str) or p.is_absolute() or p.as_posix() != name |
| or ".." in p.parts or "\\" in name or not isinstance(item, dict) |
| or not isinstance(item.get("sha256"), str) or len(item["sha256"]) != 64 |
| or type(item.get("size")) is not int or item["size"] < 0): |
| raise ValueError("unsafe or malformed file map entry") |
| actual = set() |
| for p in root.rglob("*"): |
| rel = p.relative_to(root).as_posix() |
| |
| if rel == ".cache" or rel.startswith(".cache/"): |
| if rel == ".cache" and p.is_symlink(): |
| raise ValueError("linked outer cache root") |
| if rel != ".cache" and not rel.startswith(".cache/huggingface/") and rel != ".cache/huggingface": |
| raise ValueError("unexpected outer cache entry: " + rel) |
| continue |
| if p.is_symlink(): |
| raise ValueError("symlink outside allowed outer cache: " + rel) |
| if p.is_file(): |
| actual.add(rel) |
| elif not p.is_dir(): |
| raise ValueError("unsupported entry: " + rel) |
| if actual != set(files) | {NAME}: |
| raise ValueError("file-set mismatch: missing=%s extra=%s" % (sorted(set(files) - actual), sorted(actual - set(files) - {NAME}))) |
| checkpoint = root / "checkpoint" |
| if not checkpoint.is_dir() or checkpoint.is_symlink(): |
| raise ValueError("missing or linked checkpoint directory") |
| for name, item in files.items(): |
| p = root / name |
| if p.stat().st_size != item["size"] or digest(p) != item["sha256"]: |
| raise ValueError("size or SHA-256 mismatch: " + name) |
| if files["checkpoint/axl-ablit-manifest.json"]["sha256"] != MANIFEST_SHA256: |
| raise ValueError("original checkpoint manifest changed") |
| print("PASS: outer file set and all inventory SHA-256 hashes (%d files); runtime contract not checked" % len(files)) |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("root", type=Path, help="outer download directory") |
| parser.add_argument("--outer-sha256", required=True, help="trusted 64-character published digest of outer manifest") |
| args = parser.parse_args() |
| if len(args.outer_sha256) != 64 or any(c not in "0123456789abcdef" for c in args.outer_sha256): |
| parser.error("--outer-sha256 must be a lowercase SHA-256") |
| try: |
| check(args.root, args.outer_sha256) |
| except (OSError, ValueError, KeyError, json.JSONDecodeError) as e: |
| parser.exit(1, "FAIL: %s\n" % e) |
|
|
