gpu-rent/src/gpu_rent/ready.py

"""Wait until SwarmUI HTTP is up and backend is Idle (on the VM).

Also end-of-up stack verification: every enabled service must answer.
"""

from __future__ import annotations

import json
import socket
import time
import urllib.error
import urllib.request
from collections.abc import Callable
from dataclasses import dataclass

from gpu_rent.config import Config
from gpu_rent.errors import CloudError
from gpu_rent.llm_runtime import normalize_runtime
from gpu_rent.ssh_ops import run_python, run_ssh
from gpu_rent.timing import WaitLog

Log = Callable[[str], None]

_REMOTE_POLL = r"""
import json, time, urllib.request
url = "http://127.0.0.1:7801"
deadline = time.time() + 25
while time.time() < deadline:
    try:
        req = urllib.request.Request(
            url + "/API/GetNewSession",
            data=b"{}",
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req, timeout=5) as resp:
            session = json.loads(resp.read().decode())
        sid = session.get("session_id")
        if not sid:
            time.sleep(2)
            continue
        body = json.dumps({"session_id": sid}).encode()
        req2 = urllib.request.Request(
            url + "/API/GetCurrentStatus",
            data=body,
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with urllib.request.urlopen(req2, timeout=5) as resp:
            data = json.loads(resp.read().decode())
        st = data.get("status") or {}
        be = data.get("backend_status") or {}
        waiting = int(st.get("waiting_gens") or 0)
        live = int(st.get("live_gens") or 0)
        loading = int(st.get("loading_models") or 0)
        bstat = str(be.get("status") or "unknown").lower()
        any_loading = bool(be.get("any_loading"))
        # SwarmUI: "running" = backends healthy & ready to generate.
        # "idle" = suspended / cannot generate. "loading" = still starting.
        if waiting or live or loading:
            print(f"BUSY queue w={waiting} live={live} load={loading}")
        elif bstat == "empty":
            print("BUSY backend=empty (нужен first-install Comfy)")
        elif bstat in ("loading", "some_loading") or any_loading:
            print(f"BUSY backend={bstat} (Comfy стартует)")
        elif bstat == "running":
            print("READY backend=running")
        elif bstat == "idle":
            # Suspended backends still installed — first gen wakes them.
            # Do not burn 2400s waiting for running after AllowIdle.
            print("READY backend=idle (backends present, suspended)")
        elif bstat in ("disabled", "all_disabled"):
            print(f"BUSY backend={bstat}")
        elif bstat == "errored":
            print("BUSY backend=errored")
        else:
            print(f"BUSY backend={bstat}")
        raise SystemExit(0)
    except Exception as exc:
        print(f"WAIT {exc}")
        time.sleep(3)
print("WAIT timeout-slice")
"""

# One-shot probe of configured stack endpoints on the VM (JSON line).
_REMOTE_STACK_PROBE = r'''
import json, urllib.error, urllib.request, subprocess

def http_ok(url, timeout=4.0):
    try:
        req = urllib.request.Request(url, method="GET")
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            code = getattr(resp, "status", 200) or 200
            body = resp.read(256)
            return True, f"HTTP {code} ({len(body)}b)"
    except Exception as exc:
        return False, str(exc)[:160]

def unit_active(name):
    try:
        out = subprocess.check_output(
            ["systemctl", "is-active", name],
            text=True,
            stderr=subprocess.DEVNULL,
        ).strip()
        return out
    except Exception:
        return "unknown"

checks = []
want_swarm = WANT_SWARM
want_ollama = WANT_OLLAMA

if want_swarm:
    ok, detail = http_ok("http://127.0.0.1:7801/")
    if not ok:
        # API may answer when / does not
        ok2, d2 = False, detail
        try:
            req = urllib.request.Request(
                "http://127.0.0.1:7801/API/GetNewSession",
                data=b"{}",
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=5) as resp:
                ok2 = True
                d2 = f"API session HTTP {getattr(resp, 'status', 200)}"
        except Exception as exc:
            d2 = str(exc)[:160]
        ok, detail = ok2, d2
    checks.append({
        "name": "swarmui",
        "ok": ok,
        "detail": detail,
        "unit": unit_active("swarmui"),
    })

if want_ollama:
    ok, detail = http_ok("http://127.0.0.1:11434/api/tags")
    checks.append({
        "name": "ollama",
        "ok": ok,
        "detail": detail,
        "unit": unit_active("gpu-rent-ollama"),
    })

print(json.dumps({"checks": checks}, ensure_ascii=False))
'''


