280 lines
11 KiB
Python
280 lines
11 KiB
Python
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()
|
|
self._backfilled_once = False
|
|
|
|
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
|
|
)
|
|
''')
|
|
# No schema migrations for now; keep it simple
|
|
|
|
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:
|
|
# Prefer accurate timestamps from queue_items when available
|
|
def _parse_dt(val: Any) -> Optional[datetime]:
|
|
if val is None:
|
|
return None
|
|
if isinstance(val, datetime):
|
|
return val
|
|
if isinstance(val, str):
|
|
try:
|
|
# sqlite will store Python datetimes as strings like 'YYYY-MM-DD HH:MM:SS[.ffffff]'
|
|
return datetime.fromisoformat(val)
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
created_at = None
|
|
completed_at = None
|
|
try:
|
|
cur = conn.execute(
|
|
'SELECT created_at, started_at, completed_at FROM queue_items WHERE prompt_id = ?',
|
|
(prompt_id,),
|
|
)
|
|
row = cur.fetchone()
|
|
if row:
|
|
started = _parse_dt(row['started_at'])
|
|
created = _parse_dt(row['created_at'])
|
|
completed = _parse_dt(row['completed_at'])
|
|
created_at = started or created or datetime.now()
|
|
completed_at = completed or datetime.now()
|
|
if duration_seconds is None and created_at and completed_at:
|
|
try:
|
|
duration_seconds = max(0.0, (completed_at - created_at).total_seconds())
|
|
except Exception:
|
|
duration_seconds = None
|
|
except Exception:
|
|
# Fallback to now if any error occurs
|
|
pass
|
|
|
|
if created_at is None:
|
|
created_at = datetime.now()
|
|
if completed_at is None:
|
|
completed_at = created_at
|
|
|
|
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:
|
|
# On-the-fly lightweight backfill for rows missing duration or with identical timestamps
|
|
try:
|
|
if not self._backfilled_once:
|
|
self._backfill_history_rows(conn)
|
|
self._backfilled_once = True
|
|
except Exception:
|
|
# Non-fatal if backfill fails
|
|
pass
|
|
|
|
cur = conn.execute(
|
|
'SELECT * FROM job_history ORDER BY id DESC LIMIT ?', (limit,)
|
|
)
|
|
return [dict(row) for row in cur.fetchall()]
|
|
|
|
def get_job_timestamps_and_workflow(self, prompt_id: str) -> Optional[Dict[str, Any]]:
|
|
"""Return created_at/started_at/completed_at and workflow JSON (text) for a given prompt_id."""
|
|
with self._get_conn() as conn:
|
|
cur = conn.execute(
|
|
'SELECT created_at, started_at, completed_at, workflow FROM queue_items WHERE prompt_id = ?',
|
|
(prompt_id,),
|
|
)
|
|
row = cur.fetchone()
|
|
return dict(row) if row else None
|
|
|
|
def get_average_duration_for_workflow(self, workflow_text: Optional[str], min_samples: int = 2) -> Optional[float]:
|
|
"""Compute an average historical duration for a given workflow text. Returns None if not enough data."""
|
|
if not workflow_text:
|
|
return None
|
|
with self._get_conn() as conn:
|
|
try:
|
|
cur = conn.execute(
|
|
'''
|
|
SELECT duration_seconds FROM job_history
|
|
WHERE workflow = ? AND duration_seconds IS NOT NULL AND duration_seconds > 0
|
|
ORDER BY id DESC LIMIT 20
|
|
''',
|
|
(workflow_text,),
|
|
)
|
|
vals = [float(r['duration_seconds']) for r in cur.fetchall()]
|
|
if len(vals) < min_samples:
|
|
return None
|
|
# Use median for robustness
|
|
vals.sort()
|
|
n = len(vals)
|
|
if n % 2 == 1:
|
|
return vals[n // 2]
|
|
else:
|
|
return (vals[n // 2 - 1] + vals[n // 2]) / 2.0
|
|
except Exception:
|
|
return None
|
|
|
|
def _backfill_history_rows(self, conn: sqlite3.Connection) -> None:
|
|
"""Compute and set duration_seconds/accurate timestamps for existing history rows when possible."""
|
|
try:
|
|
cur = conn.execute(
|
|
'''
|
|
SELECT j.id as jid, j.created_at as j_created, j.completed_at as j_completed, j.duration_seconds as j_dur,
|
|
j.prompt_id as pid, qi.created_at as qi_created, qi.started_at as qi_started, qi.completed_at as qi_completed
|
|
FROM job_history j
|
|
LEFT JOIN queue_items qi ON qi.prompt_id = j.prompt_id
|
|
WHERE (j.duration_seconds IS NULL OR j.duration_seconds <= 0 OR j.created_at = j.completed_at)
|
|
AND qi.completed_at IS NOT NULL
|
|
'''
|
|
)
|
|
rows = cur.fetchall()
|
|
if not rows:
|
|
return
|
|
|
|
def _parse_dt(val: Any) -> Optional[datetime]:
|
|
if val is None:
|
|
return None
|
|
if isinstance(val, datetime):
|
|
return val
|
|
if isinstance(val, str):
|
|
try:
|
|
return datetime.fromisoformat(val)
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
for r in rows:
|
|
started = _parse_dt(r['qi_started']) or _parse_dt(r['qi_created']) or _parse_dt(r['j_created']) or datetime.now()
|
|
completed = _parse_dt(r['qi_completed']) or _parse_dt(r['j_completed']) or started
|
|
try:
|
|
dur = max(0.0, (completed - started).total_seconds())
|
|
except Exception:
|
|
dur = None
|
|
conn.execute(
|
|
'UPDATE job_history SET duration_seconds = ?, created_at = ?, completed_at = ? WHERE id = ?',
|
|
(dur, started, completed, r['jid'])
|
|
)
|
|
except Exception:
|
|
# Ignore errors; not critical
|
|
return |