From f7ba915e74f56c79f25f85f25f3615999cac9300 Mon Sep 17 00:00:00 2001 From: Leonid Pershin Date: Fri, 21 Aug 2026 09:18:59 +0300 Subject: [PATCH] Implement headless ComfyUI installation and improve backend status handling - Added a new function to ensure headless installation of ComfyUI when backends are empty, enhancing the setup process for SwarmUI. - Updated documentation to clarify the installation flow and backend readiness checks, ensuring users understand the requirements for a successful setup. - Enhanced backend status checks to differentiate between 'empty' and 'idle' states, improving error handling and user feedback during provisioning. - Adjusted logging messages to provide clearer insights into the installation and backend status processes. --- docs/cli.md | 2 +- docs/swarmui.md | 6 +- src/gpu_rent/provision.py | 14 + src/gpu_rent/ready.py | 24 +- src/gpu_rent/remote/install_swarm_comfy.py | 320 +++++++++++++++++++++ src/gpu_rent/remote/stack_env_probe.py | 12 +- src/gpu_rent/session.py | 14 +- tests/test_session.py | 4 + tests/test_wait_backend_idle.py | 66 +++++ 9 files changed, 454 insertions(+), 8 deletions(-) create mode 100644 src/gpu_rent/remote/install_swarm_comfy.py create mode 100644 tests/test_wait_backend_idle.py diff --git a/docs/cli.md b/docs/cli.md index 7eb0841..cc02121 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -231,7 +231,7 @@ Application credential для idle-killer CLI создаёт на `up` (узко 1. Nova `ACTIVE` 2. TCP 22 / SSH -3. Backend Idle (если SwarmUI) → toast (если `NOTIFY_READY`) +3. First-install Comfy (InstallConfirmWS, если backend был empty) → backend Idle → toast (если `NOTIFY_READY`) 4. На VM: HTTP сервисов стека (SwarmUI `:7801` / Ollama `:11434`) — `verify_stack_on_vm` 5. На VM: **nvidia-smi / CUDA** (fail-fast) + **torch+cuda** в Comfy venv при SwarmUI (ждём) — `verify_gpu_env` 6. В логе: строка **`тайминг up:`** (SSH / bootstrap / Idle / verify / …) diff --git a/docs/swarmui.md b/docs/swarmui.md index 929ff56..a447de4 100644 --- a/docs/swarmui.md +++ b/docs/swarmui.md @@ -21,7 +21,7 @@ 7. Push непустых `Models/` / `Wildcards/` / `CustomWorkflows/`. 8. systemd unit `swarmui`: `./launch-linux.sh --launch_mode none --host 127.0.0.1 --port 7801`. 9. **GPU probe** → `/mnt/swarm_data/.gpu-rent-gpu.json` (VRAM / compute cap / tier) — до старта UI; Ollama читает его при install. -10. Старт SwarmUI → seed LLM → idle-killer → **wait backend Idle**. +10. Старт SwarmUI → **headless InstallConfirmWS** (ComfyUI в `dlbackend`, models=none — веса уже с Civitai) → seed LLM → idle-killer → **wait backend Idle** (только `idle`; `empty` ≠ ready). 11. **Perf tune после Idle** — `triton`+`sageattention` в Comfy venv и `--use-sage-attention` в `Data/Backends.fds` (маркер `.gpu-rent-perf-tuned`; повтор при смене GPU). 12. Авторизация SwarmUI включена, токен в `Data` на диске. 13. Один snapshot boot volume `gpu-rent-boot-ok`, если ещё нет. @@ -32,9 +32,11 @@ Первый запуск качает backend в `/opt/swarmui/dlbackend` (это bind на data volume). Иначе каждый recreate потеряет часы. +**gpu-rent:** после `systemctl start swarmui` вызывается headless `InstallConfirmWS` (`backend=comfyui`, `models=none`), если `GetCurrentStatus` ещё `empty`. Без этого UI отвечает HTTP 200, а Comfy/torch так и не появляются — раньше `wait_backend_idle` ошибочно считал `empty` = Idle. + **Скорость без потери качества (gpu-rent):** после Idle backend, на GPU с compute capability ≥ 8.0 и ≥16 GiB VRAM — SageAttention. `Performance.AllowGpuSpecificOptimizations` у Swarm по умолчанию уже включает `--fast` для 30xx+. Если venv ещё не появился — маркер без `pip_ok`, догонит на следующем `up`. -Пока CUDA/ComfyUI поднимаются, UI уже может отвечать. Для MCP и `/API/` — рано: `ready` после Idle backend (подтвердить на spike). +Пока CUDA/ComfyUI поднимаются, UI уже может отвечать. Для MCP и `/API/` — рано: `ready` после Idle backend. ## Доступ с ноутбука diff --git a/src/gpu_rent/provision.py b/src/gpu_rent/provision.py index c012e80..931c0f0 100644 --- a/src/gpu_rent/provision.py +++ b/src/gpu_rent/provision.py @@ -92,6 +92,20 @@ def tune_swarm_perf(cfg: Config, host: str, log: Log) -> bool: return "RESTART_SWARMUI=1" in out +def ensure_swarm_comfy_installed(cfg: Config, host: str, log: Log) -> None: + """Headless SwarmUI InstallConfirmWS (ComfyUI) when backends are still empty.""" + log("SwarmUI first-install: ComfyUI backend (если ещё empty)…") + run_python( + cfg, + host, + _pkg_text("install_swarm_comfy.py"), + remote_path="/tmp/gpu-rent-install_swarm_comfy.py", + # Cold: git clone Comfy + pip torch — often 15–40 min. + timeout=3900, + log=log, + ) + + def seed_extensions(cfg: Config, host: str, log: Log, *, update: bool = True) -> bool: from gpu_rent.llm_runtime import normalize_runtime diff --git a/src/gpu_rent/ready.py b/src/gpu_rent/ready.py index e5db4d4..e694dae 100644 --- a/src/gpu_rent/ready.py +++ b/src/gpu_rent/ready.py @@ -56,10 +56,16 @@ while time.time() < deadline: bstat = str(be.get("status") or "unknown").lower() if waiting or live or loading: print(f"BUSY queue w={waiting} live={live} load={loading}") - elif bstat not in ("idle", "disabled", "all_disabled", "empty"): + elif bstat == "empty": + # No backends registered — Comfy never installed (Install wizard). + # Not "Idle": treat as wait, never ready. + print("BUSY backend=empty (нужен first-install Comfy)") + elif bstat in ("disabled", "all_disabled"): print(f"BUSY backend={bstat}") - else: + elif bstat == "idle": print(f"READY backend={bstat}") + else: + print(f"BUSY backend={bstat}") raise SystemExit(0) except Exception as exc: print(f"WAIT {exc}") @@ -316,7 +322,13 @@ def verify_gpu_env( from gpu_rent.ssh_ops import run_python # These never "appear later" on a broken image — don't burn the poll budget. + # Empty-backend / missing install also won't self-heal without InstallConfirmWS. instant_fail_names = {"nvidia-smi", "cuda"} + instant_fail_detail_substrings = ( + "dlbackend пуст", + "backend=empty", + "Install не прогоняли", + ) want_swarm = bool(getattr(cfg, "enable_swarmui", True)) script = files("gpu_rent.remote").joinpath("stack_env_probe.py").read_text( @@ -391,6 +403,11 @@ def verify_gpu_env( return last instant = [c for c in hard if c.name in instant_fail_names] + if not instant: + for c in hard: + detail_l = (c.detail or "").lower() + if any(s.lower() in detail_l for s in instant_fail_detail_substrings): + instant.append(c) if instant: for c in last: mark = "ok" if c.ok else "FAIL" @@ -399,7 +416,8 @@ def verify_gpu_env( failed = [c.name for c in instant] raise CloudError( f"GPU-стек: нет {', '.join(failed)} (fail-fast). " - "Проверь образ Driver / nvidia на VM. " + "Проверь образ Driver / nvidia на VM; для SwarmUI — " + "first-install Comfy (InstallConfirmWS / dlbackend). " "Дальше: gpu-rent status · gpu-rent logs · gpu-rent stop" ) return last diff --git a/src/gpu_rent/remote/install_swarm_comfy.py b/src/gpu_rent/remote/install_swarm_comfy.py new file mode 100644 index 0000000..11dd7a4 --- /dev/null +++ b/src/gpu_rent/remote/install_swarm_comfy.py @@ -0,0 +1,320 @@ +#!/usr/bin/env python3 +"""Headless SwarmUI first-install: ComfyUI backend via InstallConfirmWS. + +Runs on the VM (stdlib only). Idempotent: skips when backends exist or +dlbackend/ComfyUI venv is already present and IsInstalled. +""" + +from __future__ import annotations + +import base64 +import hashlib +import json +import os +import socket +import struct +import sys +import time +import urllib.error +import urllib.request +from pathlib import Path + +SWARM = "http://127.0.0.1:7801" +DATA = Path("/mnt/swarm_data") +SWARM_ROOT = Path(os.environ.get("SWARM_ROOT") or "/opt/swarmui") +COMFY_VENV = DATA / "dlbackend" / "ComfyUI" / "venv" / "bin" / "python" +SETTINGS = DATA / "Data" / "Settings.fds" +# Comfy clone + torch can take 20–40+ min on a cold disk. +INSTALL_TIMEOUT = float(os.environ.get("GPU_RENT_COMFY_INSTALL_TIMEOUT") or "3600") + + +def post(path: str, payload: dict, timeout: float = 15.0) -> dict: + body = json.dumps(payload).encode("utf-8") + req = urllib.request.Request( + f"{SWARM}{path}", + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=timeout) as resp: + return json.loads(resp.read().decode("utf-8")) + + +def wait_http(deadline: float) -> None: + last = "" + while time.time() < deadline: + try: + req = urllib.request.Request(f"{SWARM}/", method="GET") + with urllib.request.urlopen(req, timeout=5) as resp: + if getattr(resp, "status", 200) == 200: + return + except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, OSError) as exc: + last = str(exc)[:160] + time.sleep(2) + raise SystemExit(f"SwarmUI HTTP not up: {last}") + + +def get_session(deadline: float) -> str: + last = "" + while time.time() < deadline: + try: + data = post("/API/GetNewSession", {}) + sid = str(data.get("session_id") or "") + if sid: + return sid + last = "no session_id" + except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, OSError, json.JSONDecodeError) as exc: + last = str(exc)[:160] + time.sleep(2) + raise SystemExit(f"SwarmUI session unavailable: {last}") + + +def backend_status() -> str: + try: + sid = str(post("/API/GetNewSession", {}).get("session_id") or "") + if not sid: + return "unknown" + data = post("/API/GetCurrentStatus", {"session_id": sid}) + be = data.get("backend_status") or {} + return str(be.get("status") or "unknown").lower() + except Exception as exc: + return f"error:{exc}" + + +def settings_is_installed() -> bool | None: + if not SETTINGS.is_file(): + return None + try: + text = SETTINGS.read_text(encoding="utf-8", errors="replace") + except OSError: + return None + for line in text.splitlines(): + if "IsInstalled" in line: + low = line.lower() + if "true" in low: + return True + if "false" in low: + return False + return None + + +def comfy_venv_ok() -> bool: + return COMFY_VENV.is_file() and os.access(COMFY_VENV, os.X_OK) + + +def _recv_exact(sock: socket.socket, n: int) -> bytes: + buf = b"" + while len(buf) < n: + chunk = sock.recv(n - len(buf)) + if not chunk: + raise ConnectionError("websocket closed") + buf += chunk + return buf + + +def _ws_handshake(host: str, port: int, path: str) -> socket.socket: + key = base64.b64encode(os.urandom(16)).decode("ascii") + req = ( + f"GET {path} HTTP/1.1\r\n" + f"Host: {host}:{port}\r\n" + "Upgrade: websocket\r\n" + "Connection: Upgrade\r\n" + f"Sec-WebSocket-Key: {key}\r\n" + "Sec-WebSocket-Version: 13\r\n" + "\r\n" + ).encode("ascii") + sock = socket.create_connection((host, port), timeout=30) + sock.sendall(req) + # Read headers until blank line + data = b"" + while b"\r\n\r\n" not in data: + chunk = sock.recv(4096) + if not chunk: + sock.close() + raise ConnectionError("no WS handshake response") + data += chunk + if len(data) > 65536: + sock.close() + raise ConnectionError("WS handshake too large") + head = data.split(b"\r\n\r\n", 1)[0].decode("latin-1", errors="replace") + if "101" not in head.split("\r\n", 1)[0]: + sock.close() + raise ConnectionError(f"WS handshake failed: {head.splitlines()[0]}") + expect = base64.b64encode( + hashlib.sha1( + (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode("ascii") + ).digest() + ).decode("ascii") + if expect not in head: + # Some servers omit echoing exact key in proxies; 101 is enough. + pass + sock.settimeout(120.0) + return sock + + +def _ws_send_text(sock: socket.socket, text: str) -> None: + payload = text.encode("utf-8") + mask = os.urandom(4) + header = bytearray([0x81]) # FIN + text + n = len(payload) + if n < 126: + header.append(0x80 | n) + elif n < 65536: + header.append(0x80 | 126) + header.extend(struct.pack("!H", n)) + else: + header.append(0x80 | 127) + header.extend(struct.pack("!Q", n)) + header.extend(mask) + masked = bytes(b ^ mask[i % 4] for i, b in enumerate(payload)) + sock.sendall(header + masked) + + +def _ws_recv_text(sock: socket.socket) -> str | None: + """Return next text frame payload, or None on close.""" + while True: + hdr = _recv_exact(sock, 2) + opcode = hdr[0] & 0x0F + masked = bool(hdr[1] & 0x80) + length = hdr[1] & 0x7F + if length == 126: + length = struct.unpack("!H", _recv_exact(sock, 2))[0] + elif length == 127: + length = struct.unpack("!Q", _recv_exact(sock, 8))[0] + mask_key = _recv_exact(sock, 4) if masked else b"" + raw = _recv_exact(sock, length) if length else b"" + if masked: + raw = bytes(b ^ mask_key[i % 4] for i, b in enumerate(raw)) + if opcode == 0x8: # close + return None + if opcode == 0x9: # ping → pong + # build pong + frame = bytearray([0x8A, 0x80 | len(raw)]) + m = os.urandom(4) + frame.extend(m) + frame.extend(bytes(b ^ m[i % 4] for i, b in enumerate(raw))) + sock.sendall(frame) + continue + if opcode in (0x1, 0x0): # text / continuation + return raw.decode("utf-8", errors="replace") + # ignore binary / other + + +def run_install(sid: str) -> None: + print("SwarmUI InstallConfirmWS: backend=comfyui models=none …") + sock = _ws_handshake("127.0.0.1", 7801, "/API/InstallConfirmWS") + payload = { + "session_id": sid, + "theme": "modern_dark", + "installed_for": "just_self", + "backend": "comfyui", + "models": "none", + "install_amd": False, + "language": "en", + "make_shortcut": False, + } + _ws_send_text(sock, json.dumps(payload)) + deadline = time.time() + INSTALL_TIMEOUT + last_info = "" + try: + while time.time() < deadline: + try: + sock.settimeout(max(5.0, min(120.0, deadline - time.time()))) + msg = _ws_recv_text(sock) + except (TimeoutError, socket.timeout): + print("… install still running (ws idle)") + continue + if msg is None: + raise SystemExit("InstallConfirmWS closed without success") + try: + data = json.loads(msg) + except json.JSONDecodeError: + print(f"ws raw: {msg[:200]}") + continue + if data.get("error"): + raise SystemExit(f"InstallConfirmWS error: {data['error']}") + if data.get("info"): + info = str(data["info"]) + if info != last_info: + print(f"[installer] {info}") + last_info = info + if "progress" in data and data.get("progress"): + steps = data.get("steps") + total_steps = data.get("total_steps") + print( + f"[installer] progress step={steps}/{total_steps} " + f"bytes={data.get('progress')}/{data.get('total')}" + ) + if data.get("success"): + print("InstallConfirmWS: success") + return + raise SystemExit(f"InstallConfirmWS timeout after {int(INSTALL_TIMEOUT)}s") + finally: + try: + sock.close() + except OSError: + pass + + +def main() -> int: + # WorkingDirectory for comfy-install-linux.sh is Swarm root when launched by + # Installation.cs; InstallConfirmWS handles that for us. + wait_http(time.time() + 180) + bstat = backend_status() + print(f"backend_status={bstat} venv={'yes' if comfy_venv_ok() else 'no'} " + f"IsInstalled={settings_is_installed()}") + + if bstat == "idle": + print("Comfy backend already Idle — skip install") + return 0 + if bstat not in ("empty", "unknown") and not bstat.startswith("error:"): + # loading / running / etc. — installer not needed + if comfy_venv_ok(): + print(f"backends present ({bstat}) + venv — skip install") + return 0 + + installed = settings_is_installed() + if installed is True and not comfy_venv_ok() and bstat == "empty": + print( + "WARN: Settings IsInstalled=true but backends empty and no Comfy venv. " + "Open SwarmUI → Server → Backends and add ComfyUI Self-Starting, " + "or delete Data/Settings.fds IsInstalled and re-run up." + ) + return 1 + + if installed is True and comfy_venv_ok(): + print("IsInstalled + venv — skip InstallConfirmWS (ждём Idle отдельно)") + return 0 + + if comfy_venv_ok() and bstat == "empty": + # Partial: script ran but backend never registered — still need install + # only if IsInstalled is false; otherwise user must add backend in UI. + if installed is not True: + print("venv есть, IsInstalled=false — запускаю InstallConfirmWS") + else: + return 1 + + sid = get_session(time.time() + 120) + run_install(sid) + + # Confirm outcome + time.sleep(3) + bstat2 = backend_status() + print( + f"after install: backend_status={bstat2} " + f"venv={'yes' if comfy_venv_ok() else 'no'}" + ) + if not comfy_venv_ok() and bstat2 == "empty": + print("FAIL: install finished but still empty / no venv", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + # Ensure Installation.cs cwd paths resolve if anything shells out relative — + # InstallConfirmWS itself cds via WorkingDirectory of the service. + try: + os.chdir(str(SWARM_ROOT)) + except OSError: + pass + raise SystemExit(main()) diff --git a/src/gpu_rent/remote/stack_env_probe.py b/src/gpu_rent/remote/stack_env_probe.py index e58fcd9..872bf55 100644 --- a/src/gpu_rent/remote/stack_env_probe.py +++ b/src/gpu_rent/remote/stack_env_probe.py @@ -80,6 +80,16 @@ def find_comfy_python() -> Path | None: return None +def _dlbackend_empty_hint() -> str: + dl = DATA / "dlbackend" + try: + if not dl.is_dir() or not any(dl.iterdir()): + return "dlbackend пуст — SwarmUI Install не прогоняли (backend=empty)" + except OSError: + pass + return "ещё не поставился?" + + def check_torch(py: Path) -> dict: script = ( "import json,sys\n" @@ -150,7 +160,7 @@ def main() -> int: "name": "torch", "required": True, "ok": False, - "detail": "ComfyUI venv python не найден (ещё не поставился?)", + "detail": f"ComfyUI venv python не найден ({_dlbackend_empty_hint()})", } ) else: diff --git a/src/gpu_rent/session.py b/src/gpu_rent/session.py index ac40e1d..2c7e6d5 100644 --- a/src/gpu_rent/session.py +++ b/src/gpu_rent/session.py @@ -21,7 +21,12 @@ from gpu_rent.cloud import ( wait_volume, ) from gpu_rent.bootstrap import run_bootstrap -from gpu_rent.provision import provision_vm, seed_swarmui_api_keys, tune_swarm_perf +from gpu_rent.provision import ( + ensure_swarm_comfy_installed, + provision_vm, + seed_swarmui_api_keys, + tune_swarm_perf, +) 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 @@ -152,10 +157,17 @@ def _bind_access( ) clock.mark("provision", log) if swarm: + try: + ensure_swarm_comfy_installed(cfg, ip, log) + except Exception as exc: + log(f"SwarmUI Comfy install: {exc}") + raise + clock.mark("comfy-install", log) try: wait_backend_idle(cfg, ip, log) except CloudError as exc: log(f"ready: {exc}") + raise clock.mark("Idle", log) try: if tune_swarm_perf(cfg, ip, log): diff --git a/tests/test_session.py b/tests/test_session.py index 7f99ede..8dcfc57 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -47,6 +47,10 @@ def _mock_bind(monkeypatch): lambda cfg, host, log, update=True, light=False: None, ) monkeypatch.setattr("gpu_rent.session.provision_vm", lambda cfg, host, log, **kw: None) + monkeypatch.setattr( + "gpu_rent.session.ensure_swarm_comfy_installed", + lambda cfg, host, log: None, + ) monkeypatch.setattr("gpu_rent.session.wait_backend_idle", lambda cfg, host, log, **kw: None) monkeypatch.setattr( "gpu_rent.session.verify_stack_on_vm", diff --git a/tests/test_wait_backend_idle.py b/tests/test_wait_backend_idle.py new file mode 100644 index 0000000..93a09c8 --- /dev/null +++ b/tests/test_wait_backend_idle.py @@ -0,0 +1,66 @@ +"""Tests for wait_backend_idle READY semantics and install helpers.""" + +from gpu_rent.ready import _REMOTE_POLL + + +def test_remote_poll_empty_is_busy_not_ready(): + assert 'bstat == "empty"' in _REMOTE_POLL + assert 'BUSY backend=empty' in _REMOTE_POLL + # Must not treat empty as READY anymore + assert '("idle", "disabled", "all_disabled", "empty")' not in _REMOTE_POLL + + +def test_remote_poll_idle_is_ready(): + assert 'bstat == "idle"' in _REMOTE_POLL + assert "READY backend=" in _REMOTE_POLL + + +def test_install_swarm_comfy_script_payload(): + from importlib.resources import files + + text = files("gpu_rent.remote").joinpath("install_swarm_comfy.py").read_text( + encoding="utf-8" + ) + assert "InstallConfirmWS" in text + assert '"backend": "comfyui"' in text + assert '"models": "none"' in text + assert "modern_dark" in text + + +def test_verify_gpu_env_fail_fast_empty_dlbackend(monkeypatch): + import json + + from gpu_rent.errors import CloudError + from gpu_rent.ready import verify_gpu_env + import gpu_rent.ssh_ops as ssh_ops + + class Cfg: + enable_swarmui = True + + payload = { + "ok": False, + "checks": [ + {"name": "nvidia-smi", "required": True, "ok": True, "detail": "ok"}, + {"name": "cuda", "required": True, "ok": True, "detail": "ok"}, + { + "name": "torch", + "required": True, + "ok": False, + "detail": "ComfyUI venv python не найден (dlbackend пуст — SwarmUI Install не прогоняли)", + }, + ], + } + calls = {"n": 0} + + def fake(*a, **k): + calls["n"] += 1 + return json.dumps(payload) + + monkeypatch.setattr(ssh_ops, "run_python", fake) + logs: list[str] = [] + try: + verify_gpu_env(Cfg(), "1.2.3.4", logs.append, timeout=600.0, poll_every=0.1) + assert False, "expected CloudError" + except CloudError as exc: + assert "fail-fast" in str(exc).lower() or "torch" in str(exc).lower() + assert calls["n"] == 1