| """ |
| Memory management utilities for efficient-context. |
| """ |
|
|
| import logging |
| import gc |
| import os |
| import psutil |
| from typing import Optional, Dict, Any |
| from contextlib import contextmanager |
|
|
| |
| logging.basicConfig(level=logging.INFO) |
| logger = logging.getLogger(__name__) |
|
|
| class MemoryManager: |
| """ |
| Manages memory usage for efficient context handling. |
| |
| This class provides utilities to monitor and optimize memory usage |
| when working with large language models and context on CPU. |
| """ |
| |
| def __init__( |
| self, |
| target_usage_percent: float = 80.0, |
| aggressive_cleanup: bool = False, |
| memory_monitor_interval: Optional[float] = None, |
| ): |
| """ |
| Initialize the MemoryManager. |
| |
| Args: |
| target_usage_percent: Target memory usage as percentage of available memory |
| aggressive_cleanup: Whether to perform aggressive garbage collection |
| memory_monitor_interval: Interval for memory monitoring in seconds (None to disable) |
| """ |
| self.target_usage_percent = target_usage_percent |
| self.aggressive_cleanup = aggressive_cleanup |
| self.memory_monitor_interval = memory_monitor_interval |
| self.monitor_active = False |
| |
| logger.info( |
| "MemoryManager initialized with target usage: %.1f%%", |
| target_usage_percent |
| ) |
| |
| def get_memory_usage(self) -> Dict[str, Any]: |
| """ |
| Get current memory usage statistics. |
| |
| Returns: |
| stats: Dictionary of memory usage statistics |
| """ |
| |
| process = psutil.Process(os.getpid()) |
| process_memory = process.memory_info() |
| |
| |
| system_memory = psutil.virtual_memory() |
| |
| |
| process_percent = (process_memory.rss / system_memory.total) * 100 |
| system_percent = system_memory.percent |
| |
| return { |
| "process_rss_bytes": process_memory.rss, |
| "process_vms_bytes": process_memory.vms, |
| "process_percent": process_percent, |
| "system_available_bytes": system_memory.available, |
| "system_total_bytes": system_memory.total, |
| "system_used_percent": system_percent, |
| } |
| |
| def log_memory_usage(self) -> None: |
| """Log memory usage statistics.""" |
| stats = self.get_memory_usage() |
| |
| logger.info( |
| "Memory usage: Process: %.1f%% (%.1f MB), System: %.1f%% (%.1f GB available)", |
| stats["process_percent"], |
| stats["process_rss_bytes"] / (1024 * 1024), |
| stats["system_used_percent"], |
| stats["system_available_bytes"] / (1024 * 1024 * 1024) |
| ) |
| |
| def cleanup_memory(self) -> None: |
| """Perform memory cleanup.""" |
| |
| collected = gc.collect() |
| |
| if self.aggressive_cleanup: |
| |
| collected += gc.collect() |
| |
| logger.debug("Memory cleanup: Collected %d objects", collected) |
| |
| def _check_memory_threshold(self) -> bool: |
| """ |
| Check if memory usage exceeds the target threshold. |
| |
| Returns: |
| exceeded: Whether the threshold is exceeded |
| """ |
| stats = self.get_memory_usage() |
| return stats["system_used_percent"] > self.target_usage_percent |
| |
| @contextmanager |
| def optimize_memory(self): |
| """ |
| Context manager for optimizing memory during operations. |
| |
| Example: |
| ``` |
| with memory_manager.optimize_memory(): |
| # Run memory-intensive operations |
| ``` |
| """ |
| |
| if logger.isEnabledFor(logging.DEBUG): |
| self.log_memory_usage() |
| |
| try: |
| |
| yield |
| finally: |
| |
| if self._check_memory_threshold(): |
| logger.info("Memory threshold exceeded, performing cleanup") |
| self.cleanup_memory() |
| |
| |
| if logger.isEnabledFor(logging.DEBUG): |
| self.log_memory_usage() |
|
|
