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"))
|
||||
|
||||
Reference in New Issue
Block a user