@dataclass(frozen=True)
class ServiceCheck:
    name: str
    ok: bool
    detail: str
    where: str = "vm"  # vm | local


Log = Callable[[str], None]


def collect_swarm_diagnostics(cfg: Config, host: str, log: Log) -> str:
    """Run remote swarm_diag.py; stream to log; return full report text."""
    from importlib.resources import files

    log("диагностика SwarmUI/Comfy…")
    try:
        script = files("gpu_rent.remote").joinpath("swarm_diag.py").read_text(
            encoding="utf-8"
        )
        out = run_python(
            cfg,
            host,
            script,
            remote_path="/tmp/gpu-rent-swarm_diag.py",
            timeout=90,
            log=log,
        )
        return out or ""
    except Exception as exc:
        # Best-effort fallback when full script fails
        log(f"diag script fail: {exc}")
        try:
            fb = run_ssh(
                cfg,
                host,
                "echo '=== systemctl swarmui ==='; systemctl is-active swarmui; "
                "echo; echo '=== journalctl -u swarmui -n 60 ==='; "
                "sudo -n journalctl -u swarmui -n 60 --no-pager 2>/dev/null || true",
                check=False,
                timeout=45,
            )
            for line in (fb or "").splitlines():
                log(line)
            return fb or ""
        except Exception as exc2:
            log(f"diag fallback fail: {exc2}")
            return ""


def wait_backend_idle(
    cfg: Config,
    host: str,
    log: Log,
    *,
    timeout: float = 2400.0,
    poll_every: float = 15.0,
    errored_fail_sec: float = 120.0,
    disabled_fail_sec: float = 90.0,
    loading_fail_sec: float = 900.0,
) -> None:
    """Block until SwarmUI backends are ready (status=running, no queue).

    Note: SwarmUI ``idle`` means suspended backends (cannot generate).
    Ready-to-use is ``running``. Sustained ``errored`` / ``disabled`` / long
    ``loading`` fail-fasts.
    """
    deadline = time.time() + timeout
    log("жду ready backend (running)…")
    last = ""
    errored_since: float | None = None
    disabled_since: float | None = None
    loading_since: float | None = None
    last_loading_note = 0.0
    pretty = {
        "BUSY backend=loading (Comfy стартует)": "… Comfy стартует",
        "BUSY backend=some_loading (Comfy стартует)": "… Comfy стартует (часть бэкендов)",
        "BUSY backend=empty (нужен first-install Comfy)": "… backend пуст — нужен install",
        "BUSY backend=errored": "… backend errored — смотри journalctl -u swarmui",
        "BUSY backend=disabled": "… backend disabled (пустой StartScript?)",
        "BUSY backend=all_disabled": "… все backends disabled",
        "READY backend=running": "backend ready (running)",
        "READY backend=idle (backends present, suspended)": "backend ready (idle/suspended)",
    }
    while time.time() < deadline:
        try:
            out = run_ssh(
                cfg,
                host,
                "python3 - <<'PY'\n" + _REMOTE_POLL + "\nPY",
                check=False,
                timeout=40,
            ).strip()
        except Exception as exc:
            out = f"WAIT ssh: {exc}"
        line = out.splitlines()[-1] if out else "WAIT empty"
        shown = pretty.get(line, line)
        if "BUSY backend=loading" in line or "BUSY backend=some_loading" in line:
            now = time.time()
            if loading_since is None:
                loading_since = now
            elapsed = int(now - loading_since)
            shown = f"… Comfy стартует ({elapsed // 60}м {elapsed % 60}с)"
            if now - last_loading_note >= 60:
                last_loading_note = now
                try:
                    j = run_ssh(
                        cfg,
                        host,
                        "sudo -n journalctl -u swarmui -n 8 --no-pager -o cat 2>/dev/null "
                        "| tail -n 8 || true",
                        check=False,
                        timeout=25,
                    ).strip()
                    if j:
                        for jl in j.splitlines()[-4:]:
                            log(f"  journal: {jl[:160]}")
                except Exception:
                    pass
            if elapsed >= loading_fail_sec:
                collect_swarm_diagnostics(cfg, host, log)
                raise CloudError(
                    f"backend=loading уже {elapsed} с — похоже завис "
                    "(часто FrontendVersion / pip / сеть). Диагностика выше. "
                    "Попробуй gpu-rent up снова или Server → Backends → Restart."
                )
        else:
            loading_since = None

        if shown != last:
            log(shown)
            last = shown
        if line.startswith("READY"):
            return
        if "BUSY backend=errored" in line:
            now = time.time()
            if errored_since is None:
                errored_since = now
            elif now - errored_since >= errored_fail_sec:
                collect_swarm_diagnostics(cfg, host, log)
                raise CloudError(
                    f"backend=errored уже {int(now - errored_since)} с — не заживёт само. "
                    "Диагностика выше (+ /mnt/swarm_data/.gpu-rent-last-diag.txt на VM). "
                    "Server → Backends → Restart или gpu-rent stop / up после фикса."
                )
        else:
            errored_since = None
        if "BUSY backend=disabled" in line or "BUSY backend=all_disabled" in line:
            now = time.time()
            if disabled_since is None:
                disabled_since = now
            elif now - disabled_since >= disabled_fail_sec:
                collect_swarm_diagnostics(cfg, host, log)
                raise CloudError(
                    f"backend=disabled уже {int(now - disabled_since)} с — "
                    "обычно пустой StartScript (нужен dlbackend/ComfyUI/main.py). "
                    "Диагностика выше. Server → Backends → Edit StartScript или re-run up."
                )
        else:
            disabled_since = None
        time.sleep(poll_every)
    collect_swarm_diagnostics(cfg, host, log)
    raise CloudError(
        f"backend не стал ready (running) за {int(timeout)} с. "
        "Диагностика выше (+ .gpu-rent-last-diag.txt на data volume). GPU ещё жив."
    )


