87 lines
2.7 KiB
Python
87 lines
2.7 KiB
Python
"""
|
|
Cron-скрипт: liveness probe всех моделей из DB.
|
|
Запускается внутри контейнера по расписанию (cron / supervisord).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
from datetime import datetime, timezone
|
|
|
|
DB_PATH = os.environ.get("DB_PATH", "/data/arena.db")
|
|
OMNIROUTE_URL = os.environ.get("OMNIROUTE_URL", "http://192.168.1.101:20128")
|
|
OMNIROUTE_KEY = os.environ.get("OMNIROUTE_API_KEY", "")
|
|
TIMEOUT = int(os.environ.get("CHECK_TIMEOUT", "30"))
|
|
CONCURRENCY = int(os.environ.get("CHECK_CONCURRENCY", "5"))
|
|
|
|
|
|
def check_model(model_id: str) -> tuple[bool, int | None]:
|
|
t0 = time.time()
|
|
try:
|
|
body = json.dumps({
|
|
"model": model_id,
|
|
"messages": [{"role": "user", "content": "2+2=?"}],
|
|
"max_tokens": 5,
|
|
}).encode()
|
|
req = urllib.request.Request(
|
|
f"{OMNIROUTE_URL}/v1/chat/completions",
|
|
data=body,
|
|
headers={
|
|
"Authorization": f"Bearer {OMNIROUTE_KEY}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
)
|
|
with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
|
|
resp = json.loads(r.read())
|
|
content = resp["choices"][0]["message"]["content"]
|
|
ok = bool(content.strip())
|
|
except Exception as e:
|
|
print(f" ✗ {model_id}: {e}", file=sys.stderr)
|
|
ok = False
|
|
latency = int((time.time() - t0) * 1000) if ok else None
|
|
return ok, latency
|
|
|
|
|
|
def run():
|
|
db = sqlite3.connect(DB_PATH, check_same_thread=False)
|
|
models = [r[0] for r in db.execute("SELECT id FROM models").fetchall()]
|
|
db.close()
|
|
|
|
print(f"[probe] checking {len(models)} models @ {datetime.now(timezone.utc).isoformat()}")
|
|
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
|
def do(mid):
|
|
ok, lat = check_model(mid)
|
|
return mid, ok, lat
|
|
|
|
db = sqlite3.connect(DB_PATH, check_same_thread=False)
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
|
|
with ThreadPoolExecutor(max_workers=CONCURRENCY) as ex:
|
|
futs = {ex.submit(do, m): m for m in models}
|
|
done = ok_count = 0
|
|
for f in as_completed(futs):
|
|
mid, ok, lat = f.result()
|
|
db.execute(
|
|
"INSERT INTO checks (model_id, checked_at, ok, latency_ms) VALUES (?,?,?,?)",
|
|
(mid, now, int(ok), lat),
|
|
)
|
|
done += 1
|
|
if ok:
|
|
ok_count += 1
|
|
status = "✓" if ok else "✗"
|
|
print(f" {status} {mid} ({lat}ms)" if lat else f" {status} {mid}")
|
|
|
|
db.commit()
|
|
db.close()
|
|
print(f"[probe] done: {ok_count}/{done} ok")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run()
|