137 lines
5.4 KiB
Python
137 lines
5.4 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
|
|
|
|
from usage_utils import extract_usage
|
|
|
|
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, str | None, dict]:
|
|
"""Возвращает (ok, latency_ms, error, usage).
|
|
|
|
usage извлекается ДО валидации контента: если парсинг choices падает,
|
|
потраченные токены всё равно сохраняются. Таймаут/ошибка сети без тела
|
|
ответа дают usage=None по всем полям (NULL, а не нули).
|
|
"""
|
|
t0 = time.time()
|
|
error = None
|
|
usage = {"input_tokens": None, "output_tokens": None,
|
|
"total_tokens": None, "usage_json": None}
|
|
try:
|
|
body = json.dumps({
|
|
"model": model_id,
|
|
"messages": [{"role": "user", "content": "2+2=?"}],
|
|
"max_tokens": 5,
|
|
"reasoning_effort": "none",
|
|
}).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())
|
|
# Usage — до любой валидации контента.
|
|
usage = extract_usage(resp)
|
|
# HTTP 200 может содержать только reasoning при исчерпании лимита.
|
|
# Наличие текста в любом из двух полей достаточно для liveness.
|
|
choice = resp["choices"][0]
|
|
message = choice["message"]
|
|
# Liveness: непустой текст в content ИЛИ reasoning_content.
|
|
# Usage само по себе не доказывает, что модель выдала ответ.
|
|
ok = any(isinstance(message.get(field), str) and message[field].strip()
|
|
for field in ("content", "reasoning_content"))
|
|
if not ok:
|
|
error = f"Empty content and reasoning (finish_reason={choice.get('finish_reason')})"
|
|
except Exception as e:
|
|
error = f"{type(e).__name__}: {e}"
|
|
print(f" ✗ {model_id}: {error}", file=sys.stderr)
|
|
ok = False
|
|
latency = int((time.time() - t0) * 1000)
|
|
return ok, latency, error, usage
|
|
|
|
|
|
def _ensure_usage_columns(db: sqlite3.Connection):
|
|
"""Идемпотентно добавить колонки token usage, если БД ещё со старой схемой.
|
|
|
|
Cron стартует раньше uvicorn, поэтому probe не может полагаться на то,
|
|
что init_db из main.py уже прогнал миграцию.
|
|
"""
|
|
cols = [r[1] for r in db.execute("PRAGMA table_info(checks)").fetchall()]
|
|
for col, decl in (
|
|
("input_tokens", "INTEGER"),
|
|
("output_tokens", "INTEGER"),
|
|
("total_tokens", "INTEGER"),
|
|
("usage_json", "TEXT"),
|
|
):
|
|
if col not in cols:
|
|
db.execute(f"ALTER TABLE checks ADD COLUMN {col} {decl}")
|
|
db.commit()
|
|
|
|
|
|
def run():
|
|
db = sqlite3.connect(DB_PATH, check_same_thread=False)
|
|
_ensure_usage_columns(db)
|
|
model_cols = [r[1] for r in db.execute("PRAGMA table_info(models)")]
|
|
if "active" not in model_cols:
|
|
db.execute("ALTER TABLE models ADD COLUMN active INTEGER NOT NULL DEFAULT 1")
|
|
db.commit()
|
|
models = [r[0] for r in db.execute("SELECT id FROM models WHERE active = 1").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, err, usage = check_model(mid)
|
|
return mid, ok, lat, err, usage
|
|
|
|
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, err, usage = f.result()
|
|
db.execute(
|
|
"INSERT INTO checks (model_id, checked_at, ok, latency_ms, error, "
|
|
"input_tokens, output_tokens, total_tokens, usage_json) "
|
|
"VALUES (?,?,?,?,?,?,?,?,?)",
|
|
(mid, now, int(ok), lat, err,
|
|
usage["input_tokens"], usage["output_tokens"],
|
|
usage["total_tokens"], usage["usage_json"]),
|
|
)
|
|
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()
|