def _expected_services(cfg: Config) -> tuple[bool, bool]:
    swarm = bool(getattr(cfg, "enable_swarmui", True))
    rt = normalize_runtime(getattr(cfg, "llm_runtime", "none"))
    return swarm, rt == "ollama"


def _probe_vm_once(cfg: Config, host: str) -> list[ServiceCheck]:
    want_swarm, want_ollama = _expected_services(cfg)
    script = (
        _REMOTE_STACK_PROBE.replace("WANT_SWARM", "True" if want_swarm else "False")
        .replace("WANT_OLLAMA", "True" if want_ollama else "False")
    )
    out = run_ssh(
        cfg,
        host,
        "python3 - <<'PY'\n" + script + "\nPY",
        check=False,
        timeout=60,
    ).strip()
    line = ""
    for row in reversed(out.splitlines()):
        row = row.strip()
        if row.startswith("{"):
            line = row
            break
    if not line:
        return [
            ServiceCheck("stack", False, f"нет JSON от probe: {out[-200:]}", "vm")
        ]
    try:
        data = json.loads(line)
    except json.JSONDecodeError:
        return [ServiceCheck("stack", False, f"битый JSON: {line[:200]}", "vm")]
    checks: list[ServiceCheck] = []
    for item in data.get("checks") or []:
        if not isinstance(item, dict):
            continue
        name = str(item.get("name") or "?")
        detail = str(item.get("detail") or "")
        unit = str(item.get("unit") or "")
        if unit and unit != "unknown":
            detail = f"{detail}; unit={unit}"
        checks.append(
            ServiceCheck(name=name, ok=bool(item.get("ok")), detail=detail, where="vm")
        )
    return checks


def verify_stack_on_vm(
    cfg: Config,
    host: str,
    log: Log,
    *,
    timeout: float = 300.0,
    poll_every: float = 8.0,
    raise_on_fail: bool = True,
) -> list[ServiceCheck]:
    """Poll until every enabled service answers on the VM loopback."""
    want_swarm, want_ollama = _expected_services(cfg)
    if not (want_swarm or want_ollama):
        log("проверка стека: нечего ждать (swarm off, LLM none)")
        return []

    names = []
    if want_swarm:
        names.append("SwarmUI :7801")
    if want_ollama:
        names.append("Ollama :11434")
    log(f"проверка на VM: {', '.join(names)}")

    deadline = time.time() + timeout
    last: list[ServiceCheck] = []
    wait = WaitLog(log, every=30.0)
    while time.time() < deadline:
        try:
            last = _probe_vm_once(cfg, host)
        except Exception as exc:
            last = [ServiceCheck("ssh", False, str(exc)[:200], "vm")]
        if last and all(c.ok for c in last):
            for c in last:
                log(f"  [ok] {c.name}: {c.detail}")
            log("проверка VM: всё отвечает")
            return last
        bad = ", ".join(f"{c.name}={c.detail}" for c in last if not c.ok) or "пусто"
        wait.tick(f"  … ещё нет: {bad}")
        time.sleep(poll_every)

    for c in last:
        mark = "ok" if c.ok else "FAIL"
        log(f"  [{mark}] {c.name}: {c.detail}")
    if raise_on_fail and any(not c.ok for c in last):
        failed = [c.name for c in last if not c.ok]
        raise CloudError(
            f"сервисы не ответили на VM за {int(timeout)} с: {', '.join(failed)}. "
            "GPU жив — смотри journalctl / gpu-rent logs"
        )
    return last


