model_runtime.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""
Unified model lifecycle — one resident model in VRAM, Ollama-style.
 
Owns phase transitions (idle → loading → ready → unloading), coordinates LLM
loaders, and is the single source of truth for daemon + CLI load/unload.
"""
from __future__ import annotations
 
import threading
import time
from dataclasses import dataclass
from enum import Enum
from typing import Any, Callable, Literal
 
ProgressFn = Callable[[str, float], None]
 
_lock = threading.Lock()
# Serialize load_weights / unload_weights so two builds never share the process.
_load_serial = threading.RLock()
_phase = "idle"
_target_id: str | None = None
_resident_id: str | None = None
_kind: Literal["llm"] | None = None
_error: str | None = None
_started_at: float | None = None
_loaded_at: float | None = None
_load_message: str | None = None
# Bumped on unload/cancel — in-flight builds check this and discard weights.
_load_epoch = 0
_load_epoch_at_start = 0
 
 
class ModelPhase(str, Enum):
    IDLE = "idle"
    LOADING = "loading"
    READY = "ready"
    UNLOADING = "unloading"
    ERROR = "error"
 
 
@dataclass(frozen=True)
class RuntimeSnapshot:
    phase: ModelPhase
    model_id: str | None
    target_id: str | None
    kind: Literal["llm"] | None
    error: str | None
    elapsed_s: float
    loaded_for_s: float | None
 
 
def snapshot() -> RuntimeSnapshot:
    """Thread-safe view of the in-process lifecycle state."""
    with _lock:
        now = time.monotonic()
        elapsed = (now - _started_at) if _started_at is not None else 0.0
        loaded_for = (now - _loaded_at) if _loaded_at is not None else None
        return RuntimeSnapshot(
            phase=ModelPhase(_phase),
            model_id=_resident_id,
            target_id=_target_id if _phase == "loading" else None,
            kind=_kind,
            error=_error,
            elapsed_s=elapsed,
            loaded_for_s=loaded_for,
        )
 
 
def _set_phase(
    phase: str,
    *,
    model_id: str | None = None,
    target_id: str | None = None,
    kind: Literal["llm"] | None = None,
    error: str | None = None,
) -> None:
    global _phase, _resident_id, _target_id, _kind, _error, _started_at, _loaded_at
    with _lock:
        _phase = phase
        if model_id is not None:
            _resident_id = model_id
        if target_id is not None:
            _target_id = target_id
        if kind is not None:
            _kind = kind
        _error = error
        if phase == "loading":
            _started_at = time.monotonic()
            _loaded_at = None
        elif phase == "ready":
            _loaded_at = time.monotonic()
            _error = None
        elif phase in ("idle", "unloading"):
            if phase == "idle":
                _resident_id = None
                _target_id = None
                _kind = None
                _loaded_at = None
            _error = None if phase == "idle" else _error
 
 
def reset_state() -> None:
    """Clear lifecycle markers (tests / daemon shutdown)."""
    with _lock:
        global _phase, _target_id, _resident_id, _kind, _error, _started_at, _loaded_at, _load_message
        global _load_epoch, _load_epoch_at_start
        _phase = "idle"
        _target_id = None
        _resident_id = None
        _kind = None
        _error = None
        _started_at = None
        _loaded_at = None
        _load_message = None
        _load_epoch = 0
        _load_epoch_at_start = 0
 
 
def resolve_kind(model_id: str) -> tuple[str, Literal["llm"]]:
    """Resolve slug + model family for any supported id."""
    from aquin.compute.model_loader import resolve_model_id
 
    return resolve_model_id(model_id), "llm"
 
 
def resident_from_cache() -> tuple[str | None, Literal["llm"] | None]:
    """What is actually resident in this process's VRAM caches."""
    try:
        from aquin.compute.model_loader import get_loaded_llm_id
 
        llm = get_loaded_llm_id()
        if llm:
            return llm, "llm"
    except Exception:
        pass
    return None, None
 
 
def sync_from_cache() -> RuntimeSnapshot:
    """Reconcile lifecycle state with actual VRAM caches (health checks)."""
    global _phase, _resident_id, _target_id, _kind, _error, _loaded_at
    resident, kind = resident_from_cache()
    with _lock:
        if _phase == "loading":
            pass
        elif resident:
            _resident_id = resident
            _kind = kind
            if _phase != "error":
                _phase = "ready"
            _error = None
        elif _phase == "ready":
            _phase = "idle"
            _resident_id = None
            _target_id = None
            _kind = None
            _loaded_at = None
    return snapshot()
 
 
def begin_load(model_id: str) -> tuple[str, Literal["llm"]]:
    global _load_message, _load_epoch_at_start
    slug, kind = resolve_kind(model_id)
    _set_phase("loading", target_id=slug, kind=kind, error=None)
    with _lock:
        _load_message = f"resolving {slug}"
        _load_epoch_at_start = _load_epoch
    return slug, kind
 
 
def finish_load(slug: str, kind: Literal["llm"]) -> None:
    global _load_message
    _set_phase("ready", model_id=slug, kind=kind, error=None)
    with _lock:
        _load_message = None
 
 
