Implement GPU environment verification in session management

- Added a new function `verify_gpu_env` to check GPU stack readiness, including nvidia-smi, CUDA, and torch in the Comfy virtual environment when SwarmUI is enabled. - Updated the session management to call `verify_gpu_env`, capturing GPU environment status and errors in the state notes. - Enhanced documentation in `cli.md` to reflect the new GPU environment verification process. - Added tests for `verify_gpu_env` to ensure proper functionality and error handling during GPU checks.
2026-08-21 06:58:23 +03:00
parent 09b7c36f3b
commit 3c8225a69e
7 changed files with 488 additions and 14 deletions
@@ -311,6 +311,107 @@ def _http_local(url: str, timeout: float = 3.0) -> tuple[bool, str]:
        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."""
+    from importlib.resources import files
+
+    from gpu_rent.ssh_ops import run_python
+
+    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"
+    prefix = f"import os\nos.environ['GPU_RENT_CHECK_SWARM']={swarm_flag!r}\n"
+    log(
+        "проверка GPU-стека: nvidia-smi, CUDA"
+        + (", torch в Comfy venv" if want_swarm else " (llm-only — без torch)")
+    )
+
+    deadline = time.time() + timeout
+    last: list[ServiceCheck] = []
+    while time.time() < deadline:
+        try:
+            out = run_python(
+                cfg,
+                host,
+                prefix + script,
+                remote_path="/tmp/gpu-rent-stack_env_probe.py",
+                timeout=180,
+                log=None,
+            )
+        except Exception as exc:
+            last = [ServiceCheck("gpu-env", False, str(exc)[:200], "vm")]
+            log(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")]
+            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
+
+        bad = ", ".join(f"{c.name}={c.detail}" for c in hard)
+        log(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"GPU-стек не готов за {int(timeout)} с: {', '.join(failed)}. "
+            "Нужны nvidia-smi, CUDA; для SwarmUI — torch с cuda в Comfy venv "
+            "(journalctl -u swarmui / первый старт backend)."
+        )
+    return last
+
+
 def verify_stack_local(
    cfg: Config,
    log: Log,
@@ -326,7 +427,11 @@ def verify_stack_local(
        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")
+            (
+                "ollama",
+                int(cfg.ollama_local_port),
+                f"http://127.0.0.1:{cfg.ollama_local_port}/api/tags",
+            )
        )
    if want_llama:
        p = int(cfg.llamacpp_local_port)
@@ -335,7 +440,10 @@ def verify_stack_local(
    if not targets:
        return []

-    log("проверка туннеля (localhost): " + ", ".join(f"{n}:{port}" for n, port, _ in targets))
+    log(
+        "проверка туннеля (localhost): "
+        + ", ".join(f"{n}:{port}" for n, port, _ in targets)
+    )
    deadline = time.time() + timeout
    last: list[ServiceCheck] = []
    while time.time() < deadline:
@@ -347,12 +455,9 @@ def verify_stack_local(
            if url:
                ok, detail = _http_local(url)
                if not ok and name == "llamacpp":
-                    ok, detail = _http_local(
-                        f"http://127.0.0.1:{port}/v1/models"
-                    )
+                    ok, detail = _http_local(f"http://127.0.0.1:{port}/v1/models")
                last.append(ServiceCheck(name, ok, detail, "local"))
            else:
-                # SwarmUI: TCP enough (API may need POST); try API too
                ok, detail = _http_local(f"http://127.0.0.1:{port}/")
                if not ok:
                    try:
@@ -365,9 +470,7 @@ def verify_stack_local(
                        with urllib.request.urlopen(req, timeout=4) as resp:
                            ok = True
                            detail = f"API HTTP {getattr(resp, 'status', 200)}"
-                    except Exception as exc:
-                        detail = f"TCP ok; HTTP {exc}"[:160]
-                        # TCP open is enough for swarm local check
+                    except Exception:
                        ok = True
                        detail = f"TCP :{port} open"
                last.append(ServiceCheck(name, ok, detail, "local"))
@@ -0,0 +1,170 @@
+#!/usr/bin/env python3
+"""Check NVIDIA driver / CUDA / PyTorch on the VM. Prints one JSON object."""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+COMFY_PIPS = [
+    Path("/opt/swarmui/dlbackend/comfy/venv/bin/python"),
+    Path("/opt/swarmui/dlbackend/comfy/ComfyUI/venv/bin/python"),
+]
+
+
+def run(argv: list[str], timeout: float = 30.0) -> tuple[int, str]:
+    try:
+        out = subprocess.check_output(
+            argv, text=True, stderr=subprocess.STDOUT, timeout=timeout
+        )
+        return 0, out.strip()
+    except subprocess.CalledProcessError as exc:
+        return exc.returncode, (exc.output or "").strip()
+    except (OSError, subprocess.TimeoutExpired) as exc:
+        return 1, str(exc)
+
+
+def check_nvidia() -> dict:
+    code, out = run(["nvidia-smi", "-L"])
+    if code != 0 or not out:
+        return {"ok": False, "detail": out or "nvidia-smi missing/failed"}
+    # Driver / CUDA from nvidia-smi header
+    code2, q = run(
+        [
+            "nvidia-smi",
+            "--query-gpu=driver_version,name,memory.total",
+            "--format=csv,noheader",
+        ]
+    )
+    detail = out.splitlines()[0] if out else "ok"
+    if code2 == 0 and q:
+        detail = q.splitlines()[0].strip()
+    return {"ok": True, "detail": detail}
+
+
+def check_cuda_runtime() -> dict:
+    # libcuda present?
+    code, out = run(["bash", "-lc", "ldconfig -p 2>/dev/null | grep -E 'libcuda\\.so' | head -n1"])
+    if code == 0 and out:
+        return {"ok": True, "detail": out.strip()}
+    # Fallback: nvidia-smi reports CUDA Version
+    code2, smi = run(["nvidia-smi"])
+    if code2 == 0 and "CUDA Version" in smi:
+        for line in smi.splitlines():
+            if "CUDA Version" in line:
+                return {"ok": True, "detail": line.strip()[:120]}
+    return {"ok": False, "detail": "libcuda.so не найден (драйвер/CUDA runtime)"}
+
+
+def find_comfy_python() -> Path | None:
+    for p in COMFY_PIPS:
+        if p.is_file() and os.access(p, os.X_OK):
+            return p
+    return None
+
+
+def check_torch(py: Path) -> dict:
+    script = (
+        "import json,sys\n"
+        "try:\n"
+        " import torch\n"
+        "except Exception as e:\n"
+        " print(json.dumps({'ok':False,'detail':f'import torch: {e}'})); sys.exit(0)\n"
+        "cuda=bool(torch.cuda.is_available())\n"
+        "dev=torch.cuda.get_device_name(0) if cuda else ''\n"
+        "print(json.dumps({\n"
+        " 'ok': cuda,\n"
+        " 'detail': (\n"
+        "  f'torch={torch.__version__} cuda={torch.version.cuda} '\n"
+        "  f'available={cuda} device={dev}'\n"
+        " ),\n"
+        "}))\n"
+    )
+    code, out = run([str(py), "-c", script], timeout=120.0)
+    line = ""
+    for row in reversed((out or "").splitlines()):
+        row = row.strip()
+        if row.startswith("{"):
+            line = row
+            break
+    if not line:
+        return {"ok": False, "detail": f"torch probe fail: {(out or '')[:200]}"}
+    try:
+        return json.loads(line)
+    except json.JSONDecodeError:
+        return {"ok": False, "detail": f"bad torch JSON: {line[:160]}"}
+
+
+def check_optional_pip(py: Path, names: list[str]) -> dict:
+    missing = []
+    present = []
+    for name in names:
+        code, _ = run([str(py), "-c", f"import {name}"], timeout=30.0)
+        if code == 0:
+            present.append(name)
+        else:
+            missing.append(name)
+    return {
+        "ok": True,  # optional — never hard-fail
+        "detail": (
+            f"есть: {', '.join(present) or '—'}; нет: {', '.join(missing) or '—'}"
+        ),
+        "missing": missing,
+        "present": present,
+    }
+
+
+def main() -> int:
+    want_swarm = (os.environ.get("GPU_RENT_CHECK_SWARM") or "1").strip() != "0"
+    checks: list[dict] = []
+
+    nv = check_nvidia()
+    checks.append({"name": "nvidia-smi", "required": True, **nv})
+
+    cuda = check_cuda_runtime()
+    # CUDA runtime required whenever we have a GPU session
+    checks.append({"name": "cuda", "required": True, **cuda})
+
+    if want_swarm:
+        py = find_comfy_python()
+        if py is None:
+            checks.append(
+                {
+                    "name": "torch",
+                    "required": True,
+                    "ok": False,
+                    "detail": "ComfyUI venv python не найден (ещё не поставился?)",
+                }
+            )
+        else:
+            torch_c = check_torch(py)
+            checks.append({"name": "torch", "required": True, "venv": str(py), **torch_c})
+            opt = check_optional_pip(py, ["triton", "sageattention"])
+            checks.append({"name": "triton/sage", "required": False, **opt})
+    else:
+        checks.append(
+            {
+                "name": "torch",
+                "required": False,
+                "ok": True,
+                "detail": "skip (llm-only, без Comfy venv)",
+            }
+        )
+
+    required_ok = all(c.get("ok") for c in checks if c.get("required"))
+    print(
+        json.dumps(
+            {"ok": required_ok, "checks": checks},
+            ensure_ascii=False,
+        )
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    # Always 0 — caller polls JSON `ok` (torch may appear after first Comfy start).
+    main()
+    sys.exit(0)
@@ -22,7 +22,7 @@ from gpu_rent.cloud import (
 )
 from gpu_rent.bootstrap import run_bootstrap
 from gpu_rent.provision import provision_vm, tune_swarm_perf
-from gpu_rent.ready import verify_stack_on_vm, wait_backend_idle
+from gpu_rent.ready import verify_gpu_env, verify_stack_on_vm, wait_backend_idle
 from gpu_rent.snapshot import ensure_boot_snapshot
 from gpu_rent.notify import notify_ready
 from gpu_rent.config import Config
@@ -140,6 +140,18 @@ def _bind_access(
        save_state(state)
        raise

+    try:
+        gpu_checks = verify_gpu_env(cfg, ip, log, timeout=600.0)
+        state.notes = dict(state.notes or {})
+        state.notes["gpu_env"] = [
+            {"name": c.name, "ok": c.ok, "detail": c.detail} for c in gpu_checks
+        ]
+    except CloudError as exc:
+        state.notes = dict(state.notes or {})
+        state.notes["gpu_env_error"] = str(exc)[:500]
+        save_state(state)
+        raise
+
    try:
        ensure_boot_snapshot(
            conn,