Enhance backend idle management and diagnostics in SwarmUI

- Updated the `wait_backend_idle` function to handle disabled backends, introducing a new failure condition and diagnostics collection for better error handling. - Refactored the `install_swarm_comfy` script to include a new `configure_comfy_backend` function, streamlining backend configuration and enabling better management of backend states. - Added tests to validate the new functionality for handling disabled backends and ensuring robust diagnostics, improving overall backend management.
2026-08-21 11:00:01 +03:00
parent f8bbde8ac0
commit a18074c985
3 changed files with 205 additions and 29 deletions
@@ -298,12 +298,88 @@ def list_backends(sid: str) -> dict:
        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 comfy_start_script() -> str:
+    """Relative StartScript Swarm expects (cwd=/opt/swarmui, dlbackend bind-mounted)."""
+    candidates = (
+        "dlbackend/ComfyUI/main.py",
+        "dlbackend/comfy/ComfyUI/main.py",
+        "dlbackend/comfy/main.py",
    )
+    for rel in candidates:
+        if (SWARM_ROOT / rel).is_file() or (DATA / rel).is_file():
+            return rel
+    return "dlbackend/ComfyUI/main.py"
+
+
+def edit_backend(sid: str, backend_id: int, *, title: str, settings: dict) -> dict:
+    return post(
+        "/API/EditBackend",
+        {
+            "session_id": sid,
+            "backend_id": backend_id,
+            "title": title,
+            "settings": settings,
+        },
+        timeout=120.0,
+    )
+
+
+def toggle_backend(sid: str, backend_id: int, *, enabled: bool) -> dict:
+    return post(
+        "/API/ToggleBackend",
+        {
+            "session_id": sid,
+            "backend_id": int(backend_id),
+            "enabled": bool(enabled),
+        },
+        timeout=120.0,
+    )
+
+
+def configure_comfy_backend(sid: str, backend_id: int) -> None:
+    """Set StartScript after AddNewBackend (defaults are empty → stays disabled)."""
+    script = comfy_start_script()
+    settings = {
+        "StartScript": script,
+        "ExtraArgs": "",
+        "DisableInternalArgs": False,
+        "AutoUpdate": "false",
+        "UpdateManagedNodes": "false",
+        "FrontendVersion": "LatestSwarmValidated",
+        "EnablePreviews": "true",
+        "GPU_ID": "0",
+        "OverQueue": 1,
+        "AutoRestart": True,
+    }
+    print(f"EditBackend id={backend_id} StartScript={script}", flush=True)
+    try:
+        result = edit_backend(
+            sid, backend_id, title="ComfyUI Self-Starting", settings=settings
+        )
+        print(f"EditBackend: {result}", flush=True)
+    except Exception as exc:
+        print(f"EditBackend FAIL: {exc}", flush=True)
+        raise
+    try:
+        tog = toggle_backend(sid, backend_id, enabled=True)
+        print(f"ToggleBackend enabled: {tog}", flush=True)
+    except Exception as exc:
+        print(f"ToggleBackend WARN: {exc}", flush=True)
+
+
+def _backend_entries(backends: dict) -> list[tuple[int, dict]]:
+    out: list[tuple[int, dict]] = []
+    if not isinstance(backends, dict):
+        return out
+    for key, val in backends.items():
+        if key in {"error"} or not isinstance(val, dict):
+            continue
+        try:
+            bid = int(val.get("id", key))
+        except (TypeError, ValueError):
+            continue
+        out.append((bid, val))
+    return out


 def recover_empty_backends() -> str:
@@ -312,30 +388,64 @@ def recover_empty_backends() -> str:
    if changed:
        restart_swarmui_local()
        bstat, _ = backend_status_detail()
-        if bstat not in ("empty", "unknown") and not bstat.startswith("error:"):
+        if bstat not in ("empty", "unknown", "disabled", "all_disabled") 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
+    entries = _backend_entries(backends)