def _tcp_ok(port: int, host: str = "127.0.0.1", timeout: float = 0.8) -> bool:
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(timeout)
    try:
        return sock.connect_ex((host, port)) == 0
    finally:
        sock.close()


def _http_local(url: str, timeout: float = 3.0) -> tuple[bool, str]:
    try:
        req = urllib.request.Request(url, method="GET")
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            code = getattr(resp, "status", 200) or 200
            return True, f"HTTP {code}"
    except Exception as exc:
        return False, str(exc)[:160]


def verify_gpu_env(
    cfg: Config,
    host: str,
    log: Log,
    *,
    timeout: float = 600.0,
    poll_every: float = 15.0,
    raise_on_fail: bool = True,
) -> list[ServiceCheck]:
    """nvidia-smi / CUDA / torch(+cuda) in Comfy venv when SwarmUI is on.

    Driver/CUDA missing → fail immediately (won't appear later).
    Torch/Comfy venv → poll until timeout (first Comfy start installs them).
    """
    from importlib.resources import files

    from gpu_rent.ssh_ops import run_python

    # These never "appear later" on a broken image — don't burn the poll budget.
    # Empty-backend / missing install also won't self-heal without InstallConfirmWS.
    instant_fail_names = {"nvidia-smi", "cuda"}
    instant_fail_detail_substrings = (
        "dlbackend пуст",
        "backend=empty",
        "Install не прогоняли",
        "available=false",
        "cuda=none",
        "без cuda",
    )

    want_swarm = bool(getattr(cfg, "enable_swarmui", True))
    script = files("gpu_rent.remote").joinpath("stack_env_probe.py").read_text(
        encoding="utf-8"
    )
    swarm_flag = "1" if want_swarm else "0"
    log(
        "проверка GPU-стека: nvidia-smi, CUDA"
        + (", torch в Comfy venv" if want_swarm else " (llm-only — без torch)")
    )

    deadline = time.time() + timeout
    last: list[ServiceCheck] = []
    wait = WaitLog(log, every=30.0)
    while time.time() < deadline:
        try:
            out = run_python(
                cfg,
                host,
                script,
                remote_path="/tmp/gpu-rent-stack_env_probe.py",
                timeout=180,
                log=None,
                env={"GPU_RENT_CHECK_SWARM": swarm_flag},
            )
        except Exception as exc:
            last = [ServiceCheck("gpu-env", False, str(exc)[:200], "vm")]
            wait.tick(f"  … gpu-env: {exc}")
            time.sleep(poll_every)
            continue

        data: dict = {}
        for line in reversed(out.splitlines()):
            line = line.strip()
            if line.startswith("{"):
                try:
                    data = json.loads(line)
                    break
                except json.JSONDecodeError:
                    continue
        checks_raw = data.get("checks") if isinstance(data, dict) else None
        if not isinstance(checks_raw, list):
            last = [ServiceCheck("gpu-env", False, f"нет JSON: {out[-180:]}", "vm")]
            wait.tick(f"  … gpu-env: нет JSON")
            time.sleep(poll_every)
            continue

        last = []
        for item in checks_raw:
            if not isinstance(item, dict):
                continue
            name = str(item.get("name") or "?")
            ok = bool(item.get("ok"))
            detail = str(item.get("detail") or "")
            required = bool(item.get("required", True))
            if not required:
                if not ok:
                    log(f"  [warn] {name}: {detail}")
                last.append(
                    ServiceCheck(
                        name, True, detail if ok else f"(optional) {detail}", "vm"
                    )
                )
            else:
                last.append(ServiceCheck(name, ok, detail, "vm"))

        hard = [c for c in last if not c.ok]
        if not hard:
            for c in last:
                log(f"  [ok] {c.name}: {c.detail}")
            log("проверка GPU-стека: ок")
            return last

        instant = [c for c in hard if c.name in instant_fail_names]
        if not instant:
            for c in hard:
                detail_l = (c.detail or "").lower()
                if any(s.lower() in detail_l for s in instant_fail_detail_substrings):
                    instant.append(c)
        if instant:
            for c in last:
                mark = "ok" if c.ok else "FAIL"
                log(f"  [{mark}] {c.name}: {c.detail}")
            if raise_on_fail:
                failed = [c.name for c in instant]
                raise CloudError(
                    f"GPU-стек: нет {', '.join(failed)} (fail-fast). "
                    "Проверь образ Driver / nvidia на VM; для SwarmUI — "
                    "first-install Comfy (InstallConfirmWS / dlbackend). "
                    "Дальше: gpu-rent status · gpu-rent logs · gpu-rent stop"
                )
            return last

        bad = ", ".join(f"{c.name}={c.detail}" for c in hard)
        wait.tick(f"  … ждём torch/Comfy: {bad}")
        time.sleep(poll_every)

    for c in last:
        mark = "ok" if c.ok else "FAIL"
        log(f"  [{mark}] {c.name}: {c.detail}")
    if raise_on_fail and any(not c.ok for c in last):
        failed = [c.name for c in last if not c.ok]
        raise CloudError(
            f"GPU-стек не готов за {int(timeout)} с: {', '.join(failed)}. "
            "Нужны nvidia-smi, CUDA; для SwarmUI — torch с cuda в Comfy venv "
            "(journalctl -u swarmui / первый старт backend). "
            "Дальше: gpu-rent status · gpu-rent logs · gpu-rent stop"
        )
    return last


