model-mirror-cli 0.2.0__py3-none-any.whl

model_mirror/__init__.py

@@ -0,0 +1,7 @@
+"""Mirror and validate Hugging Face model archives."""
+
+from importlib.metadata import version
+
+__all__ = ["__version__"]
+
+__version__ = version("model-mirror-cli")

model_mirror/audit.py

@@ -0,0 +1,214 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+DTYPE_SIZES = {
+    "BOOL": 1,
+    "U8": 1,
+    "I8": 1,
+    "F8_E4M3": 1,
+    "F8_E5M2": 1,
+    "U16": 2,
+    "I16": 2,
+    "F16": 2,
+    "BF16": 2,
+    "U32": 4,
+    "I32": 4,
+    "F32": 4,
+    "U64": 8,
+    "I64": 8,
+    "F64": 8,
+}
+
+
+@dataclass(slots=True)
+class SafetensorsValidation:
+    tensor_count: int
+    tensor_bytes: int
+    tensor_names: set[str]
+    warnings: list[str] = field(default_factory=list)
+
+
+@dataclass(slots=True)
+class AuditResult:
+    checked_safetensors: int = 0
+    indexed_tensors: int = 0
+    missing_files: list[str] = field(default_factory=list)
+    failures: list[str] = field(default_factory=list)
+    warnings: list[str] = field(default_factory=list)
+
+    @property
+    def ok(self) -> bool:
+        return not self.missing_files and not self.failures
+
+
+def product(values: list[int]) -> int:
+    result = 1
+    for value in values:
+        result *= value
+    return result
+
+
+def read_json(path: Path) -> dict:
+    with path.open("r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def read_safetensors_header(path: Path) -> tuple[dict, int, int]:
+    with path.open("rb") as handle:
+        size_bytes = handle.read(8)
+        if len(size_bytes) != 8:
+            raise ValueError("file is too small to contain a safetensors header")
+        header_size = int.from_bytes(size_bytes, "little")
+        if header_size <= 0:
+            raise ValueError(f"invalid safetensors header size: {header_size}")
+        header_bytes = handle.read(header_size)
+        if len(header_bytes) != header_size:
+            raise ValueError("file ended before the safetensors header was complete")
+    return json.loads(header_bytes), header_size, path.stat().st_size
+
+
+def validate_safetensors_file(path: Path) -> SafetensorsValidation:
+    warnings = []
+    header, header_size, file_size = read_safetensors_header(path)
+    data_size = file_size - 8 - header_size
+    if data_size < 0:
+        raise ValueError("file is smaller than its declared safetensors header")
+
+    tensors = {key: value for key, value in header.items() if key != "__metadata__"}
+    intervals = []
+    tensor_bytes = 0
+    for tensor_name, spec in tensors.items():
+        if not isinstance(spec, dict):
+            raise ValueError(f"tensor {tensor_name!r} has invalid metadata")
+        dtype = spec.get("dtype")
+        shape = spec.get("shape")
+        offsets = spec.get("data_offsets")
+        if not isinstance(shape, list) or not all(isinstance(item, int) and item >= 0 for item in shape):
+            raise ValueError(f"tensor {tensor_name!r} has invalid shape")
+        if not isinstance(offsets, list) or len(offsets) != 2:
+            raise ValueError(f"tensor {tensor_name!r} has invalid data_offsets")
+        start, end = offsets
+        if not isinstance(start, int) or not isinstance(end, int) or start < 0 or end < start:
+            raise ValueError(f"tensor {tensor_name!r} has invalid byte range")
+        if end > data_size:
+            raise ValueError(f"tensor {tensor_name!r} extends beyond file data")
+
+        observed_bytes = end - start
+        item_size = DTYPE_SIZES.get(dtype)
+        if item_size is None:
+            warnings.append(f"{path.name}: tensor {tensor_name!r} has unknown dtype {dtype!r}")
+        else:
+            expected_bytes = product(shape) * item_size
+            if observed_bytes != expected_bytes:
+                raise ValueError(
+                    f"tensor {tensor_name!r} byte size mismatch: "
+                    f"offsets={observed_bytes} shape*dtype={expected_bytes}"
+                )
+        tensor_bytes += observed_bytes
+        intervals.append((start, end, tensor_name))
+
+    intervals.sort()
+    previous_end = 0
+    for start, end, tensor_name in intervals:
+        if start < previous_end:
+            raise ValueError(f"tensor {tensor_name!r} overlaps previous tensor data")
+        if start > previous_end:
+            warnings.append(f"{path.name}: gap before tensor {tensor_name!r}")
+        previous_end = end
+    if previous_end < data_size:
+        warnings.append(f"{path.name}: {data_size - previous_end} trailing byte(s) after tensor data")
+
+    return SafetensorsValidation(
+        tensor_count=len(tensors),
+        tensor_bytes=tensor_bytes,
+        tensor_names=set(tensors),
+        warnings=warnings,
+    )
+
+
+def audit_model(
+    root: Path,
+    *,
+    skip_transformers: bool = False,
+    trust_remote_code: bool = False,
+    strict: bool = False,
+    require_config: bool = True,
+) -> AuditResult:
+    result = AuditResult()
+    config_path = root / "config.json"
+    if not config_path.exists():
+        if require_config:
+            result.missing_files.append("config.json")
+    else:
+        try:
+            read_json(config_path)
+        except Exception as exc:
+            result.failures.append(f"config.json: {exc}")
+
+    if not skip_transformers:
+        try:
+            from transformers import AutoConfig, AutoTokenizer
+
+            AutoConfig.from_pretrained(root, local_files_only=True, trust_remote_code=trust_remote_code)
+            try:
+                AutoTokenizer.from_pretrained(root, local_files_only=True, trust_remote_code=trust_remote_code)
+            except Exception as exc:
+                result.warnings.append(f"AutoTokenizer metadata load failed: {exc}")
+        except Exception as exc:
+            result.warnings.append(f"Transformers metadata audit skipped/failed: {exc}")
+
+    indexed_tensors: dict[str, str] = {}
+    indexed_files: set[str] = set()
+    for index_path in sorted(root.glob("*.safetensors.index.json")):
+        try:
+            index = read_json(index_path)
+            weight_map = index.get("weight_map")
+            if not isinstance(weight_map, dict) or not weight_map:
+                raise ValueError("missing or empty weight_map")
+            for tensor_name, filename in weight_map.items():
+                if not isinstance(filename, str):
+                    raise ValueError(f"weight_map entry for {tensor_name!r} is not a file name")
+                indexed_tensors[tensor_name] = filename
+                indexed_files.add(filename)
+        except Exception as exc:
+            result.failures.append(f"{index_path.name}: {exc}")
+
+    result.indexed_tensors = len(indexed_tensors)
+    for filename in sorted(indexed_files):
+        if not (root / filename).exists():
+            result.missing_files.append(filename)
+
+    actual_tensor_names = set()
+    for path in sorted(root.glob("*.safetensors")):
+        try:
+            validation = validate_safetensors_file(path)
+            result.checked_safetensors += 1
+            result.warnings.extend(validation.warnings)
+            actual_tensor_names.update(validation.tensor_names)
+
+            if indexed_tensors:
+                expected_for_shard = {
+                    tensor for tensor, filename in indexed_tensors.items() if filename == path.name
+                }
+                missing = expected_for_shard - validation.tensor_names
+                extra = validation.tensor_names - expected_for_shard
+                if missing:
+                    result.failures.append(f"{path.name}: {len(missing)} indexed tensor(s) missing")
+                if extra:
+                    message = f"{path.name}: {len(extra)} unindexed tensor(s)"
+                    if strict:
+                        result.failures.append(message)
+                    else:
+                        result.warnings.append(message)
+        except Exception as exc:
+            result.failures.append(f"{path.name}: {exc}")
+
+    missing_from_headers = set(indexed_tensors) - actual_tensor_names
+    if missing_from_headers:
+        result.failures.append(f"{len(missing_from_headers)} indexed tensor(s) missing from safetensors headers")
+
+    return result

model_mirror/checksums.py

@@ -0,0 +1,371 @@
+from __future__ import annotations
+
+import hashlib
+import json
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable, Protocol
+
+
+CHECKSUMS = ".checksums"
+MANIFEST = ".manifest"
+MANIFEST_SCHEMA = "model-mirror-manifest"
+MANIFEST_VERSION = 1
+SKIP_DIRS = {".model-mirror", ".cache", ".archive"}
+SKIP_FILES = {
+    CHECKSUMS,
+    MANIFEST,
+    ".checksums.tmp",
+    ".manifest.tmp",
+    ".verification",
+    ".verification.lock",
+    ".verification.tmp",
+}
+
+
+@dataclass(slots=True)
+class ChecksumResult:
+    checked: int = 0
+    missing: list[str] = field(default_factory=list)
+    failures: list[str] = field(default_factory=list)
+    extras: list[str] = field(default_factory=list)
+
+    @property
+    def ok(self) -> bool:
+        return not self.missing and not self.failures and not self.extras
+
+
+@dataclass(slots=True)
+class ChecksumWriteResult:
+    total: int = 0
+    hashed: int = 0
+    skipped: int = 0
+    removed: int = 0
+
+
+@dataclass(frozen=True, slots=True)
+class FileHashes:
+    sha256: str
+    git_blob_sha1: str
+
+
+@dataclass(slots=True)
+class FileHashState:
+    sha256: object
+    git_blob_sha1: object
+    bytes_hashed: int = 0
+
+    @property
+    def hashes(self) -> FileHashes:
+        return FileHashes(sha256=self.sha256.hexdigest(), git_blob_sha1=self.git_blob_sha1.hexdigest())
+
+
+class StreamingAccumulator(Protocol):
+    def update(self, data: bytes) -> None: ...
+
+    def reset(self) -> None: ...
+
+
+class HashingWriter:
+    def __init__(
+        self,
+        handle,
+        *,
+        expected_size: int,
+        hash_state: FileHashState | None = None,
+        accumulators: tuple[StreamingAccumulator, ...] = (),
+        on_progress: Callable[[int], None] | None = None,
+    ):
+        self._handle = handle
+        self._expected_size = expected_size
+        self._hash_state = hash_state or new_hash_state(expected_size)
+        self._accumulators = accumulators
+        self._on_progress = on_progress
+
+    def write(self, data: bytes) -> int:
+        written = self._handle.write(data)
+        chunk = data if written == len(data) else data[:written]
+        self._hash_state.sha256.update(chunk)
+        self._hash_state.git_blob_sha1.update(chunk)
+        for accumulator in self._accumulators:
+            accumulator.update(chunk)
+        self._hash_state.bytes_hashed += written
+        if self._on_progress is not None:
+            self._on_progress(self._hash_state.bytes_hashed)
+        return written
+
+    def tell(self) -> int:
+        return self._handle.tell()
+
+    def seek(self, offset: int, whence: int = 0) -> int:
+        return self._handle.seek(offset, whence)
+
+    def truncate(self, size: int | None = None) -> int:
+        truncated = self._handle.truncate(size)
+        if truncated == 0:
+            self._hash_state = new_hash_state(self._expected_size)
+            for accumulator in self._accumulators:
+                accumulator.reset()
+        elif truncated != self._hash_state.bytes_hashed:
+            raise OSError("cannot preserve streaming hash state after non-zero truncate")
+        return truncated
+
+    def flush(self) -> None:
+        self._handle.flush()
+
+    def fileno(self) -> int:
+        return self._handle.fileno()
+
+    @property
+    def hashes(self) -> FileHashes:
+        return self._hash_state.hashes
+
+
+def new_hash_state(expected_size: int) -> FileHashState:
+    sha256 = hashlib.sha256()
+    git_blob_sha1 = hashlib.sha1()
+    git_blob_sha1.update(f"blob {expected_size}\0".encode("ascii"))
+    return FileHashState(sha256=sha256, git_blob_sha1=git_blob_sha1)
+
+
+def hash_file_prefix(
+    path: Path,
+    *,
+    total_size: int,
+    prefix_size: int,
+    accumulators: tuple[StreamingAccumulator, ...] = (),
+    on_progress: Callable[[int], None] | None = None,
+) -> FileHashState:
+    state = new_hash_state(total_size)
+    remaining = prefix_size
+    with path.open("rb") as handle:
+        while remaining > 0:
+            chunk = handle.read(min(16 * 1024 * 1024, remaining))
+            if not chunk:
+                raise OSError(f"short read while hashing prefix: {path}")
+            state.sha256.update(chunk)
+            state.git_blob_sha1.update(chunk)
+            for accumulator in accumulators:
+                accumulator.update(chunk)
+            state.bytes_hashed += len(chunk)
+            if on_progress is not None:
+                on_progress(state.bytes_hashed)
+            remaining -= len(chunk)
+    return state
+
+
+def iter_payload_files(root: Path):
+    for path in sorted(root.rglob("*")):
+        if not path.is_file():
+            continue
+        rel = path.relative_to(root)
+        if rel.parts and rel.parts[0] in SKIP_DIRS:
+            continue
+        if rel.as_posix() in SKIP_FILES:
+            continue
+        yield path
+
+
+def file_hashes(
+    path: Path,
+    *,
+    accumulators: tuple[StreamingAccumulator, ...] = (),
+    on_progress: Callable[[int], None] | None = None,
+) -> FileHashes:
+    stat = path.stat()
+    digest = hashlib.sha256()
+    blob_digest = hashlib.sha1()
+    blob_digest.update(f"blob {stat.st_size}\0".encode("ascii"))
+    bytes_hashed = 0
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(16 * 1024 * 1024), b""):
+            digest.update(chunk)
+            blob_digest.update(chunk)
+            for accumulator in accumulators:
+                accumulator.update(chunk)
+            bytes_hashed += len(chunk)
+            if on_progress is not None:
+                on_progress(bytes_hashed)
+    return FileHashes(sha256=digest.hexdigest(), git_blob_sha1=blob_digest.hexdigest())
+
+
+def write_checksums(root: Path, *, max_workers: int = 1) -> ChecksumWriteResult:
+    manifest = load_manifest(root)
+    payload_files = list(iter_payload_files(root))
+    current_paths = {path.relative_to(root).as_posix() for path in payload_files}
+    result = ChecksumWriteResult(total=len(payload_files))
+
+    for rel in sorted(set(manifest) - current_paths):
+        manifest.pop(rel, None)
+        result.removed += 1
+
+    work: list[Path] = []
+    for path in payload_files:
+        rel = path.relative_to(root).as_posix()
+        stat = path.stat()
+        if record_is_current(manifest.get(rel), stat.st_size, stat.st_mtime_ns):
+            result.skipped += 1
+            continue
+        work.append(path)
+
+    if result.removed and not work:
+        write_manifest(root, manifest)
+    if not work:
+        return result
+
+    workers = max(1, max_workers)
+    if workers == 1:
+        for path in work:
+            row = checksum_row(root, path)
+            manifest[row["path"]] = row
+            result.hashed += 1
+            write_manifest(root, manifest)
+        return result
+
+    with ThreadPoolExecutor(max_workers=workers) as executor:
+        futures = [executor.submit(checksum_row, root, path) for path in work]
+        for future in as_completed(futures):
+            row = future.result()
+            manifest[row["path"]] = row
+            result.hashed += 1
+            write_manifest(root, manifest)
+    return result
+
+
+def record_is_current(row: dict | None, size: int, mtime_ns: int) -> bool:
+    if row is None:
+        return False
+    return (
+        row.get("size") == size
+        and row.get("mtime_ns") == mtime_ns
+        and bool(row.get("sha256"))
+        and bool(row.get("git_blob_sha1"))
+    )
+
+
+def checksum_row(root: Path, path: Path) -> dict:
+    hashes = file_hashes(path)
+    return checksum_row_from_hashes(root, path, hashes)
+
+
+def checksum_row_from_hashes(root: Path, path: Path, hashes: FileHashes) -> dict:
+    stat = path.stat()
+    return {
+        "path": path.relative_to(root).as_posix(),
+        "sha256": hashes.sha256,
+        "git_blob_sha1": hashes.git_blob_sha1,
+        "size": stat.st_size,
+        "mtime_ns": stat.st_mtime_ns,
+    }
+
+
+def write_manifest(root: Path, manifest: dict[str, dict]) -> None:
+    manifest_tmp = root / f"{MANIFEST}.tmp"
+    with manifest_tmp.open("w", encoding="utf-8", newline="\n") as handle:
+        handle.write(json.dumps({"schema": MANIFEST_SCHEMA, "version": MANIFEST_VERSION}, sort_keys=True) + "\n")
+        for rel in sorted(manifest):
+            handle.write(json.dumps(manifest[rel], sort_keys=True) + "\n")
+    manifest_tmp.replace(root / MANIFEST)
+    remove_obsolete_checksum_files(root)
+
+
+def remove_obsolete_checksum_files(root: Path) -> None:
+    for path in (root / CHECKSUMS, root / f"{CHECKSUMS}.tmp"):
+        path.unlink(missing_ok=True)
+
+
+def update_checksums(root: Path, paths: list[str], *, max_workers: int = 1) -> ChecksumWriteResult:
+    manifest = load_manifest(root)
+    result = ChecksumWriteResult(total=len(paths))
+    work: list[Path] = []
+    for rel in paths:
+        path = root / rel
+        if not path.exists() or not path.is_file():
+            if manifest.pop(rel, None) is not None:
+                result.removed += 1
+            continue
+        work.append(path)
+    if result.removed and not work:
+        write_manifest(root, manifest)
+    if not work:
+        return result
+
+    workers = max(1, max_workers)
+    if workers == 1:
+        for path in work:
+            row = checksum_row(root, path)
+            manifest[row["path"]] = row
+            result.hashed += 1
+            write_manifest(root, manifest)
+        return result
+
+    with ThreadPoolExecutor(max_workers=workers) as executor:
+        futures = [executor.submit(checksum_row, root, path) for path in work]
+        for future in as_completed(futures):
+            row = future.result()
+            manifest[row["path"]] = row
+            result.hashed += 1
+            write_manifest(root, manifest)
+    return result
+
+
+def load_manifest(root: Path) -> dict[str, dict]:
+    manifest: dict[str, dict] = {}
+    manifest_path = root / MANIFEST
+    if manifest_path.exists():
+        with manifest_path.open("r", encoding="utf-8") as handle:
+            saw_record = False
+            for line_number, line in enumerate(handle, 1):
+                if not line.strip():
+                    continue
+                try:
+                    row = json.loads(line)
+                except Exception as exc:
+                    raise ValueError(f"Malformed line {line_number} in {manifest_path}") from exc
+                if not saw_record:
+                    if not row.get("schema"):
+                        raise ValueError(f"Manifest missing header: {manifest_path}")
+                    validate_manifest_header(row, manifest_path)
+                    saw_record = True
+                    continue
+                saw_record = True
+                try:
+                    manifest[row["path"]] = row
+                except KeyError as exc:
+                    raise ValueError(f"Malformed line {line_number} in {manifest_path}") from exc
+    return manifest
+
+
+def validate_manifest_header(row: dict, path: Path) -> None:
+    if row.get("schema") != MANIFEST_SCHEMA:
+        raise ValueError(f"Unsupported manifest schema in {path}: {row.get('schema')}")
+    if row.get("version") != MANIFEST_VERSION:
+        raise ValueError(f"Unsupported manifest version in {path}: {row.get('version')}")
+
+
+def verify_checksums(root: Path, strict: bool = False) -> ChecksumResult:
+    manifest = load_manifest(root)
+    if not manifest:
+        return ChecksumResult()
+
+    result = ChecksumResult()
+    for rel, row in manifest.items():
+        path = root / rel
+        if not path.exists():
+            result.missing.append(rel)
+            continue
+        result.checked += 1
+        hashes = file_hashes(path)
+        if hashes.sha256 != row.get("sha256") or hashes.git_blob_sha1 != row.get("git_blob_sha1"):
+            result.failures.append(rel)
+
+    if strict:
+        tracked = set(manifest)
+        result.extras = [
+            path.relative_to(root).as_posix()
+            for path in iter_payload_files(root)
+            if path.relative_to(root).as_posix() not in tracked
+        ]
+
+    return result