92 lines
4.0 KiB
Python
92 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Bound a busy llama.cpp slot even if the client loses its cancellation path.
|
|
|
|
Only task IDs, monotonic timestamps and service invocation IDs are persisted.
|
|
The wrapper holds the same maintenance lock as the updater and restart helper.
|
|
"""
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import time
|
|
from urllib.request import urlopen
|
|
|
|
|
|
def observe(previous, slots, invocation, now, seconds, tokens):
|
|
old = previous if previous.get("invocation") == invocation else {}
|
|
active = {}
|
|
reasons = []
|
|
for slot in slots:
|
|
if not slot.get("is_processing"):
|
|
continue
|
|
if "id_task" not in slot:
|
|
raise ValueError("busy slot without task ID")
|
|
key = f'{slot["id"]}:{slot["id_task"]}'
|
|
started = old.get("active", {}).get(key, now)
|
|
if started > now:
|
|
started = now
|
|
active[key] = started
|
|
decoded = max((item.get("n_decoded", 0) for item in slot.get("next_token", [])), default=0)
|
|
if now - started >= seconds:
|
|
reasons.append(f"slot {slot['id']}: request exceeded {seconds}s")
|
|
if decoded > tokens:
|
|
reasons.append(f"slot {slot['id']}: output exceeded {tokens} tokens")
|
|
return {"invocation": invocation, "active": active}, reasons
|
|
|
|
|
|
def systemctl(*args):
|
|
return subprocess.run(["systemctl", "--user", *args], check=True,
|
|
capture_output=True, text=True, timeout=35).stdout.strip()
|
|
|
|
|
|
def main():
|
|
service = os.environ.get("LLAMA_SERVICE", "llama-qwen.service")
|
|
url = os.environ.get("LLAMA_URL", "http://127.0.0.1:8080").rstrip("/")
|
|
seconds = int(os.environ.get("LLAMA_MAX_REQUEST_SECONDS", "600"))
|
|
tokens = int(os.environ.get("LLAMA_MAX_OUTPUT_TOKENS", "2048"))
|
|
if seconds <= 0 or tokens <= 0:
|
|
raise ValueError("watchdog limits must be positive")
|
|
state_dir = Path(os.environ.get("XDG_STATE_HOME", str(Path.home() / ".local/state"))) / "hermes-stack-maintenance"
|
|
state_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
path = state_dir / "llama-watchdog.json"
|
|
try:
|
|
previous = json.loads(path.read_text())
|
|
except (OSError, ValueError):
|
|
previous = {}
|
|
boot = Path("/proc/sys/kernel/random/boot_id").read_text().strip()
|
|
if previous.get("boot") != boot:
|
|
previous = {}
|
|
status = systemctl("show", service, "--property=ActiveState", "--property=InvocationID")
|
|
info = dict(line.split("=", 1) for line in status.splitlines())
|
|
if info.get("ActiveState") != "active":
|
|
return # Respect explicit stops and maintenance; never start an inactive server.
|
|
invocation = info["InvocationID"]
|
|
now = time.monotonic()
|
|
try:
|
|
with urlopen(url + "/slots", timeout=5) as response:
|
|
slots = json.load(response)
|
|
if not isinstance(slots, list):
|
|
raise ValueError("invalid slot response")
|
|
state, reasons = observe(previous, slots, invocation, now, seconds, tokens)
|
|
except (OSError, ValueError, KeyError, TypeError):
|
|
# A single failed probe is not evidence of a hung model. Require a full
|
|
# request budget of consecutive failures within the same server process.
|
|
since = previous.get("unreachable_since", now) if previous.get("invocation") == invocation else now
|
|
state = {"invocation": invocation, "active": previous.get("active", {}), "unreachable_since": since}
|
|
reasons = [f"slot endpoint unreachable for {seconds}s"] if now - since >= seconds else []
|
|
last_restart = previous.get("last_restart", -1e12)
|
|
state["last_restart"] = last_restart
|
|
if reasons and now - last_restart >= max(300, seconds):
|
|
print("llama watchdog: " + "; ".join(reasons) + "; restarting " + service, flush=True)
|
|
systemctl("restart", "--no-block", service)
|
|
state = {"invocation": invocation, "active": {}, "last_restart": now}
|
|
state["boot"] = boot
|
|
pending = path.with_suffix(".pending")
|
|
pending.write_text(json.dumps(state) + "\n")
|
|
pending.chmod(0o600)
|
|
pending.replace(path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|