zeroquantz 0.1.0__py3-none-any.whl

zeroquantz/core/artifacts.py

@@ -0,0 +1,179 @@
+"""Manage on-disk artifacts: downloaded models (the HF cache) and the quantized
+models ZeroQuantz produces.
+
+Downloaded models are read from / deleted via ``huggingface_hub``'s cache API.
+Quantized outputs are tracked in a small registry (``~/.zeroquantz/quantized.json``)
+that ZeroQuantz appends to whenever it exports/quantizes, so they can be listed,
+opened, and removed later.
+"""
+
+from __future__ import annotations
+
+import json
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+from zeroquantz.utils.env import utc_timestamp
+from zeroquantz.utils.logging import get_logger
+from zeroquantz.utils.paths import paths
+
+log = get_logger(__name__)
+
+
+# ---- downloaded models (Hugging Face cache) --------------------------------
+
+
+@dataclass
+class DownloadedModel:
+    repo_id: str
+    size_bytes: int
+    nb_files: int
+    revisions: list[str]
+    path: str
+
+
+def list_downloaded_models() -> list[DownloadedModel]:
+    """List model repos in the HF cache, largest first. Safe if the cache is empty."""
+    try:
+        from huggingface_hub import scan_cache_dir
+    except ImportError:
+        return []
+    try:
+        info = scan_cache_dir()
+    except Exception as exc:  # pragma: no cover - cache-dependent
+        log.debug("scan_cache_dir failed: %s", exc)
+        return []
+    models: list[DownloadedModel] = []
+    for repo in info.repos:
+        if getattr(repo, "repo_type", "model") != "model":
+            continue
+        models.append(
+            DownloadedModel(
+                repo_id=repo.repo_id,
+                size_bytes=int(repo.size_on_disk),
+                nb_files=int(repo.nb_files),
+                revisions=[r.commit_hash for r in repo.revisions],
+                path=str(repo.repo_path),
+            )
+        )
+    models.sort(key=lambda m: m.size_bytes, reverse=True)
+    return models
+
+
+def delete_downloaded_model(repo_id: str) -> int:
+    """Delete every cached revision of ``repo_id``. Returns bytes freed."""
+    from huggingface_hub import scan_cache_dir
+
+    info = scan_cache_dir()
+    hashes: list[str] = []
+    for repo in info.repos:
+        if repo.repo_id == repo_id and getattr(repo, "repo_type", "model") == "model":
+            hashes.extend(r.commit_hash for r in repo.revisions)
+    if not hashes:
+        return 0
+    strategy = info.delete_revisions(*hashes)
+    freed = int(strategy.expected_freed_size)
+    strategy.execute()
+    return freed
+
+
+# ---- quantized outputs (ZeroQuantz registry) --------------------------------
+
+
+class QuantizedArtifact(BaseModel):
+    path: str
+    base_model: str | None = None
+    format_id: str | None = None
+    method: str | None = None
+    size_bytes: int = 0
+    created_at: str = Field(default_factory=utc_timestamp)
+
+    @property
+    def exists(self) -> bool:
+        return Path(self.path).exists()
+
+
+def dir_size(path: str | Path) -> int:
+    total = 0
+    root = Path(path)
+    if not root.exists():
+        return 0
+    for p in root.rglob("*"):
+        if p.is_file():
+            try:
+                total += p.stat().st_size
+            except OSError:
+                pass
+    return total
+
+
+class QuantizedRegistry:
+    """A JSON-backed registry of quantized model directories ZeroQuantz produced."""
+
+    def __init__(self, file: Path | None = None) -> None:
+        self.file = file or (paths().ensure().home / "quantized.json")
+
+    def _load(self) -> list[QuantizedArtifact]:
+        if not self.file.exists():
+            return []
+        try:
+            data = json.loads(self.file.read_text(encoding="utf-8"))
+            return [QuantizedArtifact.model_validate(d) for d in data]
+        except (OSError, ValueError):
+            return []
+
+    def _save(self, items: list[QuantizedArtifact]) -> None:
+        try:
+            self.file.write_text(
+                json.dumps([i.model_dump(mode="json") for i in items], indent=2),
+                encoding="utf-8",
+            )
+        except OSError as exc:  # pragma: no cover
+            log.warning("could not save quantized registry: %s", exc)
+
+    def add(
+        self,
+        path: str,
+        *,
+        base_model: str | None = None,
+        format_id: str | None = None,
+        method: str | None = None,
+    ) -> QuantizedArtifact:
+        items = [i for i in self._load() if Path(i.path) != Path(path)]
+        artifact = QuantizedArtifact(
+            path=str(Path(path).resolve()),
+            base_model=base_model,
+            format_id=format_id,
+            method=method,
+            size_bytes=dir_size(path),
+        )
+        items.append(artifact)
+        self._save(items)
+        return artifact
+
+    def list(self) -> list[QuantizedArtifact]:
+        """List tracked artifacts, pruning any whose directory is gone."""
+        items = self._load()
+        live = [i for i in items if i.exists]
+        if len(live) != len(items):
+            self._save(live)
+        # refresh sizes cheaply from disk
+        for i in live:
+            i.size_bytes = dir_size(i.path)
+        live.sort(key=lambda i: i.created_at, reverse=True)
+        return live
+
+    def delete(self, path: str, *, remove_files: bool = True) -> int:
+        target = Path(path)
+        freed = dir_size(target)
+        if remove_files and target.exists():
+            shutil.rmtree(target, ignore_errors=True)
+        self._save([i for i in self._load() if Path(i.path) != target])
+        return freed
+
+
+def default_quantized_registry() -> QuantizedRegistry:
+    return QuantizedRegistry()

