Enhance SwarmUI integration and GPU environment verification

- Updated CLI documentation to reflect the new handling of `CIVITAI_API_TOKEN`, which is now automatically passed to SwarmUI user settings during startup.
- Improved the `render_access_panel` function to include additional warnings for idle-killer failures and stack errors, enhancing user feedback.
- Introduced a new function `seed_swarmui_api_keys` to manage API key injection into SwarmUI, ensuring seamless integration with the Model Downloader.
- Enhanced GPU environment verification logic to include fail-fast checks for critical components like CUDA, improving error handling and user notifications.
- Updated tests to validate the new API key handling and access panel behavior, ensuring robustness in the integration process.
This commit is contained in:
Leonid Pershin
2026-08-21 07:09:20 +03:00
parent 1ec615c03e
commit adba4976ee
20 changed files with 657 additions and 58 deletions
+34 -5
View File
@@ -17,6 +17,7 @@ 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_ssh
from gpu_rent.timing import WaitLog
Log = Callable[[str], None]
@@ -266,6 +267,7 @@ def verify_stack_on_vm(
deadline = time.time() + timeout
last: list[ServiceCheck] = []
wait = WaitLog(log, every=30.0)
while time.time() < deadline:
try:
last = _probe_vm_once(cfg, host)
@@ -277,7 +279,7 @@ def verify_stack_on_vm(
log("проверка VM: всё отвечает")
return last
bad = ", ".join(f"{c.name}={c.detail}" for c in last if not c.ok) or "пусто"
log(f" … ещё нет: {bad}")
wait.tick(f" … ещё нет: {bad}")
time.sleep(poll_every)
for c in last:
@@ -320,11 +322,18 @@ def verify_gpu_env(
poll_every: float = 15.0,
raise_on_fail: bool = True,
) -> list[ServiceCheck]:
"""nvidia-smi / CUDA / torch(+cuda) in Comfy venv when SwarmUI is on."""
"""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.
instant_fail_names = {"nvidia-smi", "cuda"}
want_swarm = bool(getattr(cfg, "enable_swarmui", True))
script = files("gpu_rent.remote").joinpath("stack_env_probe.py").read_text(
encoding="utf-8"
@@ -338,6 +347,7 @@ def verify_gpu_env(
deadline = time.time() + timeout
last: list[ServiceCheck] = []
wait = WaitLog(log, every=30.0)
while time.time() < deadline:
try:
out = run_python(
@@ -350,7 +360,7 @@ def verify_gpu_env(
)
except Exception as exc:
last = [ServiceCheck("gpu-env", False, str(exc)[:200], "vm")]
log(f" … gpu-env: {exc}")
wait.tick(f" … gpu-env: {exc}")
time.sleep(poll_every)
continue
@@ -366,6 +376,7 @@ def verify_gpu_env(
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
@@ -395,8 +406,22 @@ def verify_gpu_env(
log("проверка GPU-стека: ок")
return last
instant = [c for c in hard if c.name in instant_fail_names]
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. "
"Дальше: gpu-rent status · gpu-rent logs · gpu-rent stop"
)
return last
bad = ", ".join(f"{c.name}={c.detail}" for c in hard)
log(f"ещё нет: {bad}")
wait.tick(f"ждём torch/Comfy: {bad}")
time.sleep(poll_every)
for c in last:
@@ -407,7 +432,8 @@ def verify_gpu_env(
raise CloudError(
f"GPU-стек не готов за {int(timeout)} с: {', '.join(failed)}. "
"Нужны nvidia-smi, CUDA; для SwarmUI — torch с cuda в Comfy venv "
"(journalctl -u swarmui / первый старт backend)."
"(journalctl -u swarmui / первый старт backend). "
"Дальше: gpu-rent status · gpu-rent logs · gpu-rent stop"
)
return last
@@ -446,6 +472,7 @@ def verify_stack_local(
)
deadline = time.time() + timeout
last: list[ServiceCheck] = []
wait = WaitLog(log, every=15.0)
while time.time() < deadline:
last = []
for name, port, url in targets:
@@ -479,6 +506,8 @@ def verify_stack_local(
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: