nodes/abdullahceylan-ac-comfyui-queue-manager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add error handling and logging

Browse Source
This commit is contained in:
Abdullah Ceylan
2025-09-27 19:40:36 +03:00
parent a94b0f1ada
commit f433da681d
12 changed files with 3975 additions and 206 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+40 -2
View File
@@ -123,8 +123,29 @@ def register_menu_extension():
def initialize_queue_manager():
"""Initialize the queue manager system."""
try:
# Initialize logging and monitoring first
from logging_config import setup_logging
from performance_monitor import setup_performance_monitoring
from health_check import setup_health_checks
# Set up logging
log_dir = current_dir / "logs"
logging_config = setup_logging(
log_dir=log_dir,
console_level=20, # INFO level
file_level=10, # DEBUG level
structured=True
)
logger = logging_config.get_logger("init")
logger.info(f"Queue Manager logging initialized (v{__version__})")
# Set up performance monitoring
perf_monitor = setup_performance_monitoring(max_history_size=1000)
logger.info("Performance monitoring initialized")
# Import and initialize core services
from api_routes import setup_api_routes
from api_routes import QueueManagerAPI
from database import SQLiteDatabase
from queue_service import QueueService
@@ -132,17 +153,34 @@ def initialize_queue_manager():
db_path = current_dir / "queue_manager.db"
database = SQLiteDatabase(str(db_path))
database.initialize()
logger.info("Database initialized")
# Initialize queue service
queue_service = QueueService(database)
logger.info("Queue service initialized")
# Set up health checks with actual services
health_manager = setup_health_checks(database=database, queue_service=queue_service)
health_manager.start_monitoring(interval=300.0) # 5 minutes
logger.info("Health monitoring started")
# Set up API routes (will be used by the web interface)
setup_api_routes(queue_service)
api = QueueManagerAPI(queue_service)
logger.info("API routes initialized")
logger.info(f"Queue Manager initialized successfully (v{__version__})")
print(f"[Queue Manager] Initialized successfully (v{__version__})")
return True
except Exception as e:
# Fall back to basic logging if initialization fails
import logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("queue_manager.init")
logger.error(f"Failed to initialize queue manager: {e}")
print(f"[Queue Manager] Failed to initialize: {e}")
return False
api_routes.py
+115 -50
View File
@@ -12,10 +12,14 @@ from typing import Any
from flask import Flask, jsonify, request, send_from_directory
from werkzeug.exceptions import BadRequest, NotFound
from error_handler import ErrorHandler, create_error_response, with_api_error_handling
from exceptions import APIError, NotFoundError, QueueManagerError, ValidationError
from logging_config import get_logger
from models import QueueFilter, QueueStatus
from queue_service import QueueService, QueueServiceError
from performance_monitor import get_performance_monitor, time_function
from queue_service import QueueService
logger = logging.getLogger(__name__)
logger = get_logger("api_routes")
class QueueManagerAPI:
@@ -42,8 +46,10 @@ class QueueManagerAPI:
def _register_routes(self) -> None:
"""Register all API routes."""
# Health check
# Health check endpoints
self.app.route('/health', methods=['GET'])(self.health_check)
self.app.route('/health/detailed', methods=['GET'])(self.detailed_health_check)
self.app.route('/health/metrics', methods=['GET'])(self.health_metrics)
# Static files
self.app.route('/', methods=['GET'])(self.serve_index)
@@ -73,64 +79,115 @@ class QueueManagerAPI:
def _register_error_handlers(self) -> None:
"""Register error handlers for the Flask app."""
@self.app.errorhandler(QueueServiceError)
def handle_queue_service_error(error: QueueServiceError) -> tuple[dict[str, Any], int]:
"""Handle queue service errors."""
logger.error(f"Queue service error: {error}")
return {
'error': 'Queue service error',
'message': str(error),
'timestamp': datetime.utcnow().isoformat()
}, 500
@self.app.errorhandler(QueueManagerError)
def handle_queue_manager_error(error: QueueManagerError) -> tuple[dict[str, Any], int]:
"""Handle queue manager errors."""
ErrorHandler.log_error(error, context={"handler": "flask_error_handler"})
# Determine HTTP status code
if isinstance(error, APIError):
status_code = error.details.get('status_code', 500)
elif isinstance(error, ValidationError):
status_code = 400
elif isinstance(error, NotFoundError):
status_code = 404
else:
status_code = 500
return create_error_response(error), status_code
@self.app.errorhandler(BadRequest)
def handle_bad_request(error: BadRequest) -> tuple[dict[str, Any], int]:
"""Handle bad request errors."""
logger.warning(f"Bad request: {error}")
return {
'error': 'Bad request',
'message': str(error.description),
'timestamp': datetime.utcnow().isoformat()
}, 400
validation_error = ValidationError(
str(error.description) or "Bad request",
details={"werkzeug_error": str(error)}
)
return create_error_response(validation_error), 400
@self.app.errorhandler(NotFound)
def handle_not_found(error: NotFound) -> tuple[dict[str, Any], int]:
"""Handle not found errors."""
return {
'error': 'Not found',
'message': 'The requested resource was not found',
'timestamp': datetime.utcnow().isoformat()
}, 404
not_found_error = NotFoundError(
"The requested resource was not found",
details={"werkzeug_error": str(error)}
)
return create_error_response(not_found_error), 404
@self.app.errorhandler(Exception)
def handle_generic_error(error: Exception) -> tuple[dict[str, Any], int]:
"""Handle generic errors."""
logger.error(f"Unexpected error: {error}", exc_info=True)
return {
'error': 'Internal server error',
'message': 'An unexpected error occurred',
'timestamp': datetime.utcnow().isoformat()
}, 500
# Convert to structured error
api_error = APIError(
"An unexpected error occurred",
status_code=500,
details={"original_error": str(error)},
cause=error
)
return create_error_response(api_error, include_traceback=False), 500
@with_api_error_handling(endpoint="/health", method="GET")
@time_function("api.health_check")
def health_check(self) -> dict[str, Any]:
"""Health check endpoint."""
# Test queue service connectivity
queue_state = self.queue_service.get_queue_state()
return {
'status': 'healthy',
'timestamp': datetime.utcnow().isoformat(),
'queue_state': queue_state.value,
'version': '1.0.0'
}
@with_api_error_handling(endpoint="/health/detailed", method="GET")
@time_function("api.detailed_health_check")
def detailed_health_check(self) -> dict[str, Any]:
"""Detailed health check endpoint with comprehensive diagnostics."""
from health_check import get_health_check_manager
import asyncio
# Get health check manager and run all checks
health_manager = get_health_check_manager()
# Run health checks in a new event loop
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
# Test queue service connectivity
queue_state = self.queue_service.get_queue_state()
return {
'status': 'healthy',
'timestamp': datetime.utcnow().isoformat(),
'queue_state': queue_state.value,
'version': '1.0.0'
}
except Exception as e:
logger.error(f"Health check failed: {e}")
return {
'status': 'unhealthy',
'timestamp': datetime.utcnow().isoformat(),
'error': str(e)
}, 503
results = loop.run_until_complete(health_manager.run_all_checks())
overall_health = health_manager.get_overall_health()
finally:
loop.close()
return overall_health
@with_api_error_handling(endpoint="/health/metrics", method="GET")
@time_function("api.health_metrics")
def health_metrics(self) -> dict[str, Any]:
"""Health metrics endpoint with performance data."""
perf_monitor = get_performance_monitor()
# Get metrics summary for the last hour
since = datetime.utcnow() - timedelta(hours=1)
metrics_summary = perf_monitor.get_metrics_summary(since)
# Get operation statistics
operation_stats = perf_monitor.get_operation_stats()
# Get health status
health_status = perf_monitor.get_health_status()
return {
'timestamp': datetime.utcnow().isoformat(),
'health_status': health_status,
'metrics_summary': metrics_summary,
'operation_stats': operation_stats
}
def serve_index(self) -> str:
"""Serve the main index.html file."""
@@ -146,15 +203,15 @@ class QueueManagerAPI:
def _validate_json_request(self) -> dict[str, Any]:
"""Validate and parse JSON request data."""
if not request.is_json:
raise BadRequest("Request must be JSON")
raise ValidationError("Request must be JSON", field="content_type")
try:
data = request.get_json()
if data is None:
raise BadRequest("Invalid JSON data")
raise ValidationError("Invalid JSON data", field="request_body")
return data
except Exception as e:
raise BadRequest(f"Failed to parse JSON: {e}")
raise ValidationError(f"Failed to parse JSON: {e}", field="request_body", cause=e)
def _parse_queue_filter(self, params: dict[str, Any]) -> QueueFilter:
"""Parse queue filter parameters from request."""
@@ -169,15 +226,21 @@ class QueueManagerAPI:
try:
filter_criteria.status = [QueueStatus(status) for status in status_values]
except ValueError as e:
raise BadRequest(f"Invalid status value: {e}")
raise ValidationError(f"Invalid status value: {e}", field="status", value=status_values)
# Parse workflow name filter
if 'workflow_name' in params:
filter_criteria.workflow_name = params['workflow_name']
workflow_name = params['workflow_name']
if not isinstance(workflow_name, str):
raise ValidationError("Workflow name must be a string", field="workflow_name", value=workflow_name)
filter_criteria.workflow_name = workflow_name
# Parse search term
if 'search' in params:
filter_criteria.search_term = params['search']
search_term = params['search']
if not isinstance(search_term, str):
raise ValidationError("Search term must be a string", field="search", value=search_term)
filter_criteria.search_term = search_term
# Parse date range
if 'date_from' in params or 'date_to' in params:
@@ -192,10 +255,12 @@ class QueueManagerAPI:
date_to = datetime.fromisoformat(params['date_to'])
if date_from and date_to:
if date_from > date_to:
raise ValidationError("date_from must be before date_to", field="date_range")
filter_criteria.date_range = (date_from, date_to)
except ValueError as e:
raise BadRequest(f"Invalid date format: {e}")
raise ValidationError(f"Invalid date format: {e}", field="date_range", cause=e)
return filter_criteria
database.py
+130 -121
View File
@@ -13,14 +13,14 @@ from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Generator
from error_handler import ErrorHandler, with_database_error_handling
from exceptions import ConnectionError, DatabaseError, SchemaError, ValidationError
from interfaces import DatabaseInterface
from logging_config import get_logger
from models import QueueConfig, QueueFilter, QueueItem, QueueStatus
from performance_monitor import get_performance_monitor, time_function
logger = logging.getLogger(__name__)
class DatabaseError(Exception):
"""Custom exception for database operations."""
logger = get_logger("database")
class SQLiteDatabase(DatabaseInterface):
@@ -95,7 +95,7 @@ class SQLiteDatabase(DatabaseInterface):
logger.debug(f"Created new database connection for thread {threading.current_thread().name}")
except sqlite3.Error as e:
logger.error(f"Failed to connect to database: {e}")
raise DatabaseError(f"Database connection failed: {e}") from e
raise ErrorHandler.handle_database_error(e, "connect", context={"db_path": str(self.db_path)})
return self._local.connection
@@ -108,47 +108,53 @@ class SQLiteDatabase(DatabaseInterface):
yield cursor
except sqlite3.Error as e:
logger.error(f"Database operation failed: {e}")
conn.rollback()
raise DatabaseError(f"Database operation failed: {e}") from e
try:
conn.rollback()
except sqlite3.Error as rollback_error:
logger.error(f"Failed to rollback transaction: {rollback_error}")
raise ErrorHandler.handle_database_error(e, "cursor_operation")
finally:
cursor.close()
try:
cursor.close()
except sqlite3.Error as close_error:
logger.warning(f"Failed to close cursor: {close_error}")
@with_database_error_handling(operation="initialize")
def initialize(self) -> bool:
"""Initialize the database and create tables if needed."""
if self._initialized:
return True
with self._connection_lock:
try:
with self._get_cursor() as cursor:
# Create all tables
cursor.executescript(self.CREATE_TABLES_SQL)
# Check and update schema version
cursor.execute("SELECT version FROM schema_version LIMIT 1")
result = cursor.fetchone()
if result is None:
# First time setup
cursor.execute("INSERT INTO schema_version (version) VALUES (?)", (self.SCHEMA_VERSION,))
logger.info(f"Database initialized with schema version {self.SCHEMA_VERSION}")
else:
current_version = result[0]
if current_version < self.SCHEMA_VERSION:
self._migrate_schema(cursor, current_version, self.SCHEMA_VERSION)
elif current_version > self.SCHEMA_VERSION:
logger.warning(f"Database schema version {current_version} is newer than expected {self.SCHEMA_VERSION}")
# Set default configuration if not exists
self._initialize_default_config(cursor)
self._initialized = True
logger.info("Database initialization completed successfully")
return True
with self._get_cursor() as cursor:
# Create all tables
cursor.executescript(self.CREATE_TABLES_SQL)
except Exception as e:
logger.error(f"Database initialization failed: {e}")
return False
# Check and update schema version
cursor.execute("SELECT version FROM schema_version LIMIT 1")
result = cursor.fetchone()
if result is None:
# First time setup
cursor.execute("INSERT INTO schema_version (version) VALUES (?)", (self.SCHEMA_VERSION,))
logger.info(f"Database initialized with schema version {self.SCHEMA_VERSION}")
else:
current_version = result[0]
if current_version < self.SCHEMA_VERSION:
self._migrate_schema(cursor, current_version, self.SCHEMA_VERSION)
elif current_version > self.SCHEMA_VERSION:
logger.warning(f"Database schema version {current_version} is newer than expected {self.SCHEMA_VERSION}")
raise SchemaError(
f"Database schema version {current_version} is newer than expected {self.SCHEMA_VERSION}",
schema_version=current_version
)
# Set default configuration if not exists
self._initialize_default_config(cursor)
self._initialized = True
logger.info("Database initialization completed successfully")
return True
def _migrate_schema(self, cursor: sqlite3.Cursor, from_version: int, to_version: int) -> None:
"""Migrate database schema from one version to another."""
@@ -171,105 +177,108 @@ class SQLiteDatabase(DatabaseInterface):
(key, json.dumps(value))
)
@with_database_error_handling(operation="INSERT", table="queue_items")
@time_function("database.create_queue_item")
def create_queue_item(self, item: QueueItem) -> bool:
"""Create a new queue item in the database."""
try:
with self._get_cursor() as cursor:
cursor.execute("""
INSERT INTO queue_items (
id, workflow_name, workflow_data, status, created_at, updated_at,
started_at, completed_at, error_message, result_data
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""", (
item.id,
item.workflow_name,
json.dumps(item.workflow_data),
item.status.value,
item.created_at.isoformat(),
item.updated_at.isoformat(),
item.started_at.isoformat() if item.started_at else None,
item.completed_at.isoformat() if item.completed_at else None,
item.error_message,
json.dumps(item.result_data) if item.result_data else None
))
logger.debug(f"Created queue item {item.id}")
return True
except Exception as e:
logger.error(f"Failed to create queue item {item.id}: {e}")
return False
# Validate item data
if not item.id:
raise ValidationError("Queue item ID cannot be empty", field="id")
if not item.workflow_name:
raise ValidationError("Workflow name cannot be empty", field="workflow_name")
with self._get_cursor() as cursor:
cursor.execute("""
INSERT INTO queue_items (
id, workflow_name, workflow_data, status, created_at, updated_at,
started_at, completed_at, error_message, result_data
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""", (
item.id,
item.workflow_name,
json.dumps(item.workflow_data),
item.status.value,
item.created_at.isoformat(),
item.updated_at.isoformat(),
item.started_at.isoformat() if item.started_at else None,
item.completed_at.isoformat() if item.completed_at else None,
item.error_message,
json.dumps(item.result_data) if item.result_data else None
))
logger.debug(f"Created queue item {item.id}")
return True
@with_database_error_handling(operation="SELECT", table="queue_items")
@time_function("database.get_queue_item")
def get_queue_item(self, item_id: str) -> QueueItem | None:
"""Retrieve a queue item by ID."""
try:
with self._get_cursor() as cursor:
cursor.execute("SELECT * FROM queue_items WHERE id = ?", (item_id,))
row = cursor.fetchone()
if row:
return self._row_to_queue_item(row)
return None
except Exception as e:
logger.error(f"Failed to get queue item {item_id}: {e}")
if not item_id:
raise ValidationError("Item ID cannot be empty", field="item_id")
with self._get_cursor() as cursor:
cursor.execute("SELECT * FROM queue_items WHERE id = ?", (item_id,))
row = cursor.fetchone()
if row:
return self._row_to_queue_item(row)
return None
@with_database_error_handling(operation="SELECT", table="queue_items")
@time_function("database.get_all_queue_items")
def get_all_queue_items(self) -> list[QueueItem]:
"""Retrieve all queue items."""
try:
with self._get_cursor() as cursor:
cursor.execute("SELECT * FROM queue_items ORDER BY created_at DESC")
rows = cursor.fetchall()
return [self._row_to_queue_item(row) for row in rows]
except Exception as e:
logger.error(f"Failed to get all queue items: {e}")
return []
with self._get_cursor() as cursor:
cursor.execute("SELECT * FROM queue_items ORDER BY created_at DESC")
rows = cursor.fetchall()
return [self._row_to_queue_item(row) for row in rows]
@with_database_error_handling(operation="UPDATE", table="queue_items")
def update_queue_item(self, item: QueueItem) -> bool:
"""Update an existing queue item."""
try:
# Update the updated_at timestamp
item.updated_at = datetime.now(timezone.utc)
if not item.id:
raise ValidationError("Queue item ID cannot be empty", field="id")
# Update the updated_at timestamp
item.updated_at = datetime.now(timezone.utc)
with self._get_cursor() as cursor:
cursor.execute("""
UPDATE queue_items SET
workflow_name = ?, workflow_data = ?, status = ?, updated_at = ?,
started_at = ?, completed_at = ?, error_message = ?, result_data = ?
WHERE id = ?
""", (
item.workflow_name,
json.dumps(item.workflow_data),
item.status.value,
item.updated_at.isoformat(),
item.started_at.isoformat() if item.started_at else None,
item.completed_at.isoformat() if item.completed_at else None,
item.error_message,
json.dumps(item.result_data) if item.result_data else None,
item.id
))
with self._get_cursor() as cursor:
cursor.execute("""
UPDATE queue_items SET
workflow_name = ?, workflow_data = ?, status = ?, updated_at = ?,
started_at = ?, completed_at = ?, error_message = ?, result_data = ?
WHERE id = ?
""", (
item.workflow_name,
json.dumps(item.workflow_data),
item.status.value,
item.updated_at.isoformat(),
item.started_at.isoformat() if item.started_at else None,
item.completed_at.isoformat() if item.completed_at else None,
item.error_message,
json.dumps(item.result_data) if item.result_data else None,
item.id
))
if cursor.rowcount > 0:
logger.debug(f"Updated queue item {item.id}")
return True
else:
logger.warning(f"Queue item {item.id} not found for update")
return False
except Exception as e:
logger.error(f"Failed to update queue item {item.id}: {e}")
return False
if cursor.rowcount > 0:
logger.debug(f"Updated queue item {item.id}")
return True
else:
logger.warning(f"Queue item {item.id} not found for update")
return False
@with_database_error_handling(operation="DELETE", table="queue_items")
def delete_queue_item(self, item_id: str) -> bool:
"""Delete a queue item by ID."""
try:
with self._get_cursor() as cursor:
cursor.execute("DELETE FROM queue_items WHERE id = ?", (item_id,))
if cursor.rowcount > 0:
logger.debug(f"Deleted queue item {item_id}")
return True
else:
logger.warning(f"Queue item {item_id} not found for deletion")
return False
except Exception as e:
logger.error(f"Failed to delete queue item {item_id}: {e}")
return False
if not item_id:
raise ValidationError("Item ID cannot be empty", field="item_id")
with self._get_cursor() as cursor:
cursor.execute("DELETE FROM queue_items WHERE id = ?", (item_id,))
if cursor.rowcount > 0:
logger.debug(f"Deleted queue item {item_id}")
return True
else:
logger.warning(f"Queue item {item_id} not found for deletion")
return False
def get_items_by_status(self, status: QueueStatus) -> list[QueueItem]:
"""Retrieve queue items by status."""
error_handler.py
+467
View File
@@ -0,0 +1,467 @@
"""
Error handling utilities for the ComfyUI Queue Manager.
Provides centralized error handling, logging, and recovery mechanisms.
"""
from __future__ import annotations
import functools
import logging
import sqlite3
import traceback
from datetime import datetime, timezone
from typing import Any, Callable, Dict, Optional, Type, TypeVar, Union
from exceptions import (
APIError,
ConnectionError,
DatabaseError,
QueueManagerError,
QueueServiceError,
SchemaError,
TimeoutError,
ValidationError,
WorkflowExecutionError,
)
logger = logging.getLogger(__name__)
T = TypeVar('T')
class ErrorHandler:
"""Centralized error handling and recovery utilities."""
@staticmethod
def handle_database_error(
error: Exception,
operation: str,
table: Optional[str] = None,
context: Optional[Dict[str, Any]] = None
) -> DatabaseError:
"""Handle and convert database errors to structured exceptions.
Args:
error: Original exception
operation: Database operation that failed
table: Database table involved
context: Additional context information
Returns:
Structured DatabaseError
"""
context = context or {}
# Handle SQLite-specific errors
if isinstance(error, sqlite3.Error):
if isinstance(error, sqlite3.OperationalError):
error_msg = str(error).lower()
if "database is locked" in error_msg:
return ConnectionError(
f"Database is locked during {operation}",
operation=operation,
table=table,
details=context,
cause=error
)
elif "no such table" in error_msg:
return SchemaError(
f"Table does not exist during {operation}",
operation=operation,
table=table,
details=context,
cause=error
)
elif "disk i/o error" in error_msg:
return DatabaseError(
f"Disk I/O error during {operation}",
error_code="DB_IO_ERROR",
operation=operation,
table=table,
details=context,
cause=error
)
else:
return DatabaseError(
f"Database operational error during {operation}: {error}",
error_code="DB_OPERATIONAL_ERROR",
operation=operation,
table=table,
details=context,
cause=error
)
elif isinstance(error, sqlite3.IntegrityError):
return ValidationError(
f"Data integrity violation during {operation}: {error}",
details={**context, "operation": operation, "table": table},
cause=error
)
elif isinstance(error, sqlite3.DatabaseError):
return DatabaseError(
f"Database error during {operation}: {error}",
error_code="DB_GENERAL_ERROR",
operation=operation,
table=table,
details=context,
cause=error
)
# Handle general exceptions
return DatabaseError(
f"Unexpected error during {operation}: {error}",
error_code="DB_UNEXPECTED_ERROR",
operation=operation,
table=table,
details=context,
cause=error
)
@staticmethod
def handle_api_error(
error: Exception,
endpoint: str,
method: str = "GET",
context: Optional[Dict[str, Any]] = None
) -> APIError:
"""Handle and convert API errors to structured exceptions.
Args:
error: Original exception
endpoint: API endpoint that failed
method: HTTP method
context: Additional context information
Returns:
Structured APIError
"""
context = context or {}
context.update({"endpoint": endpoint, "method": method})
if isinstance(error, QueueManagerError):
# Convert existing queue manager errors to API errors
return APIError(
error.message,
status_code=500,
endpoint=endpoint,
details=context,
cause=error
)
return APIError(
f"API error at {method} {endpoint}: {error}",
status_code=500,
endpoint=endpoint,
details=context,
cause=error
)
@staticmethod
def handle_workflow_error(
error: Exception,
workflow_id: Optional[str] = None,
stage: Optional[str] = None,
context: Optional[Dict[str, Any]] = None
) -> WorkflowExecutionError:
"""Handle and convert workflow execution errors.
Args:
error: Original exception
workflow_id: ID of the workflow that failed
stage: Execution stage where error occurred
context: Additional context information
Returns:
Structured WorkflowExecutionError
"""
context = context or {}
return WorkflowExecutionError(
f"Workflow execution failed: {error}",
workflow_id=workflow_id,
execution_stage=stage,
details=context,
cause=error
)
@staticmethod
def log_error(
error: Exception,
context: Optional[Dict[str, Any]] = None,
level: int = logging.ERROR
) -> None:
"""Log an error with structured information.
Args:
error: Exception to log
context: Additional context information
level: Logging level
"""
context = context or {}
if isinstance(error, QueueManagerError):
# Log structured queue manager errors
log_data = {
"error_type": error.__class__.__name__,
"error_code": error.error_code,
"message": error.message,
"details": error.details,
**context
}
logger.log(level, f"Queue Manager Error: {error.message}", extra=log_data)
if error.cause:
logger.log(level, f"Caused by: {error.cause}", exc_info=error.cause)
else:
# Log general exceptions
logger.log(level, f"Unexpected error: {error}", extra=context, exc_info=error)
def with_error_handling(
error_type: Type[QueueManagerError] = QueueManagerError,
operation: Optional[str] = None,
log_errors: bool = True,
reraise: bool = True
):
"""Decorator for automatic error handling and logging.
Args:
error_type: Type of exception to convert errors to
operation: Name of the operation for context
log_errors: Whether to log errors automatically
reraise: Whether to reraise the converted exception
"""
def decorator(func: Callable[..., T]) -> Callable[..., T]:
@functools.wraps(func)
def wrapper(*args, **kwargs) -> T:
try:
return func(*args, **kwargs)
except QueueManagerError:
# Re-raise queue manager errors as-is
raise
except Exception as e:
# Convert other exceptions
operation_name = operation or func.__name__
if error_type == DatabaseError:
converted_error = ErrorHandler.handle_database_error(
e, operation_name, context={"function": func.__name__}
)
elif error_type == APIError:
converted_error = ErrorHandler.handle_api_error(
e, operation_name, context={"function": func.__name__}
)
elif error_type == WorkflowExecutionError:
converted_error = ErrorHandler.handle_workflow_error(
e, context={"function": func.__name__}
)
else:
converted_error = error_type(
f"Error in {operation_name}: {e}",
details={"function": func.__name__},
cause=e
)
if log_errors:
ErrorHandler.log_error(converted_error)
if reraise:
raise converted_error
return converted_error
return wrapper
return decorator
def with_database_error_handling(
operation: Optional[str] = None,
table: Optional[str] = None,
log_errors: bool = True
):
"""Decorator specifically for database operations.
Args:
operation: Database operation name
table: Database table name
log_errors: Whether to log errors automatically
"""
def decorator(func: Callable[..., T]) -> Callable[..., T]:
@functools.wraps(func)
def wrapper(*args, **kwargs) -> T:
try:
return func(*args, **kwargs)
except QueueManagerError:
# Re-raise queue manager errors as-is
raise
except Exception as e:
operation_name = operation or func.__name__
converted_error = ErrorHandler.handle_database_error(
e,
operation_name,
table=table,
context={"function": func.__name__}
)
if log_errors:
ErrorHandler.log_error(converted_error)
raise converted_error
return wrapper
return decorator
def with_api_error_handling(
endpoint: Optional[str] = None,
method: str = "GET",
log_errors: bool = True
):
"""Decorator specifically for API operations.
Args:
endpoint: API endpoint name
method: HTTP method
log_errors: Whether to log errors automatically
"""
def decorator(func: Callable[..., T]) -> Callable[..., T]:
@functools.wraps(func)
def wrapper(*args, **kwargs) -> T:
try:
return func(*args, **kwargs)
except QueueManagerError:
# Re-raise queue manager errors as-is
raise
except Exception as e:
endpoint_name = endpoint or func.__name__
converted_error = ErrorHandler.handle_api_error(
e,
endpoint_name,
method=method,
context={"function": func.__name__}
)
if log_errors:
ErrorHandler.log_error(converted_error)
raise converted_error
return wrapper
return decorator
class ErrorRecovery:
"""Utilities for error recovery and retry mechanisms."""
@staticmethod
def retry_with_backoff(
func: Callable[..., T],
max_retries: int = 3,
base_delay: float = 1.0,
max_delay: float = 60.0,
backoff_factor: float = 2.0,
exceptions: tuple = (Exception,)
) -> T:
"""Retry a function with exponential backoff.
Args:
func: Function to retry
max_retries: Maximum number of retry attempts
base_delay: Initial delay between retries
max_delay: Maximum delay between retries
backoff_factor: Factor to multiply delay by each retry
exceptions: Tuple of exceptions to catch and retry on
Returns:
Result of the function call
Raises:
The last exception if all retries fail
"""
import time
last_exception = None
delay = base_delay
for attempt in range(max_retries + 1):
try:
return func()
except exceptions as e:
last_exception = e
if attempt == max_retries:
# Last attempt failed, re-raise
break
logger.warning(
f"Attempt {attempt + 1} failed, retrying in {delay:.2f}s: {e}"
)
time.sleep(delay)
delay = min(delay * backoff_factor, max_delay)
# All retries failed
if last_exception:
raise last_exception
else:
raise RuntimeError("All retry attempts failed")
@staticmethod
def safe_execute(
func: Callable[..., T],
default_value: Optional[T] = None,
log_errors: bool = True,
exceptions: tuple = (Exception,)
) -> Optional[T]:
"""Execute a function safely, returning a default value on error.
Args:
func: Function to execute
default_value: Value to return on error
log_errors: Whether to log errors
exceptions: Tuple of exceptions to catch
Returns:
Function result or default value
"""
try:
return func()
except exceptions as e:
if log_errors:
logger.error(f"Safe execution failed: {e}", exc_info=True)
return default_value
def create_error_response(
error: Exception,
include_traceback: bool = False
) -> Dict[str, Any]:
"""Create a standardized error response dictionary.
Args:
error: Exception to convert to response
include_traceback: Whether to include traceback in response
Returns:
Dictionary containing error information
"""
if isinstance(error, QueueManagerError):
response = error.to_dict()
else:
response = {
"error_type": error.__class__.__name__,
"error_code": "UNEXPECTED_ERROR",
"message": str(error),
"details": {},
"cause": None
}
response["timestamp"] = datetime.now(timezone.utc).isoformat()
if include_traceback:
response["traceback"] = traceback.format_exc()
return response
exceptions.py
+450
View File
@@ -0,0 +1,450 @@
"""
Custom exception classes for the ComfyUI Queue Manager.
Provides structured error handling across all components.
"""
from __future__ import annotations
from typing import Any, Dict, Optional
class QueueManagerError(Exception):
"""Base exception class for all queue manager errors."""
def __init__(
self,
message: str,
error_code: Optional[str] = None,
details: Optional[Dict[str, Any]] = None,
cause: Optional[Exception] = None
):
"""Initialize the exception.
Args:
message: Human-readable error message
error_code: Machine-readable error code
details: Additional error details
cause: Original exception that caused this error
"""
super().__init__(message)
self.message = message
self.error_code = error_code or self.__class__.__name__
self.details = details or {}
self.cause = cause
def to_dict(self) -> Dict[str, Any]:
"""Convert the exception to a dictionary for serialization."""
return {
"error_type": self.__class__.__name__,
"error_code": self.error_code,
"message": self.message,
"details": self.details,
"cause": str(self.cause) if self.cause else None
}
class DatabaseError(QueueManagerError):
"""Exception raised for database-related errors."""
def __init__(
self,
message: str,
operation: Optional[str] = None,
table: Optional[str] = None,
**kwargs
):
"""Initialize the database error.
Args:
message: Error message
operation: Database operation that failed (e.g., 'INSERT', 'UPDATE')
table: Database table involved in the operation
**kwargs: Additional arguments passed to parent
"""
details = kwargs.get('details', {})
if operation:
details['operation'] = operation
if table:
details['table'] = table
super().__init__(message, details=details, **kwargs)
class ConnectionError(DatabaseError):
"""Exception raised for database connection errors."""
def __init__(self, message: str, **kwargs):
super().__init__(
message,
error_code="DB_CONNECTION_ERROR",
**kwargs
)
class SchemaError(DatabaseError):
"""Exception raised for database schema-related errors."""
def __init__(self, message: str, schema_version: Optional[int] = None, **kwargs):
details = kwargs.get('details', {})
if schema_version is not None:
details['schema_version'] = schema_version
super().__init__(
message,
error_code="DB_SCHEMA_ERROR",
details=details,
**kwargs
)
class ValidationError(QueueManagerError):
"""Exception raised for data validation errors."""
def __init__(
self,
message: str,
field: Optional[str] = None,
value: Optional[Any] = None,
**kwargs
):
"""Initialize the validation error.
Args:
message: Error message
field: Field name that failed validation
value: Value that failed validation
**kwargs: Additional arguments passed to parent
"""
details = kwargs.get('details', {})
if field:
details['field'] = field
if value is not None:
details['value'] = str(value) # Convert to string for serialization
super().__init__(
message,
error_code="VALIDATION_ERROR",
details=details,
**kwargs
)
class QueueServiceError(QueueManagerError):
"""Exception raised for queue service operations."""
def __init__(
self,
message: str,
item_id: Optional[str] = None,
operation: Optional[str] = None,
**kwargs
):
"""Initialize the queue service error.
Args:
message: Error message
item_id: Queue item ID involved in the operation
operation: Queue operation that failed
**kwargs: Additional arguments passed to parent
"""
details = kwargs.get('details', {})
if item_id:
details['item_id'] = item_id
if operation:
details['operation'] = operation
super().__init__(
message,
error_code="QUEUE_SERVICE_ERROR",
details=details,
**kwargs
)
class WorkflowExecutionError(QueueManagerError):
"""Exception raised for workflow execution errors."""
def __init__(
self,
message: str,
workflow_id: Optional[str] = None,
execution_stage: Optional[str] = None,
**kwargs
):
"""Initialize the workflow execution error.
Args:
message: Error message
workflow_id: Workflow ID that failed
execution_stage: Stage of execution where error occurred
**kwargs: Additional arguments passed to parent
"""
details = kwargs.get('details', {})
if workflow_id:
details['workflow_id'] = workflow_id
if execution_stage:
details['execution_stage'] = execution_stage
super().__init__(
message,
error_code="WORKFLOW_EXECUTION_ERROR",
details=details,
**kwargs
)
class TimeoutError(QueueManagerError):
"""Exception raised for timeout-related errors."""
def __init__(
self,
message: str,
timeout_seconds: Optional[float] = None,
operation: Optional[str] = None,
**kwargs
):
"""Initialize the timeout error.
Args:
message: Error message
timeout_seconds: Timeout value that was exceeded
operation: Operation that timed out
**kwargs: Additional arguments passed to parent
"""
details = kwargs.get('details', {})
if timeout_seconds is not None:
details['timeout_seconds'] = timeout_seconds
if operation:
details['operation'] = operation
super().__init__(
message,
error_code="TIMEOUT_ERROR",
details=details,
**kwargs
)
class ConfigurationError(QueueManagerError):
"""Exception raised for configuration-related errors."""
def __init__(
self,
message: str,
config_key: Optional[str] = None,
config_value: Optional[Any] = None,
**kwargs
):
"""Initialize the configuration error.
Args:
message: Error message
config_key: Configuration key that caused the error
config_value: Configuration value that caused the error
**kwargs: Additional arguments passed to parent
"""
details = kwargs.get('details', {})
if config_key:
details['config_key'] = config_key
if config_value is not None:
details['config_value'] = str(config_value)
super().__init__(
message,
error_code="CONFIGURATION_ERROR",
details=details,
**kwargs
)
class APIError(QueueManagerError):
"""Exception raised for API-related errors."""
def __init__(
self,
message: str,
status_code: int = 500,
endpoint: Optional[str] = None,
**kwargs
):
"""Initialize the API error.
Args:
message: Error message
status_code: HTTP status code
endpoint: API endpoint that caused the error
**kwargs: Additional arguments passed to parent
"""
details = kwargs.get('details', {})
details['status_code'] = status_code
if endpoint:
details['endpoint'] = endpoint
super().__init__(
message,
error_code="API_ERROR",
details=details,
**kwargs
)
class AuthenticationError(APIError):
"""Exception raised for authentication errors."""
def __init__(self, message: str = "Authentication required", **kwargs):
super().__init__(
message,
status_code=401,
error_code="AUTHENTICATION_ERROR",
**kwargs
)
class AuthorizationError(APIError):
"""Exception raised for authorization errors."""
def __init__(self, message: str = "Access denied", **kwargs):
super().__init__(
message,
status_code=403,
error_code="AUTHORIZATION_ERROR",
**kwargs
)
class NotFoundError(APIError):
"""Exception raised when a resource is not found."""
def __init__(
self,
message: str,
resource_type: Optional[str] = None,
resource_id: Optional[str] = None,
**kwargs
):
"""Initialize the not found error.
Args:
message: Error message
resource_type: Type of resource that was not found
resource_id: ID of resource that was not found
**kwargs: Additional arguments passed to parent
"""
details = kwargs.get('details', {})
if resource_type:
details['resource_type'] = resource_type
if resource_id:
details['resource_id'] = resource_id
super().__init__(
message,
status_code=404,
error_code="NOT_FOUND_ERROR",
details=details,
**kwargs
)
class ConflictError(APIError):
"""Exception raised for resource conflicts."""
def __init__(self, message: str, **kwargs):
super().__init__(
message,
status_code=409,
error_code="CONFLICT_ERROR",
**kwargs
)
class RateLimitError(APIError):
"""Exception raised when rate limits are exceeded."""
def __init__(
self,
message: str = "Rate limit exceeded",
retry_after: Optional[int] = None,
**kwargs
):
"""Initialize the rate limit error.
Args:
message: Error message
retry_after: Seconds to wait before retrying
**kwargs: Additional arguments passed to parent
"""
details = kwargs.get('details', {})
if retry_after is not None:
details['retry_after'] = retry_after
super().__init__(
message,
status_code=429,
error_code="RATE_LIMIT_ERROR",
details=details,
**kwargs
)
class ImportExportError(QueueManagerError):
"""Exception raised for import/export operations."""
def __init__(
self,
message: str,
operation: Optional[str] = None,
file_path: Optional[str] = None,
**kwargs
):
"""Initialize the import/export error.
Args:
message: Error message
operation: Operation that failed ('import' or 'export')
file_path: File path involved in the operation
**kwargs: Additional arguments passed to parent
"""
details = kwargs.get('details', {})
if operation:
details['operation'] = operation
if file_path:
details['file_path'] = file_path
super().__init__(
message,
error_code="IMPORT_EXPORT_ERROR",
details=details,
**kwargs
)
class ArchiveError(QueueManagerError):
"""Exception raised for archive operations."""
def __init__(
self,
message: str,
operation: Optional[str] = None,
item_count: Optional[int] = None,
**kwargs
):
"""Initialize the archive error.
Args:
message: Error message
operation: Archive operation that failed
item_count: Number of items involved in the operation
**kwargs: Additional arguments passed to parent
"""
details = kwargs.get('details', {})
if operation:
details['operation'] = operation
if item_count is not None:
details['item_count'] = item_count
super().__init__(
message,
error_code="ARCHIVE_ERROR",
details=details,
**kwargs
)
health_check.py
+635
View File
@@ -0,0 +1,635 @@
"""
Health check system for the ComfyUI Queue Manager.
Provides comprehensive health monitoring and diagnostics.
"""
from __future__ import annotations
import asyncio
import logging
import threading
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone, timedelta
from enum import Enum
from typing import Any, Callable, Dict, List, Optional
from logging_config import get_logger
from performance_monitor import get_performance_monitor
logger = get_logger("health_check")
class HealthStatus(Enum):
"""Health status enumeration."""
HEALTHY = "healthy"
DEGRADED = "degraded"
UNHEALTHY = "unhealthy"
UNKNOWN = "unknown"
@dataclass
class HealthCheckResult:
"""Result of a health check."""
name: str
status: HealthStatus
message: str
details: Dict[str, Any] = field(default_factory=dict)
timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
duration: float = 0.0 # seconds
def to_dict(self) -> Dict[str, Any]:
"""Convert result to dictionary."""
return {
"name": self.name,
"status": self.status.value,
"message": self.message,
"details": self.details,
"timestamp": self.timestamp.isoformat(),
"duration": self.duration
}
class HealthCheck:
"""Base class for health checks."""
def __init__(self, name: str, timeout: float = 30.0):
"""Initialize health check.
Args:
name: Name of the health check
timeout: Timeout in seconds
"""
self.name = name
self.timeout = timeout
async def check(self) -> HealthCheckResult:
"""Perform the health check.
Returns:
HealthCheckResult
"""
start_time = time.time()
try:
# Run the actual check with timeout
result = await asyncio.wait_for(self._perform_check(), timeout=self.timeout)
duration = time.time() - start_time
result.duration = duration
return result
except asyncio.TimeoutError:
duration = time.time() - start_time
return HealthCheckResult(
name=self.name,
status=HealthStatus.UNHEALTHY,
message=f"Health check timed out after {self.timeout} seconds",
duration=duration,
details={"timeout": self.timeout}
)
except Exception as e:
duration = time.time() - start_time
return HealthCheckResult(
name=self.name,
status=HealthStatus.UNHEALTHY,
message=f"Health check failed: {e}",
duration=duration,
details={"error": str(e), "error_type": type(e).__name__}
)
async def _perform_check(self) -> HealthCheckResult:
"""Perform the actual health check logic.
This method should be overridden by subclasses.
Returns:
HealthCheckResult
"""
raise NotImplementedError("Subclasses must implement _perform_check")
class DatabaseHealthCheck(HealthCheck):
"""Health check for database connectivity and performance."""
def __init__(self, database, timeout: float = 10.0):
"""Initialize database health check.
Args:
database: Database instance to check
timeout: Timeout in seconds
"""
super().__init__("database", timeout)
self.database = database
async def _perform_check(self) -> HealthCheckResult:
"""Check database health."""
details = {}
try:
# Check if database is initialized
if not hasattr(self.database, '_initialized') or not self.database._initialized:
return HealthCheckResult(
name=self.name,
status=HealthStatus.UNHEALTHY,
message="Database not initialized",
details={"initialized": False}
)
# Test basic connectivity
start_time = time.time()
test_result = await asyncio.get_event_loop().run_in_executor(
None, self._test_database_connection
)
connection_time = time.time() - start_time
details.update(test_result)
details["connection_time"] = connection_time
# Determine status based on results
if not test_result.get("connected", False):
status = HealthStatus.UNHEALTHY
message = "Database connection failed"
elif connection_time > 5.0:
status = HealthStatus.DEGRADED
message = f"Database connection slow ({connection_time:.2f}s)"
else:
status = HealthStatus.HEALTHY
message = "Database is healthy"
return HealthCheckResult(
name=self.name,
status=status,
message=message,
details=details
)
except Exception as e:
return HealthCheckResult(
name=self.name,
status=HealthStatus.UNHEALTHY,
message=f"Database health check failed: {e}",
details={"error": str(e)}
)
def _test_database_connection(self) -> Dict[str, Any]:
"""Test database connection and basic operations."""
result = {
"connected": False,
"schema_version": None,
"table_count": 0,
"last_error": None
}
try:
# Test connection
with self.database._get_cursor() as cursor:
# Check schema version
cursor.execute("SELECT version FROM schema_version LIMIT 1")
version_row = cursor.fetchone()
if version_row:
result["schema_version"] = version_row[0]
# Count tables
cursor.execute("SELECT COUNT(*) FROM sqlite_master WHERE type='table'")
table_count_row = cursor.fetchone()
if table_count_row:
result["table_count"] = table_count_row[0]
# Test a simple query
cursor.execute("SELECT COUNT(*) FROM queue_items")
item_count_row = cursor.fetchone()
if item_count_row:
result["queue_item_count"] = item_count_row[0]
result["connected"] = True
except Exception as e:
result["last_error"] = str(e)
return result
class QueueServiceHealthCheck(HealthCheck):
"""Health check for queue service functionality."""
def __init__(self, queue_service, timeout: float = 15.0):
"""Initialize queue service health check.
Args:
queue_service: Queue service instance to check
timeout: Timeout in seconds
"""
super().__init__("queue_service", timeout)
self.queue_service = queue_service
async def _perform_check(self) -> HealthCheckResult:
"""Check queue service health."""
details = {}
try:
# Check if service is initialized
if not hasattr(self.queue_service, '_initialized') or not self.queue_service._initialized:
return HealthCheckResult(
name=self.name,
status=HealthStatus.UNHEALTHY,
message="Queue service not initialized",
details={"initialized": False}
)
# Test basic operations
start_time = time.time()
test_result = await asyncio.get_event_loop().run_in_executor(
None, self._test_queue_service_operations
)
operation_time = time.time() - start_time
details.update(test_result)
details["operation_time"] = operation_time
# Determine status
if test_result.get("errors", 0) > 0:
status = HealthStatus.UNHEALTHY
message = f"Queue service has {test_result['errors']} errors"
elif operation_time > 10.0:
status = HealthStatus.DEGRADED
message = f"Queue service operations slow ({operation_time:.2f}s)"
else:
status = HealthStatus.HEALTHY
message = "Queue service is healthy"
return HealthCheckResult(
name=self.name,
status=status,
message=message,
details=details
)
except Exception as e:
return HealthCheckResult(
name=self.name,
status=HealthStatus.UNHEALTHY,
message=f"Queue service health check failed: {e}",
details={"error": str(e)}
)
def _test_queue_service_operations(self) -> Dict[str, Any]:
"""Test queue service operations."""
result = {
"queue_state": None,
"item_count": 0,
"config_accessible": False,
"errors": 0,
"last_error": None
}
try:
# Test getting queue state
queue_state = self.queue_service.get_queue_state()
result["queue_state"] = queue_state.value
except Exception as e:
result["errors"] += 1
result["last_error"] = f"get_queue_state: {e}"
try:
# Test getting queue items
items = self.queue_service.get_queue_items()
result["item_count"] = len(items)
except Exception as e:
result["errors"] += 1
result["last_error"] = f"get_queue_items: {e}"
try:
# Test getting configuration
config = self.queue_service.get_config()
result["config_accessible"] = config is not None
except Exception as e:
result["errors"] += 1
result["last_error"] = f"get_config: {e}"
return result
class SystemResourceHealthCheck(HealthCheck):
"""Health check for system resources."""
def __init__(self, timeout: float = 5.0):
"""Initialize system resource health check.
Args:
timeout: Timeout in seconds
"""
super().__init__("system_resources", timeout)
async def _perform_check(self) -> HealthCheckResult:
"""Check system resource health."""
try:
perf_monitor = get_performance_monitor()
health_status = perf_monitor.get_health_status()
# Map performance monitor status to health status
status_mapping = {
"healthy": HealthStatus.HEALTHY,
"degraded": HealthStatus.DEGRADED,
"unhealthy": HealthStatus.UNHEALTHY
}
status = status_mapping.get(health_status["status"], HealthStatus.UNKNOWN)
return HealthCheckResult(
name=self.name,
status=status,
message=f"System resources are {health_status['status']}",
details=health_status
)
except Exception as e:
return HealthCheckResult(
name=self.name,
status=HealthStatus.UNHEALTHY,
message=f"System resource check failed: {e}",
details={"error": str(e)}
)
class HealthCheckManager:
"""Manager for running and coordinating health checks."""
def __init__(self):
"""Initialize health check manager."""
self._checks: Dict[str, HealthCheck] = {}
self._last_results: Dict[str, HealthCheckResult] = {}
self._lock = threading.RLock()
# Monitoring settings
self._monitoring_enabled = False
self._monitoring_interval = 60.0 # seconds
self._monitoring_thread: Optional[threading.Thread] = None
self._stop_monitoring = threading.Event()
logger.info("Health check manager initialized")
def register_check(self, health_check: HealthCheck) -> None:
"""Register a health check.
Args:
health_check: Health check to register
"""
with self._lock:
self._checks[health_check.name] = health_check
logger.info(f"Registered health check: {health_check.name}")
def unregister_check(self, name: str) -> None:
"""Unregister a health check.
Args:
name: Name of the health check to remove
"""
with self._lock:
self._checks.pop(name, None)
self._last_results.pop(name, None)
logger.info(f"Unregistered health check: {name}")
async def run_check(self, name: str) -> HealthCheckResult:
"""Run a specific health check.
Args:
name: Name of the health check to run
Returns:
HealthCheckResult
"""
with self._lock:
health_check = self._checks.get(name)
if not health_check:
return HealthCheckResult(
name=name,
status=HealthStatus.UNKNOWN,
message=f"Health check '{name}' not found"
)
result = await health_check.check()
with self._lock:
self._last_results[name] = result
# Log significant health changes
if result.status != HealthStatus.HEALTHY:
logger.warning(
f"Health check failed: {name}",
extra={
"check_name": name,
"status": result.status.value,
"message": result.message,
"duration": result.duration
}
)
return result
async def run_all_checks(self) -> Dict[str, HealthCheckResult]:
"""Run all registered health checks.
Returns:
Dictionary mapping check names to results
"""
with self._lock:
check_names = list(self._checks.keys())
results = {}
# Run checks concurrently
tasks = [self.run_check(name) for name in check_names]
check_results = await asyncio.gather(*tasks, return_exceptions=True)
for name, result in zip(check_names, check_results):
if isinstance(result, Exception):
results[name] = HealthCheckResult(
name=name,
status=HealthStatus.UNHEALTHY,
message=f"Health check failed with exception: {result}",
details={"exception": str(result)}
)
else:
results[name] = result
return results
def get_last_results(self) -> Dict[str, HealthCheckResult]:
"""Get the last results from all health checks.
Returns:
Dictionary mapping check names to last results
"""
with self._lock:
return self._last_results.copy()
def get_overall_health(self) -> Dict[str, Any]:
"""Get overall health status.
Returns:
Dictionary containing overall health information
"""
with self._lock:
results = self._last_results.copy()
if not results:
return {
"status": HealthStatus.UNKNOWN.value,
"message": "No health checks have been run",
"timestamp": datetime.now(timezone.utc).isoformat(),
"checks": {}
}
# Determine overall status
statuses = [result.status for result in results.values()]
if any(status == HealthStatus.UNHEALTHY for status in statuses):
overall_status = HealthStatus.UNHEALTHY
message = "One or more health checks are unhealthy"
elif any(status == HealthStatus.DEGRADED for status in statuses):
overall_status = HealthStatus.DEGRADED
message = "One or more health checks are degraded"
elif any(status == HealthStatus.UNKNOWN for status in statuses):
overall_status = HealthStatus.UNKNOWN
message = "One or more health checks have unknown status"
else:
overall_status = HealthStatus.HEALTHY
message = "All health checks are healthy"
return {
"status": overall_status.value,
"message": message,
"timestamp": datetime.now(timezone.utc).isoformat(),
"checks": {name: result.to_dict() for name, result in results.items()}
}
def start_monitoring(self, interval: float = 60.0) -> None:
"""Start continuous health monitoring.
Args:
interval: Monitoring interval in seconds
"""
if self._monitoring_thread and self._monitoring_thread.is_alive():
logger.warning("Health monitoring is already running")
return
self._monitoring_interval = interval
self._monitoring_enabled = True
self._stop_monitoring.clear()
self._monitoring_thread = threading.Thread(
target=self._monitoring_loop,
name="HealthMonitor",
daemon=True
)
self._monitoring_thread.start()
logger.info(f"Health monitoring started with {interval}s interval")
def stop_monitoring(self) -> None:
"""Stop continuous health monitoring."""
if not self._monitoring_thread or not self._monitoring_thread.is_alive():
return
self._monitoring_enabled = False
self._stop_monitoring.set()
self._monitoring_thread.join(timeout=5.0)
if self._monitoring_thread.is_alive():
logger.warning("Health monitoring thread did not stop gracefully")
else:
logger.info("Health monitoring stopped")
def _monitoring_loop(self) -> None:
"""Health monitoring loop that runs in a separate thread."""
logger.info("Health monitoring loop started")
while self._monitoring_enabled and not self._stop_monitoring.is_set():
try:
# Run health checks
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
results = loop.run_until_complete(self.run_all_checks())
# Log overall health status
overall_health = self.get_overall_health()
logger.info(
f"Health monitoring cycle completed: {overall_health['status']}",
extra={
"overall_status": overall_health["status"],
"check_count": len(results),
"healthy_checks": sum(1 for r in results.values() if r.status == HealthStatus.HEALTHY),
"degraded_checks": sum(1 for r in results.values() if r.status == HealthStatus.DEGRADED),
"unhealthy_checks": sum(1 for r in results.values() if r.status == HealthStatus.UNHEALTHY)
}
)
loop.close()
# Wait for next cycle
if self._stop_monitoring.wait(timeout=self._monitoring_interval):
break
except Exception as e:
logger.error(f"Error in health monitoring loop: {e}", exc_info=True)
time.sleep(5.0) # Brief pause before retrying
logger.info("Health monitoring loop stopped")
def close(self) -> None:
"""Close the health check manager and cleanup resources."""
self.stop_monitoring()
with self._lock:
self._checks.clear()
self._last_results.clear()
logger.info("Health check manager closed")
# Global health check manager instance
_health_check_manager: Optional[HealthCheckManager] = None
def get_health_check_manager() -> HealthCheckManager:
"""Get the global health check manager instance.
Returns:
HealthCheckManager instance
"""
global _health_check_manager
if _health_check_manager is None:
_health_check_manager = HealthCheckManager()
return _health_check_manager
def setup_health_checks(database=None, queue_service=None) -> HealthCheckManager:
"""Set up health checks for the queue manager.
Args:
database: Database instance for health checks
queue_service: Queue service instance for health checks
Returns:
HealthCheckManager instance
"""
manager = get_health_check_manager()
# Register system resource check
manager.register_check(SystemResourceHealthCheck())
# Register database check if available
if database:
manager.register_check(DatabaseHealthCheck(database))
# Register queue service check if available
if queue_service:
manager.register_check(QueueServiceHealthCheck(queue_service))
return manager
logging_config.py
+450
View File
@@ -0,0 +1,450 @@
"""
Logging configuration for the ComfyUI Queue Manager.
Provides structured logging with rotation, filtering, and monitoring capabilities.
"""
from __future__ import annotations
import json
import logging
import logging.handlers
import os
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Optional
class StructuredFormatter(logging.Formatter):
"""Custom formatter that outputs structured JSON logs."""
def __init__(self, include_extra: bool = True):
"""Initialize the structured formatter.
Args:
include_extra: Whether to include extra fields from log records
"""
super().__init__()
self.include_extra = include_extra
def format(self, record: logging.LogRecord) -> str:
"""Format a log record as structured JSON.
Args:
record: Log record to format
Returns:
JSON-formatted log string
"""
# Base log data
log_data = {
"timestamp": datetime.fromtimestamp(record.created, tz=timezone.utc).isoformat(),
"level": record.levelname,
"logger": record.name,
"message": record.getMessage(),
"module": record.module,
"function": record.funcName,
"line": record.lineno,
}
# Add thread information
if record.thread:
log_data["thread_id"] = record.thread
log_data["thread_name"] = record.threadName
# Add process information
if record.process:
log_data["process_id"] = record.process
log_data["process_name"] = record.processName
# Add exception information
if record.exc_info:
log_data["exception"] = {
"type": record.exc_info[0].__name__ if record.exc_info[0] else None,
"message": str(record.exc_info[1]) if record.exc_info[1] else None,
"traceback": self.formatException(record.exc_info) if record.exc_info else None
}
# Add extra fields if enabled
if self.include_extra:
# Get all extra attributes (those not in the standard LogRecord)
standard_attrs = {
'name', 'msg', 'args', 'levelname', 'levelno', 'pathname', 'filename',
'module', 'lineno', 'funcName', 'created', 'msecs', 'relativeCreated',
'thread', 'threadName', 'processName', 'process', 'getMessage',
'exc_info', 'exc_text', 'stack_info', 'message'
}
extra_data = {}
for key, value in record.__dict__.items():
if key not in standard_attrs and not key.startswith('_'):
# Ensure value is JSON serializable
try:
json.dumps(value)
extra_data[key] = value
except (TypeError, ValueError):
extra_data[key] = str(value)
if extra_data:
log_data["extra"] = extra_data
return json.dumps(log_data, ensure_ascii=False)
class QueueManagerFilter(logging.Filter):
"""Custom filter for queue manager logs."""
def __init__(self, component: Optional[str] = None, min_level: int = logging.INFO):
"""Initialize the filter.
Args:
component: Specific component to filter for (e.g., 'database', 'api')
min_level: Minimum log level to allow
"""
super().__init__()
self.component = component
self.min_level = min_level
def filter(self, record: logging.LogRecord) -> bool:
"""Filter log records based on criteria.
Args:
record: Log record to filter
Returns:
True if record should be logged
"""
# Check minimum level
if record.levelno < self.min_level:
return False
# Check component filter
if self.component:
# Check if the logger name contains the component
if self.component not in record.name:
return False
return True
class LoggingConfig:
"""Configuration manager for logging setup."""
DEFAULT_FORMAT = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
DEFAULT_DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
def __init__(self, log_dir: Optional[Path] = None):
"""Initialize logging configuration.
Args:
log_dir: Directory for log files. If None, uses default location.
"""
self.log_dir = log_dir or Path(__file__).parent / "logs"
self.log_dir.mkdir(parents=True, exist_ok=True)
# Configuration settings
self.console_level = logging.INFO
self.file_level = logging.DEBUG
self.structured_logging = True
self.max_file_size = 10 * 1024 * 1024 # 10MB
self.backup_count = 5
# Logger instances
self._loggers: Dict[str, logging.Logger] = {}
self._configured = False
def configure_logging(
self,
console_level: int = logging.INFO,
file_level: int = logging.DEBUG,
structured: bool = True,
max_file_size: int = 10 * 1024 * 1024,
backup_count: int = 5
) -> None:
"""Configure logging for the queue manager.
Args:
console_level: Log level for console output
file_level: Log level for file output
structured: Whether to use structured JSON logging
max_file_size: Maximum size of log files before rotation
backup_count: Number of backup files to keep
"""
self.console_level = console_level
self.file_level = file_level
self.structured_logging = structured
self.max_file_size = max_file_size
self.backup_count = backup_count
# Configure root logger
root_logger = logging.getLogger()
root_logger.setLevel(logging.DEBUG)
# Remove existing handlers to avoid duplicates
for handler in root_logger.handlers[:]:
root_logger.removeHandler(handler)
# Configure console handler
self._configure_console_handler(root_logger)
# Configure file handlers
self._configure_file_handlers(root_logger)
# Configure component-specific loggers
self._configure_component_loggers()
self._configured = True
# Log configuration completion
logger = self.get_logger("logging_config")
logger.info("Logging configuration completed", extra={
"console_level": logging.getLevelName(console_level),
"file_level": logging.getLevelName(file_level),
"structured": structured,
"log_dir": str(self.log_dir),
"max_file_size": max_file_size,
"backup_count": backup_count
})
def _configure_console_handler(self, root_logger: logging.Logger) -> None:
"""Configure console logging handler."""
console_handler = logging.StreamHandler(sys.stdout)
console_handler.setLevel(self.console_level)
if self.structured_logging:
console_formatter = StructuredFormatter(include_extra=False)
else:
console_formatter = logging.Formatter(
self.DEFAULT_FORMAT,
datefmt=self.DEFAULT_DATE_FORMAT
)
console_handler.setFormatter(console_formatter)
console_handler.addFilter(QueueManagerFilter(min_level=self.console_level))
root_logger.addHandler(console_handler)
def _configure_file_handlers(self, root_logger: logging.Logger) -> None:
"""Configure file logging handlers."""
# Main application log
main_log_file = self.log_dir / "queue_manager.log"
main_handler = logging.handlers.RotatingFileHandler(
main_log_file,
maxBytes=self.max_file_size,
backupCount=self.backup_count,
encoding='utf-8'
)
main_handler.setLevel(self.file_level)
if self.structured_logging:
main_formatter = StructuredFormatter(include_extra=True)
else:
main_formatter = logging.Formatter(
self.DEFAULT_FORMAT,
datefmt=self.DEFAULT_DATE_FORMAT
)
main_handler.setFormatter(main_formatter)
root_logger.addHandler(main_handler)
# Error-only log
error_log_file = self.log_dir / "errors.log"
error_handler = logging.handlers.RotatingFileHandler(
error_log_file,
maxBytes=self.max_file_size,
backupCount=self.backup_count,
encoding='utf-8'
)
error_handler.setLevel(logging.ERROR)
error_handler.setFormatter(main_formatter)
root_logger.addHandler(error_handler)
# Performance log
perf_log_file = self.log_dir / "performance.log"
perf_handler = logging.handlers.RotatingFileHandler(
perf_log_file,
maxBytes=self.max_file_size,
backupCount=self.backup_count,
encoding='utf-8'
)
perf_handler.setLevel(logging.INFO)
perf_handler.setFormatter(main_formatter)
perf_handler.addFilter(QueueManagerFilter(component="performance"))
root_logger.addHandler(perf_handler)
def _configure_component_loggers(self) -> None:
"""Configure component-specific loggers."""
components = [
"database",
"queue_service",
"api_routes",
"workflow_executor",
"execution_monitor",
"archive_service",
"filter_service",
"import_export_service"
]
for component in components:
logger = logging.getLogger(f"custom_nodes.comfyui-queue-manager.{component}")
logger.setLevel(logging.DEBUG)
# Inherit handlers from root logger
logger.propagate = True
def get_logger(self, name: str) -> logging.Logger:
"""Get a logger instance for a specific component.
Args:
name: Logger name (component name)
Returns:
Logger instance
"""
full_name = f"custom_nodes.comfyui-queue-manager.{name}"
if full_name not in self._loggers:
logger = logging.getLogger(full_name)
logger.setLevel(logging.DEBUG)
self._loggers[full_name] = logger
return self._loggers[full_name]
def set_log_level(self, component: str, level: int) -> None:
"""Set log level for a specific component.
Args:
component: Component name
level: Log level
"""
logger = self.get_logger(component)
logger.setLevel(level)
# Log the level change
logger.info(f"Log level changed to {logging.getLevelName(level)}", extra={
"component": component,
"new_level": logging.getLevelName(level)
})
def cleanup_old_logs(self, days_to_keep: int = 30) -> int:
"""Clean up old log files.
Args:
days_to_keep: Number of days of logs to keep
Returns:
Number of files cleaned up
"""
from datetime import timedelta
import time
cutoff_time = time.time() - (days_to_keep * 24 * 60 * 60)
cleaned_count = 0
for log_file in self.log_dir.glob("*.log*"):
try:
if log_file.stat().st_mtime < cutoff_time:
log_file.unlink()
cleaned_count += 1
except OSError as e:
logger = self.get_logger("logging_config")
logger.warning(f"Failed to clean up log file {log_file}: {e}")
if cleaned_count > 0:
logger = self.get_logger("logging_config")
logger.info(f"Cleaned up {cleaned_count} old log files", extra={
"cleaned_count": cleaned_count,
"days_to_keep": days_to_keep
})
return cleaned_count
def get_log_stats(self) -> Dict[str, Any]:
"""Get statistics about log files.
Returns:
Dictionary containing log file statistics
"""
stats = {
"log_dir": str(self.log_dir),
"total_files": 0,
"total_size": 0,
"files": []
}
for log_file in self.log_dir.glob("*.log*"):
try:
file_stat = log_file.stat()
file_info = {
"name": log_file.name,
"size": file_stat.st_size,
"modified": datetime.fromtimestamp(file_stat.st_mtime, tz=timezone.utc).isoformat()
}
stats["files"].append(file_info)
stats["total_files"] += 1
stats["total_size"] += file_stat.st_size
except OSError:
continue
return stats
# Global logging configuration instance
_logging_config: Optional[LoggingConfig] = None
def get_logging_config() -> LoggingConfig:
"""Get the global logging configuration instance.
Returns:
LoggingConfig instance
"""
global _logging_config
if _logging_config is None:
_logging_config = LoggingConfig()
return _logging_config
def setup_logging(
log_dir: Optional[Path] = None,
console_level: int = logging.INFO,
file_level: int = logging.DEBUG,
structured: bool = True
) -> LoggingConfig:
"""Set up logging for the queue manager.
Args:
log_dir: Directory for log files
console_level: Log level for console output
file_level: Log level for file output
structured: Whether to use structured JSON logging
Returns:
LoggingConfig instance
"""
global _logging_config
if log_dir:
_logging_config = LoggingConfig(log_dir)
else:
_logging_config = get_logging_config()
_logging_config.configure_logging(
console_level=console_level,
file_level=file_level,
structured=structured
)
return _logging_config
def get_logger(name: str) -> logging.Logger:
"""Get a logger instance for a component.
Args:
name: Component name
Returns:
Logger instance
"""
config = get_logging_config()
return config.get_logger(name)
performance_monitor.py
+545
View File
@@ -0,0 +1,545 @@
"""
Performance monitoring for the ComfyUI Queue Manager.
Tracks metrics, performance data, and system health.
"""
from __future__ import annotations
import functools
import logging
import psutil
import threading
import time
from collections import defaultdict, deque
from contextlib import contextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone, timedelta
from typing import Any, Callable, Dict, List, Optional, TypeVar
from logging_config import get_logger
logger = get_logger("performance_monitor")
T = TypeVar('T')
@dataclass
class PerformanceMetric:
"""Data class for performance metrics."""
name: str
value: float
unit: str
timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
tags: Dict[str, str] = field(default_factory=dict)
def to_dict(self) -> Dict[str, Any]:
"""Convert metric to dictionary."""
return {
"name": self.name,
"value": self.value,
"unit": self.unit,
"timestamp": self.timestamp.isoformat(),
"tags": self.tags
}
@dataclass
class OperationStats:
"""Statistics for a specific operation."""
operation_name: str
total_calls: int = 0
total_time: float = 0.0
min_time: float = float('inf')
max_time: float = 0.0
error_count: int = 0
last_called: Optional[datetime] = None
@property
def average_time(self) -> float:
"""Calculate average execution time."""
return self.total_time / self.total_calls if self.total_calls > 0 else 0.0
@property
def success_rate(self) -> float:
"""Calculate success rate as percentage."""
if self.total_calls == 0:
return 0.0
return ((self.total_calls - self.error_count) / self.total_calls) * 100
def update(self, execution_time: float, error: bool = False) -> None:
"""Update statistics with new execution data."""
self.total_calls += 1
self.total_time += execution_time
self.min_time = min(self.min_time, execution_time)
self.max_time = max(self.max_time, execution_time)
self.last_called = datetime.now(timezone.utc)
if error:
self.error_count += 1
def to_dict(self) -> Dict[str, Any]:
"""Convert stats to dictionary."""
return {
"operation_name": self.operation_name,
"total_calls": self.total_calls,
"total_time": self.total_time,
"average_time": self.average_time,
"min_time": self.min_time if self.min_time != float('inf') else 0.0,
"max_time": self.max_time,
"error_count": self.error_count,
"success_rate": self.success_rate,
"last_called": self.last_called.isoformat() if self.last_called else None
}
class PerformanceMonitor:
"""Performance monitoring and metrics collection."""
def __init__(self, max_history_size: int = 1000):
"""Initialize performance monitor.
Args:
max_history_size: Maximum number of metrics to keep in history
"""
self.max_history_size = max_history_size
self._lock = threading.RLock()
# Metrics storage
self._metrics_history: deque = deque(maxlen=max_history_size)
self._operation_stats: Dict[str, OperationStats] = {}
self._counters: Dict[str, int] = defaultdict(int)
self._gauges: Dict[str, float] = {}
# System monitoring
self._system_metrics_enabled = True
self._system_metrics_interval = 60.0 # seconds
self._system_monitor_thread: Optional[threading.Thread] = None
self._stop_system_monitoring = threading.Event()
# Performance thresholds
self._thresholds = {
"database_operation_time": 1.0, # seconds
"api_response_time": 2.0, # seconds
"workflow_execution_time": 300.0, # seconds
"memory_usage_percent": 80.0, # percent
"cpu_usage_percent": 90.0, # percent
}
logger.info("Performance monitor initialized", extra={
"max_history_size": max_history_size,
"system_metrics_enabled": self._system_metrics_enabled
})
def start_system_monitoring(self) -> None:
"""Start system metrics monitoring thread."""
if self._system_monitor_thread and self._system_monitor_thread.is_alive():
logger.warning("System monitoring is already running")
return
self._stop_system_monitoring.clear()
self._system_monitor_thread = threading.Thread(
target=self._system_monitoring_loop,
name="SystemMonitor",
daemon=True
)
self._system_monitor_thread.start()
logger.info("System monitoring started", extra={
"interval": self._system_metrics_interval
})
def stop_system_monitoring(self) -> None:
"""Stop system metrics monitoring thread."""
if not self._system_monitor_thread or not self._system_monitor_thread.is_alive():
return
self._stop_system_monitoring.set()
self._system_monitor_thread.join(timeout=5.0)
if self._system_monitor_thread.is_alive():
logger.warning("System monitoring thread did not stop gracefully")
else:
logger.info("System monitoring stopped")
def _system_monitoring_loop(self) -> None:
"""System monitoring loop that runs in a separate thread."""
logger.info("System monitoring loop started")
while not self._stop_system_monitoring.is_set():
try:
self._collect_system_metrics()
# Wait for interval or stop signal
if self._stop_system_monitoring.wait(timeout=self._system_metrics_interval):
break
except Exception as e:
logger.error(f"Error in system monitoring loop: {e}", exc_info=True)
time.sleep(5.0) # Brief pause before retrying
logger.info("System monitoring loop stopped")
def _collect_system_metrics(self) -> None:
"""Collect system performance metrics."""
try:
# CPU metrics
cpu_percent = psutil.cpu_percent(interval=1)
self.record_gauge("system.cpu.usage_percent", cpu_percent)
# Memory metrics
memory = psutil.virtual_memory()
self.record_gauge("system.memory.usage_percent", memory.percent)
self.record_gauge("system.memory.available_mb", memory.available / 1024 / 1024)
self.record_gauge("system.memory.used_mb", memory.used / 1024 / 1024)
# Disk metrics
disk = psutil.disk_usage('/')
self.record_gauge("system.disk.usage_percent", (disk.used / disk.total) * 100)
self.record_gauge("system.disk.free_gb", disk.free / 1024 / 1024 / 1024)
# Process metrics
process = psutil.Process()
process_memory = process.memory_info()
self.record_gauge("process.memory.rss_mb", process_memory.rss / 1024 / 1024)
self.record_gauge("process.memory.vms_mb", process_memory.vms / 1024 / 1024)
self.record_gauge("process.cpu.percent", process.cpu_percent())
self.record_gauge("process.threads.count", process.num_threads())
# Check thresholds
self._check_performance_thresholds({
"cpu_usage_percent": cpu_percent,
"memory_usage_percent": memory.percent,
})
except Exception as e:
logger.error(f"Failed to collect system metrics: {e}")
def _check_performance_thresholds(self, metrics: Dict[str, float]) -> None:
"""Check if metrics exceed performance thresholds."""
for metric_name, value in metrics.items():
threshold = self._thresholds.get(metric_name)
if threshold and value > threshold:
logger.warning(
f"Performance threshold exceeded: {metric_name}",
extra={
"metric": metric_name,
"value": value,
"threshold": threshold,
"severity": "high" if value > threshold * 1.2 else "medium"
}
)
def record_metric(self, metric: PerformanceMetric) -> None:
"""Record a performance metric.
Args:
metric: Performance metric to record
"""
with self._lock:
self._metrics_history.append(metric)
# Log metric if it's significant
if metric.name.endswith("_time") and metric.value > 1.0:
logger.info(f"Performance metric recorded: {metric.name}", extra={
"metric_name": metric.name,
"value": metric.value,
"unit": metric.unit,
"tags": metric.tags
})
def record_counter(self, name: str, increment: int = 1, tags: Optional[Dict[str, str]] = None) -> None:
"""Record a counter metric.
Args:
name: Counter name
increment: Amount to increment by
tags: Optional tags for the metric
"""
with self._lock:
self._counters[name] += increment
metric = PerformanceMetric(
name=name,
value=increment,
unit="count",
tags=tags or {}
)
self.record_metric(metric)
def record_gauge(self, name: str, value: float, tags: Optional[Dict[str, str]] = None) -> None:
"""Record a gauge metric.
Args:
name: Gauge name
value: Gauge value
tags: Optional tags for the metric
"""
with self._lock:
self._gauges[name] = value
metric = PerformanceMetric(
name=name,
value=value,
unit="gauge",
tags=tags or {}
)
self.record_metric(metric)
def record_timing(self, name: str, duration: float, tags: Optional[Dict[str, str]] = None) -> None:
"""Record a timing metric.
Args:
name: Timing metric name
duration: Duration in seconds
tags: Optional tags for the metric
"""
metric = PerformanceMetric(
name=name,
value=duration,
unit="seconds",
tags=tags or {}
)
self.record_metric(metric)
# Check timing thresholds
threshold_key = name.replace(".", "_")
threshold = self._thresholds.get(threshold_key)
if threshold and duration > threshold:
logger.warning(
f"Timing threshold exceeded: {name}",
extra={
"metric": name,
"duration": duration,
"threshold": threshold,
"tags": tags or {}
}
)
def update_operation_stats(self, operation: str, execution_time: float, error: bool = False) -> None:
"""Update statistics for an operation.
Args:
operation: Operation name
execution_time: Execution time in seconds
error: Whether the operation resulted in an error
"""
with self._lock:
if operation not in self._operation_stats:
self._operation_stats[operation] = OperationStats(operation)
self._operation_stats[operation].update(execution_time, error)
# Record timing metric
self.record_timing(f"operation.{operation}.time", execution_time, {
"operation": operation,
"status": "error" if error else "success"
})
@contextmanager
def time_operation(self, operation: str, tags: Optional[Dict[str, str]] = None):
"""Context manager for timing operations.
Args:
operation: Operation name
tags: Optional tags for the metric
"""
start_time = time.time()
error_occurred = False
try:
yield
except Exception as e:
error_occurred = True
raise
finally:
execution_time = time.time() - start_time
self.update_operation_stats(operation, execution_time, error_occurred)
# Add error status to tags
final_tags = tags or {}
final_tags["status"] = "error" if error_occurred else "success"
self.record_timing(f"operation.{operation}.time", execution_time, final_tags)
def get_operation_stats(self, operation: Optional[str] = None) -> Dict[str, Any]:
"""Get operation statistics.
Args:
operation: Specific operation name, or None for all operations
Returns:
Dictionary containing operation statistics
"""
with self._lock:
if operation:
stats = self._operation_stats.get(operation)
return stats.to_dict() if stats else {}
else:
return {
op_name: stats.to_dict()
for op_name, stats in self._operation_stats.items()
}
def get_metrics_summary(self, since: Optional[datetime] = None) -> Dict[str, Any]:
"""Get summary of recorded metrics.
Args:
since: Only include metrics since this timestamp
Returns:
Dictionary containing metrics summary
"""
with self._lock:
# Filter metrics by timestamp if specified
if since:
filtered_metrics = [
metric for metric in self._metrics_history
if metric.timestamp >= since
]
else:
filtered_metrics = list(self._metrics_history)
# Group metrics by name
metrics_by_name = defaultdict(list)
for metric in filtered_metrics:
metrics_by_name[metric.name].append(metric.value)
# Calculate summary statistics
summary = {}
for name, values in metrics_by_name.items():
if values:
summary[name] = {
"count": len(values),
"min": min(values),
"max": max(values),
"avg": sum(values) / len(values),
"total": sum(values) if name.endswith("_count") else None
}
return {
"period": {
"start": since.isoformat() if since else None,
"end": datetime.now(timezone.utc).isoformat(),
"total_metrics": len(filtered_metrics)
},
"metrics": summary,
"counters": dict(self._counters),
"gauges": dict(self._gauges)
}
def get_health_status(self) -> Dict[str, Any]:
"""Get overall system health status.
Returns:
Dictionary containing health status information
"""
with self._lock:
# Check recent errors
recent_time = datetime.now(timezone.utc) - timedelta(minutes=5)
recent_errors = sum(
1 for metric in self._metrics_history
if (metric.timestamp >= recent_time and
metric.tags.get("status") == "error")
)
# Calculate error rates
total_operations = sum(stats.total_calls for stats in self._operation_stats.values())
total_errors = sum(stats.error_count for stats in self._operation_stats.values())
overall_error_rate = (total_errors / total_operations * 100) if total_operations > 0 else 0
# Determine health status
if recent_errors > 10 or overall_error_rate > 10:
health_status = "unhealthy"
elif recent_errors > 5 or overall_error_rate > 5:
health_status = "degraded"
else:
health_status = "healthy"
return {
"status": health_status,
"timestamp": datetime.now(timezone.utc).isoformat(),
"metrics": {
"recent_errors": recent_errors,
"total_operations": total_operations,
"total_errors": total_errors,
"error_rate_percent": overall_error_rate
},
"system": {
"cpu_percent": self._gauges.get("system.cpu.usage_percent", 0),
"memory_percent": self._gauges.get("system.memory.usage_percent", 0),
"disk_percent": self._gauges.get("system.disk.usage_percent", 0)
}
}
def reset_stats(self) -> None:
"""Reset all performance statistics."""
with self._lock:
self._metrics_history.clear()
self._operation_stats.clear()
self._counters.clear()
self._gauges.clear()
logger.info("Performance statistics reset")
def close(self) -> None:
"""Close the performance monitor and cleanup resources."""
self.stop_system_monitoring()
logger.info("Performance monitor closed")
def time_function(operation_name: Optional[str] = None, tags: Optional[Dict[str, str]] = None):
"""Decorator for timing function execution.
Args:
operation_name: Name for the operation (defaults to function name)
tags: Optional tags for the metric
"""
def decorator(func: Callable[..., T]) -> Callable[..., T]:
@functools.wraps(func)
def wrapper(*args, **kwargs) -> T:
op_name = operation_name or func.__name__
with get_performance_monitor().time_operation(op_name, tags):
return func(*args, **kwargs)
return wrapper
return decorator
# Global performance monitor instance
_performance_monitor: Optional[PerformanceMonitor] = None
def get_performance_monitor() -> PerformanceMonitor:
"""Get the global performance monitor instance.
Returns:
PerformanceMonitor instance
"""
global _performance_monitor
if _performance_monitor is None:
_performance_monitor = PerformanceMonitor()
_performance_monitor.start_system_monitoring()
return _performance_monitor
def setup_performance_monitoring(max_history_size: int = 1000) -> PerformanceMonitor:
"""Set up performance monitoring.
Args:
max_history_size: Maximum number of metrics to keep in history
Returns:
PerformanceMonitor instance
"""
global _performance_monitor
if _performance_monitor:
_performance_monitor.close()
_performance_monitor = PerformanceMonitor(max_history_size)
_performance_monitor.start_system_monitoring()
return _performance_monitor
queue_service.py
+15 -16
View File
@@ -11,6 +11,8 @@ from typing import Any
from archive_service import ArchiveService
from database import SQLiteDatabase
from error_handler import ErrorHandler, with_error_handling
from exceptions import QueueServiceError, ValidationError
from filter_service import FilterService
from interfaces import DatabaseInterface, QueueServiceInterface
from models import QueueConfig, QueueFilter, QueueItem, QueueState, QueueStatus
@@ -18,10 +20,6 @@ from models import QueueConfig, QueueFilter, QueueItem, QueueState, QueueStatus
logger = logging.getLogger(__name__)
class QueueServiceError(Exception):
"""Custom exception for queue service operations."""
class QueueService(QueueServiceInterface):
"""Core queue service implementation."""
@@ -44,6 +42,7 @@ class QueueService(QueueServiceInterface):
self._initialized = True
logger.info("Queue service initialized successfully")
@with_error_handling(error_type=QueueServiceError, operation="add_workflow")
def add_workflow(
self, workflow_data: dict[str, Any], workflow_name: str = ""
) -> str:
@@ -60,10 +59,13 @@ class QueueService(QueueServiceInterface):
QueueServiceError: If the workflow cannot be added
"""
if not self._initialized:
raise QueueServiceError("Queue service not initialized")
raise QueueServiceError("Queue service not initialized", operation="add_workflow")
if not workflow_data:
raise QueueServiceError("Workflow data cannot be empty")
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__)
# Create a new queue item
item = QueueItem(
@@ -74,11 +76,12 @@ class QueueService(QueueServiceInterface):
# Save to database
if not self.database.create_queue_item(item):
raise QueueServiceError(f"Failed to create queue item {item.id}")
raise QueueServiceError(f"Failed to create queue item {item.id}", item_id=item.id, operation="create")
logger.info(f"Added workflow to queue: {item.id} - {item.workflow_name}")
return item.id
@with_error_handling(error_type=QueueServiceError, operation="get_queue_items")
def get_queue_items(self, status: QueueStatus | None = None) -> list[QueueItem]:
"""Get queue items, optionally filtered by status.
@@ -89,16 +92,12 @@ class QueueService(QueueServiceInterface):
List of queue items
"""
if not self._initialized:
raise QueueServiceError("Queue service not initialized")
raise QueueServiceError("Queue service not initialized", operation="get_queue_items")
try:
if status is not None:
return self.database.get_items_by_status(status)
else:
return self.database.get_all_queue_items()
except Exception as e:
logger.error(f"Failed to get queue items: {e}")
raise QueueServiceError(f"Failed to get queue items: {e}") from e
if status is not None:
return self.database.get_items_by_status(status)
else:
return self.database.get_all_queue_items()
def get_queue_item(self, item_id: str) -> QueueItem | None:
"""Get a specific queue item by ID.
test_error_handling.py
+499
View File
@@ -0,0 +1,499 @@
"""
Tests for error handling functionality in the ComfyUI Queue Manager.
"""
import json
import pytest
import sqlite3
import tempfile
from datetime import datetime, timezone
from pathlib import Path
from unittest.mock import Mock, patch
from database import SQLiteDatabase
from error_handler import ErrorHandler, ErrorRecovery, with_database_error_handling
from exceptions import (
APIError,
ConnectionError,
DatabaseError,
NotFoundError,
QueueManagerError,
QueueServiceError,
SchemaError,
ValidationError,
WorkflowExecutionError,
)
from models import QueueItem, QueueStatus
from queue_service import QueueService
class TestCustomExceptions:
"""Test custom exception classes."""
def test_queue_manager_error_basic(self):
"""Test basic QueueManagerError functionality."""
error = QueueManagerError("Test error")
assert error.message == "Test error"
assert error.error_code == "QueueManagerError"
assert error.details == {}
assert error.cause is None
def test_queue_manager_error_with_details(self):
"""Test QueueManagerError with additional details."""
details = {"key": "value", "number": 42}
cause = ValueError("Original error")
error = QueueManagerError(
"Test error",
error_code="CUSTOM_ERROR",
details=details,
cause=cause
)
assert error.message == "Test error"
assert error.error_code == "CUSTOM_ERROR"
assert error.details == details
assert error.cause == cause
def test_queue_manager_error_to_dict(self):
"""Test QueueManagerError serialization."""
error = QueueManagerError(
"Test error",
error_code="TEST_ERROR",
details={"field": "value"},
cause=ValueError("Cause")
)
result = error.to_dict()
assert result["error_type"] == "QueueManagerError"
assert result["error_code"] == "TEST_ERROR"
assert result["message"] == "Test error"
assert result["details"] == {"field": "value"}
assert result["cause"] == "Cause"
def test_database_error_with_operation(self):
"""Test DatabaseError with operation details."""
error = DatabaseError(
"Database failed",
operation="INSERT",
table="queue_items"
)
assert error.message == "Database failed"
assert error.details["operation"] == "INSERT"
assert error.details["table"] == "queue_items"
def test_validation_error_with_field(self):
"""Test ValidationError with field details."""
error = ValidationError(
"Invalid value",
field="workflow_name",
value="invalid"
)
assert error.message == "Invalid value"
assert error.details["field"] == "workflow_name"
assert error.details["value"] == "invalid"
def test_api_error_with_status_code(self):
"""Test APIError with status code."""
error = APIError(
"API failed",
status_code=404,
endpoint="/api/items"
)
assert error.message == "API failed"
assert error.details["status_code"] == 404
assert error.details["endpoint"] == "/api/items"
def test_workflow_execution_error(self):
"""Test WorkflowExecutionError with workflow details."""
error = WorkflowExecutionError(
"Execution failed",
workflow_id="test-123",
execution_stage="validation"
)
assert error.message == "Execution failed"
assert error.details["workflow_id"] == "test-123"
assert error.details["execution_stage"] == "validation"
class TestErrorHandler:
"""Test ErrorHandler utility functions."""
def test_handle_database_error_sqlite_locked(self):
"""Test handling of SQLite database locked error."""
original_error = sqlite3.OperationalError("database is locked")
result = ErrorHandler.handle_database_error(
original_error,
"INSERT",
"queue_items"
)
assert isinstance(result, ConnectionError)
assert "Database is locked during INSERT" in result.message
assert result.details["operation"] == "INSERT"
assert result.details["table"] == "queue_items"
assert result.cause == original_error
def test_handle_database_error_no_such_table(self):
"""Test handling of SQLite no such table error."""
original_error = sqlite3.OperationalError("no such table: missing_table")
result = ErrorHandler.handle_database_error(
original_error,
"SELECT",
"missing_table"
)
assert isinstance(result, SchemaError)
assert "Table does not exist during SELECT" in result.message
assert result.details["operation"] == "SELECT"
assert result.details["table"] == "missing_table"
def test_handle_database_error_integrity_error(self):
"""Test handling of SQLite integrity error."""
original_error = sqlite3.IntegrityError("UNIQUE constraint failed")
result = ErrorHandler.handle_database_error(
original_error,
"INSERT",
"queue_items"
)
assert isinstance(result, ValidationError)
assert "Data integrity violation during INSERT" in result.message
assert result.details["operation"] == "INSERT"
assert result.details["table"] == "queue_items"
def test_handle_database_error_generic(self):
"""Test handling of generic database error."""
original_error = Exception("Generic error")
result = ErrorHandler.handle_database_error(
original_error,
"UPDATE",
"queue_items"
)
assert isinstance(result, DatabaseError)
assert "Unexpected error during UPDATE" in result.message
assert result.details["operation"] == "UPDATE"
assert result.details["table"] == "queue_items"
def test_handle_api_error(self):
"""Test handling of API errors."""
original_error = ValueError("Invalid input")
result = ErrorHandler.handle_api_error(
original_error,
"/api/items",
"POST"
)
assert isinstance(result, APIError)
assert "API error at POST /api/items" in result.message
assert result.details["endpoint"] == "/api/items"
assert result.details["method"] == "POST"
assert result.cause == original_error
def test_handle_workflow_error(self):
"""Test handling of workflow execution errors."""
original_error = RuntimeError("Execution failed")
result = ErrorHandler.handle_workflow_error(
original_error,
workflow_id="test-123",
stage="validation"
)
assert isinstance(result, WorkflowExecutionError)
assert "Workflow execution failed" in result.message
assert result.details["workflow_id"] == "test-123"
assert result.details["execution_stage"] == "validation"
assert result.cause == original_error
class TestErrorHandlingDecorators:
"""Test error handling decorators."""
def test_with_database_error_handling_success(self):
"""Test database error handling decorator with successful operation."""
@with_database_error_handling(operation="test_op", table="test_table")
def successful_function():
return "success"
result = successful_function()
assert result == "success"
def test_with_database_error_handling_sqlite_error(self):
"""Test database error handling decorator with SQLite error."""
@with_database_error_handling(operation="test_op", table="test_table")
def failing_function():
raise sqlite3.OperationalError("database is locked")
with pytest.raises(ConnectionError) as exc_info:
failing_function()
error = exc_info.value
assert "Database is locked during test_op" in error.message
assert error.details["operation"] == "test_op"
assert error.details["table"] == "test_table"
def test_with_database_error_handling_generic_error(self):
"""Test database error handling decorator with generic error."""
@with_database_error_handling(operation="test_op")
def failing_function():
raise ValueError("Generic error")
with pytest.raises(DatabaseError) as exc_info:
failing_function()
error = exc_info.value
assert "Unexpected error during test_op" in error.message
assert error.details["operation"] == "test_op"
def test_with_database_error_handling_reraise_queue_manager_error(self):
"""Test that QueueManagerError is re-raised as-is."""
@with_database_error_handling(operation="test_op")
def failing_function():
raise ValidationError("Already a queue manager error")
with pytest.raises(ValidationError) as exc_info:
failing_function()
error = exc_info.value
assert error.message == "Already a queue manager error"
class TestErrorRecovery:
"""Test error recovery utilities."""
def test_retry_with_backoff_success_first_try(self):
"""Test retry mechanism with success on first try."""
call_count = 0
def successful_function():
nonlocal call_count
call_count += 1
return "success"
result = ErrorRecovery.retry_with_backoff(successful_function, max_retries=3)
assert result == "success"
assert call_count == 1
def test_retry_with_backoff_success_after_retries(self):
"""Test retry mechanism with success after retries."""
call_count = 0
def eventually_successful_function():
nonlocal call_count
call_count += 1
if call_count < 3:
raise ValueError("Temporary error")
return "success"
with patch('time.sleep'): # Mock sleep to speed up test
result = ErrorRecovery.retry_with_backoff(
eventually_successful_function,
max_retries=3,
base_delay=0.1
)
assert result == "success"
assert call_count == 3
def test_retry_with_backoff_all_retries_fail(self):
"""Test retry mechanism when all retries fail."""
call_count = 0
def always_failing_function():
nonlocal call_count
call_count += 1
raise ValueError(f"Error {call_count}")
with patch('time.sleep'): # Mock sleep to speed up test
with pytest.raises(ValueError) as exc_info:
ErrorRecovery.retry_with_backoff(
always_failing_function,
max_retries=2,
base_delay=0.1
)
assert "Error 3" in str(exc_info.value) # Last error
assert call_count == 3 # Initial + 2 retries
def test_safe_execute_success(self):
"""Test safe execution with successful function."""
def successful_function():
return "success"
result = ErrorRecovery.safe_execute(successful_function)
assert result == "success"
def test_safe_execute_with_error(self):
"""Test safe execution with error and default value."""
def failing_function():
raise ValueError("Error")
result = ErrorRecovery.safe_execute(failing_function, default_value="default")
assert result == "default"
def test_safe_execute_with_error_no_default(self):
"""Test safe execution with error and no default value."""
def failing_function():
raise ValueError("Error")
result = ErrorRecovery.safe_execute(failing_function)
assert result is None
class TestDatabaseErrorHandling:
"""Test error handling in database operations."""
def setup_method(self):
"""Set up test database."""
self.temp_dir = tempfile.mkdtemp()
self.db_path = Path(self.temp_dir) / "test.db"
self.database = SQLiteDatabase(self.db_path)
def teardown_method(self):
"""Clean up test database."""
if hasattr(self, 'database'):
self.database.close()
def test_database_initialization_error_handling(self):
"""Test error handling during database initialization."""
# Create a database with invalid path (directory instead of file)
invalid_path = Path(self.temp_dir) / "invalid_dir"
invalid_path.mkdir()
database = SQLiteDatabase(invalid_path)
# This should handle the error gracefully
with pytest.raises(DatabaseError):
database.initialize()
def test_create_queue_item_validation_error(self):
"""Test validation error when creating queue item."""
self.database.initialize()
# Create item with empty ID
item = QueueItem(id="", workflow_name="test")
with pytest.raises(ValidationError) as exc_info:
self.database.create_queue_item(item)
error = exc_info.value
assert "Queue item ID cannot be empty" in error.message
assert error.details["field"] == "id"
def test_get_queue_item_validation_error(self):
"""Test validation error when getting queue item."""
self.database.initialize()
with pytest.raises(ValidationError) as exc_info:
self.database.get_queue_item("")
error = exc_info.value
assert "Item ID cannot be empty" in error.message
assert error.details["field"] == "item_id"
def test_database_connection_error_handling(self):
"""Test database connection error handling."""
# Create database with invalid permissions
invalid_db = SQLiteDatabase("/root/invalid.db") # Should fail on most systems
with pytest.raises(DatabaseError):
invalid_db._get_connection()
class TestQueueServiceErrorHandling:
"""Test error handling in queue service operations."""
def setup_method(self):
"""Set up test queue service."""
self.temp_dir = tempfile.mkdtemp()
self.db_path = Path(self.temp_dir) / "test.db"
self.database = SQLiteDatabase(self.db_path)
self.database.initialize()
self.queue_service = QueueService(self.database)
def teardown_method(self):
"""Clean up test resources."""
if hasattr(self, 'queue_service'):
self.queue_service.close()
def test_add_workflow_validation_error_empty_data(self):
"""Test validation error when adding workflow with empty data."""
with pytest.raises(ValidationError) as exc_info:
self.queue_service.add_workflow({})
error = exc_info.value
assert "Workflow data cannot be empty" in error.message
assert error.details["field"] == "workflow_data"
def test_add_workflow_validation_error_invalid_type(self):
"""Test validation error when adding workflow with invalid data type."""
with pytest.raises(ValidationError) as exc_info:
self.queue_service.add_workflow("not a dict")
error = exc_info.value
assert "Workflow data must be a dictionary" in error.message
assert error.details["field"] == "workflow_data"
assert error.details["value"] == "str"
def test_queue_service_not_initialized_error(self):
"""Test error when using uninitialized queue service."""
# Create service without initializing
uninit_service = QueueService.__new__(QueueService)
uninit_service._initialized = False
with pytest.raises(QueueServiceError) as exc_info:
uninit_service.add_workflow({"test": "data"})
error = exc_info.value
assert "Queue service not initialized" in error.message
assert error.details["operation"] == "add_workflow"
class TestWorkflowExecutorErrorHandling:
"""Test error handling in workflow executor."""
def setup_method(self):
"""Set up test workflow executor."""
from workflow_executor import WorkflowExecutor
self.executor = WorkflowExecutor()
def teardown_method(self):
"""Clean up test resources."""
if hasattr(self, 'executor'):
self.executor.close()
def test_execute_workflow_validation_error_empty_data(self):
"""Test validation error when executing workflow with empty data."""
with pytest.raises(ValidationError) as exc_info:
self.executor.execute_workflow({})
error = exc_info.value
assert "Workflow data cannot be empty" in error.message
assert error.details["field"] == "workflow_data"
def test_execute_workflow_validation_error_invalid_type(self):
"""Test validation error when executing workflow with invalid data type."""
with pytest.raises(ValidationError) as exc_info:
self.executor.execute_workflow("not a dict")
error = exc_info.value
assert "Workflow data must be a dictionary" in error.message
assert error.details["field"] == "workflow_data"
assert error.details["value"] == "str"
if __name__ == "__main__":
pytest.main([__file__])
test_logging_monitoring.py
+602
View File
@@ -0,0 +1,602 @@
"""
Tests for logging and monitoring functionality in the ComfyUI Queue Manager.
"""
import asyncio
import json
import logging
import tempfile
import threading
import time
from datetime import datetime, timezone, timedelta
from pathlib import Path
from unittest.mock import Mock, patch
import pytest
from health_check import (
DatabaseHealthCheck,
HealthCheck,
HealthCheckManager,
HealthCheckResult,
HealthStatus,
QueueServiceHealthCheck,
SystemResourceHealthCheck,
)
from logging_config import LoggingConfig, QueueManagerFilter, StructuredFormatter
from performance_monitor import (
OperationStats,
PerformanceMetric,
PerformanceMonitor,
time_function,
)
class TestStructuredFormatter:
"""Test structured logging formatter."""
def test_basic_formatting(self):
"""Test basic log record formatting."""
formatter = StructuredFormatter()
# Create a log record
record = logging.LogRecord(
name="test.logger",
level=logging.INFO,
pathname="/test/path.py",
lineno=42,
msg="Test message",
args=(),
exc_info=None
)
# Format the record
formatted = formatter.format(record)
# Parse the JSON
log_data = json.loads(formatted)
assert log_data["level"] == "INFO"
assert log_data["logger"] == "test.logger"
assert log_data["message"] == "Test message"
assert log_data["line"] == 42
assert "timestamp" in log_data
def test_formatting_with_extra_fields(self):
"""Test formatting with extra fields."""
formatter = StructuredFormatter(include_extra=True)
record = logging.LogRecord(
name="test.logger",
level=logging.ERROR,
pathname="/test/path.py",
lineno=42,
msg="Test error",
args=(),
exc_info=None
)
# Add extra fields
record.user_id = "test-user"
record.operation = "test_operation"
record.custom_data = {"key": "value"}
formatted = formatter.format(record)
log_data = json.loads(formatted)
assert "extra" in log_data
assert log_data["extra"]["user_id"] == "test-user"
assert log_data["extra"]["operation"] == "test_operation"
assert log_data["extra"]["custom_data"] == {"key": "value"}
def test_formatting_with_exception(self):
"""Test formatting with exception information."""
formatter = StructuredFormatter()
try:
raise ValueError("Test exception")
except ValueError:
record = logging.LogRecord(
name="test.logger",
level=logging.ERROR,
pathname="/test/path.py",
lineno=42,
msg="Error occurred",
args=(),
exc_info=True
)
formatted = formatter.format(record)
log_data = json.loads(formatted)
assert "exception" in log_data
assert log_data["exception"]["type"] == "ValueError"
assert log_data["exception"]["message"] == "Test exception"
assert "traceback" in log_data["exception"]
class TestQueueManagerFilter:
"""Test logging filter."""
def test_level_filtering(self):
"""Test filtering by log level."""
filter_obj = QueueManagerFilter(min_level=logging.WARNING)
# Create records with different levels
info_record = logging.LogRecord(
name="test", level=logging.INFO, pathname="", lineno=0,
msg="Info message", args=(), exc_info=None
)
warning_record = logging.LogRecord(
name="test", level=logging.WARNING, pathname="", lineno=0,
msg="Warning message", args=(), exc_info=None
)
assert not filter_obj.filter(info_record)
assert filter_obj.filter(warning_record)
def test_component_filtering(self):
"""Test filtering by component."""
filter_obj = QueueManagerFilter(component="database")
# Create records with different logger names
db_record = logging.LogRecord(
name="queue_manager.database", level=logging.INFO, pathname="", lineno=0,
msg="Database message", args=(), exc_info=None
)
api_record = logging.LogRecord(
name="queue_manager.api", level=logging.INFO, pathname="", lineno=0,
msg="API message", args=(), exc_info=None
)
assert filter_obj.filter(db_record)
assert not filter_obj.filter(api_record)
class TestLoggingConfig:
"""Test logging configuration."""
def setup_method(self):
"""Set up test logging configuration."""
self.temp_dir = Path(tempfile.mkdtemp())
self.log_config = LoggingConfig(log_dir=self.temp_dir)
def test_configure_logging(self):
"""Test logging configuration setup."""
self.log_config.configure_logging(
console_level=logging.INFO,
file_level=logging.DEBUG,
structured=True
)
# Check that log files are created
assert (self.temp_dir / "queue_manager.log").exists()
assert (self.temp_dir / "errors.log").exists()
assert (self.temp_dir / "performance.log").exists()
def test_get_logger(self):
"""Test getting component loggers."""
self.log_config.configure_logging()
logger = self.log_config.get_logger("test_component")
assert logger.name == "custom_nodes.comfyui-queue-manager.test_component"
assert logger.level == logging.DEBUG
def test_set_log_level(self):
"""Test setting log level for components."""
self.log_config.configure_logging()
self.log_config.set_log_level("test_component", logging.WARNING)
logger = self.log_config.get_logger("test_component")
assert logger.level == logging.WARNING
def test_cleanup_old_logs(self):
"""Test cleaning up old log files."""
self.log_config.configure_logging()
# Create some old log files
old_log = self.temp_dir / "old.log"
old_log.touch()
# Set modification time to 31 days ago
old_time = time.time() - (31 * 24 * 60 * 60)
old_log.stat().st_mtime = old_time
cleaned_count = self.log_config.cleanup_old_logs(days_to_keep=30)
# Note: This test might not work on all systems due to file system limitations
# The cleanup_old_logs method should handle this gracefully
assert cleaned_count >= 0
def test_get_log_stats(self):
"""Test getting log file statistics."""
self.log_config.configure_logging()
stats = self.log_config.get_log_stats()
assert "log_dir" in stats
assert "total_files" in stats
assert "total_size" in stats
assert "files" in stats
assert stats["total_files"] >= 0
class TestPerformanceMetric:
"""Test performance metric data class."""
def test_metric_creation(self):
"""Test creating performance metrics."""
metric = PerformanceMetric(
name="test.metric",
value=1.5,
unit="seconds",
tags={"operation": "test"}
)
assert metric.name == "test.metric"
assert metric.value == 1.5
assert metric.unit == "seconds"
assert metric.tags["operation"] == "test"
assert isinstance(metric.timestamp, datetime)
def test_metric_to_dict(self):
"""Test converting metric to dictionary."""
metric = PerformanceMetric(
name="test.metric",
value=2.0,
unit="count"
)
result = metric.to_dict()
assert result["name"] == "test.metric"
assert result["value"] == 2.0
assert result["unit"] == "count"
assert "timestamp" in result
class TestOperationStats:
"""Test operation statistics."""
def test_stats_initialization(self):
"""Test initializing operation stats."""
stats = OperationStats("test_operation")
assert stats.operation_name == "test_operation"
assert stats.total_calls == 0
assert stats.total_time == 0.0
assert stats.error_count == 0
assert stats.average_time == 0.0
assert stats.success_rate == 0.0
def test_stats_update(self):
"""Test updating operation stats."""
stats = OperationStats("test_operation")
# Update with successful operation
stats.update(1.5, error=False)
assert stats.total_calls == 1
assert stats.total_time == 1.5
assert stats.average_time == 1.5
assert stats.min_time == 1.5
assert stats.max_time == 1.5
assert stats.error_count == 0
assert stats.success_rate == 100.0
# Update with failed operation
stats.update(2.0, error=True)
assert stats.total_calls == 2
assert stats.total_time == 3.5
assert stats.average_time == 1.75
assert stats.min_time == 1.5
assert stats.max_time == 2.0
assert stats.error_count == 1
assert stats.success_rate == 50.0
def test_stats_to_dict(self):
"""Test converting stats to dictionary."""
stats = OperationStats("test_operation")
stats.update(1.0, error=False)
result = stats.to_dict()
assert result["operation_name"] == "test_operation"
assert result["total_calls"] == 1
assert result["average_time"] == 1.0
assert result["success_rate"] == 100.0
class TestPerformanceMonitor:
"""Test performance monitoring."""
def setup_method(self):
"""Set up test performance monitor."""
self.monitor = PerformanceMonitor(max_history_size=100)
def teardown_method(self):
"""Clean up test resources."""
if hasattr(self, 'monitor'):
self.monitor.close()
def test_record_metric(self):
"""Test recording performance metrics."""
metric = PerformanceMetric("test.metric", 1.0, "seconds")
self.monitor.record_metric(metric)
# Check that metric was recorded
assert len(self.monitor._metrics_history) == 1
assert self.monitor._metrics_history[0] == metric
def test_record_counter(self):
"""Test recording counter metrics."""
self.monitor.record_counter("test.counter", increment=5)
assert self.monitor._counters["test.counter"] == 5
assert len(self.monitor._metrics_history) == 1
def test_record_gauge(self):
"""Test recording gauge metrics."""
self.monitor.record_gauge("test.gauge", 42.0)
assert self.monitor._gauges["test.gauge"] == 42.0
assert len(self.monitor._metrics_history) == 1
def test_record_timing(self):
"""Test recording timing metrics."""
self.monitor.record_timing("test.timing", 1.5)
assert len(self.monitor._metrics_history) == 1
metric = self.monitor._metrics_history[0]
assert metric.name == "test.timing"
assert metric.value == 1.5
assert metric.unit == "seconds"
def test_update_operation_stats(self):
"""Test updating operation statistics."""
self.monitor.update_operation_stats("test_op", 1.0, error=False)
stats = self.monitor.get_operation_stats("test_op")
assert stats["operation_name"] == "test_op"
assert stats["total_calls"] == 1
assert stats["average_time"] == 1.0
def test_time_operation_context_manager(self):
"""Test timing operation context manager."""
with self.monitor.time_operation("test_context"):
time.sleep(0.01) # Small delay
stats = self.monitor.get_operation_stats("test_context")
assert stats["total_calls"] == 1
assert stats["average_time"] > 0
def test_time_operation_with_error(self):
"""Test timing operation with error."""
with pytest.raises(ValueError):
with self.monitor.time_operation("test_error"):
raise ValueError("Test error")
stats = self.monitor.get_operation_stats("test_error")
assert stats["total_calls"] == 1
assert stats["error_count"] == 1
assert stats["success_rate"] == 0.0
def test_get_metrics_summary(self):
"""Test getting metrics summary."""
# Record some metrics
self.monitor.record_timing("test.timing", 1.0)
self.monitor.record_timing("test.timing", 2.0)
self.monitor.record_counter("test.counter", 5)
summary = self.monitor.get_metrics_summary()
assert "metrics" in summary
assert "test.timing" in summary["metrics"]
assert summary["metrics"]["test.timing"]["count"] == 2
assert summary["metrics"]["test.timing"]["avg"] == 1.5
def test_get_health_status(self):
"""Test getting health status."""
# Record some operations
self.monitor.update_operation_stats("test_op", 1.0, error=False)
self.monitor.update_operation_stats("test_op", 1.5, error=False)
health = self.monitor.get_health_status()
assert "status" in health
assert "metrics" in health
assert health["metrics"]["total_operations"] == 2
assert health["metrics"]["total_errors"] == 0
@patch('psutil.cpu_percent')
@patch('psutil.virtual_memory')
@patch('psutil.disk_usage')
@patch('psutil.Process')
def test_collect_system_metrics(self, mock_process, mock_disk, mock_memory, mock_cpu):
"""Test collecting system metrics."""
# Mock system metrics
mock_cpu.return_value = 50.0
mock_memory.return_value = Mock(percent=60.0, available=1024*1024*1024, used=512*1024*1024)
mock_disk.return_value = Mock(total=1024*1024*1024*100, used=1024*1024*1024*50, free=1024*1024*1024*50)
mock_process.return_value = Mock(
memory_info=Mock(rss=1024*1024*100, vms=1024*1024*200),
cpu_percent=Mock(return_value=25.0),
num_threads=Mock(return_value=10)
)
self.monitor._collect_system_metrics()
# Check that system metrics were recorded
assert "system.cpu.usage_percent" in self.monitor._gauges
assert "system.memory.usage_percent" in self.monitor._gauges
assert self.monitor._gauges["system.cpu.usage_percent"] == 50.0
assert self.monitor._gauges["system.memory.usage_percent"] == 60.0
class TestTimeFunctionDecorator:
"""Test time_function decorator."""
def setup_method(self):
"""Set up test performance monitor."""
self.monitor = PerformanceMonitor()
def teardown_method(self):
"""Clean up test resources."""
if hasattr(self, 'monitor'):
self.monitor.close()
def test_time_function_decorator(self):
"""Test timing function with decorator."""
@time_function("test_function")
def test_func():
time.sleep(0.01)
return "result"
result = test_func()
assert result == "result"
# Check that timing was recorded
from performance_monitor import get_performance_monitor
monitor = get_performance_monitor()
stats = monitor.get_operation_stats("test_function")
assert stats["total_calls"] >= 1
def test_time_function_with_error(self):
"""Test timing function with error."""
@time_function("test_error_function")
def error_func():
raise ValueError("Test error")
with pytest.raises(ValueError):
error_func()
# Check that error was recorded
from performance_monitor import get_performance_monitor
monitor = get_performance_monitor()
stats = monitor.get_operation_stats("test_error_function")
assert stats["total_calls"] >= 1
assert stats["error_count"] >= 1
class TestHealthCheck:
"""Test health check base class."""
def test_health_check_timeout(self):
"""Test health check timeout handling."""
class SlowHealthCheck(HealthCheck):
async def _perform_check(self):
await asyncio.sleep(2.0) # Longer than timeout
return HealthCheckResult(
name=self.name,
status=HealthStatus.HEALTHY,
message="Should not reach here"
)
check = SlowHealthCheck("slow_check", timeout=0.1)
# Run the check
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
result = loop.run_until_complete(check.check())
finally:
loop.close()
assert result.status == HealthStatus.UNHEALTHY
assert "timed out" in result.message.lower()
def test_health_check_exception(self):
"""Test health check exception handling."""
class FailingHealthCheck(HealthCheck):
async def _perform_check(self):
raise ValueError("Test error")
check = FailingHealthCheck("failing_check")
# Run the check
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
result = loop.run_until_complete(check.check())
finally:
loop.close()
assert result.status == HealthStatus.UNHEALTHY
assert "Test error" in result.message
class TestDatabaseHealthCheck:
"""Test database health check."""
def test_database_health_check_uninitialized(self):
"""Test health check with uninitialized database."""
mock_db = Mock()
mock_db._initialized = False
check = DatabaseHealthCheck(mock_db)
# Run the check
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
result = loop.run_until_complete(check.check())
finally:
loop.close()
assert result.status == HealthStatus.UNHEALTHY
assert "not initialized" in result.message.lower()
class TestHealthCheckManager:
"""Test health check manager."""
def setup_method(self):
"""Set up test health check manager."""
self.manager = HealthCheckManager()
def teardown_method(self):
"""Clean up test resources."""
if hasattr(self, 'manager'):
self.manager.close()
def test_register_and_unregister_check(self):
"""Test registering and unregistering health checks."""
check = SystemResourceHealthCheck()
self.manager.register_check(check)
assert "system_resources" in self.manager._checks
self.manager.unregister_check("system_resources")
assert "system_resources" not in self.manager._checks
def test_run_nonexistent_check(self):
"""Test running a non-existent health check."""
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
result = loop.run_until_complete(self.manager.run_check("nonexistent"))
finally:
loop.close()
assert result.status == HealthStatus.UNKNOWN
assert "not found" in result.message.lower()
def test_get_overall_health_no_checks(self):
"""Test getting overall health with no checks."""
health = self.manager.get_overall_health()
assert health["status"] == HealthStatus.UNKNOWN.value
assert "no health checks" in health["message"].lower()
if __name__ == "__main__":
pytest.main([__file__])
workflow_executor.py
+27 -17
View File
@@ -13,16 +13,14 @@ 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 WorkflowExecutionError(Exception):
"""Custom exception for workflow execution errors."""
class WorkflowExecutor(WorkflowExecutorInterface):
"""
Handles workflow execution integration with ComfyUI.
@@ -166,6 +164,7 @@ class WorkflowExecutor(WorkflowExecutorInterface):
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.
@@ -178,21 +177,27 @@ class WorkflowExecutor(WorkflowExecutorInterface):
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:
# 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,
}
# Execute the workflow
execution_id = self._execute_workflow_directly(workflow_data)
@@ -221,7 +226,12 @@ class WorkflowExecutor(WorkflowExecutorInterface):
QueueStatus.FAILED,
error_message=str(e)
)
raise WorkflowExecutionError(f"Workflow execution failed: {e}") from 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.