Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

hermes_cli/local_runtime/load_progress.py

@@ -0,0 +1,198 @@
+"""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