load_model_display.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Quiet surfaces for `aquin load model` / `aquin unload`."""
 
from __future__ import annotations
 
import sys
import time
from typing import Any, Callable
 
from aquin.tui.spinners import bullet_mark, spinner_frames, spinner_interval_ms
 
 
def load_model_payload(
    *,
    model: str,
    mode: str,
    daemon: bool,
    ok: bool = True,
    error: str | None = None,
    elapsed_s: float | None = None,
    vram: str | None = None,
) -> dict[str, Any]:
    out: dict[str, Any] = {
        "kind": "load-model",
        "ok": ok,
        "model": model,
        "mode": mode,
        "daemon": daemon,
    }
    if error:
        out["error"] = error
    if elapsed_s is not None:
        out["elapsed_s"] = round(elapsed_s, 1)
    if vram:
        out["vram"] = vram
    return out
 
 
def unload_payload(
    *,
    ok: bool,
    daemon_running: bool,
    stopped: bool,
    message: str,
) -> dict[str, Any]:
    return {
        "kind": "unload",
        "ok": ok,
        "daemonRunning": daemon_running,
        "stopped": stopped,
        "message": message,
    }
 
 
def print_phase(message: str, *, indent: int = 0, settled: bool = True) -> None:
    prefix = "  " * indent
    mark = bullet_mark() if settled else "-"
    print(f"{prefix}{mark} {message}", flush=True)
 
 
def _format_elapsed(seconds: float) -> str:
    s = int(seconds)
    if s < 60:
        return f"{s}s"
    return f"{s // 60}m {s % 60}s"
 
 
class LoadProgress:
    """
    Ollama-style load progress — one headline, indented phases, optional spinner.
    """
 
    def __init__(self, headline: str) -> None:
        self.headline = headline
        self._start = time.monotonic()
        self._last_phase = ""
        self._spinner_idx = 0
        self._frames = spinner_frames("dots")
        self._interval = spinner_interval_ms("dots") / 1000.0
        self._last_spin = 0.0
        print(self.headline, flush=True)
 
    def phase(self, message: str, *, indent: int = 1) -> None:
        self._last_phase = message
        elapsed = _format_elapsed(time.monotonic() - self._start)
        prefix = "  " * indent
        print(f"{prefix}{bullet_mark()} {message} · {elapsed}", flush=True)
 
    def tick(self, message: str | None = None) -> None:
        """Update in-place spinner line when on a TTY (daemon poll loop)."""
        if not sys.stdout.isatty():
            return
        now = time.monotonic()
        if now - self._last_spin < self._interval:
            return
        self._last_spin = now
        self._spinner_idx = (self._spinner_idx + 1) % len(self._frames)
        frame = self._frames[self._spinner_idx]
        text = message or self._last_phase or "loading"
        elapsed = _format_elapsed(now - self._start)
        line = f"  {frame} {text} · {elapsed}"
        sys.stdout.write(f"\r{line}\x1b[K")
        sys.stdout.flush()
 
    def done(self, message: str, *, indent: int = 0) -> None:
        if sys.stdout.isatty() and self._last_phase:
            sys.stdout.write("\r\x1b[K")
            sys.stdout.flush()
        elapsed = _format_elapsed(time.monotonic() - self._start)
        prefix = "  " * indent
        print(f"{prefix}{bullet_mark()} {message} · {elapsed}", flush=True)
 
    def elapsed(self) -> float:
        return time.monotonic() - self._start
 
 
def watch_daemon_load(
    *,
    target_model: str,
    poll: Callable[[], dict | None],
    timeout: float = 1800.0,
    on_tick: LoadProgress | None = None,
) -> dict:
    """
    Poll daemon health until target_model is ready or error/timeout.
    Returns {"ok": True, "model_id": ...} or {"ok": False, "error": ...}.
    """
    deadline = time.time() + timeout
    loading_model = target_model
    while time.time() < deadline:
        h = poll()
        if h is None:
            return {
                "ok": False,
                "error": "background engine stopped responding while loading the model",
            }
        loaded = h.get("model_id")
        model_status = h.get("model_status")
        loading_model = h.get("loading_model_id") or loading_model
        elapsed = h.get("elapsed_s")
        if on_tick:
            if model_status == "loading":
                msg = f"loading {loading_model}"
                if elapsed is not None:
                    msg = f"{msg}"
                on_tick.tick(msg)
            elif model_status == "ready" and loaded == target_model:
                on_tick.tick(f"{target_model} ready")
 
        if loaded == target_model and model_status == "ready":
            return {"ok": True, "model_id": loaded, "vram": h.get("vram")}
        if model_status == "error":
            detail = h.get("last_error") or "unknown load failure"
            failed = h.get("loading_model_id") or h.get("model_id") or target_model
            return {"ok": False, "error": f"{failed}: {detail}"}
        time.sleep(0.4)
 
    return {
        "ok": False,
        "error": f"timed out waiting for background engine to finish loading {loading_model}",
    }