""" 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