Agent/memory_manager.py

"""

Memory management utilities for the text2sql agent.

"""
import gc
import os
import time
import psutil
import threading
import weakref
from contextlib import contextmanager
from typing import Optional, Callable, Dict, Any

class MemoryManager:
    """Memory manager for monitoring and managing system memory."""
    
    def __init__(self, threshold_mb: float = 1000.0, critical_mb: float = 300.0,

                 check_interval: float = 0.5, monitor_active: bool = True):
        """Initialize the memory manager.



        Args:

            threshold_mb: Memory threshold in MB below which to take action

            critical_mb: Critical memory level in MB below which to abort operations

            check_interval: Interval in seconds for background memory checks

            monitor_active: Whether to start the background monitor thread

        """
        # Allow environment overrides to tune aggressiveness
        try:
            _env_threshold = os.environ.get("MEM_THRESHOLD_MB")
            _env_critical = os.environ.get("MEM_CRITICAL_MB")
            if _env_threshold: threshold_mb = float(_env_threshold)
            if _env_critical: critical_mb = float(_env_critical)
        except Exception:
            pass
        self.threshold_mb = threshold_mb
        self.critical_mb = critical_mb
        self.check_interval = check_interval
        self._lock = threading.RLock()
        self._monitor_thread = None
        self._keep_monitoring = False

        # Track instances with weak references when possible to avoid holding
        # strong references that prevent GC. Keep a fallback list for objects
        # that are not weak-referenceable.
        self._active_llm_refs = []  # list of weakref.ref where possible
        self._active_llm_strong = []  # fallback strong refs
        self._instance_count = 0

        # Soft cap to avoid accidentally registering many parallel LLMs
        self.max_instances = 1

        # Margin (MB) above critical required to allow new registrations
        self.registration_margin_mb = 200.0

        # Track last emergency cleanup time to avoid repeated rapid-fire releases
        self._last_emergency = 0.0
        self._emergency_cooldown = 10.0  # seconds between emergency attempts
        # Counter-based suspension flag to prevent emergency releases while in critical sections
        self._suspend_emergency_depth = 0

        # Minimum available memory (MB) at which we will STILL retain a single LLM instance
        # even if below critical, to avoid thrashing (re-init cost is high). Configurable.
        try:
            self.min_retain_mb = float(os.environ.get("MIN_RETAIN_LLM_MB", "250"))
        except Exception:
            self.min_retain_mb = 250.0

        # Start memory monitor if requested
        if monitor_active:
            self.start_monitor()

    def _prune_active_instances(self):
        """Remove dead weakrefs and return list of live instances."""
        live = []
        # Prune weak refs
        for ref in list(self._active_llm_refs):
            inst = ref()
            if inst is None:
                try:
                    self._active_llm_refs.remove(ref)
                except ValueError:
                    pass
            else:
                live.append(inst)
        # Add strong refs
        for inst in list(self._active_llm_strong):
            if inst is not None:
                live.append(inst)
        return live
    
    def register_llm_instance(self, llm_instance):
        """Register an LLM instance for tracking.



        This will refuse registration (and attempt to release the provided

        instance) if the system memory is too low. It prefers weak references

        so that the manager itself does not keep objects alive.

        """
        with self._lock:
            # Prune dead weakrefs first
            self._prune_active_instances()

            stats = self.get_memory_stats()
            # Relax policy: always allow the FIRST instance as long as we are above critical.
            # Only enforce margin for additional instances (even though max_instances is 1 today).
            if self._instance_count == 0:
                if stats['available_mb'] < self.critical_mb:
                    # Too dangerous even for first instance
                    try:
                        if hasattr(llm_instance, 'close'):
                            llm_instance.close()
                        elif hasattr(llm_instance, 'release'):
                            llm_instance.release()
                    except Exception:
                        pass
                    raise MemoryError(f"Refusing first LLM registration: below critical ({stats['available_mb']:.1f}MB < {self.critical_mb}MB)")
            else:
                # For subsequent instances keep the stricter margin
                if stats['available_mb'] < (self.critical_mb + self.registration_margin_mb):
                    try:
                        if hasattr(llm_instance, 'close'):
                            llm_instance.close()
                        elif hasattr(llm_instance, 'release'):
                            llm_instance.release()
                    except Exception:
                        pass
                    raise MemoryError(f"Refusing to register additional LLM (avail={stats['available_mb']:.1f}MB; need >= {self.critical_mb + self.registration_margin_mb:.1f}MB)")

            # Enforce max instances cap
            if self._instance_count >= self.max_instances:
                raise MemoryError(f"Refusing to register LLM instance: max_instances={self.max_instances} reached")

            # Try to create a weakref to avoid holding a strong reference
            try:
                ref = weakref.ref(llm_instance, lambda r: None)
                self._active_llm_refs.append(ref)
            except TypeError:
                # Object not weak-referenceable; store strong ref but warn
                self._active_llm_strong.append(llm_instance)

            self._instance_count += 1
            print(f"LLM instance registered (total: {self._instance_count})")
    
    def unregister_llm_instance(self, llm_instance):
        """Unregister an LLM instance from tracking."""
        with self._lock:
            # Remove from weak refs
            removed = False
            for ref in list(self._active_llm_refs):
                inst = ref()
                if inst is None:
                    # Dead ref; prune
                    try:
                        self._active_llm_refs.remove(ref)
                    except ValueError:
                        pass
                    continue
                if inst is llm_instance:
                    try:
                        self._active_llm_refs.remove(ref)
                    except ValueError:
                        pass
                    removed = True
                    break

            # Remove from strong refs
            if not removed:
                for i, instance in enumerate(list(self._active_llm_strong)):
                    if instance is llm_instance:
                        del self._active_llm_strong[i]
                        removed = True
                        break

            if removed:
                self._instance_count = max(0, self._instance_count - 1)
                print(f"LLM instance unregistered (remaining: {self._instance_count})")
    
    def get_memory_stats(self) -> Dict[str, float]:
        """Get current memory statistics."""
        vm = psutil.virtual_memory()
        return {
            "total_mb": vm.total / (1024 * 1024),
            "available_mb": vm.available / (1024 * 1024),
            "used_mb": vm.used / (1024 * 1024),
            "percent_used": vm.percent
        }
    
    def is_memory_safe(self) -> bool:
        """Check if memory is at a safe level."""
        return self.get_memory_stats()["available_mb"] >= self.threshold_mb
    
    def is_memory_critical(self) -> bool:
        """Check if memory is at a critical level."""
        stats = self.get_memory_stats()
        available = stats["available_mb"]
        is_critical = available <= self.critical_mb
        
        # Debug logging for memory state
        if is_critical:
            print(f"🚨 [MEMORY_DEBUG] Critical state: {available:.1f}MB <= {self.critical_mb}MB threshold")
        
        return is_critical
    
    def force_collect_garbage(self):
        """Force aggressive garbage collection."""
        # Run multiple collection cycles
        for i in range(3):
            gc.collect(i)
        
        # Clear any reference cycles
        if hasattr(gc, 'freeze'):
            gc.freeze()
        
        # Clear module caches if possible
        if hasattr(gc, 'clear_caches'):
            gc.clear_caches()
            
        # Log memory status after cleanup
        stats = self.get_memory_stats()
        #print(f"After garbage collection: {stats['available_mb']:.2f}MB available ({stats['percent_used']}% used)")
    
    def _monitor_memory(self):
        """Background thread for memory monitoring."""
        while self._keep_monitoring:
            stats = self.get_memory_stats()
            
            # Check for critical memory conditions
            if stats["available_mb"] <= self.critical_mb:
                #print(f"⚠️ CRITICAL MEMORY WARNING: Only {stats['available_mb']:.2f}MB available!")
                self.force_collect_garbage()
                
                # If still critical after cleanup, take more drastic measures
                if self.is_memory_critical():
                    now = time.time()
                    # Only attempt emergency release if cooldown passed
                    if now - self._last_emergency > self._emergency_cooldown:
                        # Check if we have tracked instances (prune dead refs first)
                        live = self._prune_active_instances()
                        if live:
                            # If only a single LLM instance and we still have above min_retain_mb, skip releasing to prevent constant re-init.
                            if len(live) == 1 and stats["available_mb"] > self.min_retain_mb:
                                print(f"⚠️ Critical memory but retaining single LLM (avail={stats['available_mb']:.1f}MB > {self.min_retain_mb}MB safeguard)")
                            else:
                                # Skip emergency release while suspended
                                if self._suspend_emergency_depth > 0:
                                    print(f"🚫 Emergency release suppressed (depth={self._suspend_emergency_depth}) during critical section")
                                else:
                                    self._emergency_release_llm_instances()
                        else:
                            # Log once per cooldown when no instances are registered
                            print("🚨 Emergency memory state detected but no LLM instances registered to release")
                        self._last_emergency = now
                    else:
                        # Cooldown active; skip repeated emergency attempts
                        pass
            
            # Regular check for low memory
            elif stats["available_mb"] <= self.threshold_mb:
                #print(f"⚠️ Low memory warning: {stats['available_mb']:.2f}MB available")
                self.force_collect_garbage()
            
            # Sleep for the check interval
            time.sleep(self.check_interval)
    
    def _emergency_release_llm_instances(self):
        """Emergency release of LLM instances to free memory."""
        with self._lock:
            # Build list of live instances from weak and strong refs
            instances = self._prune_active_instances()
            print(f"Attempting to release {len(instances)} LLM instances")

            for llm in instances:
                try:
                    # Try to close or release resources
                    if hasattr(llm, 'close'):
                        llm.close()
                    elif hasattr(llm, 'release'):
                        llm.release()
                    
                    # Set to None to break references
                    if hasattr(llm, '_client'):
                        llm._client = None
                    if hasattr(llm, '_model'):
                        llm._model = None
                except Exception as e:
                    print(f"Error releasing LLM instance: {str(e)}")
            
            # Clear our tracking lists
            self._active_llm_refs.clear()
            self._active_llm_strong.clear()
            self._instance_count = 0

            # Force garbage collection after cleanup and attempt to return pages
            self.force_collect_garbage()
            try:
                # Best-effort working set trim to return pages to OS (Windows)
                self.trim_working_set()
            except Exception:
                pass

    # ---------------- Public high-level cleanup APIs -----------------
    def release_all_llms(self):
        """Public wrapper to release all tracked LLM instances (non-emergency manual call)."""
        self._emergency_release_llm_instances()

    def trim_working_set(self):
        """Attempt to return unused pages to OS (best-effort, platform-specific)."""
        try:
            import platform, ctypes
            if platform.system().lower() == 'windows':
                PROCESS_SET_QUOTA = 0x0100
                PROCESS_QUERY_INFORMATION = 0x0400
                kernel32 = ctypes.windll.kernel32
                psapi = ctypes.windll.psapi
                GetCurrentProcess = kernel32.GetCurrentProcess
                hProc = GetCurrentProcess()
                psapi.EmptyWorkingSet(hProc)
            else:
                # For *nix, reading /proc/self might encourage trimming after gc
                pass
        except Exception as e:
            print(f"Working set trim not supported/failed: {e}")

    def perform_full_cleanup(self) -> Dict[str, float]:
        """Comprehensive cleanup: release LLMs, GC, trim working set, return new stats."""
        before = self.get_memory_stats()
        self.release_all_llms()
        self.force_collect_garbage()
        self.trim_working_set()
        after = self.get_memory_stats()
        delta = after['available_mb'] - before['available_mb']
        print(f"Post-query full cleanup reclaimed {delta:.2f}MB (avail: {after['available_mb']:.2f}MB)")
        return after

    # --------------- Critical section helpers ---------------
    @contextmanager
    def suspend_emergency(self, reason: str = ""):
        """Context manager to temporarily suspend emergency LLM releases.



        Use around critical sections (e.g., active LLM generation) to prevent

        the monitor thread from tearing down the only LLM mid-call, which can

        cause intermittent generation failures.

        """
        with self._lock:
            self._suspend_emergency_depth += 1
            if reason:
                print(f"⏸️ Suspended emergency releases (depth={self._suspend_emergency_depth}) reason={reason}")
        try:
            yield
        finally:
            with self._lock:
                self._suspend_emergency_depth = max(0, self._suspend_emergency_depth - 1)
                if reason:
                    print(f"▶️ Resumed emergency releases (depth={self._suspend_emergency_depth}) reason={reason}")
    
    def start_monitor(self):
        """Start the background memory monitor."""
        with self._lock:
            if self._monitor_thread is None or not self._monitor_thread.is_alive():
                self._keep_monitoring = True
                self._monitor_thread = threading.Thread(
                    target=self._monitor_memory,
                    daemon=True,
                    name="MemoryMonitorThread"
                )
                self._monitor_thread.start()
                print("Memory monitor started in background")
    
    def stop_monitor(self):
        """Stop the background memory monitor."""
        with self._lock:
            self._keep_monitoring = False
            if self._monitor_thread and self._monitor_thread.is_alive():
                self._monitor_thread.join(timeout=1.0)
                print("Memory monitor stopped")

