Smart Queue v1: GPU-aware queue autopilot + Cooldown & Pause node

.github/workflows/ci.yml

@@ -13,5 +13,5 @@ jobs:
       - uses: actions/setup-python@v5
         with:
           python-version: "3.12"
-      - run: pip install pytest
+      - run: pip install pytest pytest-asyncio aiohttp
       - run: pytest tests/ -v

.gitignore

@@ -6,3 +6,4 @@ node_modules/
 *.log
 .DS_Store
 *.sqlite3
+.pytest_cache/

__init__.py

@@ -1,6 +1,87 @@
 """Smart Queue — GPU-aware queue autopilot + cooldown/pause node for ComfyUI."""
 
-NODE_CLASS_MAPPINGS = {}
-NODE_DISPLAY_NAME_MAPPINGS = {}
+import asyncio
+import logging
+from pathlib import Path
 
-__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+try:
+    from server import PromptServer  # type: ignore[import-not-found]
+    _HAS_COMFY_SERVER = True
+except ImportError:
+    _HAS_COMFY_SERVER = False
+
+from .backend.autopilot import AutopilotSettings
+from .backend.autopilot_loop import run_autopilot_tick
+from .backend.autopilot_state import AutopilotState
+from .backend.gpu_monitor import poll_gpu_metrics
+from .backend.nodes.cooldown import SmartCooldownNode
+from .backend.persistence import init_db
+from .backend.queue_middleware import create_queue_middleware
+from .backend.queue_tracker import sync_queue_tracker
+from .backend.routes import register_routes
+
+logger = logging.getLogger(__name__)
+
+NODE_CLASS_MAPPINGS = {
+    "RubzGpuCooldownNode": SmartCooldownNode,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "RubzGpuCooldownNode": "Smart Cooldown & Pause",
+}
+WEB_DIRECTORY = "./web"
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
+
+_autopilot_state = AutopilotState()
+_autopilot_settings = AutopilotSettings()
+_seen_running: set = set()
+_seen_completed: set = set()
+
+TICK_INTERVAL_SECONDS = 5.0
+
+
+async def _async_poll_gpu_metrics():
+    # poll_gpu_metrics is a blocking subprocess call; keep it off the event loop.
+    return await asyncio.to_thread(poll_gpu_metrics)
+
+
+def _sync_queue_tracker_tick(conn) -> None:
+    global _seen_running, _seen_completed
+    running, queued = _server.prompt_queue.get_current_queue_volatile()
+    history = _server.prompt_queue.get_history()
+    _seen_running, _seen_completed = sync_queue_tracker(
+        conn, running, queued, history, _autopilot_state, _seen_running, _seen_completed
+    )
+
+
+async def _autopilot_background_loop(conn):
+    while True:
+        if _autopilot_settings.master_enabled:
+            await run_autopilot_tick(_autopilot_state, _autopilot_settings, _async_poll_gpu_metrics)
+        try:
+            # Sync sqlite3 connections are bound to the thread that created them
+            # (this loop's thread), so this must run inline, not via asyncio.to_thread.
+            _sync_queue_tracker_tick(conn)
+        except Exception:
+            logger.warning("Smart Queue: queue tracker sync failed, skipping this tick", exc_info=True)
+        await asyncio.sleep(TICK_INTERVAL_SECONDS)
+
+
+def _is_autopilot_enabled() -> bool:
+    return _autopilot_settings.master_enabled
+
+
+if _HAS_COMFY_SERVER:
+    _db_path = str(Path(__file__).parent / "smart_queue.sqlite3")
+    _conn = init_db(_db_path)
+
+    _server = PromptServer.instance
+    _server.app.middlewares.append(create_queue_middleware(_autopilot_state, _is_autopilot_enabled))
+    register_routes(_server.app, _conn, _autopilot_state, _autopilot_settings)
+
+    async def _start_autopilot_loop(app):
+        app["smart_queue_autopilot_task"] = asyncio.create_task(_autopilot_background_loop(_conn))
+
+    _server.app.on_startup.append(_start_autopilot_loop)
+
+    logger.info("[Smart Queue] Loaded — autopilot + Smart Cooldown & Pause node registered.")

backend/autopilot.py

@@ -0,0 +1,63 @@
+"""Pure rule-engine: GPU metrics + queue state -> pause/resume decision.
+
+No I/O here. This module never touches the network, the filesystem, or a
+subprocess — that keeps it trivially unit-testable and lets the fail-open
+guarantee live entirely in the caller (backend.autopilot_loop).
+"""
+
+from dataclasses import dataclass
+
+from .gpu_monitor import GpuMetrics
+
+
+@dataclass
+class AutopilotSettings:
+    master_enabled: bool = True
+
+    temp_rule_enabled: bool = True
+    pause_temp_c: float = 80.0
+    resume_temp_c: float = 72.0
+
+    vram_rule_enabled: bool = True
+    min_free_vram_mb: float = 1024.0
+
+    job_count_rule_enabled: bool = True
+    max_jobs_before_pause: int = 20
+
+    def update_from_dict(self, values: dict) -> None:
+        """Mutates in place so callers holding a reference (the background
+        loop, the middleware's is_enabled closure) see updates immediately."""
+        for key, value in values.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+
+
+@dataclass
+class Decision:
+    should_pause: bool
+    reasons: tuple[str, ...]
+
+
+def evaluate(
+    metrics: GpuMetrics,
+    jobs_since_resume: int,
+    currently_paused: bool,
+    settings: AutopilotSettings,
+) -> Decision:
+    reasons: list[str] = []
+
+    if settings.temp_rule_enabled and metrics.temp_c is not None:
+        threshold = settings.resume_temp_c if currently_paused else settings.pause_temp_c
+        if metrics.temp_c > threshold:
+            reasons.append(f"GPU at {metrics.temp_c:.0f}C, target {threshold:.0f}C")
+
+    if settings.vram_rule_enabled and metrics.vram_free_mb is not None:
+        if metrics.vram_free_mb < settings.min_free_vram_mb:
+            reasons.append(
+                f"Only {metrics.vram_free_mb:.0f}MB VRAM free, below {settings.min_free_vram_mb:.0f}MB"
+            )
+
+    if settings.job_count_rule_enabled and jobs_since_resume >= settings.max_jobs_before_pause:
+        reasons.append(f"{jobs_since_resume} jobs run since last pause, cooldown break")
+
+    return Decision(should_pause=len(reasons) > 0, reasons=tuple(reasons))

