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.
This commit is contained in:
Leonid Pershin
@@ -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())
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user