llm/scripts/check_gpu_readiness.py

from __future__ import annotations

import argparse
import datetime as dt
import json
import subprocess
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

from check_inference_endpoint import check_endpoint
from common import read_json


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_REPORT = ROOT / "reports" / "gpu-readiness.json"
GPU_PROFILES = ROOT / "config" / "gpu_profiles.json"
PROFILE_MODEL_EXPECTATIONS = {
    "vLLM": "qwen3-4b-instruct",
    "llama.cpp": "devstral-1c-q4",
}


def run_command(command: list[str], timeout: int) -> dict:
    try:
        result = subprocess.run(
            command,
            cwd=ROOT,
            text=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            timeout=timeout,
            check=False,
        )
    except subprocess.TimeoutExpired as exc:
        return {
            "command": command,
            "status": "failed",
            "returncode": None,
            "stdout": exc.stdout or "",
            "stderr": f"Timed out after {timeout}s",
        }
    return {
        "command": command,
        "status": "ok" if result.returncode == 0 else "failed",
        "returncode": result.returncode,
        "stdout": result.stdout.strip(),
        "stderr": result.stderr.strip(),
    }


def check_http_url(url: str, timeout: int, *, expected_model: str | None = None) -> dict:
    if url.rstrip("/").endswith("/v1/models"):
        base_url = url.rstrip("/")[: -len("/v1/models")]
        return check_endpoint(base_url, expected_model, timeout)

    started_at = time.perf_counter()
    result = {
        "url": url,
        "status": "failed",
        "latency_ms": None,
        "http_status": None,
        "error": None,
    }
    try:
        request = urllib.request.Request(url, method="GET")
        with urllib.request.urlopen(request, timeout=timeout) as response:
            result["http_status"] = response.status
    except (TimeoutError, OSError, urllib.error.URLError) as exc:
        result["error"] = str(exc)
        return result

    result["latency_ms"] = round((time.perf_counter() - started_at) * 1000)
    result["status"] = "ok" if 200 <= int(result["http_status"] or 0) < 500 else "failed"
    return result


def profile_checks(profile_id: str, timeout: int) -> dict[str, dict]:
    profiles = read_json(GPU_PROFILES)
    profile = profiles.get(profile_id)
    if not isinstance(profile, dict):
        available = ", ".join(sorted(str(name) for name in profiles))
        raise ValueError(f"Unknown GPU profile `{profile_id}`. Available: {available}")

    checks: dict[str, dict] = {}
    for index, item in enumerate(profile.get("wait") or [], start=1):
        if not isinstance(item, dict):
            continue
        name = str(item.get("name") or f"wait_{index}")
        url = str(item.get("url") or "")
        if not url:
            continue
        key = f"profile_{profile_id}_{index}_{name.lower().replace('.', '').replace(' ', '_')}"
        checks[key] = check_http_url(url, timeout, expected_model=PROFILE_MODEL_EXPECTATIONS.get(name))
    return checks


def main() -> int:
    parser = argparse.ArgumentParser(description="Check GPU host and inference endpoint readiness.")
    parser.add_argument("--profile", default="text", help="GPU profile from config/gpu_profiles.json to check.")
    parser.add_argument("--all-endpoints", action="store_true", help="Legacy mode: require vLLM and llama.cpp endpoints at the same time.")
    parser.add_argument("--vllm-url", default="http://docker-gpu.cin.su:8000")
    parser.add_argument("--vllm-model", default="qwen3-4b-instruct")
    parser.add_argument("--llama-url", default="http://docker-gpu.cin.su:8080")
    parser.add_argument("--llama-model", default="devstral-1c-q4")
    parser.add_argument("--ssh-target", default="docker-gpu.cin.su")
    parser.add_argument("--timeout", type=int, default=10)
    parser.add_argument("--report", type=Path, default=DEFAULT_REPORT)
    parser.add_argument("--print", action="store_true")
    args = parser.parse_args()

    checks: dict[str, dict] = {
        "ssh_preflight": run_command(
            [
                "powershell",
                "-NoProfile",
                "-ExecutionPolicy",
                "Bypass",
                "-File",
                "scripts/check_gpu_host.ps1",
                "-SshTarget",
                args.ssh_target,
                "-ConnectTimeoutSeconds",
                str(args.timeout),
            ],
            timeout=max(args.timeout * 6, 30),
        )
    }
    if args.all_endpoints:
        checks["vllm_endpoint"] = check_endpoint(args.vllm_url, args.vllm_model, args.timeout)
        checks["llama_endpoint"] = check_endpoint(args.llama_url, args.llama_model, args.timeout)
    else:
        checks.update(profile_checks(args.profile, args.timeout))

    failed = [name for name, check in checks.items() if check.get("status") != "ok"]
    report = {
        "created_at": dt.datetime.now(dt.UTC).isoformat(),
        "profile": args.profile,
        "all_endpoints": args.all_endpoints,
        "status": "failed" if failed else "ok",
        "failed_checks": failed,
        "checks": checks,
    }

    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"GPU readiness status: {report['status']}")
        print(f"Wrote report to {args.report}")
    return 0 if not failed else 1


if __name__ == "__main__":
    raise SystemExit(main())