199 lines
8.0 KiB
Python
199 lines
8.0 KiB
Python
"""Live model-load progress from the managed llama-server router.
|
|
|
|
llama-server's child processes emit per-tensor load progress
|
|
({stages, current, value}, throttled upstream to ~200ms) which the
|
|
router relays ONLY over its /models/sse stream — GET /models carries
|
|
just the coarse status string. This module owns one lazy background
|
|
watcher on that stream and keeps an in-memory snapshot other code can
|
|
poll cheaply:
|
|
|
|
get_loading_progress() -> {model_id: {"stage", "value", "percent"}}
|
|
|
|
"percent" is a composite across stages so a bar doesn't sprint 0->100
|
|
once per stage: the text model dominates load time (its weights dwarf
|
|
the mmproj/spec extras), so it gets the lion's share of the range and
|
|
the extras split the remainder.
|
|
|
|
The watcher starts on first call, reconnects with backoff (the router
|
|
bounces on model download/eject), and never raises into callers — no
|
|
router, no state file, or no SSE support (older engines) all read as
|
|
"nothing loading". Safe from any process on the machine: the endpoint
|
|
comes from the supervisor's machine-scoped state file.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import threading
|
|
import time
|
|
import urllib.request
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_TEXT_STAGE_SHARE = 0.85 # composite range share for the text model
|
|
_RECONNECT_DELAY_S = 3.0
|
|
_STALE_ENTRY_TTL_S = 120.0 # a loading entry with no events this long is dead
|
|
|
|
_lock = threading.Lock()
|
|
_watcher: threading.Thread | None = None
|
|
_snapshot: dict[str, dict] = {}
|
|
|
|
|
|
def _composite_percent(stages: list[str], current: str, value: float) -> int:
|
|
"""Map (stage, in-stage value) onto one 0-100 range, text-heavy."""
|
|
if not stages or current not in stages or len(stages) == 1:
|
|
return max(0, min(100, round(value * 100)))
|
|
extras = [s for s in stages if s != "text_model"]
|
|
extra_share = (1.0 - _TEXT_STAGE_SHARE) / len(extras) if extras else 0.0
|
|
offset = 0.0
|
|
for stage in stages:
|
|
share = _TEXT_STAGE_SHARE if stage == "text_model" else extra_share
|
|
if stage == current:
|
|
return max(0, min(100, round((offset + share * value) * 100)))
|
|
offset += share
|
|
return max(0, min(100, round(value * 100)))
|
|
|
|
|
|
def _endpoint() -> "tuple[str, str] | None":
|
|
"""(base_root, api_key) of the managed router, or None.
|
|
|
|
Resolved through the endpoint module's ownership-guarded reader, not
|
|
a raw state-file read: on the shared stable port, a foreign install's
|
|
server answers /health for anyone, and a raw read would attach this
|
|
watcher to someone else's SSE stream (or spin on 401s against it).
|
|
The guard's dead-pid check is the ownership proof."""
|
|
try:
|
|
from hermes_cli.local_runtime.endpoint import _state_endpoint
|
|
|
|
state = _state_endpoint()
|
|
if state is None:
|
|
return None
|
|
base = str(state.get("base_url", "")).rsplit("/v1", 1)[0]
|
|
return (base, str(state.get("api_key", ""))) if base else None
|
|
except Exception: # noqa: BLE001
|
|
return None
|
|
|
|
|
|
def _apply_event(model: str, event: str, data: dict) -> None:
|
|
with _lock:
|
|
status = str(data.get("status", ""))
|
|
if event in ("status_change", "model_status") and status == "loading":
|
|
progress = data.get("progress") or {}
|
|
stages = [str(s) for s in (progress.get("stages") or [])]
|
|
current = str(progress.get("current", ""))
|
|
value = progress.get("value")
|
|
entry = _snapshot.setdefault(model, {"stage": "", "value": 0.0,
|
|
"percent": 0, "ts": 0.0})
|
|
entry["ts"] = time.monotonic()
|
|
if current and isinstance(value, (int, float)):
|
|
entry["stage"] = current
|
|
entry["value"] = float(value)
|
|
entry["percent"] = _composite_percent(stages, current, float(value))
|
|
elif event in ("status_change", "model_status", "model_remove"):
|
|
# Any terminal status (loaded/unloaded/failed) ends the load.
|
|
if status != "loading":
|
|
_snapshot.pop(model, None)
|
|
|
|
|
|
def _watch() -> None:
|
|
while True:
|
|
endpoint = _endpoint()
|
|
if endpoint is None:
|
|
with _lock:
|
|
_snapshot.clear()
|
|
time.sleep(_RECONNECT_DELAY_S)
|
|
continue
|
|
base, key = endpoint
|
|
try:
|
|
req = urllib.request.Request(
|
|
f"{base}/models/sse",
|
|
headers={"Authorization": f"Bearer {key}",
|
|
"Accept": "text/event-stream"})
|
|
with urllib.request.urlopen(req, timeout=60) as r:
|
|
buf = b""
|
|
while True:
|
|
chunk = r.read1(4096) if hasattr(r, "read1") else r.read(4096)
|
|
if not chunk:
|
|
break
|
|
buf += chunk
|
|
while b"\n" in buf:
|
|
line, buf = buf.split(b"\n", 1)
|
|
text = line.decode("utf-8", "replace").strip()
|
|
if not text.startswith("data:"):
|
|
continue
|
|
try:
|
|
msg = json.loads(text[5:].strip())
|
|
_apply_event(str(msg.get("model", "")),
|
|
str(msg.get("event", "")),
|
|
msg.get("data") or {})
|
|
except (json.JSONDecodeError, TypeError):
|
|
continue
|
|
except Exception as exc: # noqa: BLE001 — watcher must never die loud
|
|
logger.debug("load-progress SSE reconnecting: %s", exc)
|
|
# Stream ended (router bounce, timeout, error): loading entries from
|
|
# the dead connection are unverifiable — drop rather than freeze.
|
|
with _lock:
|
|
_snapshot.clear()
|
|
time.sleep(_RECONNECT_DELAY_S)
|
|
|
|
|
|
def _ensure_watcher() -> None:
|
|
global _watcher
|
|
with _lock:
|
|
if _watcher is None or not _watcher.is_alive():
|
|
_watcher = threading.Thread(target=_watch, daemon=True,
|
|
name="llamacpp-load-progress")
|
|
_watcher.start()
|
|
|
|
|
|
def get_loading_progress() -> dict[str, dict]:
|
|
"""{model_id: {"stage", "value", "percent"}} for models loading right
|
|
now. Empty when nothing is loading (or nothing is knowable)."""
|
|
_ensure_watcher()
|
|
now = time.monotonic()
|
|
with _lock:
|
|
return {m: {"stage": e["stage"], "value": e["value"],
|
|
"percent": e["percent"]}
|
|
for m, e in _snapshot.items()
|
|
if now - e["ts"] < _STALE_ENTRY_TTL_S}
|
|
|
|
|
|
def get_prefill_progress(model: str) -> "dict | None":
|
|
"""{"processed": tokens} while the managed server is prompt-processing
|
|
for ``model``, or None (idle, decoding, unreachable, or foreign server).
|
|
|
|
llama-server's /slots reports ``n_prompt_tokens_processed`` climbing in
|
|
real time during prefill, but exposes no total — callers supply their
|
|
own denominator (the request's estimated token count). Busiest
|
|
processing slot wins when several are active: a parallel small request
|
|
(title generation) freezes its counter during decode while a live
|
|
prefill keeps climbing past it. One authenticated HTTP call per poll;
|
|
every failure reads as "no prefill" — this is garnish, never load-
|
|
bearing.
|
|
"""
|
|
ep = _endpoint()
|
|
if ep is None:
|
|
return None
|
|
base, key = ep
|
|
try:
|
|
from urllib.parse import quote
|
|
|
|
req = urllib.request.Request(
|
|
f"{base}/slots?model={quote(model)}",
|
|
headers={"Authorization": f"Bearer {key}"})
|
|
with urllib.request.urlopen(req, timeout=2) as r:
|
|
slots = json.loads(r.read())
|
|
except Exception: # noqa: BLE001
|
|
return None
|
|
best = 0
|
|
for slot in slots if isinstance(slots, list) else []:
|
|
if not slot.get("is_processing"):
|
|
continue
|
|
try:
|
|
processed = int(slot.get("n_prompt_tokens_processed") or 0)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
best = max(best, processed)
|
|
return {"processed": best} if best > 0 else None
|