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.
This commit is contained in:
Leonid Pershin
+112
-9
@@ -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)
|
||||
+13
-1
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user