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

model_mirror/removal.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+from .checksums import SKIP_DIRS, SKIP_FILES
+from .config import Config, REPO_TYPE_DIRS, safe_repo_path
+from .lock import LOCK_FILE, REMOVAL_MARKER
+from .state import utc_now
+
+
+REMOVALS_DIR = ".model-mirror-removals"
+REMOVAL_RECORD = REMOVAL_MARKER
+REMOVAL_SCHEMA = "model-mirror-removal"
+REMOVAL_VERSION = 1
+
+
+@dataclass(slots=True)
+class RemovalRecord:
+    repo_id: str
+    repo_type: str
+    original_path: str
+    status: str
+    exceptions: str
+    resolved_commit: str
+    checked_at_utc: str
+    check_age: str
+    payload_files: int
+    payload_size: int
+    started_at_utc: str = ""
+
+    def to_dict(self) -> dict:
+        return {
+            "schema": REMOVAL_SCHEMA,
+            "version": REMOVAL_VERSION,
+            **asdict(self),
+            "started_at_utc": self.started_at_utc or utc_now(),
+        }
+
+    @classmethod
+    def from_dict(cls, value: object, *, source: Path) -> RemovalRecord:
+        if (
+            not isinstance(value, dict)
+            or value.get("schema") != REMOVAL_SCHEMA
+            or value.get("version") != REMOVAL_VERSION
+        ):
+            raise ValueError(f"Unsupported removal record: {source}")
+        try:
+            return cls(
+                repo_id=str(value["repo_id"]),
+                repo_type=str(value["repo_type"]),
+                original_path=str(value["original_path"]),
+                status=str(value["status"]),
+                exceptions=str(value["exceptions"]),
+                resolved_commit=str(value["resolved_commit"]),
+                checked_at_utc=str(value["checked_at_utc"]),
+                check_age=str(value["check_age"]),
+                payload_files=int(value["payload_files"]),
+                payload_size=int(value["payload_size"]),
+                started_at_utc=str(value["started_at_utc"]),
+            )
+        except (KeyError, TypeError, ValueError) as exc:
+            raise ValueError(f"Malformed removal record: {source}") from exc
+
+
+def removal_path(config: Config, repo_id: str, repo_type: str) -> Path:
+    return (
+        Path(config.directory)
+        / REMOVALS_DIR
+        / REPO_TYPE_DIRS[repo_type]
+        / safe_repo_path(repo_id)
+    )
+
+
+def removal_record_path(root: Path) -> Path:
+    return root / REMOVAL_RECORD
+
+
+def stage_removal(root: Path, target: Path, record: RemovalRecord) -> Path:
+    if target.exists() or target.is_symlink():
+        raise FileExistsError(f"interrupted removal already exists: {target}")
+    target.parent.mkdir(parents=True, exist_ok=True)
+    write_removal_record(root, record)
+    root.rename(target)
+    return target
+
+
+def write_removal_record(root: Path, record: RemovalRecord) -> Path:
+    path = removal_record_path(root)
+    tmp = path.with_name(f"{path.name}.tmp")
+    tmp.write_text(json.dumps(record.to_dict(), sort_keys=True, indent=2) + "\n", encoding="utf-8")
+    tmp.replace(path)
+    return path
+
+
+def read_removal_record(root: Path) -> RemovalRecord | None:
+    path = removal_record_path(root)
+    if not path.exists():
+        return None
+    try:
+        value = json.loads(path.read_text(encoding="utf-8"))
+    except Exception as exc:
+        raise ValueError(f"Malformed removal record: {path}") from exc
+    return RemovalRecord.from_dict(value, source=path)
+
+
+def complete_removal(root: Path) -> None:
+    paths, directories = removable_paths(root)
+    payload = [path for path in paths if is_payload_path(root, path)]
+    metadata = [
+        path
+        for path in paths
+        if path not in payload and path.name not in {REMOVAL_RECORD, LOCK_FILE}
+    ]
+    for path in payload:
+        unlink_removal_path(path)
+    for path in metadata:
+        unlink_removal_path(path)
+    for path in directories:
+        path.rmdir()
+    removal_record_path(root).unlink(missing_ok=True)
+    (root / LOCK_FILE).unlink(missing_ok=True)
+    root.rmdir()
+
+
+def removable_paths(root: Path) -> tuple[list[Path], list[Path]]:
+    paths: list[Path] = []
+    directories: list[Path] = []
+    for current, dirnames, filenames in os.walk(root, topdown=False, followlinks=False):
+        current_path = Path(current)
+        for name in filenames:
+            paths.append(current_path / name)
+        for name in dirnames:
+            path = current_path / name
+            if path.is_symlink():
+                paths.append(path)
+            else:
+                directories.append(path)
+    paths.sort(key=lambda path: path.relative_to(root).as_posix())
+    directories.sort(key=lambda path: len(path.relative_to(root).parts), reverse=True)
+    return paths, directories
+
+
+def is_payload_path(root: Path, path: Path) -> bool:
+    rel = path.relative_to(root)
+    if rel.parts and rel.parts[0] in SKIP_DIRS:
+        return False
+    return rel.as_posix() not in SKIP_FILES and path.name != REMOVAL_RECORD
+
+
+def unlink_removal_path(path: Path) -> None:
+    path.unlink()
+
+
+def prune_empty_removal_parents(path: Path, *, stop: Path) -> None:
+    current = path
+    while current != stop and stop in current.parents:
+        try:
+            current.rmdir()
+        except OSError:
+            return
+        current = current.parent

model_mirror/repair.py

@@ -0,0 +1,399 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+from .audit import audit_model
+from .checksums import load_manifest, update_checksums, write_checksums
+from .config import Config, archive_path
+from .hub import (
+    HuggingFaceHub,
+    HubSnapshot,
+    cached_manifest_verifies,
+    get_snapshot,
+    read_snapshot_plan,
+    write_snapshot_plan,
+)
+from .lock import ModelBusyError, ModelLock
+from .state import AuditState, read_audit_state, state_from_results, write_audit_state
+from .verify import current_manifest_hash, metadata_blob_id, metadata_lfs_sha256, metadata_path, metadata_size, verify_remote
+
+
+@dataclass(slots=True)
+class RepairResult:
+    status: str
+    path: Path
+    paths: list[str]
+    upstream_status: str = "unknown"
+    resolved_commit: str = ""
+    upstream_commit: str = ""
+
+
+def repair(
+    config: Config,
+    repo_id: str,
+    *,
+    hub=None,
+    repo_type: str | None = None,
+    revision: str | None = None,
+    update: bool = False,
+    force_partial: bool = False,
+) -> RepairResult:
+    selected_type = repo_type or config.repo_type
+    selected_revision = revision or config.revision
+    selected_hub = hub or HuggingFaceHub(config)
+    root = archive_path(config, repo_id, selected_type)
+    initial_state = read_audit_state(root)
+    from .torrent_publication import (
+        assert_commit_update_allowed,
+        begin_maintenance,
+        finish_maintenance,
+        wait_for_maintenance_detach,
+    )
+
+    if (
+        update
+        and initial_state is not None
+        and initial_state.upstream_status == "changed"
+        and initial_state.upstream_commit
+    ):
+        assert_commit_update_allowed(root, initial_state.upstream_commit)
+
+    maintenance = False
+    if (
+        initial_state is not None
+        and not initial_state.clean
+        and bool(initial_state.repair_paths)
+        and not update
+    ):
+        with ModelLock(root, "torrent-maintenance", repo_id, selected_type):
+            maintenance = begin_maintenance(root) is not None
+        if maintenance:
+            wait_for_maintenance_detach(root)
+
+    result = None
+    try:
+        with ModelLock(root, "repair", repo_id, selected_type):
+            result = repair_locked(
+                config,
+                repo_id,
+                selected_hub,
+                selected_type,
+                selected_revision,
+                root,
+                update=update,
+                force_partial=force_partial,
+            )
+            if maintenance:
+                finish_maintenance(
+                    root,
+                    healthy=result.status in {"complete", "repaired"},
+                )
+            return result
+    except Exception:
+        if maintenance:
+            try:
+                with ModelLock(root, "torrent-maintenance-failed", repo_id, selected_type):
+                    finish_maintenance(root, healthy=False)
+            except ModelBusyError:
+                pass
+        raise
+
+
+def repair_locked(
+    config: Config,
+    repo_id: str,
+    selected_hub,
+    selected_type: str,
+    selected_revision: str,
+    root: Path,
+    *,
+    update: bool,
+    force_partial: bool,
+) -> RepairResult:
+    state = read_audit_state(root)
+    if state is None:
+        return RepairResult("verify-required", root, [])
+    if state.offline_only:
+        return repair_result("offline-only", root, [], state)
+    if update and state.upstream_status == "changed" and state.upstream_commit:
+        return update_to_upstream(
+            config,
+            repo_id,
+            selected_hub,
+            selected_type,
+            selected_revision,
+            root,
+            state,
+        )
+    pinned_snapshot = read_snapshot_plan(root)
+    if (
+        state.clean
+        and state.resolved_commit
+        and pinned_snapshot is not None
+        and pinned_snapshot.resolved_commit != state.resolved_commit
+    ):
+        return reconcile_snapshot_plan(
+            config,
+            repo_id,
+            selected_hub,
+            selected_type,
+            selected_revision,
+            root,
+            state,
+        )
+    if state.clean:
+        return repair_result("complete", root, [], state)
+
+    paths = sorted(set(state.repair_paths))
+    if not paths:
+        return repair_result("no-repair-paths", root, [], state)
+
+    requested_revision = state.requested_revision or selected_revision
+    target_revision = state.resolved_commit
+    if not target_revision:
+        return verification_incomplete_result(root, paths, state)
+    snapshot = snapshot_for_requested_revision(
+        get_snapshot(selected_hub, repo_id, selected_type, target_revision),
+        requested_revision,
+    )
+    if (
+        config.checksum
+        and not force_partial
+        and missing_manifest_paths(root, snapshot.files, ignored_paths=set(paths))
+    ):
+        return verification_incomplete_result(root, paths, state)
+    root.mkdir(parents=True, exist_ok=True)
+    for rel in paths:
+        target = root / rel
+        if target.exists() and target.is_file():
+            target.unlink()
+
+    download_snapshot(selected_hub, snapshot, root, allow_patterns=paths)
+    checksums_available = cached_manifest_verifies(root, snapshot.files)
+    if config.checksum and not checksums_available:
+        update_checksums(root, paths, max_workers=config.checksum_workers)
+        checksums_available = True
+
+    final_state = derive_state(
+        config,
+        repo_id,
+        selected_hub,
+        selected_type,
+        state.requested_revision or requested_revision,
+        root,
+        snapshot=snapshot,
+        resolved_commit=target_revision,
+        upstream_commit=state.upstream_commit,
+        cached=checksums_available,
+        from_manifest=checksums_available,
+        persist=False,
+    )
+    persist_repair_state(root, snapshot, final_state)
+    status = repair_status(final_state, success="repaired")
+    return repair_result(status, root, paths, final_state)
+
+
+def verification_incomplete_result(root: Path, paths: list[str], state: AuditState) -> RepairResult:
+    return RepairResult(
+        "verification-incomplete",
+        root,
+        paths,
+        state.upstream_status,
+        state.resolved_commit,
+        state.upstream_commit,
+    )
+
+
+def update_to_upstream(
+    config: Config,
+    repo_id: str,
+    selected_hub,
+    selected_type: str,
+    selected_revision: str,
+    root: Path,
+    state: AuditState,
+) -> RepairResult:
+    requested_revision = state.requested_revision or selected_revision
+    target_revision = state.upstream_commit
+    root.mkdir(parents=True, exist_ok=True)
+    snapshot = snapshot_for_requested_revision(
+        get_snapshot(selected_hub, repo_id, selected_type, target_revision),
+        requested_revision,
+    )
+    download_snapshot(selected_hub, snapshot, root)
+    checksums_available = cached_manifest_verifies(root, snapshot.files)
+    if config.checksum and not checksums_available:
+        write_checksums(root, max_workers=config.checksum_workers)
+        checksums_available = cached_manifest_verifies(root, snapshot.files)
+    final_state = derive_state(
+        config,
+        repo_id,
+        selected_hub,
+        selected_type,
+        requested_revision,
+        root,
+        snapshot=snapshot,
+        resolved_commit=target_revision,
+        upstream_commit=target_revision,
+        cached=checksums_available,
+        from_manifest=checksums_available,
+        persist=False,
+    )
+    persist_repair_state(root, snapshot, final_state)
+    status = repair_status(final_state, success="updated")
+    return repair_result(status, root, [], final_state)
+
+
+def reconcile_snapshot_plan(
+    config: Config,
+    repo_id: str,
+    selected_hub,
+    selected_type: str,
+    selected_revision: str,
+    root: Path,
+    state: AuditState,
+) -> RepairResult:
+    requested_revision = state.requested_revision or selected_revision
+    snapshot = snapshot_for_requested_revision(
+        get_snapshot(selected_hub, repo_id, selected_type, state.resolved_commit),
+        requested_revision,
+    )
+    checksums_available = cached_manifest_verifies(root, snapshot.files)
+    final_state = derive_state(
+        config,
+        repo_id,
+        selected_hub,
+        selected_type,
+        requested_revision,
+        root,
+        snapshot=snapshot,
+        resolved_commit=state.resolved_commit,
+        upstream_commit=state.upstream_commit,
+        cached=checksums_available,
+        from_manifest=checksums_available,
+        persist=False,
+    )
+    persist_repair_state(root, snapshot, final_state)
+    status = repair_status(final_state, success="repaired")
+    return repair_result(status, root, [], final_state)
+
+
+def snapshot_for_requested_revision(snapshot: HubSnapshot, requested_revision: str) -> HubSnapshot:
+    return HubSnapshot(
+        repo_id=snapshot.repo_id,
+        repo_type=snapshot.repo_type,
+        requested_revision=requested_revision,
+        resolved_commit=snapshot.resolved_commit,
+        files=snapshot.files,
+    )
+
+
+def persist_repair_state(root: Path, snapshot: HubSnapshot, state: AuditState) -> None:
+    if state.clean:
+        write_snapshot_plan(root, snapshot)
+    write_audit_state(root, state)
+
+
+def repair_status(state: AuditState, *, success: str) -> str:
+    if state.clean:
+        return success
+    if state.status == "incomplete":
+        return "verification-incomplete"
+    return "incomplete"
+
+
+def download_snapshot(selected_hub, snapshot: HubSnapshot, root: Path, *, allow_patterns: list[str] | None = None) -> None:
+    if hasattr(selected_hub, "download_snapshot"):
+        selected_hub.download_snapshot(snapshot, root, allow_patterns=allow_patterns)
+        return
+    selected_hub.snapshot_download(
+        snapshot.repo_id,
+        snapshot.repo_type,
+        snapshot.resolved_commit,
+        root,
+        allow_patterns=allow_patterns,
+    )
+
+
+def missing_manifest_paths(root: Path, metadata, *, ignored_paths: set[str]) -> list[str]:
+    manifest = load_manifest(root)
+    missing = []
+    for item in metadata:
+        rel = metadata_path(item)
+        if rel in ignored_paths:
+            continue
+        if metadata_lfs_sha256(item) is None and metadata_blob_id(item) is None:
+            continue
+        path = root / rel
+        if not path.exists() or not path.is_file():
+            continue
+        expected_size = metadata_size(item)
+        stat = path.stat()
+        if expected_size is not None and stat.st_size != expected_size:
+            continue
+        if metadata_lfs_sha256(item) is not None:
+            hash_key = "sha256"
+        else:
+            hash_key = "git_blob_sha1"
+        if current_manifest_hash(manifest, rel, stat.st_size, stat.st_mtime_ns, hash_key) is None:
+            missing.append(rel)
+    return missing
+
+
+def repair_result(
+    status: str,
+    root: Path,
+    paths: list[str],
+    state: AuditState,
+) -> RepairResult:
+    return RepairResult(
+        status,
+        root,
+        paths,
+        state.upstream_status,
+        state.resolved_commit,
+        state.upstream_commit,
+    )
+
+
+def derive_state(
+    config: Config,
+    repo_id: str,
+    hub,
+    repo_type: str,
+    requested_revision: str,
+    root: Path,
+    *,
+    snapshot: HubSnapshot | None = None,
+    resolved_commit: str | None = None,
+    upstream_commit: str | None = None,
+    cached: bool = False,
+    from_manifest: bool = False,
+    persist: bool = True,
+) -> AuditState:
+    if snapshot is None:
+        snapshot = get_snapshot(hub, repo_id, repo_type, resolved_commit or requested_revision)
+    metadata = snapshot.files
+    remote_result = verify_remote(root, metadata, cached=cached, from_manifest=from_manifest)
+    expected_paths = {metadata_path(item) for item in metadata}
+    gguf_only = "config.json" not in expected_paths and any(
+        path.lower().endswith(".gguf") for path in expected_paths
+    )
+    audit_result = (
+        audit_model(root, skip_transformers=True, require_config=not gguf_only)
+        if repo_type == "model"
+        else None
+    )
+    state = state_from_results(
+        repo_id,
+        repo_type,
+        requested_revision,
+        remote_result,
+        audit_result,
+        resolved_commit=resolved_commit or snapshot.resolved_commit,
+        upstream_commit=upstream_commit or snapshot.resolved_commit,
+    )
+    if persist:
+        write_audit_state(root, state)
+    return state

