memory_manager.py

"""
Memory Manager for Large-Scale Document Processing
Author: Martin Bacigal, 01/2025 @ https://procureai.tech
License: MIT License
"""

import gc
import logging
import asyncio
import psutil
import os
from typing import Dict, Any, List, Optional, Callable
from pathlib import Path
from dataclasses import dataclass
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
import multiprocessing
from queue import Queue
import threading
import time
from functools import wraps
import weakref

@dataclass
class MemoryConfig:
    """Memory management configuration"""
    memory_limit_mb: int = 4096
    memory_check_interval: int = 5  # seconds
    gc_threshold_percent: float = 80.0
    batch_size_adjustment: bool = True
    min_batch_size: int = 1
    max_batch_size: int = 50
    process_pool_size: int = multiprocessing.cpu_count()
    enable_memory_profiling: bool = False

class MemoryMonitor:
    """Monitor system memory usage"""

    def __init__(self, config: MemoryConfig):
        self.config = config
        self.process = psutil.Process()
        self.monitoring = False
        self.monitor_thread = None
        self.callbacks: List[Callable] = []
        self._last_gc_time = time.time()

    def start_monitoring(self):
        """Start memory monitoring in background thread"""
        self.monitoring = True
        self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True)
        self.monitor_thread.start()
        logging.info("Memory monitoring started")

    def stop_monitoring(self):
        """Stop memory monitoring"""
        self.monitoring = False
        if self.monitor_thread:
            self.monitor_thread.join(timeout=1)
        logging.info("Memory monitoring stopped")

    def _monitor_loop(self):
        """Main monitoring loop"""
        while self.monitoring:
            try:
                memory_info = self.get_memory_info()

                # Check if we need to trigger garbage collection
                if memory_info['percent'] > self.config.gc_threshold_percent:
                    self._trigger_gc()

                # Notify callbacks
                for callback in self.callbacks:
                    callback(memory_info)

            except Exception as e:
                logging.error(f"Error in memory monitor: {e}")

            time.sleep(self.config.memory_check_interval)

    def get_memory_info(self) -> Dict[str, Any]:
        """Get current memory usage information"""
        memory_info = self.process.memory_info()
        system_memory = psutil.virtual_memory()

        return {
            'process_rss_mb': memory_info.rss / 1024 / 1024,
            'process_vms_mb': memory_info.vms / 1024 / 1024,
            'system_available_mb': system_memory.available / 1024 / 1024,
            'system_total_mb': system_memory.total / 1024 / 1024,
            'percent': system_memory.percent,
            'process_percent': (memory_info.rss / system_memory.total) * 100
        }

    def _trigger_gc(self):
        """Trigger garbage collection if needed"""
        current_time = time.time()
        if current_time - self._last_gc_time > 30:  # Don't GC more than once per 30 seconds
            logging.info("Triggering garbage collection due to high memory usage")
            gc.collect()
            self._last_gc_time = current_time

    def add_callback(self, callback: Callable):
        """Add a callback for memory updates"""
        self.callbacks.append(callback)

    def remove_callback(self, callback: Callable):
        """Remove a callback"""
        if callback in self.callbacks:
            self.callbacks.remove(callback)

class BatchManager:
    """Manages batching for memory-efficient processing"""

    def __init__(self, config: MemoryConfig, memory_monitor: MemoryMonitor):
        self.config = config
        self.memory_monitor = memory_monitor
        self.current_batch_size = config.max_batch_size // 2
        self._adjustment_history = []

    def get_optimal_batch_size(self) -> int:
        """Calculate optimal batch size based on memory usage"""
        if not self.config.batch_size_adjustment:
            return self.current_batch_size

        memory_info = self.memory_monitor.get_memory_info()
        memory_percent = memory_info['process_percent']

        # Adjust batch size based on memory usage
        if memory_percent > 70:
            # High memory usage - reduce batch size
            self.current_batch_size = max(
                self.config.min_batch_size,
                int(self.current_batch_size * 0.8)
            )
        elif memory_percent < 40:
            # Low memory usage - increase batch size
            self.current_batch_size = min(
                self.config.max_batch_size,
                int(self.current_batch_size * 1.2)
            )

        # Track adjustment history
        self._adjustment_history.append({
            'timestamp': time.time(),
            'batch_size': self.current_batch_size,
            'memory_percent': memory_percent
        })

        # Keep only last 100 entries
        if len(self._adjustment_history) > 100:
            self._adjustment_history = self._adjustment_history[-100:]

        logging.debug(f"Optimal batch size: {self.current_batch_size} (Memory: {memory_percent:.1f}%)")
        return self.current_batch_size

    def create_batches(self, items: List[Any]) -> List[List[Any]]:
        """Create memory-efficient batches from items"""
        batch_size = self.get_optimal_batch_size()
        batches = []

        for i in range(0, len(items), batch_size):
            batch = items[i:i + batch_size]
            batches.append(batch)

            # Force garbage collection between large batches
            if len(batch) > self.config.max_batch_size // 2:
                gc.collect()

        return batches

class ProcessingPool:
    """Manages process pool for memory isolation"""

    def __init__(self, config: MemoryConfig):
        self.config = config
        self.executor = None
        self._active_futures = weakref.WeakSet()

    def __enter__(self):
        self.executor = ProcessPoolExecutor(max_workers=self.config.process_pool_size)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.executor:
            self.executor.shutdown(wait=True)
            self.executor = None

    def submit(self, fn, *args, **kwargs):
        """Submit a task to the process pool"""
        if not self.executor:
            raise RuntimeError("ProcessingPool not initialized")

        future = self.executor.submit(fn, *args, **kwargs)
        self._active_futures.add(future)
        return future

    def map(self, fn, iterable, timeout=None):
        """Map function over iterable using process pool"""
        if not self.executor:
            raise RuntimeError("ProcessingPool not initialized")

        return self.executor.map(fn, iterable, timeout=timeout)

def memory_efficient_processor(memory_limit_mb: int = 4096):
    """Decorator for memory-efficient processing functions"""
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            # Create memory monitor for this processing task
            config = MemoryConfig(memory_limit_mb=memory_limit_mb)
            monitor = MemoryMonitor(config)

            try:
                monitor.start_monitoring()

                # Check memory before processing
                memory_info = monitor.get_memory_info()
                if memory_info['process_rss_mb'] > memory_limit_mb:
                    raise MemoryError(f"Memory usage ({memory_info['process_rss_mb']:.1f}MB) exceeds limit ({memory_limit_mb}MB)")

                # Run the function
                result = await func(*args, **kwargs)

                # Force garbage collection after processing
                gc.collect()

                return result

            finally:
                monitor.stop_monitoring()

        return wrapper
    return decorator

class StreamingFileProcessor:
    """Process large files in streaming fashion to minimize memory usage"""

    def __init__(self, chunk_size: int = 1024 * 1024):  # 1MB chunks
        self.chunk_size = chunk_size

    async def process_file_streaming(
        self,
        file_path: Path,
        processor_func: Callable[[bytes], Any]
    ) -> List[Any]:
        """Process file in chunks to minimize memory usage"""
        results = []

        async with aiofiles.open(file_path, 'rb') as f:
            while True:
                chunk = await f.read(self.chunk_size)
                if not chunk:
                    break

                # Process chunk
                result = await asyncio.to_thread(processor_func, chunk)
                results.append(result)

                # Allow other tasks to run
                await asyncio.sleep(0)

        return results

