402 lines
15 KiB
Python
402 lines
15 KiB
Python
"""
|
|
Execution monitoring system for ComfyUI Queue Manager.
|
|
Monitors workflow execution status and provides real-time updates.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
import time
|
|
from datetime import datetime, timezone, timedelta
|
|
from typing import Any, Callable, Dict, List, Optional
|
|
|
|
from models import QueueStatus
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ExecutionMonitor:
|
|
"""
|
|
Monitors workflow execution status and provides real-time updates.
|
|
|
|
This class tracks running workflows, monitors their progress, and provides
|
|
callbacks for status changes and completion events.
|
|
"""
|
|
|
|
def __init__(self, queue_service=None, workflow_executor=None):
|
|
"""Initialize the execution monitor.
|
|
|
|
Args:
|
|
queue_service: Optional queue service for status updates
|
|
workflow_executor: Optional workflow executor for execution monitoring
|
|
"""
|
|
self.queue_service = queue_service
|
|
self.workflow_executor = workflow_executor
|
|
|
|
self._monitoring_thread: Optional[threading.Thread] = None
|
|
self._stop_monitoring = threading.Event()
|
|
self._lock = threading.RLock()
|
|
|
|
# Monitoring state
|
|
self._monitored_workflows: Dict[str, Dict[str, Any]] = {}
|
|
self._status_callbacks: Dict[str, Callable] = {}
|
|
self._monitoring_interval = 1.0 # seconds
|
|
self._is_monitoring = False
|
|
|
|
logger.info("Execution monitor initialized")
|
|
|
|
def start_monitoring(self) -> bool:
|
|
"""Start the execution monitoring thread.
|
|
|
|
Returns:
|
|
True if monitoring was started successfully
|
|
"""
|
|
with self._lock:
|
|
if self._is_monitoring:
|
|
logger.warning("Execution monitoring is already running")
|
|
return True
|
|
|
|
try:
|
|
self._stop_monitoring.clear()
|
|
self._monitoring_thread = threading.Thread(
|
|
target=self._monitoring_loop,
|
|
name="ExecutionMonitor",
|
|
daemon=True
|
|
)
|
|
self._monitoring_thread.start()
|
|
self._is_monitoring = True
|
|
|
|
logger.info("Execution monitoring started")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to start execution monitoring: {e}")
|
|
return False
|
|
|
|
def stop_monitoring(self) -> bool:
|
|
"""Stop the execution monitoring thread.
|
|
|
|
Returns:
|
|
True if monitoring was stopped successfully
|
|
"""
|
|
with self._lock:
|
|
if not self._is_monitoring:
|
|
return True
|
|
|
|
try:
|
|
self._stop_monitoring.set()
|
|
|
|
if self._monitoring_thread and self._monitoring_thread.is_alive():
|
|
self._monitoring_thread.join(timeout=5.0)
|
|
|
|
if self._monitoring_thread.is_alive():
|
|
logger.warning("Monitoring thread did not stop gracefully")
|
|
|
|
self._is_monitoring = False
|
|
self._monitoring_thread = None
|
|
|
|
logger.info("Execution monitoring stopped")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to stop execution monitoring: {e}")
|
|
return False
|
|
|
|
def is_monitoring(self) -> bool:
|
|
"""Check if execution monitoring is active."""
|
|
return self._is_monitoring
|
|
|
|
def set_monitoring_interval(self, interval: float) -> None:
|
|
"""Set the monitoring interval in seconds.
|
|
|
|
Args:
|
|
interval: Monitoring interval in seconds
|
|
"""
|
|
if interval <= 0:
|
|
raise ValueError("Monitoring interval must be positive")
|
|
|
|
self._monitoring_interval = interval
|
|
logger.info(f"Monitoring interval set to {interval} seconds")
|
|
|
|
def add_workflow_to_monitor(self, workflow_id: str, workflow_info: Dict[str, Any]) -> None:
|
|
"""Add a workflow to the monitoring system.
|
|
|
|
Args:
|
|
workflow_id: Unique identifier for the workflow
|
|
workflow_info: Information about the workflow
|
|
"""
|
|
with self._lock:
|
|
# Preserve existing added_at if provided, otherwise use current time
|
|
added_at = workflow_info.get("added_at", datetime.now(timezone.utc))
|
|
|
|
self._monitored_workflows[workflow_id] = {
|
|
**workflow_info,
|
|
"added_at": added_at,
|
|
"last_checked": None,
|
|
"check_count": 0,
|
|
}
|
|
|
|
logger.debug(f"Added workflow to monitor: {workflow_id}")
|
|
|
|
def remove_workflow_from_monitor(self, workflow_id: str) -> bool:
|
|
"""Remove a workflow from the monitoring system.
|
|
|
|
Args:
|
|
workflow_id: Identifier of the workflow to remove
|
|
|
|
Returns:
|
|
True if the workflow was removed
|
|
"""
|
|
with self._lock:
|
|
removed = self._monitored_workflows.pop(workflow_id, None) is not None
|
|
|
|
if removed:
|
|
logger.debug(f"Removed workflow from monitor: {workflow_id}")
|
|
|
|
return removed
|
|
|
|
def register_status_callback(self, callback_id: str, callback: Callable) -> None:
|
|
"""Register a callback for status change events.
|
|
|
|
Args:
|
|
callback_id: Unique identifier for the callback
|
|
callback: Function to call on status changes
|
|
"""
|
|
self._status_callbacks[callback_id] = callback
|
|
logger.debug(f"Registered status callback: {callback_id}")
|
|
|
|
def unregister_status_callback(self, callback_id: str) -> None:
|
|
"""Unregister a status callback.
|
|
|
|
Args:
|
|
callback_id: Identifier of the callback to remove
|
|
"""
|
|
self._status_callbacks.pop(callback_id, None)
|
|
logger.debug(f"Unregistered status callback: {callback_id}")
|
|
|
|
def get_monitored_workflows(self) -> Dict[str, Dict[str, Any]]:
|
|
"""Get information about all monitored workflows.
|
|
|
|
Returns:
|
|
Dictionary mapping workflow IDs to their information
|
|
"""
|
|
with self._lock:
|
|
return {wf_id: info.copy() for wf_id, info in self._monitored_workflows.items()}
|
|
|
|
def get_workflow_status(self, workflow_id: str) -> QueueStatus | None:
|
|
"""Get the current status of a monitored workflow.
|
|
|
|
Args:
|
|
workflow_id: The workflow ID
|
|
|
|
Returns:
|
|
Current status or None if not monitored
|
|
"""
|
|
with self._lock:
|
|
if workflow_id not in self._monitored_workflows:
|
|
return None
|
|
|
|
workflow_info = self._monitored_workflows[workflow_id]
|
|
return workflow_info.get("status")
|
|
|
|
def _monitoring_loop(self) -> None:
|
|
"""Main monitoring loop that runs in a separate thread."""
|
|
logger.info("Execution monitoring loop started")
|
|
|
|
try:
|
|
while not self._stop_monitoring.is_set():
|
|
try:
|
|
self._check_workflow_statuses()
|
|
|
|
# Wait for the monitoring interval or stop signal
|
|
if self._stop_monitoring.wait(timeout=self._monitoring_interval):
|
|
break
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in monitoring loop: {e}")
|
|
# Continue monitoring even if there's an error
|
|
time.sleep(1.0)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Fatal error in monitoring loop: {e}")
|
|
|
|
finally:
|
|
logger.info("Execution monitoring loop stopped")
|
|
|
|
def _check_workflow_statuses(self) -> None:
|
|
"""Check the status of all monitored workflows."""
|
|
with self._lock:
|
|
workflows_to_check = list(self._monitored_workflows.items())
|
|
|
|
for workflow_id, workflow_info in workflows_to_check:
|
|
try:
|
|
self._check_single_workflow_status(workflow_id, workflow_info)
|
|
except Exception as e:
|
|
logger.error(f"Error checking workflow {workflow_id}: {e}")
|
|
|
|
def _check_single_workflow_status(self, workflow_id: str, workflow_info: Dict[str, Any]) -> None:
|
|
"""Check the status of a single workflow.
|
|
|
|
Args:
|
|
workflow_id: The workflow ID
|
|
workflow_info: Current workflow information
|
|
"""
|
|
old_status = workflow_info.get("status")
|
|
new_status = None
|
|
|
|
# Update check information
|
|
now = datetime.now(timezone.utc)
|
|
workflow_info["last_checked"] = now
|
|
workflow_info["check_count"] = workflow_info.get("check_count", 0) + 1
|
|
|
|
# Check status from workflow executor if available
|
|
if self.workflow_executor:
|
|
try:
|
|
new_status = self.workflow_executor.get_workflow_status(workflow_id)
|
|
except Exception as e:
|
|
logger.error(f"Failed to get workflow status from executor: {e}")
|
|
|
|
# Check status from queue service if available and executor didn't provide status
|
|
if new_status is None and self.queue_service:
|
|
try:
|
|
queue_item_id = workflow_info.get("queue_item_id")
|
|
if queue_item_id:
|
|
queue_item = self.queue_service.get_queue_item(queue_item_id)
|
|
if queue_item:
|
|
new_status = queue_item.status
|
|
except Exception as e:
|
|
logger.error(f"Failed to get workflow status from queue service: {e}")
|
|
|
|
# Update status if it changed
|
|
if new_status and new_status != old_status:
|
|
workflow_info["status"] = new_status
|
|
workflow_info["status_changed_at"] = now
|
|
|
|
# Notify callbacks
|
|
self._notify_status_callbacks(workflow_id, old_status, new_status, workflow_info)
|
|
|
|
# Remove completed/failed workflows from monitoring after a delay
|
|
if new_status in (QueueStatus.COMPLETED, QueueStatus.FAILED):
|
|
self._schedule_workflow_removal(workflow_id)
|
|
|
|
def _notify_status_callbacks(
|
|
self,
|
|
workflow_id: str,
|
|
old_status: QueueStatus | None,
|
|
new_status: QueueStatus,
|
|
workflow_info: Dict[str, Any]
|
|
) -> None:
|
|
"""Notify all registered callbacks about a status change.
|
|
|
|
Args:
|
|
workflow_id: The workflow ID
|
|
old_status: Previous status
|
|
new_status: New status
|
|
workflow_info: Current workflow information
|
|
"""
|
|
callback_data = {
|
|
"workflow_id": workflow_id,
|
|
"old_status": old_status.value if old_status else None,
|
|
"new_status": new_status.value,
|
|
"workflow_info": workflow_info.copy(),
|
|
"timestamp": datetime.now(timezone.utc).isoformat(),
|
|
}
|
|
|
|
for callback_id, callback in list(self._status_callbacks.items()):
|
|
try:
|
|
callback(callback_data)
|
|
except Exception as e:
|
|
logger.error(f"Error in status callback {callback_id}: {e}")
|
|
|
|
def _schedule_workflow_removal(self, workflow_id: str, delay_seconds: float = 60.0) -> None:
|
|
"""Schedule a workflow for removal from monitoring after a delay.
|
|
|
|
Args:
|
|
workflow_id: The workflow ID
|
|
delay_seconds: Delay before removal in seconds
|
|
"""
|
|
def remove_after_delay():
|
|
time.sleep(delay_seconds)
|
|
self.remove_workflow_from_monitor(workflow_id)
|
|
|
|
removal_thread = threading.Thread(
|
|
target=remove_after_delay,
|
|
name=f"RemoveWorkflow-{workflow_id}",
|
|
daemon=True
|
|
)
|
|
removal_thread.start()
|
|
|
|
def get_monitoring_statistics(self) -> Dict[str, Any]:
|
|
"""Get statistics about the monitoring system.
|
|
|
|
Returns:
|
|
Dictionary containing monitoring statistics
|
|
"""
|
|
with self._lock:
|
|
stats = {
|
|
"is_monitoring": self._is_monitoring,
|
|
"monitoring_interval": self._monitoring_interval,
|
|
"monitored_workflows": len(self._monitored_workflows),
|
|
"registered_callbacks": len(self._status_callbacks),
|
|
"status_counts": {},
|
|
}
|
|
|
|
# Count workflows by status
|
|
for workflow_info in self._monitored_workflows.values():
|
|
status = workflow_info.get("status")
|
|
if status:
|
|
status_value = status.value if hasattr(status, 'value') else str(status)
|
|
stats["status_counts"][status_value] = stats["status_counts"].get(status_value, 0) + 1
|
|
|
|
return stats
|
|
|
|
def cleanup_old_workflows(self, max_age_hours: int = 24) -> int:
|
|
"""Clean up old workflow monitoring records.
|
|
|
|
Args:
|
|
max_age_hours: Maximum age in hours for workflow records
|
|
|
|
Returns:
|
|
Number of workflows cleaned up
|
|
"""
|
|
cutoff_time = datetime.now(timezone.utc) - timedelta(hours=max_age_hours)
|
|
cleaned_count = 0
|
|
|
|
with self._lock:
|
|
workflows_to_remove = []
|
|
|
|
for workflow_id, workflow_info in self._monitored_workflows.items():
|
|
added_at = workflow_info.get("added_at")
|
|
if added_at and added_at < cutoff_time:
|
|
# Only remove completed or failed workflows
|
|
status = workflow_info.get("status")
|
|
if status in (QueueStatus.COMPLETED, QueueStatus.FAILED):
|
|
workflows_to_remove.append(workflow_id)
|
|
|
|
for workflow_id in workflows_to_remove:
|
|
del self._monitored_workflows[workflow_id]
|
|
cleaned_count += 1
|
|
|
|
if cleaned_count > 0:
|
|
logger.info(f"Cleaned up {cleaned_count} old workflow monitoring records")
|
|
|
|
return cleaned_count
|
|
|
|
def close(self) -> None:
|
|
"""Close the execution monitor and clean up resources."""
|
|
self.stop_monitoring()
|
|
|
|
with self._lock:
|
|
self._monitored_workflows.clear()
|
|
self._status_callbacks.clear()
|
|
|
|
logger.info("Execution monitor closed")
|
|
|
|
def __enter__(self):
|
|
"""Context manager entry."""
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""Context manager exit."""
|
|
self.close() |