def fail_load(model_id: str, error: str) -> None:
    global _load_message
    _set_phase("error", model_id=model_id, error=error)
    with _lock:
        _load_message = None
 
 
def begin_unload() -> None:
    snap = snapshot()
    _set_phase("unloading", model_id=snap.model_id, kind=snap.kind)
 
 
def finish_unload() -> None:
    global _load_message
    _set_phase("idle")
    with _lock:
        _load_message = None
 
 
def request_load_cancel() -> None:
    """Mark any in-flight build as cancelled (checked after from_pretrained)."""
    global _load_epoch
    with _lock:
        _load_epoch += 1
 
 
def is_load_cancelled() -> bool:
    """True if unload/cancel happened after this load began."""
    with _lock:
        return _load_epoch != _load_epoch_at_start
 
 
def _emit(progress: ProgressFn | None, message: str) -> None:
    global _load_message
    with _lock:
        _load_message = message
    if progress:
        snap = snapshot()
        progress(message, snap.elapsed_s)
 
 
def load_weights(
    model_id: str,
    *,
    progress: ProgressFn | None = None,
) -> str:
    """
    Load model_id into this process. Returns resolved slug. Updates lifecycle state.
 
    Process-wide serialized — concurrent builds on MPS exhaust unified memory.
    """
    with _load_serial:
        slug, kind = begin_load(model_id)
        try:
            _emit(progress, f"resolving {slug}")
            from aquin.compute.model_loader import clear_sae_cache, load_model
 
            clear_sae_cache()
            if is_load_cancelled():
                raise RuntimeError(f"Model load cancelled ({slug})")
            _emit(progress, f"loading model weights")
            load_model(slug)
 
            if is_load_cancelled():
                from aquin.compute.model_loader import clear_llm_models, clear_sae_cache
 
                clear_llm_models()
                clear_sae_cache()
                finish_unload()
                raise RuntimeError(f"Model load cancelled ({slug})")
 
            from aquin.compute.model_loader import _save_active_model
 
            _save_active_model(slug)
            finish_load(slug, kind)
            _emit(progress, f"{slug} ready")
            return slug
        except Exception as exc:
            msg = str(exc) or exc.__class__.__name__
            if "cancelled" in msg.lower():
                finish_unload()
            else:
                fail_load(slug, msg)
            raise
 
 
def unload_weights(*, progress: ProgressFn | None = None, clear_active: bool = True) -> None:
    """Drop all resident weights from VRAM in this process."""
    request_load_cancel()
    with _load_serial:
        begin_unload()
        try:
            _emit(progress, "releasing VRAM")
            from aquin.compute.model_loader import clear_active_model_file, clear_llm_models, clear_sae_cache
 
            clear_llm_models()
            clear_sae_cache()
            if clear_active:
                clear_active_model_file()
        finally:
            finish_unload()
 
 
def release_foreign_daemon() -> bool:
    """
    Optionally unload a background engine so this process can claim VRAM.
 
    Default is **never** — the desktop app keeps a resident model in the daemon.
    Headless CLI tools that try load_model in a short-lived process used to call
    model_daemon.unload() here, which looked like the model “randomly disappearing”.
 
    Opt-in only: set AQUIN_CLAIM_DAEMON_VRAM=1 when you intentionally want an
    exclusive in-process load that frees the background engine first.
    """
    import os
 
    if os.environ.get("AQUIN_DAEMON") == "1":
        return False
    claim = (os.environ.get("AQUIN_CLAIM_DAEMON_VRAM") or "").strip().lower()
    if claim not in ("1", "true", "yes", "on"):
        return False
    try:
        from aquin.engine import model_daemon
 
        if model_daemon.is_running():
            model_daemon.unload()
            return True
    except Exception:
        pass
    return False
 
 
def vram_line() -> str | None:
    """Short VRAM summary for status output."""
    try:
        from aquin.compute.vram_guard import accelerator_vram_gib
        from aquin.compute.device import resolve_compute_device
 
        info = accelerator_vram_gib()
        dev = resolve_compute_device()
        if info:
            free_g, total_g = info
            return f"{free_g:.1f} / {total_g:.1f} GiB free ({dev})"
        if dev == "mps":
            return "Metal unified memory"
    except Exception:
        pass
    return None
 
 
def health_payload(*, daemon: bool, port: int) -> dict[str, Any]:
    """JSON health block for the local engine server."""
    sync_from_cache()
    snap = snapshot()
    resident, _ = resident_from_cache()
    with _lock:
        load_message = _load_message
    return {
        "status": "ok",
        "port": port,
        "daemon": daemon,
        "model_id": resident,
        "model_status": snap.phase.value,
        "loading_model_id": snap.target_id if snap.phase == ModelPhase.LOADING else None,
        "load_message": load_message if snap.phase == ModelPhase.LOADING else None,
        "model_kind": snap.kind,
        "last_error": snap.error,
        "elapsed_s": round(snap.elapsed_s, 1) if snap.phase == ModelPhase.LOADING else None,
        "loaded_for_s": round(snap.loaded_for_s, 1) if snap.loaded_for_s is not None else None,
        "vram": vram_line(),
    }