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.

src/gpu_rent/remote/stack_env_probe.py

@@ -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)