additory 0.1.0a1__py3-none-any.whl

additory/core/memory_manager.py

@@ -0,0 +1,547 @@
+# memory_manager.py
+# Centralized memory management for enhanced expressions system
+
+import gc
+import os
+import psutil
+import threading
+import time
+from typing import Dict, List, Optional, Callable, Any
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta
+from contextlib import contextmanager
+
+from .logging import log_info, log_warning
+
+
+@dataclass
+class MemorySnapshot:
+    """Memory usage snapshot at a point in time"""
+    timestamp: datetime
+    process_memory_mb: float
+    arrow_allocated_bytes: int
+    python_objects_count: int
+    gc_collections: Dict[int, int]
+    custom_metrics: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class MemoryThresholds:
+    """Memory usage thresholds for cleanup triggers"""
+    process_memory_mb: float = 500.0  # Process memory threshold
+    arrow_memory_bytes: int = 100 * 1024 * 1024  # 100MB Arrow memory
+    python_objects_count: int = 100000  # Python object count
+    cleanup_interval_seconds: float = 30.0  # Periodic cleanup interval
+
+
+class MemoryLeakDetector:
+    """Detects potential memory leaks by tracking memory growth patterns"""
+    
+    def __init__(self, window_size: int = 10, growth_threshold: float = 1.5):
+        self.window_size = window_size
+        self.growth_threshold = growth_threshold
+        self.snapshots: List[MemorySnapshot] = []
+        self.leak_warnings = 0
+    
+    def add_snapshot(self, snapshot: MemorySnapshot):
+        """Add a memory snapshot for leak detection"""
+        self.snapshots.append(snapshot)
+        
+        # Keep only the last window_size snapshots
+        if len(self.snapshots) > self.window_size:
+            self.snapshots.pop(0)
+        
+        # Check for potential leaks
+        if len(self.snapshots) >= self.window_size:
+            self._check_for_leaks()
+    
+    def _check_for_leaks(self):
+        """Check if memory usage shows signs of leaking"""
+        if len(self.snapshots) < 2:
+            return
+        
+        # Calculate memory growth rate
+        first_snapshot = self.snapshots[0]
+        last_snapshot = self.snapshots[-1]
+        
+        memory_growth = last_snapshot.process_memory_mb / max(first_snapshot.process_memory_mb, 1.0)  # Avoid division by zero
+        time_span = (last_snapshot.timestamp - first_snapshot.timestamp).total_seconds()
+        
+        # Check for sustained memory growth (relaxed time threshold for testing)
+        if memory_growth > self.growth_threshold and time_span > 10:  # 10 seconds minimum
+            self.leak_warnings += 1
+            log_warning(f"[memory_manager] Potential memory leak detected: "
+                       f"{memory_growth:.2f}x growth over {time_span:.1f}s")
+    
+    def get_leak_status(self) -> Dict[str, Any]:
+        """Get current leak detection status"""
+        return {
+            "leak_warnings": self.leak_warnings,
+            "snapshots_count": len(self.snapshots),
+            "monitoring_window": self.window_size,
+            "growth_threshold": self.growth_threshold
+        }
+
+
+class MemoryManager:
+    """Centralized memory management for enhanced expressions system"""
+    
+    def __init__(self):
+        # Configuration
+        self.thresholds = MemoryThresholds()
+        self.monitoring_enabled = True
+        self.auto_cleanup_enabled = True
+        
+        # State tracking
+        self.snapshots: List[MemorySnapshot] = []
+        self.cleanup_callbacks: List[Callable[[], None]] = []
+        self.leak_detector = MemoryLeakDetector()
+        
+        # Statistics
+        self.stats = {
+            "total_cleanups": 0,
+            "forced_cleanups": 0,
+            "auto_cleanups": 0,
+            "memory_freed_mb": 0.0,
+            "last_cleanup": None
+        }
+        
+        # Background monitoring
+        self._monitoring_thread = None
+        self._stop_monitoring = threading.Event()
+        
+        # Process reference for memory monitoring
+        try:
+            self.process = psutil.Process(os.getpid())
+        except Exception as e:
+            log_warning(f"[memory_manager] Failed to initialize process monitor: {e}")
+            self.process = None
+        
+        log_info("[memory_manager] Memory Manager initialized")
+    
+    def start_monitoring(self, interval_seconds: float = 30.0):
+        """Start background memory monitoring"""
+        if self._monitoring_thread and self._monitoring_thread.is_alive():
+            log_warning("[memory_manager] Monitoring already running")
+            return
+        
+        self.thresholds.cleanup_interval_seconds = interval_seconds
+        self._stop_monitoring.clear()
+        
+        self._monitoring_thread = threading.Thread(
+            target=self._monitoring_loop,
+            daemon=True,
+            name="MemoryMonitor"
+        )
+        self._monitoring_thread.start()
+        
+        log_info(f"[memory_manager] Started background monitoring (interval: {interval_seconds}s)")
+    
+    def stop_monitoring(self):
+        """Stop background memory monitoring"""
+        if self._monitoring_thread and self._monitoring_thread.is_alive():
+            self._stop_monitoring.set()
+            self._monitoring_thread.join(timeout=5.0)
+            log_info("[memory_manager] Stopped background monitoring")
+    
+    def _monitoring_loop(self):
+        """Background monitoring loop"""
+        while not self._stop_monitoring.is_set():
+            try:
+                if self.monitoring_enabled:
+                    # Take memory snapshot
+                    snapshot = self.take_snapshot()
+                    
+                    # Check for cleanup needs
+                    if self.auto_cleanup_enabled and self._should_cleanup(snapshot):
+                        self.cleanup_if_needed()
+                
+                # Wait for next interval
+                self._stop_monitoring.wait(self.thresholds.cleanup_interval_seconds)
+                
+            except Exception as e:
+                log_warning(f"[memory_manager] Monitoring loop error: {e}")
+                time.sleep(5)  # Brief pause before retrying
+    
+    def register_cleanup_callback(self, callback: Callable[[], None]):
+        """Register a callback function for memory cleanup"""
+        self.cleanup_callbacks.append(callback)
+        log_info(f"[memory_manager] Registered cleanup callback: {callback.__name__}")
+    
+    def unregister_cleanup_callback(self, callback: Callable[[], None]):
+        """Unregister a cleanup callback"""
+        if callback in self.cleanup_callbacks:
+            self.cleanup_callbacks.remove(callback)
+            log_info(f"[memory_manager] Unregistered cleanup callback: {callback.__name__}")
+    
+    def take_snapshot(self) -> MemorySnapshot:
+        """Take a snapshot of current memory usage"""
+        try:
+            # Process memory
+            process_memory_mb = 0.0
+            if self.process:
+                memory_info = self.process.memory_info()
+                process_memory_mb = memory_info.rss / 1024 / 1024
+            
+            # Arrow memory (if available)
+            arrow_allocated_bytes = 0
+            try:
+                import pyarrow as pa
+                arrow_allocated_bytes = pa.default_memory_pool().bytes_allocated()
+            except ImportError:
+                pass
+            
+            # Python objects
+            python_objects_count = len(gc.get_objects())
+            
+            # GC statistics
+            gc_collections = {i: gc.get_count()[i] for i in range(3)}
+            
+            snapshot = MemorySnapshot(
+                timestamp=datetime.now(),
+                process_memory_mb=process_memory_mb,
+                arrow_allocated_bytes=arrow_allocated_bytes,
+                python_objects_count=python_objects_count,
+                gc_collections=gc_collections
+            )
+            
+            # Store snapshot
+            self.snapshots.append(snapshot)
+            
+            # Keep only recent snapshots (last 100)
+            if len(self.snapshots) > 100:
+                self.snapshots.pop(0)
+            
+            # Add to leak detector
+            self.leak_detector.add_snapshot(snapshot)
+            
+            return snapshot
+            
+        except Exception as e:
+            log_warning(f"[memory_manager] Failed to take memory snapshot: {e}")
+            return MemorySnapshot(
+                timestamp=datetime.now(),
+                process_memory_mb=0.0,
+                arrow_allocated_bytes=0,
+                python_objects_count=0,
+                gc_collections={}
+            )
+    
+    def _should_cleanup(self, snapshot: MemorySnapshot) -> bool:
+        """Check if cleanup is needed based on thresholds"""
+        return (
+            snapshot.process_memory_mb > self.thresholds.process_memory_mb or
+            snapshot.arrow_allocated_bytes > self.thresholds.arrow_memory_bytes or
+            snapshot.python_objects_count > self.thresholds.python_objects_count
+        )
+    
+    def cleanup_if_needed(self) -> bool:
+        """Perform cleanup if memory usage exceeds thresholds"""
+        snapshot = self.take_snapshot()
+        
+        if not self._should_cleanup(snapshot):
+            return False
+        
+        log_info(f"[memory_manager] Auto cleanup triggered - "
+                f"Memory: {snapshot.process_memory_mb:.1f}MB, "
+                f"Arrow: {snapshot.arrow_allocated_bytes} bytes, "
+                f"Objects: {snapshot.python_objects_count}")
+        
+        memory_before = snapshot.process_memory_mb
+        self._perform_cleanup()
+        
+        # Take another snapshot to measure cleanup effectiveness
+        after_snapshot = self.take_snapshot()
+        memory_freed = max(0, memory_before - after_snapshot.process_memory_mb)
+        
+        self.stats["auto_cleanups"] += 1
+        self.stats["memory_freed_mb"] += memory_freed
+        self.stats["last_cleanup"] = datetime.now()
+        
+        log_info(f"[memory_manager] Auto cleanup completed - "
+                f"Freed: {memory_freed:.1f}MB")
+        
+        return True
+    
+    def force_cleanup(self) -> Dict[str, float]:
+        """Force immediate memory cleanup"""
+        log_info("[memory_manager] Forcing memory cleanup")
+        
+        before_snapshot = self.take_snapshot()
+        memory_before = before_snapshot.process_memory_mb
+        
+        self._perform_cleanup()
+        
+        after_snapshot = self.take_snapshot()
+        memory_after = after_snapshot.process_memory_mb
+        memory_freed = max(0, memory_before - memory_after)
+        
+        self.stats["forced_cleanups"] += 1
+        self.stats["memory_freed_mb"] += memory_freed
+        self.stats["last_cleanup"] = datetime.now()
+        
+        cleanup_stats = {
+            "memory_before_mb": memory_before,
+            "memory_after_mb": memory_after,
+            "memory_freed_mb": memory_freed,
+            "arrow_freed_bytes": before_snapshot.arrow_allocated_bytes - after_snapshot.arrow_allocated_bytes
+        }
+        
+        log_info(f"[memory_manager] Forced cleanup completed - "
+                f"Freed: {memory_freed:.1f}MB")
+        
+        return cleanup_stats
+    
+    def _perform_cleanup(self):
+        """Perform the actual cleanup operations"""
+        # Call registered cleanup callbacks
+        for callback in self.cleanup_callbacks:
+            try:
+                callback()
+            except Exception as e:
+                log_warning(f"[memory_manager] Cleanup callback failed: {e}")
+        
+        # Force garbage collection
+        collected = gc.collect()
+        
+        # Additional cleanup for specific libraries
+        self._cleanup_arrow_memory()
+        self._cleanup_polars_memory()
+        
+        self.stats["total_cleanups"] += 1
+        
+        log_info(f"[memory_manager] Cleanup performed - "
+                f"GC collected: {collected} objects")
+    
+    def _cleanup_arrow_memory(self):
+        """Cleanup Arrow-specific memory"""
+        try:
+            import pyarrow as pa
+            # Force Arrow memory pool cleanup
+            pool = pa.default_memory_pool()
+            allocated_before = pool.bytes_allocated()
+            
+            # Trigger garbage collection to free Arrow objects
+            gc.collect()
+            
+            allocated_after = pool.bytes_allocated()
+            freed = allocated_before - allocated_after
+            
+            if freed > 0:
+                log_info(f"[memory_manager] Arrow cleanup freed {freed} bytes")
+                
+        except ImportError:
+            pass
+        except Exception as e:
+            log_warning(f"[memory_manager] Arrow cleanup failed: {e}")
+    
+    def _cleanup_polars_memory(self):
+        """Cleanup Polars-specific memory"""
+        try:
+            import polars as pl
+            # Polars cleanup is mostly handled by Rust's memory management
+            # But we can clear any cached data
+            gc.collect()
+            
+        except ImportError:
+            pass
+        except Exception as e:
+            log_warning(f"[memory_manager] Polars cleanup failed: {e}")
+    
+    def get_memory_stats(self) -> Dict[str, Any]:
+        """Get comprehensive memory statistics"""
+        current_snapshot = self.take_snapshot()
+        
+        # Calculate memory trends
+        memory_trend = "stable"
+        if len(self.snapshots) >= 10:
+            recent_memory = sum(s.process_memory_mb for s in self.snapshots[-5:]) / min(5, len(self.snapshots))
+            older_memory = sum(s.process_memory_mb for s in self.snapshots[-10:-5]) / max(1, min(5, len(self.snapshots) - 5))
+            
+            if recent_memory > older_memory * 1.1:
+                memory_trend = "increasing"
+            elif recent_memory < older_memory * 0.9:
+                memory_trend = "decreasing"
+        elif len(self.snapshots) >= 2:
+            # Simple trend for fewer snapshots
+            if self.snapshots[-1].process_memory_mb > self.snapshots[0].process_memory_mb * 1.1:
+                memory_trend = "increasing"
+            elif self.snapshots[-1].process_memory_mb < self.snapshots[0].process_memory_mb * 0.9:
+                memory_trend = "decreasing"
+        
+        return {
+            "current": {
+                "process_memory_mb": current_snapshot.process_memory_mb,
+                "arrow_allocated_bytes": current_snapshot.arrow_allocated_bytes,
+                "python_objects_count": current_snapshot.python_objects_count,
+                "gc_collections": current_snapshot.gc_collections
+            },
+            "thresholds": {
+                "process_memory_mb": self.thresholds.process_memory_mb,
+                "arrow_memory_bytes": self.thresholds.arrow_memory_bytes,
+                "python_objects_count": self.thresholds.python_objects_count
+            },
+            "trends": {
+                "memory_trend": memory_trend,
+                "snapshots_count": len(self.snapshots),
+                "monitoring_enabled": self.monitoring_enabled
+            },
+            "cleanup_stats": self.stats.copy(),
+            "leak_detection": self.leak_detector.get_leak_status()
+        }
+    
+    def get_memory_usage_mb(self) -> float:
+        """Get current memory usage in MB"""
+        if self.process:
+            try:
+                return self.process.memory_info().rss / 1024 / 1024
+            except Exception:
+                pass
+        return 0.0
+    
+    def set_thresholds(self, **kwargs):
+        """Update memory thresholds"""
+        for key, value in kwargs.items():
+            if hasattr(self.thresholds, key):
+                setattr(self.thresholds, key, value)
+                log_info(f"[memory_manager] Updated threshold {key} = {value}")
+            else:
+                log_warning(f"[memory_manager] Unknown threshold: {key}")
+    
+    def enable_monitoring(self):
+        """Enable memory monitoring"""
+        self.monitoring_enabled = True
+        log_info("[memory_manager] Memory monitoring enabled")
+    
+    def disable_monitoring(self):
+        """Disable memory monitoring"""
+        self.monitoring_enabled = False
+        log_info("[memory_manager] Memory monitoring disabled")
+    
+    def enable_auto_cleanup(self):
+        """Enable automatic cleanup"""
+        self.auto_cleanup_enabled = True
+        log_info("[memory_manager] Auto cleanup enabled")
+    
+    def disable_auto_cleanup(self):
+        """Disable automatic cleanup"""
+        self.auto_cleanup_enabled = False
+        log_info("[memory_manager] Auto cleanup disabled")
+    
+    @contextmanager
+    def memory_context(self, description: str = "operation"):
+        """Context manager for monitoring memory usage during operations"""
+        before_snapshot = self.take_snapshot()
+        start_time = datetime.now()
+        
+        log_info(f"[memory_manager] Starting {description} - "
+                f"Memory: {before_snapshot.process_memory_mb:.1f}MB")
+        
+        try:
+            yield self
+        finally:
+            after_snapshot = self.take_snapshot()
+            duration = (datetime.now() - start_time).total_seconds()
+            memory_delta = after_snapshot.process_memory_mb - before_snapshot.process_memory_mb
+            
+            log_info(f"[memory_manager] Completed {description} - "
+                    f"Duration: {duration:.2f}s, "
+                    f"Memory delta: {memory_delta:+.1f}MB")
+            
+            # Auto cleanup if needed
+            if self.auto_cleanup_enabled and memory_delta > 50:  # 50MB increase
+                self.cleanup_if_needed()
+    
+    def reset_stats(self):
+        """Reset memory management statistics"""
+        self.stats = {
+            "total_cleanups": 0,
+            "forced_cleanups": 0,
+            "auto_cleanups": 0,
+            "memory_freed_mb": 0.0,
+            "last_cleanup": None
+        }
+        self.leak_detector = MemoryLeakDetector()
+        log_info("[memory_manager] Statistics reset")
+    
+    def __del__(self):
+        """Cleanup when memory manager is destroyed"""
+        try:
+            self.stop_monitoring()
+        except Exception:
+            pass
+
+
+# Global memory manager instance
+_global_memory_manager = None
+
+
+def get_memory_manager() -> MemoryManager:
+    """Get the global memory manager instance"""
+    global _global_memory_manager
+    if _global_memory_manager is None:
+        _global_memory_manager = MemoryManager()
+    return _global_memory_manager
+
+
+def cleanup_memory():
+    """Convenience function for forcing memory cleanup"""
+    return get_memory_manager().force_cleanup()
+
+
+def get_memory_stats():
+    """Convenience function for getting memory statistics"""
+    return get_memory_manager().get_memory_stats()
+
+
+def memory_context(description: str = "operation"):
+    """Convenience function for memory monitoring context"""
+    return get_memory_manager().memory_context(description)
+
+
+def memory_profile(description: str = None):
+    """
+    Decorator for profiling memory usage of functions.
+    
+    Args:
+        description: Optional description for the operation
+        
+    Example:
+        @memory_profile("data processing")
+        def process_data(df):
+            return df.with_columns(...)
+    """
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            func_description = description or f"{func.__module__}.{func.__name__}"
+            with memory_context(func_description):
+                return func(*args, **kwargs)
+        return wrapper
+    return decorator
+
+
+def track_memory_usage(func):
+    """
+    Simple decorator to track memory usage of a function.
+    
+    Example:
+        @track_memory_usage
+        def expensive_operation():
+            # ... do work
+            pass
+    """
+    def wrapper(*args, **kwargs):
+        manager = get_memory_manager()
+        before = manager.get_memory_usage_mb()
+        
+        try:
+            result = func(*args, **kwargs)
+            return result
+        finally:
+            after = manager.get_memory_usage_mb()
+            delta = after - before
+            
+            from .logging import log_info
+            log_info(f"[memory_profile] {func.__name__}: {delta:+.1f}MB")
+    
+    return wrapper