-    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"
+    if not entries:
+        print("ListBackends empty — AddNewBackend comfyui_selfstart", flush=True)
+        try:
+            result = add_comfy_selfstart(sid)
+            print(f"AddNewBackend: {result}", flush=True)
+            bid = int(result.get("id", 0))
+        except Exception as exc:
+            print(f"AddNewBackend FAIL: {exc}", flush=True)
+            return "empty"
+        try:
+            configure_comfy_backend(sid, bid)
+        except Exception:
+            return "empty"
+    else:
+        # Backends exist but may lack StartScript / be disabled.
+        for bid, meta in entries:
+            settings = meta.get("settings") or {}
+            script = str(settings.get("StartScript") or "").strip()
+            enabled = bool(meta.get("enabled", True))
+            status = str(meta.get("status") or "").lower()
+            if not script or status in {"disabled", "errored", "waiting"}:
+                print(
+                    f"fix backend id={bid} status={status} StartScript={script!r}",
+                    flush=True,
+                )
+                try:
+                    configure_comfy_backend(sid, bid)
+                except Exception as exc:
+                    print(f"configure FAIL id={bid}: {exc}", flush=True)

-    # Default StartScript may be dlbackend/comfy/...; our tree is dlbackend/ComfyUI.
-    time.sleep(2)
-    bstat, _ = backend_status_detail()
-    return bstat
+    deadline = time.time() + 180
+    last = "disabled"
+    while time.time() < deadline:
+        time.sleep(5)
+        last, msg = backend_status_detail()
+        extra = f" — {msg[:80]}" if msg else ""
+        print(f"recover poll: backend_status={last}{extra}", flush=True)
+        if last in ("running", "idle", "loading", "some_loading"):
+            return last
+        if last in ("disabled", "all_disabled"):
+            # One more enable attempt
+            sid2 = get_session(time.time() + 30)
+            for bid, _meta in _backend_entries(list_backends(sid2)):
+                try:
+                    toggle_backend(sid2, bid, enabled=True)
+                except Exception:
+                    pass
+    return last


 def run_diagnostics() -> None:
@@ -637,7 +747,9 @@ def main() -> int:
    if bstat == "idle":
        print("backends present (idle/suspended) + skip install")
        return 0
-    if bstat not in ("empty", "unknown", "errored") and not bstat.startswith("error:"):
+    # disabled = empty StartScript / not usable — must recover, not skip.
+    need_recover = bstat in ("empty", "disabled", "all_disabled", "unknown")
+    if not need_recover and bstat != "errored" and not bstat.startswith("error:"):
        if comfy_venv_ok():
            print(f"backends present ({bstat}) + venv — skip install")
            return 0
@@ -652,22 +764,31 @@ def main() -> int:
        run_diagnostics()
        return 1

-    if installed is True and comfy_venv_ok() and bstat not in ("empty", "errored"):
+    if (
+        installed is True
+        and comfy_venv_ok()
+        and not need_recover
+        and bstat != "errored"
+        and not bstat.startswith("error:")
+    ):
        print("IsInstalled + venv — skip InstallConfirmWS (ждём ready отдельно)")
        return 0

-    if comfy_venv_ok() and bstat == "empty":
+    if comfy_venv_ok() and need_recover:
        print(
-            "backend empty при venv/IsInstalled — recover (sanitize FDS / AddNewBackend)…",
+            f"backend {bstat} при venv — recover (sanitize FDS / StartScript / AddNewBackend)…",
            flush=True,
        )
        bstat = recover_empty_backends()
-        if bstat not in ("empty", "unknown") and not bstat.startswith("error:"):
+        # disabled = empty StartScript / not usable — not a successful recover
+        ok_stats = ("running", "idle", "loading", "some_loading")
+        if bstat in ok_stats:
            print(f"recovered empty → {bstat}")
            return 0
+        print(f"recover incomplete → {bstat}", flush=True)
        if installed is True:
            print(
-                "WARN: всё ещё empty после recover — пробую InstallConfirmWS",
+                "WARN: всё ещё не ready после recover — пробую InstallConfirmWS",
                flush=True,
            )
        else: