utils/memory_optimizer.py

"""
Memory Optimization Utilities

This module provides utilities for managing memory usage in the translation service,
including model cleanup, garbage collection, and memory monitoring.
"""

import gc
import logging
import psutil
import torch
from typing import Dict, Any, Optional
from contextlib import contextmanager
import threading
import time

logger = logging.getLogger(__name__)


class MemoryOptimizer:
    """Memory optimization manager for the translation service"""
    
    def __init__(self, max_memory_percent: float = 85.0):
        self.max_memory_percent = max_memory_percent
        self.model_cache: Dict[str, Any] = {}
        self.last_cleanup = time.time()
        self.cleanup_interval = 300  # 5 minutes
        self.lock = threading.Lock()
        
        logger.info(f"Memory optimizer initialized with {max_memory_percent}% threshold")
    
    def get_memory_usage(self) -> Dict[str, float]:
        """Get current memory usage statistics"""
        process = psutil.Process()
        memory_info = process.memory_info()
        
        return {
            "rss_mb": memory_info.rss / 1024 / 1024,  # Resident Set Size
            "vms_mb": memory_info.vms / 1024 / 1024,  # Virtual Memory Size
            "percent": process.memory_percent(),
            "available_mb": psutil.virtual_memory().available / 1024 / 1024
        }
    
    def should_cleanup(self) -> bool:
        """Check if memory cleanup is needed"""
        memory_usage = self.get_memory_usage()
        return memory_usage["percent"] > self.max_memory_percent
    
    def cleanup_models(self, keep_essential: bool = True):
        """Clean up loaded models to free memory"""
        with self.lock:
            try:
                # Clear PyTorch cache
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                    torch.cuda.synchronize()
                
                # Clear model cache (keep essential models if requested)
                if keep_essential:
                    # Keep only the most commonly used models
                    essential_models = ["nllb", "m2m"]
                    self.model_cache = {
                        k: v for k, v in self.model_cache.items() 
                        if k in essential_models
                    }
                else:
                    self.model_cache.clear()
                
                # Force garbage collection
                gc.collect()
                
                self.last_cleanup = time.time()
                
                memory_after = self.get_memory_usage()
                logger.info(f"Memory cleanup completed. Usage: {memory_after['percent']:.1f}%")
                
            except Exception as e:
                logger.error(f"Error during memory cleanup: {e}")
    
    def auto_cleanup_if_needed(self):
        """Automatically cleanup if memory usage is high or enough time has passed"""
        current_time = time.time()
        
        if (self.should_cleanup() or 
            current_time - self.last_cleanup > self.cleanup_interval):
            self.cleanup_models()
    
    def cache_model(self, model_name: str, model: Any):
        """Cache a model with memory management"""
        with self.lock:
            # Check if we need to cleanup before adding new model
            if self.should_cleanup():
                self.cleanup_models(keep_essential=True)
            
            self.model_cache[model_name] = model
            logger.debug(f"Cached model: {model_name}")
    
    def get_cached_model(self, model_name: str) -> Optional[Any]:
        """Get a cached model"""
        with self.lock:
            return self.model_cache.get(model_name)
    
    def remove_model(self, model_name: str):
        """Remove a specific model from cache"""
        with self.lock:
            if model_name in self.model_cache:
                del self.model_cache[model_name]
                logger.debug(f"Removed model from cache: {model_name}")
    
    @contextmanager
    def memory_context(self, context_name: str):
        """Context manager for memory-intensive operations"""
        memory_before = self.get_memory_usage()
        logger.debug(f"Starting {context_name}. Memory: {memory_before['percent']:.1f}%")
        
        try:
            yield
        finally:
            memory_after = self.get_memory_usage()
            memory_diff = memory_after['percent'] - memory_before['percent']
            
            logger.debug(f"Completed {context_name}. Memory change: {memory_diff:+.1f}%")
            
            # Auto cleanup if memory usage increased significantly
            if memory_diff > 5.0:  # More than 5% increase
                self.auto_cleanup_if_needed()
    
    def optimize_for_translation(self, text_length: int):
        """Optimize memory based on translation task requirements"""
        # For long texts, ensure we have enough memory
        if text_length > 10000:  # Long text
            self.auto_cleanup_if_needed()
        
        # For very long texts, be more aggressive
        if text_length > 50000:  # Very long text
            self.cleanup_models(keep_essential=True)
    
    def get_memory_status(self) -> Dict[str, Any]:
        """Get comprehensive memory status"""
        memory_usage = self.get_memory_usage()
        
        return {
            "total_gb": round(memory_usage.get("total", 0) / (1024**3), 2),
            "available_gb": round(memory_usage.get("available", 0) / (1024**3), 2),
            "used_gb": round(memory_usage.get("used", 0) / (1024**3), 2),
            "percent": memory_usage.get("percent", 0),
            "loaded_models_count": len(self.model_cache),
            "loaded_models": list(self.model_cache.keys()),
            "memory_usage": memory_usage,
            "last_cleanup": self.last_cleanup,
            "should_cleanup": self.should_cleanup(),
            "torch_cuda_available": torch.cuda.is_available(),
            "torch_cuda_memory": self._get_torch_cuda_memory() if torch.cuda.is_available() else None
        }
    
    def _get_torch_cuda_memory(self) -> Dict[str, float]:
        """Get PyTorch CUDA memory information"""
        if not torch.cuda.is_available():
            return {}
        
        return {
            "allocated_mb": torch.cuda.memory_allocated() / 1024 / 1024,
            "cached_mb": torch.cuda.memory_reserved() / 1024 / 1024,
            "max_allocated_mb": torch.cuda.max_memory_allocated() / 1024 / 1024
        }


# Global memory optimizer instance
memory_optimizer = MemoryOptimizer()


def optimize_memory_for_task(task_type: str, **kwargs):
    """Decorator to optimize memory for specific tasks"""
    def decorator(func):
        def wrapper(*args, **func_kwargs):
            with memory_optimizer.memory_context(f"{task_type}_{func.__name__}"):
                # Pre-task optimization
                if task_type == "translation":
                    text_length = kwargs.get('text_length', 0)
                    memory_optimizer.optimize_for_translation(text_length)
                
                return func(*args, **func_kwargs)
        return wrapper
    return decorator


def cleanup_memory():
    """Convenience function for manual memory cleanup"""
    memory_optimizer.cleanup_models()


def get_memory_status() -> Dict[str, Any]:
    """Get current memory status"""
    return memory_optimizer.get_memory_status()