def verify_stack_local(
    cfg: Config,
    log: Log,
    *,
    timeout: float = 60.0,
    poll_every: float = 2.0,
    raise_on_fail: bool = True,
) -> list[ServiceCheck]:
    """After tunnel: local ports + light HTTP for enabled services."""
    want_swarm, want_ollama = _expected_services(cfg)
    targets: list[tuple[str, int, str | None]] = []
    if want_swarm:
        targets.append(("swarmui", int(cfg.swarmui_local_port), None))
    if want_ollama:
        targets.append(
            (
                "ollama",
                int(cfg.ollama_local_port),
                f"http://127.0.0.1:{cfg.ollama_local_port}/api/tags",
            )
        )

    if not targets:
        return []

    log(
        "проверка туннеля (localhost): "
        + ", ".join(f"{n}:{port}" for n, port, _ in targets)
    )
    deadline = time.time() + timeout
    last: list[ServiceCheck] = []
    wait = WaitLog(log, every=15.0)
    while time.time() < deadline:
        last = []
        for name, port, url in targets:
            if not _tcp_ok(port):
                last.append(ServiceCheck(name, False, f"порт {port} закрыт", "local"))
                continue
            if url:
                ok, detail = _http_local(url)
                last.append(ServiceCheck(name, ok, detail, "local"))
            else:
                ok, detail = _http_local(f"http://127.0.0.1:{port}/")
                if not ok:
                    try:
                        req = urllib.request.Request(
                            f"http://127.0.0.1:{port}/API/GetNewSession",
                            data=b"{}",
                            headers={"Content-Type": "application/json"},
                            method="POST",
                        )
                        with urllib.request.urlopen(req, timeout=4) as resp:
                            ok = True
                            detail = f"API HTTP {getattr(resp, 'status', 200)}"
                    except Exception as exc:
                        ok = False
                        detail = f"HTTP/API fail (TCP open): {str(exc)[:100]}"
                last.append(ServiceCheck(name, ok, detail, "local"))
        if last and all(c.ok for c in last):
            for c in last:
                log(f"  [ok] localhost {c.name}: {c.detail}")
            log("проверка туннеля: всё доступно")
            return last
        bad = ", ".join(f"{c.name}={c.detail}" for c in last if not c.ok) or "пусто"
        wait.tick(f"  … localhost ещё нет: {bad}")
        time.sleep(poll_every)

    for c in last:
        mark = "ok" if c.ok else "FAIL"
        log(f"  [{mark}] localhost {c.name}: {c.detail}")
    if raise_on_fail and any(not c.ok for c in last):
        failed = [c.name for c in last if not c.ok]
        raise CloudError(
            f"туннель поднят, но локально не отвечает: {', '.join(failed)}"
        )
    return last