Leonid Pershin
f8bbde8ac0
- Improved the `install_swarm_comfy` function to handle empty backend states more effectively, introducing recovery mechanisms and enhanced logging for better visibility. - Updated the `tune_swarm_perf` function to always sanitize backend FDS corruption, ensuring consistent performance tuning. - Added new tests to validate the functionality of backend recovery and FDS sanitization, ensuring robustness in backend management.
718 lines
25 KiB
Python
718 lines
25 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 re
|
||
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"
|
||
BACKENDS_FDS = DATA / "Data" / "Backends.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 backend_status_detail() -> tuple[str, str]:
|
||
"""Return (status, message) from GetCurrentStatus."""
|
||
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(),
|
||
str(be.get("message") or "").strip(),
|
||
)
|
||
except Exception as exc:
|
||
return f"error:{exc}", ""
|
||
|
||
|
||
def restart_backends(sid: str, *, which: str = "all") -> dict:
|
||
"""POST /API/RestartBackends — recovers ERRORED Comfy self-start."""
|
||
return post(
|
||
"/API/RestartBackends",
|
||
{"session_id": sid, "backend": which},
|
||
timeout=180.0,
|
||
)
|
||
|
||
|
||
def recover_errored_backends(*, wait_sec: float = 180.0) -> str:
|
||
"""RestartBackends then poll until not errored (or timeout). Returns last status."""
|
||
bstat, msg = backend_status_detail()
|
||
print(
|
||
f"backend errored — RestartBackends(all)"
|
||
+ (f" ({msg[:120]})" if msg else ""),
|
||
flush=True,
|
||
)
|
||
sid = get_session(time.time() + 60)
|
||
try:
|
||
result = restart_backends(sid)
|
||
print(f"RestartBackends: {result}", flush=True)
|
||
except (urllib.error.URLError, urllib.error.HTTPError, TimeoutError, OSError, json.JSONDecodeError) as exc:
|
||
print(f"RestartBackends FAIL: {exc}", flush=True)
|
||
return bstat
|
||
|
||
deadline = time.time() + wait_sec
|
||
last = "errored"
|
||
while time.time() < deadline:
|
||
time.sleep(8)
|
||
last, msg2 = backend_status_detail()
|
||
extra = f" — {msg2[:100]}" if msg2 else ""
|
||
print(f"after restart: backend_status={last}{extra}", flush=True)
|
||
if last != "errored" and not last.startswith("error:"):
|
||
return last
|
||
return last
|
||
|
||
|
||
def sanitize_backends_fds() -> bool:
|
||
"""Repair ExtraArgs: \\x --flag (FDS empty + appended flag). Return True if changed."""
|
||
path = DATA / "Data" / "Backends.fds"
|
||
if not path.is_file():
|
||
return False
|
||
try:
|
||
text = path.read_text(encoding="utf-8")
|
||
except OSError:
|
||
return False
|
||
new, n = re.subn(r"^(\s*ExtraArgs:\s*)\\x(\s+)", r"\1", text, flags=re.M)
|
||
if not n:
|
||
return False
|
||
path.write_text(new, encoding="utf-8")
|
||
print(f"sanitized Backends.fds ExtraArgs \\x corruption ({n} line(s))", flush=True)
|
||
return True
|
||
|
||
|
||
def restart_swarmui_local() -> None:
|
||
print("systemctl restart swarmui (после fix Backends.fds)", flush=True)
|
||
try:
|
||
subprocess.run(
|
||
["sudo", "-n", "systemctl", "restart", "swarmui"],
|
||
check=False,
|
||
timeout=120,
|
||
)
|
||
except (OSError, subprocess.TimeoutExpired) as exc:
|
||
print(f"WARN restart swarmui: {exc}", flush=True)
|
||
wait_http(time.time() + 180)
|
||
|
||
|
||
def list_backends(sid: str) -> dict:
|
||
try:
|
||
return post(
|
||
"/API/ListBackends",
|
||
{"session_id": sid, "nonreal": False, "full_data": False},
|
||
timeout=30.0,
|
||
)
|
||
except Exception as exc:
|
||
return {"error": str(exc)}
|
||
|
||
|
||
def add_comfy_selfstart(sid: str) -> dict:
|
||
return post(
|
||
"/API/AddNewBackend",
|
||
{"session_id": sid, "type_id": "comfyui_selfstart"},
|
||
timeout=60.0,
|
||
)
|
||
|
||
|
||
def recover_empty_backends() -> str:
|
||
"""When API says empty but venv/IsInstalled exist — fix FDS and/or AddNewBackend."""
|
||
changed = sanitize_backends_fds()
|
||
if changed:
|
||
restart_swarmui_local()
|
||
bstat, _ = backend_status_detail()
|
||
if bstat not in ("empty", "unknown") and not bstat.startswith("error:"):
|
||
print(f"after sanitize: backend_status={bstat}", flush=True)
|
||
return bstat
|
||
|
||
sid = get_session(time.time() + 60)
|
||
backends = list_backends(sid)
|
||
if isinstance(backends, dict) and any(
|
||
isinstance(v, dict) and v.get("type") for v in backends.values()
|
||
):
|
||
bstat, _ = backend_status_detail()
|
||
return bstat
|
||
|
||
print("ListBackends empty — AddNewBackend comfyui_selfstart", flush=True)
|
||
try:
|
||
result = add_comfy_selfstart(sid)
|
||
print(f"AddNewBackend: {result}", flush=True)
|
||
except Exception as exc:
|
||
print(f"AddNewBackend FAIL: {exc}", flush=True)
|
||
return "empty"
|
||
|
||
# Default StartScript may be dlbackend/comfy/...; our tree is dlbackend/ComfyUI.
|
||
time.sleep(2)
|
||
bstat, _ = backend_status_detail()
|
||
return bstat
|
||
|
||
|
||
def run_diagnostics() -> None:
|
||
"""Best-effort: prefer uploaded swarm_diag.py, else journalctl snippet."""
|
||
diag = Path("/tmp/gpu-rent-swarm_diag.py")
|
||
if diag.is_file():
|
||
try:
|
||
subprocess.run(
|
||
[sys.executable, str(diag)],
|
||
check=False,
|
||
timeout=90,
|
||
)
|
||
return
|
||
except (OSError, subprocess.TimeoutExpired) as exc:
|
||
print(f"diag script fail: {exc}", flush=True)
|
||
print("=== fallback journalctl -u swarmui -n 60 ===", flush=True)
|
||
try:
|
||
out = subprocess.check_output(
|
||
["journalctl", "-u", "swarmui", "-n", "60", "--no-pager"],
|
||
text=True,
|
||
stderr=subprocess.STDOUT,
|
||
timeout=30,
|
||
errors="replace",
|
||
)
|
||
print(out, flush=True)
|
||
except (OSError, subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
|
||
print(f"journalctl fail: {exc}", flush=True)
|
||
|
||
|
||
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, bmsg = backend_status_detail()
|
||
print(
|
||
f"backend_status={bstat} venv={'yes' if comfy_venv_ok() else 'no'} "
|
||
f"IsInstalled={settings_is_installed()}"
|
||
+ (f" msg={bmsg[:160]}" if bmsg else ""),
|
||
flush=True,
|
||
)
|
||
|
||
# Errored backends are registered but dead — restart, don't skip as "present".
|
||
if bstat == "errored":
|
||
bstat = recover_errored_backends(wait_sec=180.0)
|
||
if bstat == "errored" or bstat.startswith("error:"):
|
||
print(
|
||
"FAIL: backend still errored after RestartBackends — collecting diag",
|
||
file=sys.stderr,
|
||
)
|
||
run_diagnostics()
|
||
return 1
|
||
if bstat in ("running", "idle", "loading", "some_loading") and comfy_venv_ok():
|
||
print(f"recovered to {bstat} — skip InstallConfirmWS")
|
||
return 0
|
||
|
||
# Backends already registered (healthy / loading / suspended) + venv → skip.
|
||
if bstat == "idle":
|
||
print("backends present (idle/suspended) + skip install")
|
||
return 0
|
||
if bstat not in ("empty", "unknown", "errored") 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."
|
||
)
|
||
run_diagnostics()
|
||
return 1
|
||
|
||
if installed is True and comfy_venv_ok() and bstat not in ("empty", "errored"):
|
||
print("IsInstalled + venv — skip InstallConfirmWS (ждём ready отдельно)")
|
||
return 0
|
||
|
||
if comfy_venv_ok() and bstat == "empty":
|
||
print(
|
||
"backend empty при venv/IsInstalled — recover (sanitize FDS / AddNewBackend)…",
|
||
flush=True,
|
||
)
|
||
bstat = recover_empty_backends()
|
||
if bstat not in ("empty", "unknown") and not bstat.startswith("error:"):
|
||
print(f"recovered empty → {bstat}")
|
||
return 0
|
||
if installed is True:
|
||
print(
|
||
"WARN: всё ещё empty после recover — пробую InstallConfirmWS",
|
||
flush=True,
|
||
)
|
||
else:
|
||
print("venv есть, IsInstalled=false — запускаю InstallConfirmWS")
|
||
|
||
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 bstat2 == "errored":
|
||
bstat2 = recover_errored_backends(wait_sec=120.0)
|
||
if bstat2 == "errored":
|
||
print("FAIL: install done but backend errored", file=sys.stderr)
|
||
run_diagnostics()
|
||
return 1
|
||
if not comfy_venv_ok() and bstat2 == "empty":
|
||
print("FAIL: install finished but still empty / no venv", file=sys.stderr)
|
||
run_diagnostics()
|
||
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())
|