Implement persistent queue functionality
This commit is contained in:
Simon Cordt
+17
@@ -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']
|
||||
@@ -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()]
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user