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

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

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.
 
 ## Доступ с ноутбука
 

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
 

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

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

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:

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

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",

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