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

model_mirror/lock.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+import errno
+import fcntl
+import os
+import socket
+from dataclasses import dataclass
+from pathlib import Path
+
+import yaml
+
+from .state import utc_now
+
+
+LOCK_FILE = ".verification.lock"
+REMOVAL_MARKER = ".model-mirror-removal.json"
+
+
+class ModelBusyError(RuntimeError):
+    def __init__(self, root: Path, info: dict):
+        self.root = root
+        self.info = info
+        detail = lock_label(info)
+        super().__init__(f"model mirror is busy: {root} ({detail})")
+
+
+@dataclass(slots=True)
+class ModelLock:
+    root: Path
+    command: str
+    repo_id: str
+    repo_type: str = "model"
+    wait: bool = False
+    handle: object | None = None
+
+    def __enter__(self):
+        self.root.mkdir(parents=True, exist_ok=True)
+        path = lock_path(self.root)
+        self.handle = path.open("a+", encoding="utf-8")
+        flags = fcntl.LOCK_EX if self.wait else fcntl.LOCK_EX | fcntl.LOCK_NB
+        try:
+            fcntl.flock(self.handle.fileno(), flags)
+        except OSError as exc:
+            if exc.errno not in {errno.EACCES, errno.EAGAIN}:
+                raise
+            self.handle.seek(0)
+            info = yaml.safe_load(self.handle.read()) or {}
+            raise ModelBusyError(self.root, info) from exc
+
+        marker = self.root / REMOVAL_MARKER
+        if self.command != "remove" and (marker.exists() or marker.is_symlink()):
+            info = {
+                "command": "remove-pending",
+                "repo_id": self.repo_id,
+                "repo_type": self.repo_type,
+            }
+            fcntl.flock(self.handle.fileno(), fcntl.LOCK_UN)
+            self.handle.close()
+            self.handle = None
+            raise ModelBusyError(self.root, info)
+
+        info = {
+            "pid": os.getpid(),
+            "host": socket.gethostname(),
+            "command": self.command,
+            "repo_id": self.repo_id,
+            "repo_type": self.repo_type,
+            "started_at_utc": utc_now(),
+        }
+        self.handle.seek(0)
+        self.handle.truncate()
+        self.handle.write(yaml.safe_dump(info, sort_keys=True))
+        self.handle.flush()
+        os.fsync(self.handle.fileno())
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        if self.handle is None:
+            return False
+        self.handle.seek(0)
+        self.handle.truncate()
+        self.handle.flush()
+        os.fsync(self.handle.fileno())
+        fcntl.flock(self.handle.fileno(), fcntl.LOCK_UN)
+        self.handle.close()
+        self.handle = None
+        return False
+
+
+def lock_path(root: Path) -> Path:
+    return root / LOCK_FILE
+
+
+def read_active_lock(root: Path) -> dict | None:
+    path = lock_path(root)
+    if not path.exists():
+        return None
+    with path.open("r", encoding="utf-8") as handle:
+        try:
+            fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except OSError as exc:
+            if exc.errno not in {errno.EACCES, errno.EAGAIN}:
+                raise
+            handle.seek(0)
+            return yaml.safe_load(handle.read()) or {}
+        else:
+            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+            return None
+
+
+def lock_label(info: dict | None) -> str:
+    if not info:
+        return "lock held"
+    parts = []
+    for key in ("command", "pid", "host", "started_at_utc"):
+        value = info.get(key)
+        if value:
+            parts.append(f"{key}={value}")
+    return ", ".join(parts) if parts else "lock held"

