From a622265bb52f31f39dcbade0807e4317d22d6c52 Mon Sep 17 00:00:00 2001 From: Leonid Pershin Date: Fri, 21 Aug 2026 09:28:45 +0300 Subject: [PATCH] Enhance ComfyUI installation script with progress tracking and status detection - Added functions for formatting byte sizes, displaying progress bars, and calculating directory sizes to improve user feedback during installation. - Implemented a stage detection mechanism to provide real-time updates on the installation process of ComfyUI. - Enhanced logging to capture recent journal lines related to ComfyUI installation, improving visibility into the installation status. - Updated tests to verify the presence of new functionality in the installation script, ensuring robustness and reliability. --- src/gpu_rent/remote/install_swarm_comfy.py | 216 +++++++++++++++++++-- tests/test_wait_backend_idle.py | 3 + 2 files changed, 205 insertions(+), 14 deletions(-) diff --git a/src/gpu_rent/remote/install_swarm_comfy.py b/src/gpu_rent/remote/install_swarm_comfy.py index 11dd7a4..849d093 100644 --- a/src/gpu_rent/remote/install_swarm_comfy.py +++ b/src/gpu_rent/remote/install_swarm_comfy.py @@ -13,7 +13,9 @@ import json import os import socket import struct +import subprocess import sys +import threading import time import urllib.error import urllib.request @@ -22,12 +24,131 @@ 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" +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( @@ -201,7 +322,7 @@ def _ws_recv_text(sock: socket.socket) -> str | None: def run_install(sid: str) -> None: - print("SwarmUI InstallConfirmWS: backend=comfyui models=none …") + print("SwarmUI InstallConfirmWS: backend=comfyui models=none …", flush=True) sock = _ws_handshake("127.0.0.1", 7801, "/API/InstallConfirmWS") payload = { "session_id": sid, @@ -216,44 +337,111 @@ def run_install(sid: str) -> None: _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(120.0, deadline - time.time()))) + sock.settimeout(max(5.0, min(30.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]}") + 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(f"[installer] {info}") + print(flush=True) # break progress \r line + print(f"[installer] {info}", flush=True) 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 "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"): - print("InstallConfirmWS: 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: diff --git a/tests/test_wait_backend_idle.py b/tests/test_wait_backend_idle.py index 93a09c8..ba87d50 100644 --- a/tests/test_wait_backend_idle.py +++ b/tests/test_wait_backend_idle.py @@ -25,6 +25,9 @@ def test_install_swarm_comfy_script_payload(): assert '"backend": "comfyui"' in text assert '"models": "none"' in text assert "modern_dark" in text + assert "detect_stage" in text + assert "dlbackend=" in text + assert 'end="\\r"' in text or "end=\"\\r\"" in text def test_verify_gpu_env_fail_fast_empty_dlbackend(monkeypatch):