121 lines
3.9 KiB
Python
121 lines
3.9 KiB
Python
"""Ensure the oMLX server is up for a run.
|
|
|
|
oMLX runs as a managed background service (auto-starts on login) and loads models
|
|
on-demand with LRU eviction, so there's usually nothing to do but health-check. For the
|
|
nightly launchd job we still `omlx start` defensively and warm the active model so the
|
|
first agent turn isn't paying cold-load latency.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import time
|
|
|
|
import httpx
|
|
from loguru import logger
|
|
|
|
from config.settings import get_settings
|
|
|
|
_OMLX_CANDIDATES = ["omlx", os.path.expanduser("~/.omlx/bin/omlx")]
|
|
|
|
|
|
def _omlx_bin() -> str | None:
|
|
for cand in _OMLX_CANDIDATES:
|
|
path = shutil.which(cand) if os.path.sep not in cand else (cand if os.path.exists(cand) else None)
|
|
if path:
|
|
return path
|
|
return None
|
|
|
|
|
|
def server_up() -> bool:
|
|
s = get_settings()
|
|
try:
|
|
r = httpx.get(
|
|
f"{s.llm_base_url}/models",
|
|
headers={"Authorization": f"Bearer {s.llm_api_key}"},
|
|
timeout=5,
|
|
)
|
|
return r.status_code == 200
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def start_server() -> None:
|
|
if server_up():
|
|
logger.info("oMLX already up at {}", get_settings().llm_base_url)
|
|
return
|
|
bin_ = _omlx_bin()
|
|
if not bin_:
|
|
logger.warning("omlx CLI not found; cannot auto-start. Start the oMLX app manually.")
|
|
return
|
|
proc = subprocess.run([bin_, "start"], capture_output=True, text=True, timeout=120)
|
|
logger.info("omlx start -> {} {}", proc.returncode, (proc.stdout + proc.stderr).strip()[:300])
|
|
|
|
|
|
def wait_healthy(timeout_s: int = 120) -> bool:
|
|
deadline = time.time() + timeout_s
|
|
while time.time() < deadline:
|
|
if server_up():
|
|
return True
|
|
time.sleep(3)
|
|
return False
|
|
|
|
|
|
def warm_model() -> None:
|
|
"""Fire a tiny completion so the active model is resident before the first agent turn."""
|
|
from llm.client import build_llm
|
|
|
|
try:
|
|
build_llm().invoke("ok")
|
|
logger.info("warmed model")
|
|
except Exception as e:
|
|
logger.warning("warm_model failed (non-fatal): {}", e)
|
|
|
|
|
|
def clear_kv_cache(hot: bool = True) -> None:
|
|
"""Wipe oMLX's SSD prefix/KV cache (and, by default, the in-memory hot cache).
|
|
|
|
oMLX prefix-caches every prompt into an SSD-backed paged KV store so repeated prefixes
|
|
skip prefill. Within a single agent that's a big win, but across personas the prefixes
|
|
diverge, so a finished agent's blocks are dead weight that just keeps eating SSD (and the
|
|
hot cache keeps eating RAM) until the disk fills or the prefill memory guard trips. Call
|
|
this at each agent boundary to keep peak SSD/RAM use to roughly one agent's context.
|
|
|
|
Auth: oMLX's admin endpoints want a session cookie, minted by POSTing the main API key to
|
|
/admin/api/login; httpx.Client keeps that cookie for the follow-up clears. Fully non-fatal
|
|
— a cache-clear hiccup must never abort a run.
|
|
"""
|
|
s = get_settings()
|
|
base = f"http://{s.llm_host}:{s.llm_port}"
|
|
try:
|
|
with httpx.Client(base_url=base, timeout=30) as c:
|
|
c.post(
|
|
"/admin/api/login",
|
|
json={"api_key": s.llm_api_key, "remember": False},
|
|
).raise_for_status()
|
|
ssd = c.post("/admin/api/ssd-cache/clear")
|
|
ssd.raise_for_status()
|
|
msg = f"cleared oMLX SSD KV cache ({ssd.json().get('total_deleted', '?')} blocks)"
|
|
if hot:
|
|
h = c.post("/admin/api/hot-cache/clear")
|
|
h.raise_for_status()
|
|
msg += f", hot cache reclaimed {h.json().get('bytes_reclaimed', 0) / 1e6:.0f}MB"
|
|
logger.info(msg)
|
|
except Exception as e:
|
|
logger.warning("clear_kv_cache failed (non-fatal): {}", e)
|
|
|
|
|
|
def ensure_ready(warm: bool = True) -> bool:
|
|
start_server()
|
|
ok = wait_healthy(120)
|
|
if ok and warm:
|
|
warm_model()
|
|
logger.info("oMLX ready: {}", ok)
|
|
return ok
|
|
|
|
|
|
if __name__ == "__main__":
|
|
ensure_ready()
|