model_mirror/state.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+
+import yaml
+
+
+VERIFICATION_FILE = ".verification"
+
+
+@dataclass(slots=True)
+class VerificationState:
+    status: str
+    repo_id: str
+    repo_type: str = "model"
+    requested_revision: str = "main"
+    resolved_commit: str = ""
+    upstream_commit: str = ""
+    upstream_status: str = "unknown"
+    offline_only: bool = False
+    repair_paths: list[str] = field(default_factory=list)
+    issues: list[str] = field(default_factory=list)
+    checked_at_utc: str = ""
+
+    @property
+    def clean(self) -> bool:
+        return self.status == "clean"
+
+
+def audit_state_path(root: Path) -> Path:
+    return verification_state_path(root)
+
+
+def verification_state_path(root: Path) -> Path:
+    return root / VERIFICATION_FILE
+
+
+def utc_now() -> str:
+    return datetime.now(timezone.utc).isoformat(timespec="seconds")
+
+
+def read_verification_state(root: Path) -> VerificationState | None:
+    path = verification_state_path(root)
+    if not path.exists():
+        return None
+    data = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
+    if not isinstance(data, dict):
+        raise ValueError(f"Verification state must contain a YAML mapping: {path}")
+    return VerificationState(
+        status=str(data.get("status", "dirty")),
+        repo_id=str(data.get("repo_id", "")),
+        repo_type=str(data.get("repo_type", "model")),
+        requested_revision=str(data.get("requested_revision", data.get("revision", "main"))),
+        resolved_commit=str(data.get("resolved_commit", "")),
+        upstream_commit=str(data.get("upstream_commit", "")),
+        upstream_status=str(data.get("upstream_status", "unknown")),
+        offline_only=bool(data.get("offline_only", False)),
+        repair_paths=sorted(str(item) for item in data.get("repair_paths", []) or []),
+        issues=[str(item) for item in data.get("issues", []) or []],
+        checked_at_utc=str(data.get("checked_at_utc", "")),
+    )
+
+
+def write_verification_state(root: Path, state: VerificationState) -> Path:
+    path = verification_state_path(root)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    payload = {
+        "status": state.status,
+        "repo_id": state.repo_id,
+        "repo_type": state.repo_type,
+        "requested_revision": state.requested_revision,
+        "resolved_commit": state.resolved_commit,
+        "upstream_commit": state.upstream_commit,
+        "upstream_status": state.upstream_status,
+        "offline_only": state.offline_only,
+        "repair_paths": sorted(set(state.repair_paths)),
+        "issues": state.issues,
+        "checked_at_utc": state.checked_at_utc or utc_now(),
+    }
+    tmp = path.with_name(f"{path.name}.tmp")
+    tmp.write_text(yaml.safe_dump(payload, sort_keys=True), encoding="utf-8")
+    tmp.replace(path)
+    return path
+
+
+def repair_paths_from_results(remote_result, audit_result=None) -> list[str]:
+    paths = set()
+    paths.update(getattr(remote_result, "missing", []))
+    paths.update(getattr(remote_result, "size_mismatches", []))
+    paths.update(getattr(remote_result, "hash_mismatches", []))
+    if audit_result is not None:
+        paths.update(getattr(audit_result, "missing_files", []))
+        for failure in getattr(audit_result, "failures", []):
+            candidate = str(failure).split(":", 1)[0]
+            if candidate.endswith((".json", ".safetensors", ".bin", ".txt", ".model")):
+                paths.add(candidate)
+    return sorted(paths)
+
+
+def state_from_results(
+    repo_id: str,
+    repo_type: str,
+    requested_revision: str,
+    remote_result,
+    audit_result=None,
+    *,
+    resolved_commit: str = "",
+    upstream_commit: str = "",
+    offline_only: bool = False,
+) -> VerificationState:
+    issues = []
+    for name in ("missing", "size_mismatches", "hash_mismatches", "cached_hash_missing", "extras"):
+        values = getattr(remote_result, name, [])
+        issues.extend(f"{name}: {value}" for value in values)
+    if audit_result is not None:
+        issues.extend(f"missing: {value}" for value in getattr(audit_result, "missing_files", []))
+        issues.extend(f"audit: {value}" for value in getattr(audit_result, "failures", []))
+
+    ok = getattr(remote_result, "ok", False) and (audit_result is None or getattr(audit_result, "ok", False))
+    repair_paths = repair_paths_from_results(remote_result, audit_result)
+    cache_incomplete = bool(getattr(remote_result, "cached_hash_missing", [])) and not repair_paths
+    return VerificationState(
+        status="clean" if ok else "incomplete" if cache_incomplete else "dirty",
+        repo_id=repo_id,
+        repo_type=repo_type,
+        requested_revision=requested_revision,
+        resolved_commit=resolved_commit,
+        upstream_commit=upstream_commit,
+        upstream_status=upstream_status(resolved_commit, upstream_commit),
+        offline_only=offline_only,
+        repair_paths=[] if ok else repair_paths,
+        issues=[] if ok else issues,
+        checked_at_utc=utc_now(),
+    )
+
+
+def upstream_status(resolved_commit: str, upstream_commit: str) -> str:
+    if not upstream_commit:
+        return "unknown"
+    if not resolved_commit:
+        return "unknown"
+    return "current" if resolved_commit == upstream_commit else "changed"
+
+
+# Backward-compatible internal aliases while the implementation is being split up.
+AuditState = VerificationState
+read_audit_state = read_verification_state
+write_audit_state = write_verification_state