Leonid Pershin
1785ab369c
- Updated the idle-killer logic to treat SwarmUI `empty` and `disabled` states as busy, preventing unnecessary idle time during provisioning. - Enhanced the `wait_backend_idle` function to recognize suspended backends as ready, improving resource utilization and user feedback. - Refined the `install_swarm_comfy` script to skip installation when backends are already present, streamlining the setup process. - Improved the `resolve_llm_runtime` function to prioritize live configuration over stale state notes, ensuring accurate runtime detection. - Added tests to validate the new backend status handling and idle management logic, ensuring robustness and reliability.
518 lines
18 KiB
Python
518 lines
18 KiB
Python
#!/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 subprocess
|
||
import sys
|
||
import threading
|
||
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")
|
||
DLBACKEND = DATA / "dlbackend"
|
||
COMFY_ROOT = DLBACKEND / "ComfyUI"
|
||
COMFY_VENV = COMFY_ROOT / "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 fmt_bytes(n: float) -> str:
|
||
n = float(n)
|
||
for unit, div in (("GB", 1024**3), ("MB", 1024**2), ("KB", 1024), ("B", 1)):
|
||
if n >= div or unit == "B":
|
||
if unit == "B":
|
||
return f"{int(n)}B"
|
||
return f"{n / div:.1f}{unit}"
|
||
return f"{n:.0f}B"
|
||
|
||
|
||
def progress_bar(done: int, total: int | None, *, width: int = 22) -> str:
|
||
if not total or total <= 0:
|
||
return ""
|
||
pct = min(100.0, 100.0 * done / total)
|
||
filled = min(width, max(0, int(width * done / total)))
|
||
bar = "#" * filled + "-" * (width - filled)
|
||
return f"[{bar}] {pct:5.1f}% "
|
||
|
||
|
||
def dir_size_bytes(root: Path) -> int:
|
||
if not root.is_dir():
|
||
return 0
|
||
try:
|
||
out = subprocess.check_output(
|
||
["du", "-sb", str(root)],
|
||
text=True,
|
||
stderr=subprocess.DEVNULL,
|
||
timeout=15,
|
||
)
|
||
return int(out.split()[0])
|
||
except (subprocess.CalledProcessError, FileNotFoundError, OSError, ValueError, IndexError, subprocess.TimeoutExpired):
|
||
pass
|
||
total = 0
|
||
try:
|
||
for dirpath, _dirnames, filenames in os.walk(root):
|
||
for name in filenames:
|
||
try:
|
||
total += (Path(dirpath) / name).stat().st_size
|
||
except OSError:
|
||
continue
|
||
except OSError:
|
||
return total
|
||
return total
|
||
|
||
|
||
def detect_stage() -> str:
|
||
"""Best-effort stage while Linux comfy-install-linux.sh runs (no WS bytes)."""
|
||
if not COMFY_ROOT.is_dir():
|
||
return "ожидание clone ComfyUI"
|
||
if not (COMFY_ROOT / ".git").is_dir() and not (COMFY_ROOT / "main.py").is_file():
|
||
return "git clone ComfyUI…"
|
||
if not COMFY_VENV.is_file():
|
||
if (COMFY_ROOT / "venv").is_dir():
|
||
return "создание venv…"
|
||
return "ComfyUI clone ok → venv"
|
||
# Prefer filesystem check — importing torch every second is too heavy.
|
||
torch_dirs = list((COMFY_ROOT / "venv").glob("lib/python*/site-packages/torch"))
|
||
if not torch_dirs:
|
||
# pip may still be downloading into cache / tmp
|
||
return "pip install torch (CUDA)…"
|
||
req = COMFY_ROOT / "requirements.txt"
|
||
if req.is_file():
|
||
# Heuristic: if common deps missing, still on requirements.
|
||
missing = False
|
||
for pkg in ("aiohttp", "einops", "safetensors"):
|
||
if not list((COMFY_ROOT / "venv").glob(f"lib/python*/site-packages/{pkg}*")):
|
||
missing = True
|
||
break
|
||
if missing:
|
||
return "pip install requirements…"
|
||
return "Comfy deps почти готовы…"
|
||
|
||
|
||
def _journal_comfy_lines(since_sec: float = 90.0) -> list[str]:
|
||
"""Recent swarmui journal lines that look like comfy/pip progress."""
|
||
try:
|
||
out = subprocess.check_output(
|
||
[
|
||
"journalctl",
|
||
"-u",
|
||
"swarmui",
|
||
"--no-pager",
|
||
"-o",
|
||
"cat",
|
||
f"--since={int(max(1, since_sec))} seconds ago",
|
||
],
|
||
text=True,
|
||
stderr=subprocess.DEVNULL,
|
||
timeout=8,
|
||
)
|
||
except (subprocess.CalledProcessError, FileNotFoundError, OSError, subprocess.TimeoutExpired):
|
||
return []
|
||
keys = (
|
||
"comfy",
|
||
"torch",
|
||
"Downloading",
|
||
"install",
|
||
"Requirement",
|
||
"Collecting",
|
||
"Cloning",
|
||
"Making venv",
|
||
"Installation completed",
|
||
"pip",
|
||
)
|
||
lines: list[str] = []
|
||
for raw in out.splitlines():
|
||
line = raw.strip()
|
||
if not line or len(line) < 8:
|
||
continue
|
||
low = line.lower()
|
||
if any(k.lower() in low for k in keys):
|
||
# Strip noisy prefixes
|
||
if "[ComfyUI Install" in line or "STDOUT" in line or "pip" in low or "torch" in low:
|
||
lines.append(line[-180:])
|
||
return lines[-6:]
|
||
|
||
|
||
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 …", flush=True)
|
||
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 = ""
|
||
last_journal: set[str] = set()
|
||
stop = threading.Event()
|
||
t0 = time.monotonic()
|
||
size0 = dir_size_bytes(DLBACKEND)
|
||
last_size = size0
|
||
last_size_t = t0
|
||
# WS byte progress (Windows 7z / model downloads); Linux script usually silent.
|
||
ws_prog: dict = {"done": 0, "total": 0, "steps": 0, "total_steps": 0, "t": 0.0}
|
||
|
||
def emit_progress_line(*, final: bool = False) -> None:
|
||
nonlocal last_size, last_size_t
|
||
now = time.monotonic()
|
||
elapsed = max(now - t0, 0.001)
|
||
size = dir_size_bytes(DLBACKEND)
|
||
dt = max(now - last_size_t, 0.001)
|
||
speed = (size - last_size) / dt if size >= last_size else 0.0
|
||
last_size, last_size_t = size, now
|
||
stage = detect_stage()
|
||
grown = max(0, size - size0)
|
||
parts = [
|
||
f"[installer] {stage}",
|
||
f"dlbackend={fmt_bytes(size)}",
|
||
f"+{fmt_bytes(grown)}",
|
||
f"{fmt_bytes(speed)}/s",
|
||
f"{int(elapsed)}s",
|
||
]
|
||
done = int(ws_prog.get("done") or 0)
|
||
total = int(ws_prog.get("total") or 0)
|
||
if done > 0 or total > 0:
|
||
bar = progress_bar(done, total if total > 0 else None)
|
||
steps = ws_prog.get("steps")
|
||
total_steps = ws_prog.get("total_steps")
|
||
parts.insert(
|
||
1,
|
||
f"step {steps}/{total_steps} {bar}{fmt_bytes(done)}"
|
||
+ (f"/{fmt_bytes(total)}" if total else ""),
|
||
)
|
||
line = " · ".join(parts)
|
||
if final:
|
||
print(line, flush=True)
|
||
else:
|
||
print(line, end="\r", flush=True)
|
||
|
||
def monitor() -> None:
|
||
while not stop.wait(1.0):
|
||
emit_progress_line()
|
||
# Surface interesting journal lines as real log lines (newline).
|
||
for jline in _journal_comfy_lines(120.0):
|
||
if jline in last_journal:
|
||
continue
|
||
last_journal.add(jline)
|
||
# End in-place progress then print journal snippet.
|
||
print(flush=True)
|
||
print(f" … {jline}", flush=True)
|
||
|
||
mon = threading.Thread(target=monitor, name="comfy-install-progress", daemon=True)
|
||
mon.start()
|
||
try:
|
||
while time.time() < deadline:
|
||
try:
|
||
sock.settimeout(max(5.0, min(30.0, deadline - time.time())))
|
||
msg = _ws_recv_text(sock)
|
||
except (TimeoutError, socket.timeout):
|
||
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]}", flush=True)
|
||
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(flush=True) # break progress \r line
|
||
print(f"[installer] {info}", flush=True)
|
||
last_info = info
|
||
if "progress" in data:
|
||
try:
|
||
ws_prog["done"] = int(data.get("progress") or 0)
|
||
ws_prog["total"] = int(data.get("total") or 0)
|
||
ws_prog["steps"] = int(data.get("steps") or 0)
|
||
ws_prog["total_steps"] = int(data.get("total_steps") or 0)
|
||
ws_prog["t"] = time.monotonic()
|
||
except (TypeError, ValueError):
|
||
pass
|
||
# Byte tick from Swarm (rare on Linux) — refresh bar now.
|
||
if ws_prog["done"] or ws_prog["total"]:
|
||
emit_progress_line()
|
||
if data.get("success"):
|
||
stop.set()
|
||
print(flush=True)
|
||
emit_progress_line(final=True)
|
||
print("InstallConfirmWS: success", flush=True)
|
||
return
|
||
raise SystemExit(f"InstallConfirmWS timeout after {int(INSTALL_TIMEOUT)}s")
|
||
finally:
|
||
stop.set()
|
||
try:
|
||
sock.close()
|
||
except OSError:
|
||
pass
|
||
mon.join(timeout=2.0)
|
||
|
||
|
||
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()}")
|
||
|
||
# Backends already registered (any non-empty status) + venv → skip InstallConfirmWS.
|
||
if bstat == "idle":
|
||
print("backends present (idle/suspended) + skip install")
|
||
return 0
|
||
if bstat not in ("empty", "unknown") and not bstat.startswith("error:"):
|
||
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 (ждём ready отдельно)")
|
||
return 0
|
||
|
||
if comfy_venv_ok() and bstat == "empty":
|
||
if installed is not True:
|
||
print("venv есть, IsInstalled=false — запускаю InstallConfirmWS")
|
||
else:
|
||
return 1
|
||
|
||
marker = DATA / ".gpu-rent-comfy-installing"
|
||
try:
|
||
marker.write_text(f"{int(time.time())}\n", encoding="utf-8")
|
||
except OSError:
|
||
pass
|
||
try:
|
||
sid = get_session(time.time() + 120)
|
||
run_install(sid)
|
||
finally:
|
||
try:
|
||
marker.unlink(missing_ok=True)
|
||
except OSError:
|
||
pass
|
||
|
||
# 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())
|