""" Workflow execution integration for ComfyUI Queue Manager. Handles intercepting, monitoring, and executing workflows through ComfyUI. """ from __future__ import annotations import json import logging import threading import time import uuid from datetime import datetime, timezone from typing import Any, Callable, Dict, Optional from error_handler import ErrorHandler, with_error_handling from exceptions import ValidationError, WorkflowExecutionError from interfaces import WorkflowExecutorInterface from models import QueueStatus logger = logging.getLogger(__name__) class WorkflowExecutor(WorkflowExecutorInterface): """ Handles workflow execution integration with ComfyUI. This class provides the bridge between the queue manager and ComfyUI's workflow execution system, intercepting executions and monitoring status. """ def __init__(self, queue_service=None): """Initialize the workflow executor. Args: queue_service: Optional queue service for automatic queue management """ if queue_service is None: # Create a default database and queue service try: from database import SQLiteDatabase from queue_service import QueueService from pathlib import Path db_path = Path(__file__).parent / "queue_manager.db" database = SQLiteDatabase(str(db_path)) database.initialize() queue_service = QueueService(database) except Exception as e: logger.error(f"Failed to create default queue service: {e}") queue_service = None self.queue_service = queue_service self._running_workflows: Dict[str, Dict[str, Any]] = {} self._execution_callbacks: Dict[str, Callable] = {} self._lock = threading.RLock() self._interceptor_enabled = True # ComfyUI integration hooks self._original_execute = None self._comfyui_server = None logger.info("Workflow executor initialized") def enable_interceptor(self) -> None: """Enable workflow execution interception.""" self._interceptor_enabled = True logger.info("Workflow execution interceptor enabled") def disable_interceptor(self) -> None: """Disable workflow execution interception.""" self._interceptor_enabled = False logger.info("Workflow execution interceptor disabled") def is_interceptor_enabled(self) -> bool: """Check if the interceptor is enabled.""" return self._interceptor_enabled def register_execution_callback(self, callback_id: str, callback: Callable) -> None: """Register a callback for execution events. Args: callback_id: Unique identifier for the callback callback: Function to call on execution events """ with self._lock: self._execution_callbacks[callback_id] = callback logger.debug(f"Registered execution callback: {callback_id}") def unregister_execution_callback(self, callback_id: str) -> None: """Unregister an execution callback. Args: callback_id: Identifier of the callback to remove """ with self._lock: self._execution_callbacks.pop(callback_id, None) logger.debug(f"Unregistered execution callback: {callback_id}") def _notify_callbacks(self, event_type: str, workflow_id: str, data: Dict[str, Any]) -> None: """Notify all registered callbacks of an execution event. Args: event_type: Type of event (started, completed, failed, etc.) workflow_id: ID of the workflow data: Additional event data """ with self._lock: callbacks = list(self._execution_callbacks.values()) for callback in callbacks: try: callback(event_type, workflow_id, data) except Exception as e: logger.error(f"Error in execution callback: {e}") def intercept_workflow_execution(self, workflow_data: Dict[str, Any]) -> str: """Intercept a workflow execution and add it to the queue. Args: workflow_data: The workflow data to execute Returns: The workflow execution ID """ if not self._interceptor_enabled: return self._execute_workflow_directly(workflow_data) workflow_id = str(uuid.uuid4()) try: # Extract workflow name from data workflow_name = self._extract_workflow_name(workflow_data) # Add to queue if queue service is available if self.queue_service: queue_item_id = self.queue_service.add_workflow( workflow_data=workflow_data, workflow_name=workflow_name ) # Store the mapping between workflow ID and queue item ID with self._lock: self._running_workflows[workflow_id] = { "queue_item_id": queue_item_id, "workflow_data": workflow_data, "workflow_name": workflow_name, "status": QueueStatus.PENDING, "started_at": None, "completed_at": None, "error_message": None, "result_data": None, } logger.info(f"Intercepted workflow execution: {workflow_id} -> queue item: {queue_item_id}") else: # Execute directly if no queue service return self._execute_workflow_directly(workflow_data) except Exception as e: logger.error(f"Failed to intercept workflow execution: {e}") # Fall back to direct execution return self._execute_workflow_directly(workflow_data) return workflow_id @with_error_handling(error_type=WorkflowExecutionError, operation="execute_workflow") def execute_workflow(self, workflow_data: Dict[str, Any]) -> Dict[str, Any]: """Execute a workflow and return the results. Args: workflow_data: The workflow data to execute Returns: Dictionary containing execution results Raises: WorkflowExecutionError: If execution fails """ if not workflow_data: raise ValidationError("Workflow data cannot be empty", field="workflow_data") if not isinstance(workflow_data, dict): raise ValidationError("Workflow data must be a dictionary", field="workflow_data", value=type(workflow_data).__name__) workflow_id = str(uuid.uuid4()) # Create workflow info for tracking with self._lock: self._running_workflows[workflow_id] = { "workflow_data": workflow_data, "workflow_name": self._extract_workflow_name(workflow_data), "status": QueueStatus.RUNNING, "started_at": datetime.now(timezone.utc), "completed_at": None, "error_message": None, "result_data": None, } try: # Execute the workflow execution_id = self._execute_workflow_directly(workflow_data) # Create result data result = { "status": "completed", "execution_time": 0.1, "outputs": {}, "executed_at": datetime.now(timezone.utc).isoformat(), "execution_id": execution_id, } # Update status to completed self._update_workflow_status( workflow_id, QueueStatus.COMPLETED, result_data=result ) return result except Exception as e: # Update status to failed self._update_workflow_status( workflow_id, QueueStatus.FAILED, error_message=str(e) ) raise WorkflowExecutionError( f"Workflow execution failed: {e}", workflow_id=workflow_id, execution_stage="execution", cause=e ) def _execute_workflow_directly(self, workflow_data: Dict[str, Any]) -> str: """Execute a workflow directly through ComfyUI. Args: workflow_data: The workflow data to execute Returns: Workflow execution ID """ # This is a placeholder for actual ComfyUI integration # In a real implementation, this would interface with ComfyUI's execution system logger.info("Executing workflow directly through ComfyUI") # Simulate workflow execution time.sleep(0.1) # Simulate processing time # Return a workflow ID for direct execution return str(uuid.uuid4()) def _extract_workflow_name(self, workflow_data: Dict[str, Any]) -> str: """Extract a meaningful name from workflow data. Args: workflow_data: The workflow data Returns: A descriptive name for the workflow """ # Try to extract name from various possible locations if isinstance(workflow_data, dict): # Check for explicit name if "name" in workflow_data: return workflow_data["name"] # Check for title or description if "title" in workflow_data: return workflow_data["title"] # Check for workflow metadata if "workflow" in workflow_data and isinstance(workflow_data["workflow"], dict): workflow = workflow_data["workflow"] if "name" in workflow: return workflow["name"] if "title" in workflow: return workflow["title"] # Generate name based on node count or other characteristics if "nodes" in workflow_data: node_count = len(workflow_data["nodes"]) return f"Workflow with {node_count} nodes" # Default name with timestamp return f"Workflow {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" def _update_workflow_status( self, workflow_id: str, status: QueueStatus, error_message: Optional[str] = None, result_data: Optional[Dict[str, Any]] = None ) -> None: """Update the status of a workflow execution. Args: workflow_id: The workflow ID status: The new status error_message: Optional error message for failed workflows result_data: Optional result data for completed workflows """ with self._lock: if workflow_id not in self._running_workflows: logger.warning(f"Workflow {workflow_id} not found in running workflows") return workflow_info = self._running_workflows[workflow_id] old_status = workflow_info["status"] workflow_info["status"] = status # Update timestamps now = datetime.now(timezone.utc) if status == QueueStatus.RUNNING and old_status == QueueStatus.PENDING: workflow_info["started_at"] = now elif status in (QueueStatus.COMPLETED, QueueStatus.FAILED): if not workflow_info["started_at"]: workflow_info["started_at"] = now workflow_info["completed_at"] = now # Update error message and result data if error_message: workflow_info["error_message"] = error_message if result_data: workflow_info["result_data"] = result_data # Update queue service if available if self.queue_service and "queue_item_id" in workflow_info: try: self.queue_service.update_item_status( workflow_info["queue_item_id"], status, error_message=error_message, result_data=result_data ) except Exception as e: logger.error(f"Failed to update queue item status: {e}") # Notify callbacks self._notify_callbacks("status_changed", workflow_id, { "old_status": old_status.value if old_status else None, "new_status": status.value, "error_message": error_message, "result_data": result_data, }) logger.info(f"Updated workflow {workflow_id} status: {old_status} -> {status}") def is_workflow_running(self, workflow_id: str) -> bool: """Check if a workflow is currently running. Args: workflow_id: The workflow ID to check Returns: True if the workflow is running """ with self._lock: if workflow_id not in self._running_workflows: return False return self._running_workflows[workflow_id]["status"] == QueueStatus.RUNNING def cancel_workflow(self, workflow_id: str) -> bool: """Cancel a running workflow. Args: workflow_id: The workflow ID to cancel Returns: True if the workflow was cancelled successfully """ with self._lock: if workflow_id not in self._running_workflows: logger.warning(f"Cannot cancel workflow {workflow_id}: not found") return False workflow_info = self._running_workflows[workflow_id] if workflow_info["status"] not in (QueueStatus.PENDING, QueueStatus.RUNNING): logger.warning(f"Cannot cancel workflow {workflow_id}: not in cancellable state") return False try: # Update status to failed with cancellation message self._update_workflow_status( workflow_id, QueueStatus.FAILED, error_message="Workflow cancelled by user" ) # Notify callbacks self._notify_callbacks("cancelled", workflow_id, {}) logger.info(f"Cancelled workflow: {workflow_id}") return True except Exception as e: logger.error(f"Failed to cancel workflow {workflow_id}: {e}") return False def get_workflow_status(self, workflow_id: str) -> QueueStatus | None: """Get the status of a workflow execution. Args: workflow_id: The workflow ID Returns: The current status or None if not found """ with self._lock: if workflow_id not in self._running_workflows: return None return self._running_workflows[workflow_id]["status"] def get_running_workflows(self) -> Dict[str, Dict[str, Any]]: """Get information about all running workflows. Returns: Dictionary mapping workflow IDs to their information """ with self._lock: return { wf_id: info.copy() for wf_id, info in self._running_workflows.items() if info["status"] == QueueStatus.RUNNING } def get_workflow_info(self, workflow_id: str) -> Dict[str, Any] | None: """Get detailed information about a workflow. Args: workflow_id: The workflow ID Returns: Dictionary containing workflow information or None if not found """ with self._lock: if workflow_id not in self._running_workflows: return None return self._running_workflows[workflow_id].copy() def cleanup_completed_workflows(self, max_age_hours: int = 24) -> int: """Clean up completed workflow records older than specified age. Args: max_age_hours: Maximum age in hours for completed workflows Returns: Number of workflows cleaned up """ from datetime import timedelta cutoff_time = datetime.now(timezone.utc) - timedelta(hours=max_age_hours) cleaned_count = 0 with self._lock: workflows_to_remove = [] for workflow_id, info in self._running_workflows.items(): if info["status"] in (QueueStatus.COMPLETED, QueueStatus.FAILED): completed_at = info.get("completed_at") if completed_at and completed_at < cutoff_time: workflows_to_remove.append(workflow_id) for workflow_id in workflows_to_remove: del self._running_workflows[workflow_id] cleaned_count += 1 if cleaned_count > 0: logger.info(f"Cleaned up {cleaned_count} completed workflows") return cleaned_count def get_execution_statistics(self) -> Dict[str, Any]: """Get statistics about workflow executions. Returns: Dictionary containing execution statistics """ with self._lock: stats = { "total_workflows": len(self._running_workflows), "status_counts": {}, "average_execution_time": 0.0, "total_execution_time": 0.0, } execution_times = [] for info in self._running_workflows.values(): status = info["status"].value stats["status_counts"][status] = stats["status_counts"].get(status, 0) + 1 # Calculate execution time for completed workflows if info["started_at"] and info["completed_at"]: exec_time = (info["completed_at"] - info["started_at"]).total_seconds() execution_times.append(exec_time) if execution_times: stats["average_execution_time"] = sum(execution_times) / len(execution_times) stats["total_execution_time"] = sum(execution_times) return stats def close(self) -> None: """Clean up resources and close the executor.""" with self._lock: self._running_workflows.clear() self._execution_callbacks.clear() logger.info("Workflow executor closed") def __enter__(self): """Context manager entry.""" return self def __exit__(self, exc_type, exc_val, exc_tb): """Context manager exit.""" self.close()