model_daemon.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""
Client + lifecycle for the persistent model daemon (aquin.engine.daemon_main).
 
The daemon keeps one model resident in VRAM across separate `aquin` invocations.
This module is the thin client the CLI uses to start it, check it, switch models,
warm SAEs, dispatch tools, and stop it. Everything here is best-effort: if the
daemon is unavailable, callers fall back to loading in-process.
"""
from __future__ import annotations
 
import json
import os
import subprocess
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any
 
from aquin.engine.engine_info import (
    DEFAULT_ENGINE_HOST,
    resolve_engine_port,
    set_daemon_state,
)
 
_HOST = DEFAULT_ENGINE_HOST
_LOG_PATH = Path.home() / ".aquin" / "daemon.log"
 
# Generous ceiling for a large model load / long tool run.
DEFAULT_TIMEOUT = 1800
FAST_TIMEOUT = 0.2
 
 
def _port() -> int:
    return resolve_engine_port()
 
 
def _base() -> str:
    return f"http://{_HOST}:{_port()}"
 
 
def _url(path: str) -> str:
    return f"{_base()}{path}"
 
 
def _get(path: str, timeout: float = 1.5) -> tuple[int | None, Any]:
    try:
        with urllib.request.urlopen(_url(path), timeout=timeout) as resp:
            raw = resp.read().decode() or "{}"
            return resp.status, json.loads(raw)
    except Exception:
        return None, None
 
 
def _post(path: str, body: dict | None, timeout: float = DEFAULT_TIMEOUT) -> tuple[int | None, Any]:
    data = json.dumps(body or {}).encode()
    req = urllib.request.Request(
        _url(path), data=data, headers={"Content-Type": "application/json"}, method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            raw = resp.read().decode() or "{}"
            return resp.status, json.loads(raw)
    except urllib.error.HTTPError as exc:
        try:
            return exc.code, json.loads(exc.read().decode() or "{}")
        except Exception:
            return exc.code, None
    except Exception:
        return None, None
 
 
def health(timeout: float = 1.5) -> dict | None:
    """Return the daemon health payload, or None if no daemon is answering."""
    status, data = _get("/health", timeout=timeout)
    if status == 200 and isinstance(data, dict) and data.get("daemon"):
        return data
    return None
 
 
def health_fast() -> dict | None:
    """Quick probe for status screens — fails fast when daemon is down."""
    return health(timeout=FAST_TIMEOUT)
 
 
def is_running() -> bool:
    return health() is not None
 
 
def loaded_model() -> str | None:
    h = health()
    return h.get("model_id") if h else None
 
 
def _port_occupied() -> bool:
    """Something is listening on the port (may be a non-daemon chat server)."""
    status, _ = _get("/health", timeout=0.6)
    return status is not None
 
 
def _write_state(pid: int) -> None:
    try:
        set_daemon_state(
            status="running",
            pid=pid,
            port=_port(),
            started_at=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            model_id=None,
        )
    except Exception:
        pass
 
 
def _clear_state() -> None:
    try:
        set_daemon_state(status="stopped", pid=None, started_at=None, model_id=None)
    except Exception:
        pass
 
 
def _spawn() -> bool:
    """Spawn the daemon as a detached background process. Returns spawn success."""
    try:
        _LOG_PATH.parent.mkdir(parents=True, exist_ok=True)
        logf = open(_LOG_PATH, "ab")
    except Exception:
        logf = subprocess.DEVNULL  # type: ignore[assignment]
 
    creationflags = 0
    popen_kwargs: dict[str, Any] = {}
    if os.name == "nt":
        # DETACHED_PROCESS (0x00000008) + new process group so it outlives this shell.
        creationflags = subprocess.CREATE_NEW_PROCESS_GROUP | 0x00000008
    else:
        popen_kwargs["start_new_session"] = True
 
    port = _port()
    try:
        proc = subprocess.Popen(
            [sys.executable, "-m", "aquin.engine.daemon_main", "--port", str(port)],
            stdout=logf,
            stderr=logf,
            stdin=subprocess.DEVNULL,
            creationflags=creationflags,
            close_fds=True,
            **popen_kwargs,
        )
    except Exception:
        return False
    _write_state(proc.pid)
    return True
 
 
def ensure_running(wait_seconds: float = 40.0) -> bool:
    """
    Make sure the daemon HTTP server is up. Returns True if it is (or becomes)
    reachable. Does NOT load a model; use switch_model for that.
    """
    if is_running():
        return True
    # Another server (e.g. `aquin chat`) may hold the port; do not fight it.
    if _port_occupied():
        return False
    if not _spawn():
        return False
    deadline = time.time() + wait_seconds
    while time.time() < deadline:
        if is_running():
            return True
        time.sleep(0.5)
    return False
 
 
def switch_model(
    model_id: str,
    timeout: float = DEFAULT_TIMEOUT,
    *,
    on_tick: Any = None,
) -> dict | None:
    """Load model_id in the daemon, evicting any previously resident model."""
    status, data = _post("/model/switch", {"model_id": model_id}, timeout=10.0)
    if status != 200 or not isinstance(data, dict):
        return {
            "ok": False,
            "error": "background engine did not respond to model switch request",
        }
    if not data.get("ok"):
        return data
    if data.get("status") == "ready":
        h = health(timeout=2.0) or {}
        return {"ok": True, "model_id": data.get("model_id") or model_id, "vram": h.get("vram")}
 
    from aquin.load_model_display import watch_daemon_load
 
    target_model = data.get("model_id") or model_id
    return watch_daemon_load(
        target_model=target_model,
        poll=lambda: health(timeout=2.0),
        timeout=timeout,
        on_tick=on_tick,
    )
 
 
def warm_sae(model_id: str, layer: int, timeout: float = 600) -> dict | None:
    status, data = _post("/sae/load", {"model_id": model_id, "layer": layer}, timeout=timeout)
    if status == 200 and isinstance(data, dict):
        return data
    return None
 
 
def dispatch(name: str, args: dict, ctx: dict, timeout: float = DEFAULT_TIMEOUT) -> dict | None:
    """Run a tool in the daemon. Returns {"ok": bool, "result"|"error": ...} or None."""
    status, data = _post("/dispatch", {"name": name, "args": args, "ctx": ctx}, timeout=timeout)
    if status == 200 and isinstance(data, dict):
        return data
    return None
 
 
def unload() -> bool:
    """Free the daemon's VRAM (keeps the daemon process alive)."""
    status, _ = _post("/unload", {}, timeout=120)
    return status == 200
 
 
def prompt(
    text: str,
    *,
    model_id: str | None = None,
    max_new_tokens: int = 200,
    temperature: float = 0.7,
    timeout: float = DEFAULT_TIMEOUT,
) -> dict | None:
    """Generate a completion from the daemon's resident model."""
    body: dict[str, Any] = {
        "prompt": text,
        "max_new_tokens": max_new_tokens,
        "temperature": temperature,
    }
    if model_id:
        body["model_id"] = model_id
    status, data = _post("/prompt", body, timeout=timeout)
    if status == 200 and isinstance(data, dict):
        return data
    return None
 
 
def stop() -> bool:
    """Ask the daemon to shut down and clear local state."""
    status, _ = _post("/shutdown", {}, timeout=10)
    _clear_state()
    return status == 200