280 lines
8.9 KiB
Python
Executable File
280 lines
8.9 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Optimized queue consumer with integrated performance monitoring.
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import time
|
|
import signal
|
|
import argparse
|
|
import multiprocessing
|
|
import logging
|
|
from pathlib import Path
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import threading
|
|
|
|
# Configure logging
|
|
logger = logging.getLogger('app')
|
|
|
|
# Add project root directory to Python path
|
|
project_root = Path(__file__).parent.parent
|
|
sys.path.insert(0, str(project_root))
|
|
|
|
from task_queue.config import huey
|
|
from task_queue.manager import queue_manager
|
|
from task_queue.integration_tasks import process_files_async, cleanup_project_async
|
|
from huey.consumer import Consumer
|
|
from utils.settings import REDIS_URL
|
|
|
|
|
|
class OptimizedQueueConsumer:
|
|
"""Optimized queue consumer with integrated performance monitoring."""
|
|
|
|
def __init__(self, worker_type: str = "threads", workers: int = 2):
|
|
self.huey = huey
|
|
self.worker_type = worker_type
|
|
self.workers = workers
|
|
self.running = False
|
|
self.consumer = None
|
|
self.processed_count = 0
|
|
self.start_time = None
|
|
|
|
# Performance monitoring
|
|
self.performance_stats = {
|
|
'tasks_processed': 0,
|
|
'tasks_failed': 0,
|
|
'avg_processing_time': 0,
|
|
'start_time': None,
|
|
'last_activity': None
|
|
}
|
|
|
|
# Register signal handlers
|
|
signal.signal(signal.SIGINT, self._signal_handler)
|
|
signal.signal(signal.SIGTERM, self._signal_handler)
|
|
|
|
def _signal_handler(self, signum, frame):
|
|
"""Signal handler for graceful shutdown."""
|
|
logger.info(f"\nReceived signal {signum}, shutting down queue consumer...")
|
|
self.running = False
|
|
if self.consumer:
|
|
self.consumer.stop()
|
|
|
|
def setup_optimizations(self):
|
|
"""Set up performance optimizations."""
|
|
# Set environment variables
|
|
env_vars = {
|
|
'PYTHONUNBUFFERED': '1',
|
|
'PYTHONDONTWRITEBYTECODE': '1',
|
|
}
|
|
|
|
for key, value in env_vars.items():
|
|
os.environ[key] = value
|
|
|
|
# Optimize huey configuration
|
|
if hasattr(huey, 'immediate'):
|
|
huey.immediate = False
|
|
|
|
# Adjust based on worker type
|
|
if self.worker_type == "threads":
|
|
# Thread pool optimization
|
|
if hasattr(huey, 'worker_type'):
|
|
huey.worker_type = 'threads'
|
|
|
|
# Set thread pool size
|
|
if hasattr(huey, 'always_eager'):
|
|
huey.always_eager = False
|
|
|
|
logger.info("Queue consumer optimization setup complete:")
|
|
logger.info(f"- Worker type: {self.worker_type}")
|
|
logger.info(f"- Worker count: {self.workers}")
|
|
|
|
def monitor_performance(self):
|
|
"""Performance monitoring thread."""
|
|
while self.running:
|
|
time.sleep(30) # Output statistics every 30 seconds
|
|
|
|
if self.start_time:
|
|
elapsed = time.time() - self.start_time
|
|
rate = self.performance_stats['tasks_processed'] / max(1, elapsed)
|
|
|
|
logger.info(f"\n[Performance Stats]")
|
|
logger.info(f"- Uptime: {elapsed:.1f}s")
|
|
logger.info(f"- Tasks processed: {self.performance_stats['tasks_processed']}")
|
|
logger.info(f"- Failed tasks: {self.performance_stats['tasks_failed']}")
|
|
logger.info(f"- Average processing rate: {rate:.2f} tasks/s")
|
|
|
|
if self.performance_stats['avg_processing_time'] > 0:
|
|
logger.info(f"- Average processing time: {self.performance_stats['avg_processing_time']:.2f}s")
|
|
|
|
def start(self):
|
|
"""Start the queue consumer."""
|
|
logger.info("=" * 60)
|
|
logger.info("Optimized queue consumer starting")
|
|
logger.info("=" * 60)
|
|
|
|
# Apply optimizations
|
|
self.setup_optimizations()
|
|
|
|
logger.info(f"Backend: Redis ({REDIS_URL})")
|
|
logger.info("Press Ctrl+C to stop the consumer")
|
|
|
|
self.running = True
|
|
self.start_time = time.time()
|
|
self.performance_stats['start_time'] = self.start_time
|
|
|
|
# Start performance monitoring thread
|
|
monitor_thread = threading.Thread(target=self.monitor_performance, daemon=True)
|
|
monitor_thread.start()
|
|
|
|
try:
|
|
# Create consumer
|
|
self.consumer = Consumer(
|
|
self.huey,
|
|
workers=self.workers,
|
|
worker_type=self.worker_type,
|
|
max_delay=60.0, # Maximum delay
|
|
check_delay=1.0, # Check interval
|
|
periodic=True, # Enable periodic tasks
|
|
)
|
|
|
|
logger.info("Queue consumer started, waiting for tasks...")
|
|
|
|
# Start the consumer
|
|
self.consumer.run()
|
|
|
|
except KeyboardInterrupt:
|
|
logger.info("\nReceived keyboard interrupt signal")
|
|
except Exception as e:
|
|
logger.error(f"Queue consumer runtime error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
finally:
|
|
self.shutdown()
|
|
|
|
def shutdown(self):
|
|
"""Shut down the queue consumer."""
|
|
logger.info("\nShutting down queue consumer...")
|
|
self.running = False
|
|
|
|
if self.consumer:
|
|
try:
|
|
self.consumer.stop()
|
|
logger.info("Queue consumer stopped")
|
|
except Exception as e:
|
|
logger.error(f"Error stopping queue consumer: {e}")
|
|
|
|
# Output final statistics
|
|
if self.start_time:
|
|
elapsed = time.time() - self.start_time
|
|
logger.info(f"\n[Final Stats]")
|
|
logger.info(f"- Total uptime: {elapsed:.1f}s")
|
|
logger.info(f"- Total tasks processed: {self.performance_stats['tasks_processed']}")
|
|
logger.info(f"- Total failed tasks: {self.performance_stats['tasks_failed']}")
|
|
|
|
if self.performance_stats['tasks_processed'] > 0:
|
|
rate = self.performance_stats['tasks_processed'] / elapsed
|
|
logger.info(f"- Average processing rate: {rate:.2f} tasks/s")
|
|
|
|
|
|
def calculate_optimal_workers():
|
|
"""Calculate the optimal number of worker threads."""
|
|
cpu_count = multiprocessing.cpu_count()
|
|
|
|
# Based on CPU core count and system resources
|
|
if cpu_count <= 2:
|
|
return 2
|
|
elif cpu_count <= 4:
|
|
return 4
|
|
else:
|
|
return min(8, cpu_count)
|
|
|
|
|
|
def check_queue_status():
|
|
"""Check queue status."""
|
|
try:
|
|
stats = queue_manager.get_queue_stats()
|
|
|
|
logger.info("\n[Queue Status]")
|
|
if isinstance(stats, dict):
|
|
if 'total_tasks' in stats:
|
|
logger.info(f"- Total tasks: {stats['total_tasks']}")
|
|
if 'pending_tasks' in stats:
|
|
logger.info(f"- Pending tasks: {stats['pending_tasks']}")
|
|
if 'scheduled_tasks' in stats:
|
|
logger.info(f"- Scheduled tasks: {stats['scheduled_tasks']}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to get queue status: {e}")
|
|
|
|
|
|
def main():
|
|
"""Main entry point."""
|
|
parser = argparse.ArgumentParser(description="Optimized queue consumer")
|
|
parser.add_argument(
|
|
"--workers",
|
|
type=int,
|
|
default=calculate_optimal_workers(),
|
|
help=f"Number of worker threads (default: {calculate_optimal_workers()})"
|
|
)
|
|
parser.add_argument(
|
|
"--worker-type",
|
|
type=str,
|
|
default="threads",
|
|
choices=["threads", "greenlets", "gevent"],
|
|
help="Worker type (default: threads)"
|
|
)
|
|
parser.add_argument(
|
|
"--check-status",
|
|
action="store_true",
|
|
help="Check queue status and exit"
|
|
)
|
|
parser.add_argument(
|
|
"--profile",
|
|
type=str,
|
|
default="balanced",
|
|
choices=["low_memory", "balanced", "high_performance"],
|
|
help="Performance profile"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Apply performance profile
|
|
if args.profile == "low_memory":
|
|
os.environ['PYTHONOPTIMIZE'] = '1'
|
|
if args.workers > 2:
|
|
args.workers = 2
|
|
logger.info(f"Low memory mode: adjusted worker count to {args.workers}")
|
|
elif args.profile == "high_performance":
|
|
if args.workers < 4:
|
|
args.workers = 4
|
|
logger.info(f"High performance mode: adjusted worker count to {args.workers}")
|
|
|
|
# Check queue status
|
|
if args.check_status:
|
|
check_queue_status()
|
|
return
|
|
|
|
# Check environment
|
|
try:
|
|
import psutil
|
|
memory = psutil.virtual_memory()
|
|
logger.info("[System Info]")
|
|
logger.info(f"- CPU cores: {multiprocessing.cpu_count()}")
|
|
logger.info(f"- Available memory: {memory.available / (1024**3):.1f}GB")
|
|
logger.info(f"- Memory usage: {memory.percent:.1f}%")
|
|
except ImportError:
|
|
logger.info("[Tip] Install psutil to display system info: pip install psutil")
|
|
|
|
# Create and start the queue consumer
|
|
consumer = OptimizedQueueConsumer(
|
|
worker_type=args.worker_type,
|
|
workers=args.workers
|
|
)
|
|
|
|
consumer.start()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|