YG_FT/scripts/ops_diagnostics.py

from __future__ import annotations

import json
import os
import sys
import time
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen


def _json_request(url: str, token: str | None = None, timeout: float = 5.0) -> dict[str, Any]:
    headers = {"Accept": "application/json"}
    if token:
        headers["X-Compute-Token"] = token
    request = Request(url, headers=headers)
    started = time.perf_counter()
    try:
        with urlopen(request, timeout=timeout) as response:
            body = response.read().decode("utf-8", errors="replace")
            data = json.loads(body) if body else {}
            return {
                "ok": 200 <= response.status < 300,
                "status": response.status,
                "duration_ms": int((time.perf_counter() - started) * 1000),
                "data": data,
            }
    except HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        return {
            "ok": False,
            "status": exc.code,
            "duration_ms": int((time.perf_counter() - started) * 1000),
            "error": body or str(exc),
        }
    except (URLError, TimeoutError, OSError) as exc:
        return {
            "ok": False,
            "duration_ms": int((time.perf_counter() - started) * 1000),
            "error": str(exc),
        }


def _postgres_check(database_url: str | None) -> dict[str, Any]:
    if not database_url:
        return {"ok": False, "skipped": True, "error": "DATABASE_URL is not set"}
    try:
        import psycopg
    except ImportError:
        return {"ok": False, "skipped": True, "error": "psycopg is not installed"}
    started = time.perf_counter()
    try:
        with psycopg.connect(database_url.replace("postgresql+psycopg://", "postgresql://"), connect_timeout=5) as conn:
            with conn.cursor() as cursor:
                cursor.execute(
                    """
                    SELECT table_name
                    FROM information_schema.tables
                    WHERE table_schema='public'
                      AND table_name IN ('users','fine_tune_tasks','compute_nodes','compute_jobs','gpu_allocations')
                    ORDER BY table_name
                    """
                )
                tables = [row[0] for row in cursor.fetchall()]
        return {"ok": len(tables) >= 5, "duration_ms": int((time.perf_counter() - started) * 1000), "tables": tables}
    except Exception as exc:  # noqa: BLE001 - diagnostics should report instead of crashing
        return {"ok": False, "duration_ms": int((time.perf_counter() - started) * 1000), "error": str(exc)}


def main() -> int:
    route_prefix = os.getenv("ROUTE_PREFIX", "/modelTF").rstrip("/")
    app_base_url = os.getenv("APP_BASE_URL", "http://localhost:17861").rstrip("/")
    compute_base_url = os.getenv("COMPUTE_BASE_URL", "").rstrip("/")
    compute_token = os.getenv("COMPUTE_SERVICE_TOKEN")
    checks: list[dict[str, Any]] = [
        {"name": "backend health", **_json_request(f"{app_base_url}{route_prefix}/health")},
        {"name": "postgres schema", **_postgres_check(os.getenv("DATABASE_URL"))},
    ]
    if compute_base_url:
        checks.extend(
            [
                {
                    "name": "compute health",
                    **_json_request(f"{compute_base_url}{route_prefix}/v1/compute/health", compute_token),
                },
                {
                    "name": "compute gpus",
                    **_json_request(f"{compute_base_url}{route_prefix}/compute/resources/gpus", compute_token),
                },
                {
                    "name": "compute jobs",
                    **_json_request(f"{compute_base_url}{route_prefix}/compute/jobs", compute_token),
                },
            ]
        )
    else:
        checks.append({"name": "compute health", "ok": False, "skipped": True, "error": "COMPUTE_BASE_URL is not set"})
    report = {"ok": all(item.get("ok") or item.get("skipped") for item in checks), "checks": checks}
    print(json.dumps(report, ensure_ascii=False, indent=2))
    return 0 if report["ok"] else 1


if __name__ == "__main__":
    sys.exit(main())