model_mirror/mirror.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+from .checksums import write_checksums
+from .config import Config, archive_path
+from .hub import HuggingFaceHub, cached_manifest_verifies, compatible_snapshot_plan, get_snapshot, write_snapshot_plan
+from .lock import ModelLock
+from .repair import derive_state
+from .state import VerificationState, read_verification_state, write_verification_state
+from .verify import verify_remote
+
+
+@dataclass(slots=True)
+class MirrorResult:
+    status: str
+    path: Path
+    files: int
+
+
+def mirror(
+    config: Config,
+    repo_id: str,
+    *,
+    hub=None,
+    repo_type: str | None = None,
+    revision: str | None = None,
+    force: bool = False,
+    verify_after: bool = True,
+    stall_timeout_seconds: int | None = None,
+) -> MirrorResult:
+    selected_type = repo_type or config.repo_type
+    selected_revision = revision or config.revision
+    selected_hub = hub or HuggingFaceHub(config)
+    destination = archive_path(config, repo_id, selected_type)
+    with ModelLock(destination, "mirror", repo_id, selected_type):
+        existing_state = read_verification_state(destination)
+        if existing_state is None:
+            write_verification_state(
+                destination,
+                VerificationState(
+                    status="in_progress",
+                    repo_id=repo_id,
+                    repo_type=selected_type,
+                    requested_revision=selected_revision,
+                    issues=["mirror started"],
+                ),
+            )
+        return mirror_locked(
+            config,
+            repo_id,
+            selected_hub,
+            selected_type,
+            selected_revision,
+            destination,
+            existing_state=existing_state,
+            force=force,
+            verify_after=verify_after,
+            stall_timeout_seconds=stall_timeout_seconds,
+        )
+
+
+def mirror_locked(
+    config: Config,
+    repo_id: str,
+    selected_hub,
+    selected_type: str,
+    selected_revision: str,
+    destination: Path,
+    *,
+    existing_state: VerificationState | None,
+    force: bool,
+    verify_after: bool,
+    stall_timeout_seconds: int | None,
+) -> MirrorResult:
+    snapshot = select_mirror_snapshot(
+        selected_hub,
+        repo_id,
+        selected_type,
+        selected_revision,
+        destination,
+        existing_state=existing_state,
+        force=force,
+    )
+    from .torrent_publication import assert_commit_update_allowed, load_fenced_publication
+
+    assert_commit_update_allowed(destination, snapshot.resolved_commit)
+    metadata = snapshot.files
+
+    if not force and verify_remote(destination, metadata, check_hashes=False).ok:
+        if not verify_after:
+            return MirrorResult("complete", destination, len(metadata))
+        if (
+            existing_state is not None
+            and existing_state.clean
+            and existing_state.resolved_commit == snapshot.resolved_commit
+        ):
+            return MirrorResult("complete", destination, len(metadata))
+        checksums_written = cached_manifest_verifies(destination, metadata)
+        if config.checksum and not checksums_written:
+            write_checksums(destination, max_workers=config.checksum_workers)
+            checksums_written = True
+        state = derive_state(
+            config,
+            repo_id,
+            selected_hub,
+            selected_type,
+            selected_revision,
+            destination,
+            snapshot=snapshot,
+            upstream_commit=snapshot.resolved_commit,
+            cached=False,
+            from_manifest=checksums_written,
+        )
+        return MirrorResult("complete" if state.clean else "downloaded-unverified", destination, len(metadata))
+
+    fenced = load_fenced_publication(destination)
+    if fenced is not None:
+        raise RuntimeError(
+            f"published payload cannot be replaced by mirror: {fenced[0].publication_id}; "
+            "use same-commit repair or retire the publication"
+        )
+
+    write_verification_state(
+        destination,
+        VerificationState(
+            status="in_progress",
+            repo_id=repo_id,
+            repo_type=selected_type,
+            requested_revision=selected_revision,
+            resolved_commit=snapshot.resolved_commit,
+            upstream_commit=snapshot.resolved_commit,
+            upstream_status="current",
+            issues=["mirror in progress"],
+        ),
+    )
+    destination.mkdir(parents=True, exist_ok=True)
+    write_snapshot_plan(destination, snapshot)
+    download_snapshot(selected_hub, snapshot, destination, stall_timeout_seconds=stall_timeout_seconds)
+    checksums_written = cached_manifest_verifies(destination, metadata)
+    if config.checksum and not checksums_written:
+        write_checksums(destination, max_workers=config.checksum_workers)
+        checksums_written = True
+    if verify_after:
+        state = derive_state(
+            config,
+            repo_id,
+            selected_hub,
+            selected_type,
+            selected_revision,
+            destination,
+            snapshot=snapshot,
+            upstream_commit=snapshot.resolved_commit,
+            cached=False,
+            from_manifest=checksums_written,
+        )
+        return MirrorResult("downloaded" if state.clean else "downloaded-unverified", destination, len(metadata))
+    write_verification_state(
+        destination,
+        VerificationState(
+            status="dirty",
+            repo_id=repo_id,
+            repo_type=selected_type,
+            requested_revision=selected_revision,
+            resolved_commit=snapshot.resolved_commit,
+            upstream_commit=snapshot.resolved_commit,
+            upstream_status="current",
+            issues=["verification skipped"],
+        ),
+    )
+    return MirrorResult("downloaded", destination, len(metadata))
+
+
+def select_mirror_snapshot(
+    selected_hub,
+    repo_id: str,
+    selected_type: str,
+    selected_revision: str,
+    destination: Path,
+    *,
+    existing_state: VerificationState | None,
+    force: bool,
+):
+    if not force and existing_state is not None and existing_state.status == "in_progress":
+        frozen = compatible_snapshot_plan(
+            destination,
+            repo_id=repo_id,
+            repo_type=selected_type,
+            requested_revision=selected_revision,
+        )
+        if frozen is not None:
+            return frozen
+    return get_snapshot(selected_hub, repo_id, selected_type, selected_revision)
+
+
+def download_snapshot(selected_hub, snapshot, destination: Path, *, stall_timeout_seconds: int | None = None) -> None:
+    if hasattr(selected_hub, "download_snapshot"):
+        selected_hub.download_snapshot(snapshot, destination, stall_timeout_seconds=stall_timeout_seconds)
+        return
+    selected_hub.snapshot_download(snapshot.repo_id, snapshot.repo_type, snapshot.resolved_commit, destination)

