85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import datetime as dt
|
|
import json
|
|
import sys
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
DEFAULT_REPORT = ROOT / "reports" / "inference-endpoint-check.json"
|
|
|
|
|
|
def fetch_models(base_url: str, timeout: int) -> tuple[list[str], int]:
|
|
started_at = time.perf_counter()
|
|
request = urllib.request.Request(f"{base_url.rstrip('/')}/v1/models", method="GET")
|
|
with urllib.request.urlopen(request, timeout=timeout) as response:
|
|
payload = json.loads(response.read().decode("utf-8"))
|
|
latency_ms = round((time.perf_counter() - started_at) * 1000)
|
|
rows = payload.get("data") if isinstance(payload, dict) else []
|
|
if not isinstance(rows, list):
|
|
return [], latency_ms
|
|
return [str(row.get("id")) for row in rows if isinstance(row, dict) and row.get("id")], latency_ms
|
|
|
|
|
|
def check_endpoint(base_url: str, expected_model: str | None, timeout: int) -> dict:
|
|
result = {
|
|
"base_url": base_url,
|
|
"expected_model": expected_model,
|
|
"status": "failed",
|
|
"models": [],
|
|
"available": None,
|
|
"latency_ms": None,
|
|
"error": None,
|
|
}
|
|
try:
|
|
models, latency_ms = fetch_models(base_url, timeout)
|
|
except (TimeoutError, OSError, urllib.error.URLError, ValueError) as exc:
|
|
result["error"] = str(exc)
|
|
return result
|
|
|
|
result["models"] = models
|
|
result["latency_ms"] = latency_ms
|
|
result["available"] = expected_model in models if expected_model else None
|
|
result["status"] = "ok" if models and (expected_model is None or result["available"]) else "failed"
|
|
return result
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Check an OpenAI-compatible inference endpoint.")
|
|
parser.add_argument("--base-url", default="http://docker-gpu.cin.su:8000")
|
|
parser.add_argument("--expected-model")
|
|
parser.add_argument("--timeout", type=int, default=10)
|
|
parser.add_argument("--report", type=Path, default=DEFAULT_REPORT)
|
|
parser.add_argument("--print", action="store_true")
|
|
parser.add_argument("--no-report", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
result = check_endpoint(args.base_url, args.expected_model, args.timeout)
|
|
report = {
|
|
"created_at": dt.datetime.now(dt.UTC).isoformat(),
|
|
"status": result["status"],
|
|
"checks": [result],
|
|
}
|
|
|
|
if not args.no_report:
|
|
args.report.parent.mkdir(parents=True, exist_ok=True)
|
|
args.report.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
|
|
|
if args.print:
|
|
print(json.dumps(report, ensure_ascii=False, indent=2))
|
|
else:
|
|
print(f"Endpoint status: {result['status']} ({args.base_url})")
|
|
if result["error"]:
|
|
print(result["error"], file=sys.stderr)
|
|
|
|
return 0 if result["status"] == "ok" else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|