class MemoryEfficientDocumentProcessor:
    """Main class for memory-efficient document processing"""

    def __init__(
        self,
        memory_config: MemoryConfig,
        process_func: Callable
    ):
        self.config = memory_config
        self.process_func = process_func
        self.memory_monitor = MemoryMonitor(memory_config)
        self.batch_manager = BatchManager(memory_config, self.memory_monitor)
        self._processed_count = 0
        self._failed_count = 0

    async def process_documents(
        self,
        file_paths: List[Path],
        progress_callback: Optional[Callable] = None
    ) -> List[Dict[str, Any]]:
        """Process documents with memory management"""
        self.memory_monitor.start_monitoring()
        results = []

        try:
            # Create batches
            batches = self.batch_manager.create_batches(file_paths)
            total_batches = len(batches)

            logging.info(f"Processing {len(file_paths)} files in {total_batches} batches")

            for batch_idx, batch in enumerate(batches):
                # Check memory before processing batch
                memory_info = self.memory_monitor.get_memory_info()
                if memory_info['process_rss_mb'] > self.config.memory_limit_mb * 0.9:
                    logging.warning("Approaching memory limit, forcing garbage collection")
                    gc.collect()

                    # Re-check after GC
                    memory_info = self.memory_monitor.get_memory_info()
                    if memory_info['process_rss_mb'] > self.config.memory_limit_mb:
                        raise MemoryError("Memory limit exceeded")

                # Process batch
                batch_results = await self._process_batch(batch)
                results.extend(batch_results)

                # Update progress
                if progress_callback:
                    progress_callback(
                        processed=self._processed_count,
                        total=len(file_paths),
                        batch=batch_idx + 1,
                        total_batches=total_batches
                    )

                # Clean up after each batch
                del batch_results
                gc.collect()

                # Small delay to allow system to recover
                await asyncio.sleep(0.1)

            return results

        finally:
            self.memory_monitor.stop_monitoring()

    async def _process_batch(self, file_paths: List[Path]) -> List[Dict[str, Any]]:
        """Process a single batch of files"""
        results = []

        # Use process pool for memory isolation
        with ProcessingPool(self.config) as pool:
            # Submit all files in batch
            futures = []
            for file_path in file_paths:
                future = pool.submit(self._process_single_file, file_path)
                futures.append((file_path, future))

            # Collect results
            for file_path, future in futures:
                try:
                    result = future.result(timeout=self.config.memory_check_interval * 60)
                    results.append(result)
                    self._processed_count += 1
                except Exception as e:
                    logging.error(f"Error processing {file_path}: {e}")
                    results.append({
                        'file_path': str(file_path),
                        'error': str(e),
                        'status': 'failed'
                    })
                    self._failed_count += 1

        return results

    def _process_single_file(self, file_path: Path) -> Dict[str, Any]:
        """Process a single file (runs in separate process)"""
        try:
            # This runs in a separate process for memory isolation
            return self.process_func(file_path)
        except Exception as e:
            logging.error(f"Error in process_single_file for {file_path}: {e}")
            raise

# Utility functions
def get_system_memory_info() -> Dict[str, Any]:
    """Get system memory information"""
    memory = psutil.virtual_memory()
    swap = psutil.swap_memory()

    return {
        'total_mb': memory.total / 1024 / 1024,
        'available_mb': memory.available / 1024 / 1024,
        'percent': memory.percent,
        'swap_total_mb': swap.total / 1024 / 1024,
        'swap_used_mb': swap.used / 1024 / 1024,
        'swap_percent': swap.percent
    }

def estimate_file_memory_usage(file_path: Path, multiplier: float = 3.0) -> float:
    """Estimate memory usage for processing a file"""
    file_size_mb = file_path.stat().st_size / 1024 / 1024
    # Estimate based on file size and expected overhead
    return file_size_mb * multiplier

def calculate_optimal_workers(memory_limit_mb: int, avg_file_size_mb: float) -> int:
    """Calculate optimal number of workers based on memory constraints"""
    system_info = get_system_memory_info()
    available_memory = min(memory_limit_mb, system_info['available_mb'] * 0.8)

    # Estimate memory per worker (file size + overhead)
    memory_per_worker = avg_file_size_mb * 3  # 3x multiplier for processing overhead

    # Calculate workers
    optimal_workers = max(1, int(available_memory / memory_per_worker))

    # Cap at CPU count
    return min(optimal_workers, multiprocessing.cpu_count())

# Import aiofiles if available
try:
    import aiofiles
except ImportError:
    logging.warning("aiofiles not available, streaming file processing will be limited")