model_mirror/progress.py

@@ -0,0 +1,260 @@
+from __future__ import annotations
+
+import json
+import os
+import threading
+import time
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+
+
+PROGRESS_DIR = ".model-mirror"
+PROGRESS_FILE = "progress.json"
+PROGRESS_SCHEMA = "model-mirror-progress"
+PROGRESS_VERSION = 1
+DEFAULT_STALL_TIMEOUT_SECONDS = 600
+DEFAULT_STALL_RETRIES = 3
+SKIP_DIRS = {".model-mirror", ".cache", ".archive"}
+
+
+@dataclass(frozen=True, slots=True)
+class ProgressEntry:
+    path: str
+    stage: str
+    bytes_done: int
+    bytes_total: int | None
+    updated_at_utc: str
+    idle_seconds: int | None
+    stalled: bool
+    source: str
+    rate_bytes_per_second: float | None = None
+
+
+@dataclass(frozen=True, slots=True)
+class ProgressSnapshot:
+    entries: list[ProgressEntry]
+    source: str
+
+    @property
+    def active(self) -> bool:
+        return bool(self.entries)
+
+    @property
+    def stalled_count(self) -> int:
+        return sum(1 for entry in self.entries if entry.stalled)
+
+    @property
+    def any_stalled(self) -> bool:
+        return self.stalled_count > 0
+
+
+class ProgressRecorder:
+    def __init__(
+        self,
+        root: Path,
+        *,
+        min_interval_seconds: float = 5.0,
+        min_bytes: int = 64 * 1024 * 1024,
+    ):
+        self.root = root
+        self.min_interval_seconds = min_interval_seconds
+        self.min_bytes = min_bytes
+        self._active: dict[str, dict] = {}
+        self._lock = threading.Lock()
+
+    def track(self, path: str, *, total: int | None, stage: str, bytes_done: int = 0) -> FileProgress:
+        progress = FileProgress(self, path, total=total)
+        progress.update(bytes_done, stage=stage, force=True)
+        return progress
+
+    def update(
+        self,
+        path: str,
+        *,
+        total: int | None,
+        stage: str,
+        bytes_done: int,
+        rate_bytes_per_second: float | None,
+    ) -> None:
+        now = utc_now_iso()
+        with self._lock:
+            entry = self._active.get(path)
+            if entry is None:
+                entry = {
+                    "path": path,
+                    "bytes_total": total,
+                    "started_at_utc": now,
+                }
+            entry.update(
+                {
+                    "stage": stage,
+                    "bytes_done": bytes_done,
+                    "bytes_total": total,
+                    "updated_at_utc": now,
+                    "rate_bytes_per_second": rate_bytes_per_second,
+                }
+            )
+            self._active[path] = entry
+            self._write_locked()
+
+    def finish(self, path: str) -> None:
+        with self._lock:
+            self._active.pop(path, None)
+            if self._active:
+                self._write_locked()
+            else:
+                progress_path(self.root).unlink(missing_ok=True)
+
+    def _write_locked(self) -> None:
+        path = progress_path(self.root)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        payload = {
+            "schema": PROGRESS_SCHEMA,
+            "version": PROGRESS_VERSION,
+            "pid": os.getpid(),
+            "updated_at_utc": utc_now_iso(),
+            "active_files": self._active,
+        }
+        tmp = path.with_name(f"{path.name}.tmp")
+        tmp.write_text(json.dumps(payload, sort_keys=True, indent=2) + "\n", encoding="utf-8")
+        tmp.replace(path)
+
+
+class FileProgress:
+    def __init__(self, recorder: ProgressRecorder, path: str, *, total: int | None):
+        self.recorder = recorder
+        self.path = path
+        self.total = total
+        self.stage = "starting"
+        self.last_emit_time: float | None = None
+        self.last_emit_bytes: int | None = None
+
+    def update(self, bytes_done: int, *, stage: str | None = None, force: bool = False) -> None:
+        now = time.monotonic()
+        selected_stage = stage or self.stage
+        should_emit = force or self.last_emit_time is None
+        if not should_emit and now - self.last_emit_time >= self.recorder.min_interval_seconds:
+            should_emit = True
+        if (
+            not should_emit
+            and self.last_emit_bytes is not None
+            and bytes_done - self.last_emit_bytes >= self.recorder.min_bytes
+        ):
+            should_emit = True
+        if not should_emit:
+            self.stage = selected_stage
+            return
+
+        rate = None
+        if self.last_emit_time is not None and self.last_emit_bytes is not None:
+            elapsed = now - self.last_emit_time
+            if elapsed > 0:
+                rate = max(0, bytes_done - self.last_emit_bytes) / elapsed
+        self.last_emit_time = now
+        self.last_emit_bytes = bytes_done
+        self.stage = selected_stage
+        self.recorder.update(
+            self.path,
+            total=self.total,
+            stage=selected_stage,
+            bytes_done=bytes_done,
+            rate_bytes_per_second=rate,
+        )
+
+    def finish(self) -> None:
+        self.recorder.finish(self.path)
+
+
+def progress_path(root: Path) -> Path:
+    return root / PROGRESS_DIR / PROGRESS_FILE
+
+
+def progress_snapshot(
+    root: Path,
+    *,
+    stall_timeout_seconds: int = DEFAULT_STALL_TIMEOUT_SECONDS,
+    now: datetime | None = None,
+    scan_incomplete: bool = True,
+) -> ProgressSnapshot:
+    now = now or datetime.now(timezone.utc)
+    entries = progress_file_entries(root, stall_timeout_seconds=stall_timeout_seconds, now=now)
+    if entries:
+        return ProgressSnapshot(entries=entries, source="heartbeat")
+    if not scan_incomplete:
+        return ProgressSnapshot(entries=[], source="heartbeat")
+    return ProgressSnapshot(
+        entries=incomplete_file_entries(root, stall_timeout_seconds=stall_timeout_seconds, now=now),
+        source="partial-file",
+    )
+
+
+def progress_file_entries(root: Path, *, stall_timeout_seconds: int, now: datetime) -> list[ProgressEntry]:
+    path = progress_path(root)
+    if not path.exists():
+        return []
+    data = json.loads(path.read_text(encoding="utf-8"))
+    if data.get("schema") != PROGRESS_SCHEMA or data.get("version") != PROGRESS_VERSION:
+        return []
+    entries = []
+    for rel, item in sorted((data.get("active_files") or {}).items()):
+        updated_at = str(item.get("updated_at_utc") or "")
+        idle = idle_seconds(updated_at, now)
+        entries.append(
+            ProgressEntry(
+                path=str(item.get("path") or rel),
+                stage=str(item.get("stage") or "unknown"),
+                bytes_done=int(item.get("bytes_done") or 0),
+                bytes_total=item.get("bytes_total"),
+                updated_at_utc=updated_at,
+                idle_seconds=idle,
+                stalled=is_stalled(idle, stall_timeout_seconds),
+                source="heartbeat",
+                rate_bytes_per_second=item.get("rate_bytes_per_second"),
+            )
+        )
+    return entries
+
+
+def incomplete_file_entries(root: Path, *, stall_timeout_seconds: int, now: datetime) -> list[ProgressEntry]:
+    entries = []
+    for path in sorted(root.rglob("*.incomplete")):
+        if not path.is_file():
+            continue
+        rel = path.relative_to(root)
+        if rel.parts and rel.parts[0] in SKIP_DIRS:
+            continue
+        updated = datetime.fromtimestamp(path.stat().st_mtime, timezone.utc).isoformat(timespec="seconds")
+        idle = idle_seconds(updated, now)
+        rel_path = rel.as_posix()[: -len(".incomplete")]
+        entries.append(
+            ProgressEntry(
+                path=rel_path,
+                stage="partial",
+                bytes_done=path.stat().st_size,
+                bytes_total=None,
+                updated_at_utc=updated,
+                idle_seconds=idle,
+                stalled=is_stalled(idle, stall_timeout_seconds),
+                source="partial-file",
+            )
+        )
+    return entries
+
+
+def idle_seconds(updated_at_utc: str, now: datetime) -> int | None:
+    try:
+        updated = datetime.fromisoformat(updated_at_utc)
+    except ValueError:
+        return None
+    if updated.tzinfo is None:
+        updated = updated.replace(tzinfo=timezone.utc)
+    return max(0, int((now - updated).total_seconds()))
+
+
+def is_stalled(idle: int | None, stall_timeout_seconds: int) -> bool:
+    return stall_timeout_seconds > 0 and idle is not None and idle >= stall_timeout_seconds
+
+
+def utc_now_iso() -> str:
+    return datetime.now(timezone.utc).isoformat(timespec="seconds")