# Create a global memory manager instance
memory_manager = MemoryManager(
    threshold_mb=500.0,   # 500MB — warning/action threshold (reduced from 1000MB)
    critical_mb=300.0,    # 300MB — BSOD-protection hard floor (unchanged)
    check_interval=1.0,   # Check every second
    monitor_active=True   # Start monitoring immediately
)

class TimeoutManager:
    """Manages operation timeouts with proper cleanup."""
    
    @staticmethod
    def run_with_timeout(func: Callable, timeout: float, *args, **kwargs) -> Any:
        """

        Run a function with a timeout.

        

        Args:

            func: Function to run

            timeout: Timeout in seconds

            *args, **kwargs: Arguments to pass to the function

            

        Returns:

            The result of the function

            

        Raises:

            TimeoutError: If the function times out

        """
        result = None
        exception = None
        completed = threading.Event()
        
        # Thread function
        def worker():
            nonlocal result, exception
            try:
                result = func(*args, **kwargs)
            except Exception as e:
                exception = e
            finally:
                completed.set()
        
        # Start the worker thread
        thread = threading.Thread(target=worker, daemon=True)
        thread.start()
        
        # Wait for completion or timeout
        if not completed.wait(timeout):
            # Force garbage collection before raising timeout
            memory_manager.force_collect_garbage()
            
            # Raise timeout error
            raise TimeoutError(f"Operation timed out after {timeout} seconds")
        
        # If there was an exception, raise it
        if exception is not None:
            raise exception
        
        return result

def safe_check_memory(threshold_mb: Optional[float] = None) -> bool:
    """

    Safely check if there's enough memory available.

    Also performs garbage collection if memory is low.

    

    Args:

        threshold_mb: Optional threshold in MB (uses memory manager's threshold if None)

        

    Returns:

        True if memory is safe, False if it's below threshold

    """
    # Use memory manager's threshold if none provided
    if threshold_mb is None:
        threshold_mb = memory_manager.threshold_mb
    
    # Get current memory stats
    stats = memory_manager.get_memory_stats()
    
    # Log memory status
    print(f"Memory status: {stats['available_mb']:.2f}MB available of {stats['total_mb']:.2f}MB total")
    
    # If memory is below threshold, run garbage collection
    if stats['available_mb'] < threshold_mb:
        memory_manager.force_collect_garbage()
        
        # Get updated stats after garbage collection
        stats = memory_manager.get_memory_stats()
        
        # If still below threshold, return False
        if stats['available_mb'] < threshold_mb:
            print(f"⚠️ Low memory warning: Only {stats['available_mb']:.2f}MB available, need at least {threshold_mb:.2f}MB")
            return False
    
    return True