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

@@ -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"):
+            if "progress" in data:
-                steps = data.get("steps")
+                try:
-                total_steps = data.get("total_steps")
+                    ws_prog["done"] = int(data.get("progress") or 0)
-                print(
+                    ws_prog["total"] = int(data.get("total") or 0)
-                    f"[installer] progress step={steps}/{total_steps} "
+                    ws_prog["steps"] = int(data.get("steps") or 0)
-                    f"bytes={data.get('progress')}/{data.get('total')}"
+                    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:

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