latest
This commit is contained in:
@@ -0,0 +1,411 @@
|
||||
"""
|
||||
Memory-Aware Queue Management System
|
||||
====================================
|
||||
|
||||
Provides intelligent queue management based on memory usage and file sizes
|
||||
rather than simple task count limits. Supports multiple users with fair
|
||||
queuing and capacity management.
|
||||
|
||||
Features:
|
||||
- Memory-based queue limits (not just task count)
|
||||
- Fair queuing across multiple users
|
||||
- Upload capacity checking with user feedback
|
||||
- Graceful degradation under load
|
||||
- Service-specific memory tracking
|
||||
"""
|
||||
|
||||
import os
|
||||
import time
|
||||
import json
|
||||
import uuid
|
||||
import logging
|
||||
from typing import Dict, List, Optional, Any, Tuple
|
||||
from dataclasses import dataclass, asdict
|
||||
from enum import Enum
|
||||
import redis
|
||||
from .redis_manager import get_redis_manager
|
||||
import psutil
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class QueueStatus(Enum):
|
||||
ACCEPTING = "accepting" # Normal operation
|
||||
BUSY = "busy" # High load, warn users
|
||||
OVERLOADED = "overloaded" # Reject new uploads
|
||||
MAINTENANCE = "maintenance" # Manual override
|
||||
|
||||
@dataclass
|
||||
class MemoryConfig:
|
||||
"""Memory configuration for queue management."""
|
||||
max_total_memory_mb: int = 2048 # 2GB total queue memory
|
||||
max_user_memory_mb: int = 512 # 512MB per user
|
||||
max_file_size_mb: int = 100 # 100MB max file size
|
||||
memory_warning_threshold: float = 0.8 # Warn at 80%
|
||||
memory_reject_threshold: float = 0.95 # Reject at 95%
|
||||
|
||||
@dataclass
|
||||
class QueuedFile:
|
||||
"""Represents a file waiting in the queue."""
|
||||
file_id: str
|
||||
user_id: str
|
||||
filename: str
|
||||
size_bytes: int
|
||||
mime_type: str
|
||||
cabinet_id: str
|
||||
priority: int = 1
|
||||
queued_at: float = 0
|
||||
estimated_processing_time: int = 300 # seconds
|
||||
memory_estimate_mb: float = 0
|
||||
|
||||
def __post_init__(self):
|
||||
if self.queued_at == 0:
|
||||
self.queued_at = time.time()
|
||||
|
||||
# Estimate memory usage (rough heuristic)
|
||||
self.memory_estimate_mb = self._estimate_memory_usage()
|
||||
|
||||
def _estimate_memory_usage(self) -> float:
|
||||
"""Estimate memory usage for this file during processing."""
|
||||
base_mb = self.size_bytes / (1024 * 1024)
|
||||
|
||||
# Processing multipliers based on operations
|
||||
if self.mime_type == 'application/pdf':
|
||||
# PDF: original + extracted text + images + thumbnails
|
||||
return base_mb * 3.5
|
||||
elif self.mime_type.startswith('image/'):
|
||||
# Images: original + resized variants + OCR text
|
||||
return base_mb * 2.5
|
||||
else:
|
||||
# Other docs: original + PDF conversion + processing
|
||||
return base_mb * 4.0
|
||||
|
||||
class MemoryAwareQueue:
|
||||
"""Memory-aware queue management system."""
|
||||
|
||||
def __init__(self, environment: str = "dev"):
|
||||
self.redis_manager = get_redis_manager(environment)
|
||||
self.redis_client = self.redis_manager.client
|
||||
self.config = self._load_config()
|
||||
|
||||
# Redis keys
|
||||
self.upload_queue_key = "upload_queue"
|
||||
self.processing_memory_key = "processing_memory"
|
||||
self.user_quota_key = "user_quotas"
|
||||
self.system_status_key = "system_status"
|
||||
|
||||
logger.info(f"🧠 Memory-aware queue initialized (max: {self.config.max_total_memory_mb}MB)")
|
||||
|
||||
def _load_config(self) -> MemoryConfig:
|
||||
"""Load memory configuration from environment."""
|
||||
return MemoryConfig(
|
||||
max_total_memory_mb=int(os.getenv('QUEUE_MAX_MEMORY_MB', '2048')),
|
||||
max_user_memory_mb=int(os.getenv('QUEUE_MAX_USER_MEMORY_MB', '512')),
|
||||
max_file_size_mb=int(os.getenv('MAX_FILE_SIZE_MB', '100')),
|
||||
memory_warning_threshold=float(os.getenv('MEMORY_WARNING_THRESHOLD', '0.8')),
|
||||
memory_reject_threshold=float(os.getenv('MEMORY_REJECT_THRESHOLD', '0.95'))
|
||||
)
|
||||
|
||||
def check_upload_capacity(self, user_id: str, file_size_bytes: int,
|
||||
mime_type: str) -> Tuple[bool, str, Dict[str, Any]]:
|
||||
"""
|
||||
Check if system can accept a new upload.
|
||||
|
||||
Returns:
|
||||
(can_accept, message, queue_info)
|
||||
"""
|
||||
|
||||
# Create temporary QueuedFile to estimate memory
|
||||
temp_file = QueuedFile(
|
||||
file_id="temp",
|
||||
user_id=user_id,
|
||||
filename="temp",
|
||||
size_bytes=file_size_bytes,
|
||||
mime_type=mime_type,
|
||||
cabinet_id="temp"
|
||||
)
|
||||
|
||||
file_memory_mb = temp_file.memory_estimate_mb
|
||||
|
||||
# Check file size limit
|
||||
if file_size_bytes > (self.config.max_file_size_mb * 1024 * 1024):
|
||||
return False, f"File too large (max: {self.config.max_file_size_mb}MB)", {}
|
||||
|
||||
# Get current memory usage
|
||||
current_memory = self._get_current_memory_usage()
|
||||
user_memory = self._get_user_memory_usage(user_id)
|
||||
|
||||
# Check user quota
|
||||
if user_memory + file_memory_mb > self.config.max_user_memory_mb:
|
||||
return False, f"User quota exceeded (limit: {self.config.max_user_memory_mb}MB)", {
|
||||
'user_current': user_memory,
|
||||
'user_limit': self.config.max_user_memory_mb
|
||||
}
|
||||
|
||||
# Check system capacity
|
||||
total_after = current_memory + file_memory_mb
|
||||
max_memory = self.config.max_total_memory_mb
|
||||
|
||||
if total_after > (max_memory * self.config.memory_reject_threshold):
|
||||
queue_info = self._get_queue_info()
|
||||
return False, "System overloaded. Please try again later.", {
|
||||
'current_memory': current_memory,
|
||||
'max_memory': max_memory,
|
||||
'utilization': current_memory / max_memory,
|
||||
'queue_position': queue_info['total_queued'] + 1
|
||||
}
|
||||
|
||||
# Calculate wait time estimate
|
||||
wait_estimate = self._estimate_wait_time(user_id)
|
||||
|
||||
status = "ready"
|
||||
message = "Upload accepted"
|
||||
|
||||
if total_after > (max_memory * self.config.memory_warning_threshold):
|
||||
status = "busy"
|
||||
message = f"System busy. Estimated wait: {wait_estimate // 60}m {wait_estimate % 60}s"
|
||||
|
||||
return True, message, {
|
||||
'status': status,
|
||||
'estimated_wait_seconds': wait_estimate,
|
||||
'memory_usage': {
|
||||
'current': current_memory,
|
||||
'after_upload': total_after,
|
||||
'limit': max_memory,
|
||||
'utilization': total_after / max_memory
|
||||
},
|
||||
'user_quota': {
|
||||
'used': user_memory,
|
||||
'after_upload': user_memory + file_memory_mb,
|
||||
'limit': self.config.max_user_memory_mb
|
||||
}
|
||||
}
|
||||
|
||||
def enqueue_file(self, file_id: str, user_id: str, filename: str,
|
||||
size_bytes: int, mime_type: str, cabinet_id: str,
|
||||
priority: int = 1) -> Dict[str, Any]:
|
||||
"""
|
||||
Add file to upload queue.
|
||||
|
||||
Returns:
|
||||
Queue information including position and estimated wait time
|
||||
"""
|
||||
|
||||
queued_file = QueuedFile(
|
||||
file_id=file_id,
|
||||
user_id=user_id,
|
||||
filename=filename,
|
||||
size_bytes=size_bytes,
|
||||
mime_type=mime_type,
|
||||
cabinet_id=cabinet_id,
|
||||
priority=priority
|
||||
)
|
||||
|
||||
# Serialize and add to Redis queue (priority queue: higher priority = lower score)
|
||||
score = time.time() - (priority * 1000000) # Priority affects score significantly
|
||||
|
||||
self.redis_client.zadd(
|
||||
self.upload_queue_key,
|
||||
{json.dumps(asdict(queued_file)): score}
|
||||
)
|
||||
|
||||
# Update user quota tracking
|
||||
self._update_user_quota(user_id, queued_file.memory_estimate_mb, increment=True)
|
||||
|
||||
# Get queue position and wait estimate
|
||||
position = self._get_queue_position(file_id)
|
||||
wait_estimate = self._estimate_wait_time(user_id)
|
||||
|
||||
logger.info(f"📋 Queued file {file_id} for user {user_id} (pos: {position}, wait: {wait_estimate}s)")
|
||||
|
||||
return {
|
||||
'queued': True,
|
||||
'file_id': file_id,
|
||||
'queue_position': position,
|
||||
'estimated_wait_seconds': wait_estimate,
|
||||
'memory_estimate_mb': queued_file.memory_estimate_mb
|
||||
}
|
||||
|
||||
def dequeue_next_file(self, service_name: str) -> Optional[QueuedFile]:
|
||||
"""
|
||||
Get next file from queue for processing.
|
||||
|
||||
Args:
|
||||
service_name: The service requesting work (for capacity management)
|
||||
"""
|
||||
|
||||
# Check if service has capacity
|
||||
service_memory = self._get_service_memory_usage(service_name)
|
||||
service_limit = self._get_service_memory_limit(service_name)
|
||||
|
||||
if service_memory >= service_limit:
|
||||
logger.debug(f"Service {service_name} at capacity ({service_memory}/{service_limit}MB)")
|
||||
return None
|
||||
|
||||
# Get next item from priority queue (lowest score first)
|
||||
items = self.redis_client.zrange(self.upload_queue_key, 0, 0, withscores=True)
|
||||
|
||||
if not items:
|
||||
return None
|
||||
|
||||
file_data_json, score = items[0]
|
||||
file_data = json.loads(file_data_json)
|
||||
queued_file = QueuedFile(**file_data)
|
||||
|
||||
# Check if this file would exceed service memory limit
|
||||
if service_memory + queued_file.memory_estimate_mb > service_limit:
|
||||
# Skip this file for now, try smaller ones later
|
||||
logger.debug(f"File {queued_file.file_id} too large for {service_name} capacity")
|
||||
return None
|
||||
|
||||
# Remove from queue
|
||||
self.redis_client.zrem(self.upload_queue_key, file_data_json)
|
||||
|
||||
# Update tracking
|
||||
self._update_user_quota(queued_file.user_id, queued_file.memory_estimate_mb, increment=False)
|
||||
self._update_service_memory(service_name, queued_file.memory_estimate_mb, increment=True)
|
||||
|
||||
logger.info(f"🎯 Dequeued file {queued_file.file_id} for {service_name} processing")
|
||||
|
||||
return queued_file
|
||||
|
||||
def complete_processing(self, service_name: str, file_id: str, memory_used_mb: float):
|
||||
"""Mark file processing as complete and free memory."""
|
||||
self._update_service_memory(service_name, memory_used_mb, increment=False)
|
||||
logger.info(f"✅ Completed processing {file_id} in {service_name} (freed {memory_used_mb}MB)")
|
||||
|
||||
def _get_current_memory_usage(self) -> float:
|
||||
"""Get current total memory usage across all services."""
|
||||
services = ['docling', 'tika', 'llm', 'document_analysis']
|
||||
total = 0
|
||||
|
||||
for service in services:
|
||||
service_key = f"{self.processing_memory_key}:{service}"
|
||||
memory = float(self.redis_client.get(service_key) or 0)
|
||||
total += memory
|
||||
|
||||
return total
|
||||
|
||||
def _get_user_memory_usage(self, user_id: str) -> float:
|
||||
"""Get current memory usage for a specific user."""
|
||||
user_key = f"{self.user_quota_key}:{user_id}"
|
||||
return float(self.redis_client.get(user_key) or 0)
|
||||
|
||||
def _get_service_memory_usage(self, service_name: str) -> float:
|
||||
"""Get current memory usage for a service."""
|
||||
service_key = f"{self.processing_memory_key}:{service_name}"
|
||||
return float(self.redis_client.get(service_key) or 0)
|
||||
|
||||
def _get_service_memory_limit(self, service_name: str) -> float:
|
||||
"""Get memory limit for a service."""
|
||||
# Service-specific memory limits as percentage of total
|
||||
limits = {
|
||||
'docling': 0.4, # 40% for Docling (memory-intensive)
|
||||
'tika': 0.2, # 20% for Tika
|
||||
'llm': 0.3, # 30% for LLM processing
|
||||
'document_analysis': 0.1 # 10% for document analysis
|
||||
}
|
||||
|
||||
percentage = limits.get(service_name, 0.1)
|
||||
return self.config.max_total_memory_mb * percentage
|
||||
|
||||
def _update_user_quota(self, user_id: str, memory_mb: float, increment: bool):
|
||||
"""Update user memory quota tracking."""
|
||||
user_key = f"{self.user_quota_key}:{user_id}"
|
||||
|
||||
if increment:
|
||||
self.redis_client.incrbyfloat(user_key, memory_mb)
|
||||
else:
|
||||
current = float(self.redis_client.get(user_key) or 0)
|
||||
new_value = max(0, current - memory_mb)
|
||||
self.redis_client.set(user_key, new_value)
|
||||
|
||||
# Set expiration for cleanup
|
||||
self.redis_client.expire(user_key, 86400) # 24 hours
|
||||
|
||||
def _update_service_memory(self, service_name: str, memory_mb: float, increment: bool):
|
||||
"""Update service memory usage tracking."""
|
||||
service_key = f"{self.processing_memory_key}:{service_name}"
|
||||
|
||||
if increment:
|
||||
self.redis_client.incrbyfloat(service_key, memory_mb)
|
||||
else:
|
||||
current = float(self.redis_client.get(service_key) or 0)
|
||||
new_value = max(0, current - memory_mb)
|
||||
self.redis_client.set(service_key, new_value)
|
||||
|
||||
# Set expiration for cleanup
|
||||
self.redis_client.expire(service_key, 3600) # 1 hour
|
||||
|
||||
def _get_queue_position(self, file_id: str) -> int:
|
||||
"""Get position of file in queue."""
|
||||
items = self.redis_client.zrange(self.upload_queue_key, 0, -1)
|
||||
for i, item in enumerate(items):
|
||||
file_data = json.loads(item)
|
||||
if file_data['file_id'] == file_id:
|
||||
return i + 1
|
||||
return 0
|
||||
|
||||
def _estimate_wait_time(self, user_id: str) -> int:
|
||||
"""Estimate wait time for user's next file."""
|
||||
# Simple estimation based on queue position and average processing time
|
||||
queue_size = self.redis_client.zcard(self.upload_queue_key)
|
||||
avg_processing_time = 300 # 5 minutes average
|
||||
|
||||
return int(queue_size * avg_processing_time * 0.5) # Assume parallel processing
|
||||
|
||||
def _get_queue_info(self) -> Dict[str, Any]:
|
||||
"""Get comprehensive queue information."""
|
||||
total_queued = self.redis_client.zcard(self.upload_queue_key)
|
||||
current_memory = self._get_current_memory_usage()
|
||||
max_memory = self.config.max_total_memory_mb
|
||||
|
||||
return {
|
||||
'total_queued': total_queued,
|
||||
'memory_usage': {
|
||||
'current_mb': current_memory,
|
||||
'max_mb': max_memory,
|
||||
'utilization': current_memory / max_memory if max_memory > 0 else 0
|
||||
},
|
||||
'status': self._determine_system_status(current_memory, max_memory)
|
||||
}
|
||||
|
||||
def _determine_system_status(self, current_memory: float, max_memory: float) -> str:
|
||||
"""Determine current system status based on memory usage."""
|
||||
utilization = current_memory / max_memory if max_memory > 0 else 0
|
||||
|
||||
if utilization >= self.config.memory_reject_threshold:
|
||||
return "overloaded"
|
||||
elif utilization >= self.config.memory_warning_threshold:
|
||||
return "busy"
|
||||
else:
|
||||
return "ready"
|
||||
|
||||
def get_system_status(self) -> Dict[str, Any]:
|
||||
"""Get comprehensive system status for monitoring."""
|
||||
queue_info = self._get_queue_info()
|
||||
|
||||
# Service-specific info
|
||||
services = {}
|
||||
for service_name in ['docling', 'tika', 'llm', 'document_analysis']:
|
||||
services[service_name] = {
|
||||
'memory_used_mb': self._get_service_memory_usage(service_name),
|
||||
'memory_limit_mb': self._get_service_memory_limit(service_name),
|
||||
'utilization': self._get_service_memory_usage(service_name) / self._get_service_memory_limit(service_name)
|
||||
}
|
||||
|
||||
return {
|
||||
'status': queue_info['status'],
|
||||
'queue': queue_info,
|
||||
'services': services,
|
||||
'config': asdict(self.config)
|
||||
}
|
||||
|
||||
# Convenience functions
|
||||
def get_memory_queue(environment: str = "dev") -> MemoryAwareQueue:
|
||||
"""Get memory-aware queue instance."""
|
||||
return MemoryAwareQueue(environment)
|
||||
|
||||
def check_upload_capacity(user_id: str, file_size: int, mime_type: str, environment: str = "dev") -> Tuple[bool, str, Dict]:
|
||||
"""Quick capacity check for upload."""
|
||||
queue = get_memory_queue(environment)
|
||||
return queue.check_upload_capacity(user_id, file_size, mime_type)
|
||||
Reference in New Issue
Block a user