From daa8a665b422bba1a8c8512dbe8338586423ae39 Mon Sep 17 00:00:00 2001 From: Simon Cordt Date: Wed, 24 Sep 2025 00:13:29 +0200 Subject: [PATCH] Implement persistent queue functionality --- __init__.py | 17 ++++ server/database.py | 149 ++++++++++++++++++++++++++++++++ server/queue_manager.py | 186 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 352 insertions(+) create mode 100644 __init__.py create mode 100644 server/database.py create mode 100644 server/queue_manager.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..025d060 --- /dev/null +++ b/__init__.py @@ -0,0 +1,17 @@ +import logging + +# Expose empty node mappings to satisfy ComfyUI V1 loader +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +# Optional: serve web assets if present +WEB_DIRECTORY = "./web" + +# Install persistent queue hooks as soon as this extension is imported. +try: + from .server.queue_manager import queue_manager + queue_manager.initialize() +except Exception as e: + logging.debug(f"ComfyUI-PersistentQueue initialization failed: {e}") + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] \ No newline at end of file diff --git a/server/database.py b/server/database.py new file mode 100644 index 0000000..c799326 --- /dev/null +++ b/server/database.py @@ -0,0 +1,149 @@ +import os +import sqlite3 +import json +from datetime import datetime +from typing import List, Optional, Dict, Any + +import folder_paths + +class QueueDatabase: + def __init__(self, db_path: Optional[str] = None): + # Default to ComfyUI user directory to ensure write permissions and persistence across updates + if db_path is None: + user_dir = folder_paths.get_user_directory() + os.makedirs(user_dir, exist_ok=True) + db_path = os.path.join(user_dir, "persistent_queue.sqlite3") + self.db_path = db_path + self._init_database() + + def _get_conn(self) -> sqlite3.Connection: + conn = sqlite3.connect(self.db_path) + conn.row_factory = sqlite3.Row + # Improve concurrency and durability for multi-threaded usage + conn.execute("PRAGMA journal_mode=WAL;") + conn.execute("PRAGMA synchronous=NORMAL;") + return conn + + def _init_database(self): + """Create tables if they don't exist""" + with self._get_conn() as conn: + conn.execute(''' + CREATE TABLE IF NOT EXISTS queue_items ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + prompt_id TEXT UNIQUE, + workflow TEXT, + priority INTEGER DEFAULT 0, + status TEXT DEFAULT 'pending', + created_at TIMESTAMP, + started_at TIMESTAMP, + completed_at TIMESTAMP, + error TEXT + ) + ''') + + conn.execute(''' + CREATE TABLE IF NOT EXISTS job_history ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + prompt_id TEXT, + workflow TEXT, + outputs TEXT, + duration_seconds REAL, + created_at TIMESTAMP, + completed_at TIMESTAMP, + status TEXT + ) + ''') + + def add_job(self, prompt_id: str, workflow: dict, priority: int = 0) -> None: + """Add a job to the persistent queue""" + with self._get_conn() as conn: + conn.execute( + ''' + INSERT OR IGNORE INTO queue_items (prompt_id, workflow, priority, created_at) + VALUES (?, ?, ?, ?) + ''', + (prompt_id, json.dumps(workflow), priority, datetime.now()), + ) + + def remove_job(self, prompt_id: str) -> None: + with self._get_conn() as conn: + conn.execute('DELETE FROM queue_items WHERE prompt_id = ?', (prompt_id,)) + + def get_job(self, prompt_id: str) -> Optional[Dict[str, Any]]: + with self._get_conn() as conn: + cur = conn.execute('SELECT * FROM queue_items WHERE prompt_id = ?', (prompt_id,)) + row = cur.fetchone() + return dict(row) if row else None + + def get_pending_jobs(self) -> List[Dict[str, Any]]: + """Get all pending jobs ordered by priority (higher first), then created_at""" + with self._get_conn() as conn: + cursor = conn.execute( + ''' + SELECT * FROM queue_items + WHERE status = 'pending' + ORDER BY priority DESC, created_at ASC + ''' + ) + return [dict(row) for row in cursor.fetchall()] + + def update_job_status(self, prompt_id: str, status: str, error: Optional[str] = None) -> None: + """Update job status and timestamps""" + with self._get_conn() as conn: + if status == 'running': + conn.execute( + ''' + UPDATE queue_items + SET status = ?, started_at = ? + WHERE prompt_id = ? + ''', + (status, datetime.now(), prompt_id), + ) + elif status in ('completed', 'failed'): + conn.execute( + ''' + UPDATE queue_items + SET status = ?, completed_at = ?, error = ? + WHERE prompt_id = ? + ''', + (status, datetime.now(), error, prompt_id), + ) + + def update_job_priority(self, prompt_id: str, new_priority: int) -> None: + with self._get_conn() as conn: + conn.execute('UPDATE queue_items SET priority = ? WHERE prompt_id = ?', (new_priority, prompt_id)) + + def add_history( + self, + prompt_id: str, + workflow: dict, + outputs: Optional[dict], + status: str, + duration_seconds: Optional[float] = None, + ) -> None: + with self._get_conn() as conn: + now = datetime.now() + created_at = now + completed_at = now + conn.execute( + ''' + INSERT INTO job_history (prompt_id, workflow, outputs, duration_seconds, created_at, completed_at, status) + VALUES (?, ?, ?, ?, ?, ?, ?) + ''', + ( + prompt_id, + json.dumps(workflow) if workflow is not None else None, + json.dumps(outputs) if outputs is not None else None, + duration_seconds, + created_at, + completed_at, + status, + ), + ) + + def list_history(self, limit: int = 50) -> List[Dict[str, Any]]: + with self._get_conn() as conn: + cur = conn.execute( + 'SELECT * FROM job_history ORDER BY id DESC LIMIT ?', (limit,) + ) + return [dict(row) for row in cur.fetchall()] \ No newline at end of file diff --git a/server/queue_manager.py b/server/queue_manager.py new file mode 100644 index 0000000..187c585 --- /dev/null +++ b/server/queue_manager.py @@ -0,0 +1,186 @@ +import asyncio +import json +import logging +import time +from typing import Optional, Any, Dict, Tuple + +from aiohttp import web + +from .database import QueueDatabase + +class PersistentQueueManager: + def __init__(self): + self.db = QueueDatabase() + self.paused = False + self.current_job = None + self._installed = False + self._original_send_sync = None + self._original_queue_get = None + self._original_task_done = None + + def initialize(self): + """Install hooks and API routes after PromptServer is created.""" + if self._installed: + return + + # Import here to avoid circular dependencies + from server import PromptServer + import execution + + server_instance = PromptServer.instance + + # Register on-prompt handler to persist incoming prompts + server_instance.add_on_prompt_handler(self._on_prompt) + + # Monkey-patch PromptQueue.get to support pause and mark jobs running + if self._original_queue_get is None: + self._original_queue_get = execution.PromptQueue.get + + def get_wrapper(q_self, timeout=None): + if self.paused: + time.sleep(0.1) + return None + result = self._original_queue_get(q_self, timeout=timeout) + if result is not None: + try: + item, item_id = result + # item: (number, prompt_id, prompt, extra_data, outputs_to_execute) + prompt_id = item[1] + self.db.update_job_status(prompt_id, 'running') + except Exception as e: + logging.debug(f"PersistentQueue get_wrapper mark running failed: {e}") + return result + + execution.PromptQueue.get = get_wrapper + + # Monkey-patch PromptQueue.task_done to persist results/history + if self._original_task_done is None: + self._original_task_done = execution.PromptQueue.task_done + + def task_done_wrapper(q_self, item_id, history_result, status): + # Capture before original pops it + item = q_self.currently_running.get(item_id) + self._original_task_done(q_self, item_id, history_result, status) + try: + if item is not None: + prompt_id = item[1] + prompt = item[2] + status_str = status.status_str if status is not None else 'success' + completed = (status.completed if status is not None else True) + new_state = 'completed' if completed else 'failed' + self.db.update_job_status(prompt_id, new_state, error=None if completed else status_str) + self.db.add_history( + prompt_id=prompt_id, + workflow=prompt, + outputs=history_result.get('outputs', {}), + status=status_str, + duration_seconds=None, + ) + except Exception as e: + logging.debug(f"PersistentQueue task_done wrapper persist failed: {e}") + + execution.PromptQueue.task_done = task_done_wrapper + + # Add API routes + self._install_api_routes() + + # Restore pending jobs on startup + self._schedule_restore_pending_jobs() + + self._installed = True + + def _on_prompt(self, json_data: Dict[str, Any]) -> Dict[str, Any]: + try: + if "prompt" in json_data: + prompt = json_data["prompt"] + prompt_id = str(json_data.get("prompt_id")) if json_data.get("prompt_id") else None + if prompt_id is None: + import uuid + prompt_id = uuid.uuid4().hex + json_data["prompt_id"] = prompt_id + # Priority scaffold (0 default) + priority = 0 + self.db.add_job(prompt_id, prompt, priority=priority) + except Exception as e: + logging.debug(f"PersistentQueue on_prompt persist failed: {e}") + return json_data + + def pause_queue(self): + """Pause queue execution""" + self.paused = True + + def resume_queue(self): + """Resume queue execution""" + self.paused = False + + def reorder_job(self, prompt_id: str, new_priority: int): + """Change job priority""" + self.db.update_job_priority(prompt_id, new_priority) + + def _schedule_restore_pending_jobs(self): + try: + from server import PromptServer + loop = PromptServer.instance.loop + loop.create_task(self._restore_pending_jobs_async()) + except Exception as e: + logging.debug(f"PersistentQueue schedule restore failed: {e}") + + async def _restore_pending_jobs_async(self): + from server import PromptServer + import execution + server_instance = PromptServer.instance + + pending_jobs = self.db.get_pending_jobs() + for job in pending_jobs: + try: + prompt_id = job["prompt_id"] + workflow_json = job["workflow"] + prompt = json.loads(workflow_json) if isinstance(workflow_json, str) else workflow_json + valid, err, outputs_to_execute, node_errors = await execution.validate_prompt(prompt_id, prompt, None) + if valid: + number = server_instance.number + server_instance.number += 1 + extra_data = {} + server_instance.prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute)) + else: + self.db.update_job_status(prompt_id, 'failed', error=(err or {}).get('message') if isinstance(err, dict) else str(err)) + except Exception as e: + logging.debug(f"PersistentQueue restore job failed: {e}") + + # API Routes + async def _api_get_pqueue(self, request: web.Request) -> web.Response: + from server import PromptServer + running, queued = PromptServer.instance.prompt_queue.get_current_queue_volatile() + return web.json_response({ + "paused": self.paused, + "db_pending": self.db.get_pending_jobs(), + "queue_running": running, + "queue_pending": queued, + }) + + async def _api_get_history(self, request: web.Request) -> web.Response: + limit = int(request.rel_url.query.get("limit", "50")) + return web.json_response({"history": self.db.list_history(limit=limit)}) + + async def _api_pause(self, request: web.Request) -> web.Response: + self.pause_queue() + return web.json_response({"paused": True}) + + async def _api_resume(self, request: web.Request) -> web.Response: + self.resume_queue() + return web.json_response({"paused": False}) + + def _install_api_routes(self): + try: + from server import PromptServer + app = PromptServer.instance.app + app.add_routes([ + web.get('/api/pqueue', self._api_get_pqueue), + web.get('/api/pqueue/history', self._api_get_history), + web.post('/api/pqueue/pause', self._api_pause), + web.post('/api/pqueue/resume', self._api_resume), + ]) + except Exception as e: + logging.debug(f"PersistentQueue add routes failed: {e}") + +queue_manager = PersistentQueueManager() \ No newline at end of file