backend/autopilot_loop.py

@@ -0,0 +1,36 @@
+"""Background polling: one tick = poll GPU metrics, evaluate rules, apply decision.
+
+Fail-open: if metrics_provider raises, the tick is a no-op — it never leaves
+or forces the queue into a paused state due to an internal error.
+"""
+
+import logging
+from typing import Awaitable, Callable
+
+from .autopilot import AutopilotSettings, evaluate
+from .autopilot_state import AutopilotState
+from .gpu_monitor import GpuMetrics
+
+logger = logging.getLogger(__name__)
+
+MetricsProvider = Callable[[], Awaitable[GpuMetrics]]
+
+
+async def run_autopilot_tick(
+    state: AutopilotState,
+    settings: AutopilotSettings,
+    metrics_provider: MetricsProvider,
+) -> None:
+    try:
+        metrics = await metrics_provider()
+    except Exception:
+        logger.warning("Smart Queue: GPU metrics poll failed, skipping this tick", exc_info=True)
+        return
+
+    decision = evaluate(
+        metrics,
+        jobs_since_resume=state.jobs_since_resume,
+        currently_paused=state.is_paused,
+        settings=settings,
+    )
+    state.apply(decision)

backend/autopilot_state.py

@@ -0,0 +1,36 @@
+"""In-memory autopilot state, mutated by the polling loop and read by the middleware."""
+
+from .autopilot import Decision
+
+
+class AutopilotState:
+    def __init__(self) -> None:
+        self.is_paused: bool = False
+        self.jobs_since_resume: int = 0
+        self.last_reasons: tuple[str, ...] = ()
+        self.manual_paused: bool = False
+
+    def apply(self, decision: Decision) -> None:
+        was_paused = self.is_paused
+        self.is_paused = decision.should_pause
+        self.last_reasons = decision.reasons
+        if was_paused and not self.is_paused:
+            self.jobs_since_resume = 0
+
+    def record_job_started(self) -> None:
+        self.jobs_since_resume += 1
+
+    def set_manual_pause(self, paused: bool) -> None:
+        self.manual_paused = paused
+
+    @property
+    def effective_paused(self) -> bool:
+        """What the middleware actually gates on: autopilot OR a manual pause."""
+        return self.is_paused or self.manual_paused
+
+    @property
+    def effective_reasons(self) -> tuple[str, ...]:
+        reasons = self.last_reasons
+        if self.manual_paused:
+            reasons = ("Manually paused",) + reasons
+        return reasons

backend/continue_registry.py

@@ -0,0 +1,26 @@
+"""In-memory registry of pending `wait_for_click` continues, keyed by prompt_id.
+
+Bridges the synchronous node execution thread (which blocks in `wait_for_continue`)
+and the aiohttp route handler that receives the click (`signal_continue`).
+"""
+
+import threading
+
+_lock = threading.Lock()
+_events: dict[str, threading.Event] = {}
+
+
+def wait_for_continue(prompt_id: str, poll_interval: float = 0.25) -> None:
+    with _lock:
+        event = _events.setdefault(prompt_id, threading.Event())
+    try:
+        event.wait()
+    finally:
+        with _lock:
+            _events.pop(prompt_id, None)
+
+
+def signal_continue(prompt_id: str) -> None:
+    with _lock:
+        event = _events.setdefault(prompt_id, threading.Event())
+    event.set()

backend/gpu_monitor.py

@@ -0,0 +1,48 @@
+"""Polls NVIDIA GPU metrics via nvidia-smi. No third-party dependencies."""
+
+import subprocess
+from dataclasses import dataclass
+
+
+@dataclass
+class GpuMetrics:
+    temp_c: float | None
+    vram_used_mb: float | None
+    vram_total_mb: float | None
+    util_pct: float | None
+
+    @property
+    def vram_free_mb(self) -> float | None:
+        if self.vram_used_mb is None or self.vram_total_mb is None:
+            return None
+        return self.vram_total_mb - self.vram_used_mb
+
+
+_EMPTY = GpuMetrics(temp_c=None, vram_used_mb=None, vram_total_mb=None, util_pct=None)
+
+
+def poll_gpu_metrics(timeout: float = 5.0) -> GpuMetrics:
+    try:
+        result = subprocess.run(
+            [
+                "nvidia-smi",
+                "--query-gpu=temperature.gpu,memory.used,memory.total,utilization.gpu",
+                "--format=csv,noheader,nounits",
+            ],
+            capture_output=True,
+            text=True,
+            timeout=timeout,
+        )
+        if result.returncode != 0:
+            return _EMPTY
+
+        first_line = result.stdout.strip().splitlines()[0]
+        temp_s, used_s, total_s, util_s = (p.strip() for p in first_line.split(","))
+        return GpuMetrics(
+            temp_c=float(temp_s),
+            vram_used_mb=float(used_s),
+            vram_total_mb=float(total_s),
+            util_pct=float(util_s),
+        )
+    except Exception:
+        return _EMPTY