zeroquantz/core/context.py

@@ -0,0 +1,127 @@
+"""The application context — shared state for one interactive/CLI invocation.
+
+Holds the detected hardware, the backend/runtime registries, the app config, and
+the active :class:`Session`. Command handlers read and mutate this; the CLI and
+TUI both build one and hand it to the dispatcher.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+from zeroquantz.core.session import (
+    Session,
+    SessionRepository,
+    default_session_repository,
+)
+from zeroquantz.hardware.detector import HardwareDetector
+from zeroquantz.quantization.registry import BackendRegistry, default_registry
+from zeroquantz.runtimes.registry import RuntimeRegistry, default_runtime_registry
+from zeroquantz.utils.config import AppConfig, load_config
+
+if TYPE_CHECKING:
+    from zeroquantz.hardware.capabilities import HardwareProfile
+    from zeroquantz.optimization.constraints import OptimizationGoal
+    from zeroquantz.models.metadata import ModelProfile
+
+
+@dataclass
+class AppContext:
+    """Everything a command needs to run."""
+
+    hardware: HardwareProfile
+    session: Session
+    repo: SessionRepository
+    backends: BackendRegistry
+    runtimes: RuntimeRegistry
+    config: AppConfig = field(default_factory=AppConfig)
+
+    # transient, per-invocation artifacts (never persisted)
+    candidates: list = field(default_factory=list)
+    sensitivity: object | None = None
+    last_quant_result: object | None = None
+
+    # ---- convenience accessors ---------------------------------------------
+
+    @property
+    def goal(self) -> OptimizationGoal:
+        return self.session.goal
+
+    @property
+    def model_profile(self) -> ModelProfile | None:
+        return self.session.model_profile
+
+    @property
+    def model_id(self) -> str | None:
+        return self.session.model_id
+
+    def save(self) -> None:
+        """Persist the session if it is a named (non-scratch) one."""
+        if not self.session.name.startswith("scratch-"):
+            self.repo.save(self.session)
+
+    # ---- construction -------------------------------------------------------
+
+    @classmethod
+    def create(
+        cls,
+        *,
+        session_name: str | None = None,
+        continue_last: bool = False,
+        repo: SessionRepository | None = None,
+        config: AppConfig | None = None,
+        detect_hardware: bool = True,
+    ) -> AppContext:
+        repo = repo or default_session_repository()
+        config = config or load_config()
+
+        session = cls._resolve_session(repo, session_name, continue_last)
+
+        # Apply config-level default runtime/objective to a fresh goal.
+        if session.goal.is_empty() and (config.default_runtime or config.objective):
+            from zeroquantz.optimization.constraints import Objective, OptimizationGoal
+
+            session.goal = OptimizationGoal(
+                runtime=config.default_runtime,
+                objective=Objective(config.objective),
+            )
+
+        hardware = (
+            HardwareDetector.detect() if detect_hardware else _empty_hardware()
+        )
+        # Keep a hardware snapshot on the session for display + reports.
+        session.hardware = hardware
+
+        return cls(
+            hardware=hardware,
+            session=session,
+            repo=repo,
+            backends=default_registry(),
+            runtimes=default_runtime_registry(),
+            config=config,
+        )
+
+    @staticmethod
+    def _resolve_session(
+        repo: SessionRepository, session_name: str | None, continue_last: bool
+    ) -> Session:
+        if continue_last:
+            latest = repo.latest()
+            return latest or Session.new("default")
+        if session_name:
+            if repo.exists(session_name):
+                return repo.load(session_name)
+            return Session.new(session_name)
+        # Ephemeral scratch session (not auto-saved).
+        import os
+
+        return Session.new(f"scratch-{os.getpid()}")
+
+
+def _empty_hardware() -> HardwareProfile:
+    from zeroquantz.hardware.capabilities import HardwareProfile, derive_precision_support
+
+    return HardwareProfile(
+        precision_support=derive_precision_support(None, cuda_available=False)
+    )

zeroquantz/core/events.py

@@ -0,0 +1,30 @@
+"""Lightweight progress reporting shared by long-running operations.
+
+Operations accept an optional ``progress(stage: str, fraction: float)`` callback.
+The CLI wires it to a Rich progress bar; the TUI to a widget; tests pass a
+collector or nothing at all.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+ProgressCallback = Callable[[str, float], None]
+
+
+def null_progress(stage: str, fraction: float) -> None:  # noqa: ARG001
+    """A no-op progress sink."""
+
+
+class ProgressCollector:
+    """Records progress updates — handy in tests and for replaying to a UI."""
+
+    def __init__(self) -> None:
+        self.updates: list[tuple[str, float]] = []
+
+    def __call__(self, stage: str, fraction: float) -> None:
+        self.updates.append((stage, fraction))
+
+    @property
+    def last(self) -> tuple[str, float] | None:
+        return self.updates[-1] if self.updates else None

zeroquantz/core/exceptions.py

@@ -0,0 +1,105 @@
+"""ZeroQuantz's exception hierarchy.
+
+Errors are meant to be *actionable*. Every :class:`ZeroQuantzError` carries an
+optional ``detail`` paragraph and a list of concrete ``suggestions`` (commands to
+run, extras to install). The CLI and TUI render these instead of a bare
+traceback, turning "RuntimeError" into "here is exactly what went wrong and what
+to try next".
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+
+class ZeroQuantzError(Exception):
+    """Base class for all expected ZeroQuantz failures.
+
+    Args:
+        message: One-line summary of what failed.
+        detail: Optional longer explanation (may span several lines).
+        suggestions: Concrete next steps, e.g. ``["/recommend", "pip install ..."]``.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        detail: str | None = None,
+        suggestions: Sequence[str] | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.message = message
+        self.detail = detail
+        self.suggestions: list[str] = list(suggestions or [])
+
+    def format(self) -> str:
+        """Render a multi-line, human-readable representation."""
+        lines = [self.message]
+        if self.detail:
+            lines.append("")
+            lines.append(self.detail)
+        if self.suggestions:
+            lines.append("")
+            lines.append("Try:")
+            lines.extend(f"  {s}" for s in self.suggestions)
+        return "\n".join(lines)
+
+
+class DependencyError(ZeroQuantzError):
+    """A required optional dependency is not installed."""
+
+    @classmethod
+    def for_extra(cls, package: str, extra: str, *, purpose: str) -> DependencyError:
+        return cls(
+            f"'{package}' is required to {purpose}, but it is not installed.",
+            detail=(
+                f"This capability lives behind the '{extra}' extra so that the base "
+                "install stays lightweight."
+            ),
+            suggestions=[f"pip install 'zeroquantz[{extra}]'", f"pip install {package}"],
+        )
+
+
+class HardwareError(ZeroQuantzError):
+    """Something about the detected hardware blocks the requested operation."""
+
+
+class UnsupportedConfigurationError(ZeroQuantzError):
+    """A quantization/runtime configuration is not supported on this setup."""
+
+
+class ModelInspectionError(ZeroQuantzError):
+    """Model metadata could not be read."""
+
+
+class ModelLoadError(ZeroQuantzError):
+    """A model's weights could not be loaded."""
+
+
+class BackendUnavailableError(ZeroQuantzError):
+    """A quantization backend was requested but is not available."""
+
+
+class QuantizationError(ZeroQuantzError):
+    """Quantization failed to run."""
+
+
+class RuntimeCompatibilityError(ZeroQuantzError):
+    """A quantized model is incompatible with the requested runtime."""
+
+
+class VerificationError(ZeroQuantzError):
+    """Verification could not be completed."""
+
+
+class SessionError(ZeroQuantzError):
+    """A session could not be loaded, saved, or mutated."""
+
+
+class ConfigError(ZeroQuantzError):
+    """The on-disk configuration file is invalid."""
+
+
+class CommandError(ZeroQuantzError):
+    """An interactive command was malformed or could not be dispatched."""

