Teknium
436 lines
18 KiB
Python
436 lines
18 KiB
Python
"""Supervision of one llama-server in router mode.
|
|
|
|
The router process is ours (restart with backoff on crash); router children are its problem — child
|
|
failures surface via GET /models exit_code, never auto-retried here.
|
|
|
|
Readiness rules (each learned the hard way on real hardware): - health-200 is NOT readiness; every
|
|
readiness claim requires a touch generation (temp-0, expected token, generous budget,
|
|
reasoning_content scanned). - Always dial 127.0.0.1 — resolving localhost adds ~2s per request on
|
|
Windows via IPv6 fallback.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import secrets
|
|
import socket
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
from hermes_cli.local_runtime.binaries import server_binary, runtimes_root
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
TOUCH_PROMPT = "Reply with exactly one word: the capital of France."
|
|
TOUCH_EXPECT = "paris"
|
|
_RESTART_BACKOFF_S = (1, 5, 15, 60)
|
|
|
|
|
|
def state_path() -> Path:
|
|
"""Endpoint state for other Hermes processes (provider resolution reads
|
|
this to route llamacpp-alias requests at the managed server)."""
|
|
return runtimes_root() / "server.json"
|
|
|
|
|
|
def _free_port() -> int:
|
|
with socket.socket() as s:
|
|
s.bind(("127.0.0.1", 0))
|
|
return s.getsockname()[1]
|
|
|
|
|
|
# Default port for the managed server, chosen once and reused across
|
|
# restarts. Sessions persist the resolved base_url; an ephemeral port
|
|
# would strand every resumed session on a dead endpoint after each
|
|
# restart. Deliberately NOT 8080 so we never collide with a user's own
|
|
# llama-server/Ollama-adjacent stack.
|
|
_DEFAULT_PORT = 18434
|
|
|
|
|
|
def _stable_port() -> int:
|
|
"""The stable default port, falling back to an ephemeral one only when something else already
|
|
listens there (and it isn't a leftover managed server, which stop() would have cleaned up).
|
|
"""
|
|
try:
|
|
with socket.socket() as s:
|
|
s.bind(("127.0.0.1", _DEFAULT_PORT))
|
|
return _DEFAULT_PORT
|
|
except OSError:
|
|
logger.warning(
|
|
"port %d busy; managed llama-server falling back to an ephemeral "
|
|
"port — existing sessions may need a model re-pick", _DEFAULT_PORT)
|
|
return _free_port()
|
|
|
|
|
|
def _stable_api_key() -> str:
|
|
"""One key for the life of the install, persisted beside the runtimes.
|
|
|
|
Endpoint identity must survive restarts as a UNIT — sessions persist the resolved base_url +
|
|
api_key, so a per-boot key strands every resumed session on HTTP 401 exactly the way a per-boot
|
|
port would strand them on connection errors.
|
|
"""
|
|
key_path = runtimes_root() / ".api_key"
|
|
try:
|
|
existing = key_path.read_text(encoding="utf-8").strip()
|
|
if len(existing) >= 16:
|
|
return existing
|
|
except OSError:
|
|
pass
|
|
key = secrets.token_urlsafe(24)
|
|
try:
|
|
key_path.parent.mkdir(parents=True, exist_ok=True)
|
|
key_path.write_text(key, encoding="utf-8")
|
|
except OSError as exc:
|
|
logger.warning("could not persist api key (%s); sessions will need "
|
|
"a re-pick after restart", exc)
|
|
return key
|
|
|
|
|
|
class LlamaServerSupervisor:
|
|
"""Own one llama-server router process for the life of a Hermes session."""
|
|
|
|
def __init__(self, install_dir: Path, models_dir: Path, *,
|
|
models_max: int = 4, port: int | None = None,
|
|
extra_args: list[str] | None = None,
|
|
log_path: Path | None = None,
|
|
preset_path: Path | None = None):
|
|
self.install_dir = Path(install_dir)
|
|
self.models_dir = Path(models_dir)
|
|
self.models_max = models_max
|
|
self.port = port or _stable_port()
|
|
self.api_key = _stable_api_key()
|
|
self.extra_args = list(extra_args or [])
|
|
self.log_path = log_path or (self.models_dir.parent / "logs" / "llama-server.log")
|
|
self.preset_path = preset_path
|
|
self.proc: subprocess.Popen | None = None
|
|
self.primary_model: str | None = None
|
|
self._restarts = 0
|
|
self._stopping = False
|
|
self._watchdog: threading.Thread | None = None
|
|
self._log_handle = None
|
|
self._idle_since: dict[str, float] = {}
|
|
|
|
# ── endpoints ────────────────────────────────────────────
|
|
|
|
@property
|
|
def base_url(self) -> str:
|
|
return f"http://127.0.0.1:{self.port}/v1"
|
|
|
|
def _url(self, route: str) -> str:
|
|
return f"http://127.0.0.1:{self.port}{route}"
|
|
|
|
def _open(self, route: str, body: dict | None = None, timeout_s: int = 30,
|
|
*, json_type: bool = True):
|
|
headers = {"Authorization": f"Bearer {self.api_key}"}
|
|
if json_type:
|
|
headers["Content-Type"] = "application/json"
|
|
req = urllib.request.Request(
|
|
self._url(route),
|
|
data=json.dumps(body).encode() if body is not None else None,
|
|
headers=headers,
|
|
)
|
|
return urllib.request.urlopen(req, timeout=timeout_s)
|
|
|
|
def _request(self, route: str, body: dict | None = None, timeout_s: int = 30) -> dict:
|
|
with self._open(route, body, timeout_s) as r:
|
|
raw = r.read()
|
|
return json.loads(raw) if raw else {}
|
|
|
|
# ── lifecycle ────────────────────────────────────────────
|
|
|
|
def _spawn(self) -> None:
|
|
exe = server_binary(self.install_dir)
|
|
cmd = [
|
|
str(exe),
|
|
"--host", "127.0.0.1",
|
|
"--port", str(self.port),
|
|
"--api-key", self.api_key,
|
|
"--models-dir", str(self.models_dir),
|
|
"--models-max", str(self.models_max),
|
|
# The residency contract at the layer that sees every message:
|
|
# a chat request to a staged-but-unloaded model loads it (slow
|
|
# first token) instead of failing with 'model not found' —
|
|
# without this flag, chat after an eject is a bare 400/404.
|
|
"--models-autoload",
|
|
"--metrics", # opt-in flag; supervisor telemetry needs it
|
|
"--slots", # /slots endpoint is also opt-in; is_idle reads it
|
|
"--no-webui",
|
|
"--jinja",
|
|
# Direct I/O on model load: bypasses the page cache, so a
|
|
# multi-GB load doesn't evict half the OS cache — measured
|
|
# faster loads on NVMe, and our router bounces (download/
|
|
# delete/activate) reload models often enough to care.
|
|
"-dio",
|
|
]
|
|
if self.preset_path and self.preset_path.exists():
|
|
cmd += ["--models-preset", str(self.preset_path)]
|
|
cmd += self.extra_args
|
|
self.log_path.parent.mkdir(parents=True, exist_ok=True)
|
|
if self._log_handle is not None:
|
|
# The crash-restart loop calls _spawn repeatedly; without
|
|
# closing the prior handle each restart leaks one fd.
|
|
try:
|
|
self._log_handle.close()
|
|
except Exception: # noqa: BLE001 — best-effort
|
|
pass
|
|
self._log_handle = open(self.log_path, "a", encoding="utf-8", errors="replace")
|
|
self._log_handle.write(f"\n# spawn: {cmd}\n")
|
|
self._log_handle.flush()
|
|
# list-args, never a shell: spaced paths (user homes) must survive.
|
|
self.proc = subprocess.Popen(cmd, stdout=self._log_handle,
|
|
stderr=subprocess.STDOUT, cwd=str(exe.parent))
|
|
logger.info("llama-server router spawned pid=%s port=%s", self.proc.pid, self.port)
|
|
# State goes down at SPAWN, not after health: endpoint resolution
|
|
# treats a live-pid-but-not-yet-healthy server as "starting" rather
|
|
# than "unconfigured", so a readiness probe racing the boot doesn't
|
|
# throw the app back to onboarding (observed on first restart test).
|
|
self._write_state()
|
|
|
|
def start(self, timeout_s: int = 120) -> None:
|
|
self._stopping = False
|
|
self._spawn()
|
|
self._wait_health(timeout_s)
|
|
self._write_state()
|
|
self._watchdog = threading.Thread(target=self._watch, daemon=True,
|
|
name="llamacpp-supervisor")
|
|
self._watchdog.start()
|
|
|
|
def _write_state(self) -> None:
|
|
path = state_path()
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps({
|
|
"base_url": self.base_url,
|
|
"api_key": self.api_key,
|
|
"pid": self.proc.pid if self.proc else None,
|
|
}), encoding="utf-8")
|
|
|
|
def _wait_health(self, timeout_s: int) -> None:
|
|
deadline = time.monotonic() + timeout_s
|
|
while time.monotonic() < deadline:
|
|
if self.proc and self.proc.poll() is not None:
|
|
raise RuntimeError(
|
|
f"llama-server exited rc={self.proc.returncode} during startup "
|
|
f"(log: {self.log_path})")
|
|
try:
|
|
with urllib.request.urlopen(self._url("/health"), timeout=3) as r:
|
|
if r.status == 200:
|
|
return
|
|
except (urllib.error.URLError, OSError, TimeoutError):
|
|
pass
|
|
time.sleep(1)
|
|
raise TimeoutError(f"llama-server not healthy after {timeout_s}s (log: {self.log_path})")
|
|
|
|
def _watch(self) -> None:
|
|
"""Restart the router (not its children) on crash, with backoff."""
|
|
while not self._stopping:
|
|
proc = self.proc
|
|
if proc is None:
|
|
return
|
|
rc = proc.poll()
|
|
if rc is None:
|
|
time.sleep(2)
|
|
continue
|
|
if self._stopping:
|
|
return
|
|
backoff = _RESTART_BACKOFF_S[min(self._restarts, len(_RESTART_BACKOFF_S) - 1)]
|
|
logger.warning("llama-server exited rc=%s; restart #%s in %ss",
|
|
rc, self._restarts + 1, backoff)
|
|
time.sleep(backoff)
|
|
self._restarts += 1
|
|
try:
|
|
self._reap_orphaned_children()
|
|
self._spawn()
|
|
self._wait_health(120)
|
|
if self.primary_model:
|
|
self.ensure_model_ready(self.primary_model)
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.error("llama-server restart failed: %s", exc)
|
|
|
|
def stop(self) -> None:
|
|
self._stopping = True
|
|
state_path().unlink(missing_ok=True)
|
|
if self.proc and self.proc.poll() is None:
|
|
self._terminate_tree(self.proc)
|
|
if self._log_handle:
|
|
self._log_handle.close()
|
|
self._log_handle = None
|
|
|
|
@staticmethod
|
|
def _terminate_tree(proc: subprocess.Popen) -> None:
|
|
"""Terminate the router AND its model children.
|
|
|
|
The router spawns one child llama-server per loaded model, each holding gigabytes of VRAM.
|
|
Terminating only the router (on Windows, TerminateProcess with no cleanup) orphans them: the
|
|
port goes quiet but the weights stay resident. Enumerate children FIRST (the parent must be
|
|
alive to walk them), then terminate all together, escalating to kill for stragglers.
|
|
"""
|
|
children: list = []
|
|
try:
|
|
import psutil
|
|
|
|
children = psutil.Process(proc.pid).children(recursive=True)
|
|
except Exception: # noqa: BLE001 — no psutil view; still stop the router
|
|
children = []
|
|
proc.terminate()
|
|
for child in children:
|
|
try:
|
|
child.terminate()
|
|
except Exception: # noqa: BLE001
|
|
pass
|
|
try:
|
|
proc.wait(timeout=15)
|
|
except subprocess.TimeoutExpired:
|
|
proc.kill()
|
|
for child in children:
|
|
try:
|
|
if child.is_running():
|
|
child.kill()
|
|
except Exception: # noqa: BLE001
|
|
pass
|
|
|
|
def _reap_orphaned_children(self) -> None:
|
|
"""Kill model children orphaned by a router crash, before respawn.
|
|
|
|
A crashed router can't clean up its children, and a dead parent can't be walked — so match
|
|
by identity instead: any process running OUR llama-server binary whose parent is gone is an
|
|
orphan of a previous router. Their VRAM must come back before the new router loads models
|
|
next to the ghosts.
|
|
"""
|
|
try:
|
|
import psutil
|
|
|
|
exe = str(server_binary(self.install_dir))
|
|
except Exception: # noqa: BLE001
|
|
return
|
|
for p in psutil.process_iter(["exe", "ppid"]):
|
|
try:
|
|
if p.info.get("exe") != exe:
|
|
continue
|
|
if self.proc is not None and p.pid == self.proc.pid:
|
|
continue
|
|
ppid = p.info.get("ppid") or 0
|
|
if ppid and psutil.pid_exists(ppid):
|
|
continue
|
|
logger.warning("reaping orphaned llama-server child pid=%s", p.pid)
|
|
p.kill()
|
|
except (psutil.NoSuchProcess, psutil.AccessDenied):
|
|
continue
|
|
|
|
# ── model management (router endpoints) ──────────────────
|
|
|
|
def models(self) -> dict:
|
|
"""{model_id: status_value} from GET /models."""
|
|
data = self._request("/models")
|
|
return {m["id"]: m.get("status", {}).get("value", "unknown")
|
|
for m in data.get("data", [])}
|
|
|
|
def load_model(self, model_id: str, timeout_s: int = 600) -> None:
|
|
self._request("/models/load", {"model": model_id}, timeout_s=timeout_s)
|
|
|
|
def unload_model(self, model_id: str) -> None:
|
|
"""Free the child's VRAM now. Route existence verified empirically on b10290 (POST
|
|
/models/unload; bogus name -> 400 'model is not found'). Momentary action: never touches
|
|
primary_model — the declaration is durable, an eject is not (residency design).
|
|
"""
|
|
self._request("/models/unload", {"model": model_id}, timeout_s=120)
|
|
deadline = time.monotonic() + 15
|
|
while time.monotonic() < deadline:
|
|
try:
|
|
if self.models().get(model_id) not in ("loaded", "ready", "unloading"):
|
|
return
|
|
except Exception: # noqa: BLE001
|
|
return
|
|
time.sleep(0.3)
|
|
|
|
# ── idle residency (non-primary models) ──────────────────
|
|
|
|
# A model that has gone quiet gets its VRAM back after this long. A
|
|
# constant, not a knob: long enough that an active conversation never
|
|
# trips it, short enough that a wandered-off session frees ~20 GiB
|
|
# within the hour. No exemptions (residency v2): demand reloads
|
|
# anything the user comes back to.
|
|
IDLE_UNLOAD_S = 15 * 60
|
|
|
|
def sweep_idle(self, now: float | None = None) -> list[str]:
|
|
"""Unload models idle past IDLE_UNLOAD_S. Returns the model ids
|
|
unloaded. Idle means no busy slots and no queued work, tracked
|
|
per model across calls; a model seen busy resets its clock."""
|
|
now = time.monotonic() if now is None else now
|
|
unloaded: list[str] = []
|
|
try:
|
|
statuses = self.models()
|
|
except Exception: # noqa: BLE001
|
|
return unloaded
|
|
for model_id, status in statuses.items():
|
|
if status not in ("loaded", "ready") or not self.is_idle(model_id):
|
|
self._idle_since.pop(model_id, None)
|
|
continue
|
|
first_idle = self._idle_since.setdefault(model_id, now)
|
|
if now - first_idle >= self.IDLE_UNLOAD_S:
|
|
try:
|
|
self.unload_model(model_id)
|
|
self._idle_since.pop(model_id, None)
|
|
unloaded.append(model_id)
|
|
logger.info("idle-unloaded %s (idle %ds)", model_id,
|
|
int(now - first_idle))
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.warning("idle unload of %s failed: %s", model_id, exc)
|
|
return unloaded
|
|
|
|
def touch_generate(self, model_id: str, timeout_s: int = 300) -> bool:
|
|
"""The readiness proof. Generous budget + reasoning_content scan — small token budgets
|
|
false-fail reasoning models, which spend their first tokens thinking.
|
|
"""
|
|
try:
|
|
resp = self._request("/v1/chat/completions", {
|
|
"model": model_id,
|
|
"messages": [{"role": "user", "content": TOUCH_PROMPT}],
|
|
"max_tokens": 512, "temperature": 0,
|
|
}, timeout_s=timeout_s)
|
|
msg = resp["choices"][0]["message"]
|
|
blob = (msg.get("content") or "") + " " + (msg.get("reasoning_content") or "")
|
|
return TOUCH_EXPECT in blob.lower()
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.warning("touch generation failed for %s: %s", model_id, exc)
|
|
return False
|
|
|
|
def ensure_model_ready(self, model_id: str, timeout_s: int = 600) -> bool:
|
|
"""Load if needed, then prove readiness with a touch generation."""
|
|
status = self.models().get(model_id)
|
|
if status is None:
|
|
raise KeyError(f"model {model_id} not present in models dir")
|
|
if status not in ("loaded", "ready"):
|
|
self.load_model(model_id, timeout_s=timeout_s)
|
|
return self.touch_generate(model_id)
|
|
|
|
# ── telemetry ────────────────────────────────────────────
|
|
|
|
def is_idle(self, model_id: str | None = None) -> bool:
|
|
"""No processing requests and no busy slots. Router quirk: /slots and /metrics are per-child
|
|
and require ?model= (bare calls 400), and no KV-usage metric exists. With ``model_id``
|
|
checks that one child; without, every loaded child.
|
|
"""
|
|
try:
|
|
if model_id is not None:
|
|
loaded = [model_id]
|
|
else:
|
|
loaded = [m for m, status in self.models().items()
|
|
if status in ("loaded", "ready")]
|
|
for mid in loaded:
|
|
slots = self._request(f"/slots?model={mid}")
|
|
if any(s.get("is_processing") for s in slots):
|
|
return False
|
|
with self._open(f"/metrics?model={mid}", timeout_s=10, json_type=False) as r:
|
|
text = r.read().decode()
|
|
for line in text.splitlines():
|
|
if line.startswith("llamacpp:requests_processing"):
|
|
if float(line.split()[-1]) != 0.0:
|
|
return False
|
|
return True
|
|
except Exception: # noqa: BLE001
|
|
return False
|