305 lines
11 KiB
Python
305 lines
11 KiB
Python
"""
|
|
Main ComfyUI custom node for the Queue Manager.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import uuid
|
|
from datetime import datetime
|
|
from typing import Any, Optional
|
|
|
|
from models import QueueStatus
|
|
|
|
|
|
class QueueManagerNode:
|
|
"""
|
|
ComfyUI custom node for queue management functionality.
|
|
This node serves as the entry point for the queue manager system.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
"""Define the input types for the node."""
|
|
return {
|
|
"required": {},
|
|
"optional": {
|
|
"workflow_name": ("STRING", {"default": "Workflow"}),
|
|
"auto_queue": ("BOOLEAN", {"default": True}),
|
|
"priority": ("INT", {"default": 0, "min": -10, "max": 10}),
|
|
"tags": ("STRING", {"default": ""}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "BOOLEAN")
|
|
RETURN_NAMES = ("queue_item_id", "queued_successfully")
|
|
FUNCTION = "process"
|
|
CATEGORY = "Queue Management"
|
|
DESCRIPTION = (
|
|
"Manages workflow execution queues with persistence and control features. "
|
|
"Automatically captures and queues workflows for later execution."
|
|
)
|
|
|
|
# ComfyUI node metadata
|
|
OUTPUT_NODE = False
|
|
|
|
def __init__(self):
|
|
"""Initialize the queue manager node."""
|
|
self._queue_service = None
|
|
self._workflow_executor = None
|
|
self._execution_monitor = None
|
|
|
|
@property
|
|
def queue_service(self):
|
|
"""Lazy initialization of queue service."""
|
|
if self._queue_service is None:
|
|
try:
|
|
from database import SQLiteDatabase
|
|
from queue_service import QueueService
|
|
from pathlib import Path
|
|
|
|
# Initialize database and queue service
|
|
db_path = Path(__file__).parent / "queue_manager.db"
|
|
database = SQLiteDatabase(str(db_path))
|
|
database.initialize()
|
|
self._queue_service = QueueService(database)
|
|
|
|
# Initialize workflow execution components
|
|
self._initialize_workflow_execution()
|
|
|
|
except Exception as e:
|
|
print(f"[Queue Manager] Failed to initialize queue service: {e}")
|
|
# Don't set to None, leave it as None to indicate failure
|
|
return None
|
|
return self._queue_service
|
|
|
|
def _initialize_workflow_execution(self):
|
|
"""Initialize workflow execution integration components."""
|
|
try:
|
|
from workflow_executor import WorkflowExecutor
|
|
from execution_monitor import ExecutionMonitor
|
|
from workflow_interceptor import get_workflow_interceptor
|
|
|
|
# Initialize workflow executor
|
|
self._workflow_executor = WorkflowExecutor(self._queue_service)
|
|
|
|
# Initialize execution monitor
|
|
self._execution_monitor = ExecutionMonitor(
|
|
queue_service=self._queue_service,
|
|
workflow_executor=self._workflow_executor
|
|
)
|
|
|
|
# Set up workflow interceptor
|
|
interceptor = get_workflow_interceptor()
|
|
interceptor.set_workflow_executor(self._workflow_executor)
|
|
|
|
# Start execution monitoring
|
|
self._execution_monitor.start_monitoring()
|
|
|
|
# Enable workflow interception
|
|
interceptor.enable()
|
|
|
|
print("[Queue Manager] Workflow execution integration initialized")
|
|
|
|
except Exception as e:
|
|
print(f"[Queue Manager] Failed to initialize workflow execution: {e}")
|
|
# Continue without workflow execution integration
|
|
|
|
def process(
|
|
self,
|
|
workflow_name: str = "Workflow",
|
|
auto_queue: bool = True,
|
|
priority: int = 0,
|
|
tags: str = ""
|
|
) -> tuple[str, bool]:
|
|
"""
|
|
Process the node execution and optionally add to queue.
|
|
|
|
Args:
|
|
workflow_name: Name for the workflow
|
|
auto_queue: Whether to automatically add to queue
|
|
priority: Priority level for queue processing
|
|
tags: Comma-separated tags for the workflow
|
|
|
|
Returns:
|
|
Tuple containing (queue_item_id, queued_successfully)
|
|
"""
|
|
queue_item_id = str(uuid.uuid4())
|
|
queued_successfully = False
|
|
|
|
if auto_queue and self.queue_service:
|
|
try:
|
|
# Create workflow data structure
|
|
workflow_data = {
|
|
"name": workflow_name,
|
|
"priority": priority,
|
|
"tags": [tag.strip() for tag in tags.split(",") if tag.strip()],
|
|
"created_at": datetime.now().isoformat(),
|
|
"node_id": queue_item_id,
|
|
}
|
|
|
|
# Add to queue
|
|
actual_id = self.queue_service.add_workflow(workflow_data)
|
|
if actual_id:
|
|
queue_item_id = actual_id
|
|
queued_successfully = True
|
|
print(f"[Queue Manager] Added workflow '{workflow_name}' to queue: {queue_item_id}")
|
|
else:
|
|
print(f"[Queue Manager] Failed to add workflow '{workflow_name}' to queue")
|
|
|
|
except Exception as e:
|
|
print(f"[Queue Manager] Error adding workflow to queue: {e}")
|
|
|
|
elif not auto_queue:
|
|
print(f"[Queue Manager] Auto-queue disabled for workflow '{workflow_name}'")
|
|
|
|
return (queue_item_id, queued_successfully)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
"""
|
|
Determine if the node needs to be re-executed.
|
|
|
|
For queue management, we want to execute every time to capture workflows.
|
|
"""
|
|
# Return NaN to always execute (following official documentation)
|
|
return float("nan")
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, **kwargs):
|
|
"""Validate the inputs to the node."""
|
|
# Basic validation following official documentation patterns
|
|
workflow_name = kwargs.get("workflow_name", "")
|
|
priority = kwargs.get("priority", 0)
|
|
|
|
# Validate workflow name is not empty if provided
|
|
if workflow_name and not isinstance(workflow_name, str):
|
|
return "workflow_name must be a string"
|
|
|
|
# Validate priority is within range
|
|
if not isinstance(priority, int) or priority < -10 or priority > 10:
|
|
return "priority must be an integer between -10 and 10"
|
|
|
|
return True
|
|
|
|
def get_queue_status(self) -> dict[str, Any]:
|
|
"""
|
|
Get current queue status information.
|
|
|
|
Returns:
|
|
Dictionary containing queue status information
|
|
"""
|
|
if not self.queue_service:
|
|
return {"error": "Queue service not available"}
|
|
|
|
try:
|
|
items = self.queue_service.get_queue_items()
|
|
status_counts = {}
|
|
|
|
for item in items:
|
|
status = item.status.value
|
|
status_counts[status] = status_counts.get(status, 0) + 1
|
|
|
|
# Check if pause functionality is available
|
|
is_paused = False
|
|
if hasattr(self.queue_service, 'is_paused'):
|
|
is_paused = self.queue_service.is_paused()
|
|
|
|
return {
|
|
"total_items": len(items),
|
|
"status_counts": status_counts,
|
|
"is_paused": is_paused,
|
|
}
|
|
except Exception as e:
|
|
return {"error": f"Failed to get queue status: {e}"}
|
|
|
|
def get_workflow_executor(self):
|
|
"""Get the workflow executor instance."""
|
|
# Ensure queue service is initialized
|
|
_ = self.queue_service
|
|
return self._workflow_executor
|
|
|
|
def get_execution_monitor(self):
|
|
"""Get the execution monitor instance."""
|
|
# Ensure queue service is initialized
|
|
_ = self.queue_service
|
|
return self._execution_monitor
|
|
|
|
def execute_workflow_directly(self, workflow_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Execute a workflow directly through the workflow executor.
|
|
|
|
Args:
|
|
workflow_data: The workflow data to execute
|
|
|
|
Returns:
|
|
Dictionary containing execution results
|
|
"""
|
|
if not self._workflow_executor:
|
|
raise RuntimeError("Workflow executor not initialized")
|
|
|
|
try:
|
|
return self._workflow_executor.execute_workflow(workflow_data)
|
|
except Exception as e:
|
|
print(f"[Queue Manager] Failed to execute workflow: {e}")
|
|
raise
|
|
|
|
def get_execution_statistics(self) -> dict[str, Any]:
|
|
"""Get statistics about workflow executions.
|
|
|
|
Returns:
|
|
Dictionary containing execution statistics
|
|
"""
|
|
stats = {}
|
|
|
|
if self._workflow_executor:
|
|
stats["executor"] = self._workflow_executor.get_execution_statistics()
|
|
|
|
if self._execution_monitor:
|
|
stats["monitor"] = self._execution_monitor.get_monitoring_statistics()
|
|
|
|
return stats
|
|
|
|
def cleanup_execution_resources(self) -> dict[str, int]:
|
|
"""Clean up old execution resources.
|
|
|
|
Returns:
|
|
Dictionary containing cleanup counts
|
|
"""
|
|
cleanup_counts = {}
|
|
|
|
if self._workflow_executor:
|
|
try:
|
|
count = self._workflow_executor.cleanup_completed_workflows()
|
|
cleanup_counts["executor_workflows"] = count
|
|
except Exception as e:
|
|
print(f"[Queue Manager] Failed to cleanup executor workflows: {e}")
|
|
|
|
if self._execution_monitor:
|
|
try:
|
|
count = self._execution_monitor.cleanup_old_workflows()
|
|
cleanup_counts["monitor_workflows"] = count
|
|
except Exception as e:
|
|
print(f"[Queue Manager] Failed to cleanup monitor workflows: {e}")
|
|
|
|
return cleanup_counts
|
|
|
|
@classmethod
|
|
def get_node_info(cls) -> dict[str, Any]:
|
|
"""
|
|
Get information about this node for ComfyUI.
|
|
|
|
Returns:
|
|
Dictionary containing node information
|
|
"""
|
|
return {
|
|
"name": "QueueManagerNode",
|
|
"display_name": "Queue Manager",
|
|
"category": "Queue Management",
|
|
"description": cls.DESCRIPTION,
|
|
"version": "1.0.0",
|
|
"author": "ComfyUI Queue Manager",
|
|
"input_types": cls.INPUT_TYPES(),
|
|
"return_types": cls.RETURN_TYPES,
|
|
"return_names": cls.RETURN_NAMES,
|
|
}
|