zeroquantz/core/session.py

@@ -0,0 +1,202 @@
+"""Interactive session state, persistence, and undo/history.
+
+A :class:`Session` bundles the model, hardware snapshot, optimization goal, the
+selected strategy/plan, and measured results, plus an immutable history of
+configuration snapshots. Undo/checkout operate on that history (config/session
+state only — never on multi-hour quantization artifacts). Persistence goes
+through a :class:`SessionRepository` abstraction so the JSON store can be swapped
+later.
+"""
+
+from __future__ import annotations
+
+import json
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+from zeroquantz.benchmark.runner import BenchmarkResult
+from zeroquantz.core.exceptions import SessionError
+from zeroquantz.hardware.capabilities import HardwareProfile
+from zeroquantz.models.metadata import ModelProfile
+from zeroquantz.optimization.constraints import OptimizationGoal
+from zeroquantz.optimization.planner import MixedPrecisionPlan
+from zeroquantz.utils.env import utc_timestamp
+from zeroquantz.utils.paths import paths
+from zeroquantz.verification.report import VerificationReport
+
+
+class SessionSnapshot(BaseModel):
+    """Immutable record of the mutable config state at a point in time."""
+
+    label: str
+    timestamp: str = Field(default_factory=utc_timestamp)
+    model_id: str | None = None
+    goal: OptimizationGoal = Field(default_factory=OptimizationGoal)
+    selected_method: str | None = None
+
+
+class Session(BaseModel):
+    """A named, persistable interactive session."""
+
+    name: str
+    created_at: str = Field(default_factory=utc_timestamp)
+    updated_at: str = Field(default_factory=utc_timestamp)
+
+    model_id: str | None = None
+    goal: OptimizationGoal = Field(default_factory=OptimizationGoal)
+    selected_method: str | None = None
+    selected_format_id: str | None = None
+    model_local_path: str | None = None
+
+    # cached artifacts (recomputed lazily where possible)
+    model_profile: ModelProfile | None = None
+    hardware: HardwareProfile | None = None
+    plan: MixedPrecisionPlan | None = None
+    last_benchmark: BenchmarkResult | None = None
+    last_verification: VerificationReport | None = None
+
+    operation_log: list[str] = Field(default_factory=list)
+    history: list[SessionSnapshot] = Field(default_factory=list)
+    head: int = 0
+
+    # ---- construction -------------------------------------------------------
+
+    @classmethod
+    def new(cls, name: str) -> Session:
+        session = cls(name=name)
+        session.history = [SessionSnapshot(label="baseline")]
+        session.head = 0
+        return session
+
+    # ---- history / undo -----------------------------------------------------
+
+    def record(self, label: str) -> None:
+        """Push a new snapshot of the current config state and log the operation."""
+        # Truncate any redo tail before appending (new branch of history).
+        self.history = self.history[: self.head + 1]
+        self.history.append(
+            SessionSnapshot(
+                label=label,
+                model_id=self.model_id,
+                goal=self.goal,
+                selected_method=self.selected_method,
+            )
+        )
+        self.head = len(self.history) - 1
+        self.operation_log.append(label)
+        self.updated_at = utc_timestamp()
+
+    def undo(self) -> str | None:
+        """Revert to the previous snapshot. Returns the label reverted *from*."""
+        if self.head <= 0:
+            return None
+        reverted_from = self.history[self.head].label
+        self.head -= 1
+        self._apply(self.history[self.head])
+        self.updated_at = utc_timestamp()
+        return reverted_from
+
+    def checkout(self, index: int) -> SessionSnapshot:
+        if not 0 <= index < len(self.history):
+            raise SessionError(
+                f"No history entry {index}.",
+                detail=f"Valid range is 0..{len(self.history) - 1}.",
+                suggestions=["/history"],
+            )
+        self.head = index
+        self._apply(self.history[index])
+        self.updated_at = utc_timestamp()
+        return self.history[index]
+
+    def _apply(self, snapshot: SessionSnapshot) -> None:
+        self.model_id = snapshot.model_id
+        self.goal = snapshot.goal
+        self.selected_method = snapshot.selected_method
+
+    def history_view(self) -> list[tuple[int, str, bool]]:
+        return [(i, s.label, i == self.head) for i, s in enumerate(self.history)]
+
+
+# ---- persistence -----------------------------------------------------------
+
+
+class SessionRepository(ABC):
+    """Storage abstraction for sessions (JSON today, could be SQLite later)."""
+
+    @abstractmethod
+    def save(self, session: Session) -> None: ...
+
+    @abstractmethod
+    def load(self, name: str) -> Session: ...
+
+    @abstractmethod
+    def list(self) -> list[str]: ...
+
+    @abstractmethod
+    def exists(self, name: str) -> bool: ...
+
+    @abstractmethod
+    def delete(self, name: str) -> None: ...
+
+    @abstractmethod
+    def latest(self) -> Session | None: ...
+
+
+class JsonSessionRepository(SessionRepository):
+    """Stores each session as ``<dir>/<name>.json``."""
+
+    def __init__(self, directory: Path | None = None) -> None:
+        self.directory = directory or paths().ensure().sessions
+        self.directory.mkdir(parents=True, exist_ok=True)
+
+    def _path(self, name: str) -> Path:
+        safe = "".join(c for c in name if c.isalnum() or c in "-_.") or "session"
+        return self.directory / f"{safe}.json"
+
+    def save(self, session: Session) -> None:
+        session.updated_at = utc_timestamp()
+        try:
+            self._path(session.name).write_text(
+                json.dumps(session.model_dump(mode="json"), indent=2), encoding="utf-8"
+            )
+        except OSError as exc:
+            raise SessionError(
+                f"Could not save session '{session.name}'.", detail=str(exc)
+            ) from exc
+
+    def load(self, name: str) -> Session:
+        path = self._path(name)
+        if not path.exists():
+            raise SessionError(
+                f"No saved session named '{name}'.",
+                suggestions=["/sessions"],
+            )
+        try:
+            data = json.loads(path.read_text(encoding="utf-8"))
+            return Session.model_validate(data)
+        except (OSError, ValueError) as exc:
+            raise SessionError(
+                f"Session '{name}' is corrupted or unreadable.", detail=str(exc)
+            ) from exc
+
+    def list(self) -> list[str]:
+        return sorted(p.stem for p in self.directory.glob("*.json"))
+
+    def exists(self, name: str) -> bool:
+        return self._path(name).exists()
+
+    def delete(self, name: str) -> None:
+        self._path(name).unlink(missing_ok=True)
+
+    def latest(self) -> Session | None:
+        files = list(self.directory.glob("*.json"))
+        if not files:
+            return None
+        newest = max(files, key=lambda p: p.stat().st_mtime)
+        return self.load(newest.stem)
+
+
+def default_session_repository() -> SessionRepository:
+    return JsonSessionRepository()