runbooks 1.1.4__py3-none-any.whl → 1.1.5__py3-none-any.whl

runbooks/common/memory_optimization.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """
 Memory Optimization Framework for CloudOps-Runbooks
- 
+
 🎯 SRE Automation Specialist Implementation
 Following proven systematic delegation patterns for memory management and optimization.
 
@@ -29,10 +29,11 @@ import tracemalloc
 
 try:
     import psutil
+
     PSUTIL_AVAILABLE = True
 except ImportError:
     PSUTIL_AVAILABLE = False
-    
+
 from runbooks.common.rich_utils import (
     console,
     print_header,
@@ -40,7 +41,7 @@ from runbooks.common.rich_utils import (
     print_warning,
     print_error,
     create_table,
-    STATUS_INDICATORS
+    STATUS_INDICATORS,
 )
 
 logger = logging.getLogger(__name__)
@@ -49,6 +50,7 @@ logger = logging.getLogger(__name__)
 @dataclass
 class MemorySnapshot:
     """Memory usage snapshot for monitoring"""
+
     timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
     memory_mb: float = 0.0
     peak_memory_mb: float = 0.0
@@ -58,9 +60,10 @@ class MemorySnapshot:
     operation_context: Optional[str] = None
 
 
-@dataclass  
+@dataclass
 class MemoryOptimizationMetrics:
     """Memory optimization performance metrics"""
+
     operation_name: str
     start_memory_mb: float = 0.0
     peak_memory_mb: float = 0.0
@@ -75,43 +78,45 @@ class MemoryOptimizationMetrics:
 class MemoryOptimizer:
     """
     Enterprise memory optimization system for large-scale CloudOps operations
-    
+
     Provides:
     - Real-time memory monitoring with alerting thresholds
-    - Automatic garbage collection optimization  
+    - Automatic garbage collection optimization
     - Memory-efficient data processing patterns
     - Resource cleanup automation
     - Memory leak detection and prevention
     """
-    
-    def __init__(self, 
-                 warning_threshold_mb: float = 1024,
-                 critical_threshold_mb: float = 2048,
-                 monitoring_interval_seconds: float = 5.0):
+
+    def __init__(
+        self,
+        warning_threshold_mb: float = 1024,
+        critical_threshold_mb: float = 2048,
+        monitoring_interval_seconds: float = 5.0,
+    ):
         """
         Initialize memory optimizer
-        
+
         Args:
             warning_threshold_mb: Memory warning threshold in MB
-            critical_threshold_mb: Memory critical threshold in MB  
+            critical_threshold_mb: Memory critical threshold in MB
             monitoring_interval_seconds: Memory monitoring check interval
         """
         self.warning_threshold_mb = warning_threshold_mb
         self.critical_threshold_mb = critical_threshold_mb
         self.monitoring_interval_seconds = monitoring_interval_seconds
-        
+
         # Memory tracking
         self.snapshots: List[MemorySnapshot] = []
         self.metrics: List[MemoryOptimizationMetrics] = []
         self.monitoring_active = False
         self.monitoring_thread: Optional[threading.Thread] = None
-        
+
         # Process handle if psutil available
         self.process = psutil.Process() if PSUTIL_AVAILABLE else None
-        
+
         # Memory management strategies
         self._setup_gc_optimization()
-        
+
         # Weak reference tracking for leak detection
         self._tracked_objects: weakref.WeakSet = weakref.WeakSet()
 
@@ -119,74 +124,78 @@ class MemoryOptimizer:
         """Configure garbage collection optimization"""
         # Optimize garbage collection thresholds for large-scale operations
         gc.set_threshold(700, 10, 10)  # More aggressive collection
-        
+
         # Enable automatic garbage collection
         if not gc.isenabled():
             gc.enable()
-            
+
         logger.debug("Memory optimization: Garbage collection configured")
 
     @contextmanager
     def optimize_memory_usage(self, operation_name: str, enable_monitoring: bool = True):
         """
         Context manager for memory-optimized operation execution
-        
+
         Args:
             operation_name: Name of operation being optimized
             enable_monitoring: Whether to enable real-time monitoring
         """
         # Initialize metrics
         metrics = MemoryOptimizationMetrics(operation_name=operation_name)
-        
+
         # Start monitoring if requested
         if enable_monitoring:
             self.start_memory_monitoring()
-        
+
         # Record initial memory state
         initial_snapshot = self._take_memory_snapshot(operation_name)
         metrics.start_memory_mb = initial_snapshot.memory_mb
-        
+
         try:
-            console.log(f"[dim]🧠 Starting memory-optimized: {operation_name} (current: {initial_snapshot.memory_mb:.1f}MB)[/]")
-            
+            console.log(
+                f"[dim]🧠 Starting memory-optimized: {operation_name} (current: {initial_snapshot.memory_mb:.1f}MB)[/]"
+            )
+
             yield metrics
-            
+
             # Record final state and calculate savings
             final_snapshot = self._take_memory_snapshot(f"{operation_name}_end")
             metrics.end_memory_mb = final_snapshot.memory_mb
-            metrics.peak_memory_mb = max(s.peak_memory_mb for s in self.snapshots[-10:]) if self.snapshots else final_snapshot.memory_mb
-            
+            metrics.peak_memory_mb = (
+                max(s.peak_memory_mb for s in self.snapshots[-10:]) if self.snapshots else final_snapshot.memory_mb
+            )
+
             # Calculate memory efficiency
             if metrics.start_memory_mb > 0:
                 memory_change = metrics.end_memory_mb - metrics.start_memory_mb
                 if memory_change < 0:
                     metrics.memory_saved_mb = abs(memory_change)
                     metrics.optimization_techniques_applied.append("memory_cleanup")
-            
+
             metrics.success = True
             self._log_memory_results(metrics)
-            
+
         except Exception as e:
             # Handle memory issues during operation
             error_snapshot = self._take_memory_snapshot(f"{operation_name}_error")
             metrics.end_memory_mb = error_snapshot.memory_mb
             metrics.success = False
-            
+
             # Force garbage collection on error
             collected = gc.collect()
             metrics.gc_collections_triggered += collected
             metrics.optimization_techniques_applied.append("error_gc_cleanup")
-            
+
             print_error(f"Memory optimization failed for {operation_name}", e)
             raise
-            
+
         finally:
             # Stop monitoring and store metrics
             if enable_monitoring:
                 self.stop_memory_monitoring()
-                
+
             self.metrics.append(metrics)
-            
+
             # Cleanup operation-specific resources
             self._cleanup_operation_resources()
 
@@ -194,113 +203,118 @@ class MemoryOptimizer:
         """Start background memory monitoring"""
         if self.monitoring_active:
             return
-            
+
         self.monitoring_active = True
-        
+
         def monitor_memory():
             while self.monitoring_active:
                 try:
                     snapshot = self._take_memory_snapshot("monitoring")
                     self.snapshots.append(snapshot)
-                    
+
                     # Check thresholds and alert
                     self._check_memory_thresholds(snapshot)
-                    
+
                     # Limit snapshot history to prevent memory growth
                     if len(self.snapshots) > 1000:
                         self.snapshots = self.snapshots[-500:]  # Keep recent 500
-                    
+
                     time.sleep(self.monitoring_interval_seconds)
-                    
+
                 except Exception as e:
                     logger.debug(f"Memory monitoring error: {e}")
                     break
-        
+
         self.monitoring_thread = threading.Thread(target=monitor_memory, daemon=True)
         self.monitoring_thread.start()
-        
+
         logger.debug("Memory monitoring started")
 
     def stop_memory_monitoring(self):
         """Stop background memory monitoring"""
         self.monitoring_active = False
-        
+
         if self.monitoring_thread and self.monitoring_thread.is_alive():
             self.monitoring_thread.join(timeout=5.0)
-            
+
         logger.debug("Memory monitoring stopped")
 
     def _take_memory_snapshot(self, operation_context: str) -> MemorySnapshot:
         """Take comprehensive memory usage snapshot"""
         snapshot = MemorySnapshot(operation_context=operation_context)
-        
+
         if self.process and PSUTIL_AVAILABLE:
             try:
                 memory_info = self.process.memory_info()
                 snapshot.memory_mb = memory_info.rss / (1024 * 1024)
-                snapshot.peak_memory_mb = getattr(memory_info, 'peak_wset', memory_info.rss) / (1024 * 1024)
-                
+                snapshot.peak_memory_mb = getattr(memory_info, "peak_wset", memory_info.rss) / (1024 * 1024)
+
                 # Memory percentage if available
                 try:
                     snapshot.memory_percent = self.process.memory_percent()
                 except:
                     snapshot.memory_percent = 0.0
-                    
+
             except Exception as e:
                 logger.debug(f"Failed to get process memory info: {e}")
         else:
             # Fallback using tracemalloc if available
             try:
                 import tracemalloc
+
                 if tracemalloc.is_tracing():
                     current, peak = tracemalloc.get_traced_memory()
                     snapshot.memory_mb = current / (1024 * 1024)
                     snapshot.peak_memory_mb = peak / (1024 * 1024)
             except:
                 pass
-        
+
         # GC statistics
         snapshot.gc_collections = tuple(gc.get_count())
-        
+
         # Active objects count
         try:
             snapshot.active_objects = len(gc.get_objects())
         except:
             snapshot.active_objects = 0
-            
+
         return snapshot
 
     def _check_memory_thresholds(self, snapshot: MemorySnapshot):
         """Check memory usage against thresholds and alert if needed"""
         if snapshot.memory_mb > self.critical_threshold_mb:
-            console.log(f"[red]🚨 CRITICAL: Memory usage {snapshot.memory_mb:.1f}MB exceeds critical threshold {self.critical_threshold_mb}MB[/red]")
-            
+            console.log(
+                f"[red]🚨 CRITICAL: Memory usage {snapshot.memory_mb:.1f}MB exceeds critical threshold {self.critical_threshold_mb}MB[/red]"
+            )
+
             # Force aggressive garbage collection
             collected = self._force_garbage_collection()
             console.log(f"[yellow]🗑️  Emergency GC collected {collected} objects[/yellow]")
-            
+
         elif snapshot.memory_mb > self.warning_threshold_mb:
-            console.log(f"[yellow]⚠️ WARNING: Memory usage {snapshot.memory_mb:.1f}MB exceeds warning threshold {self.warning_threshold_mb}MB[/yellow]")
+            console.log(
+                f"[yellow]⚠️ WARNING: Memory usage {snapshot.memory_mb:.1f}MB exceeds warning threshold {self.warning_threshold_mb}MB[/yellow]"
+            )
 
     def _force_garbage_collection(self) -> int:
         """Force comprehensive garbage collection"""
         total_collected = 0
-        
+
         # Multiple GC passes for thorough cleanup
         for generation in range(3):
             collected = gc.collect(generation)
             total_collected += collected
-            
+
         # Additional cleanup
         gc.collect()  # Final full collection
-        
+
         return total_collected
 
     def _cleanup_operation_resources(self):
         """Clean up operation-specific resources"""
         # Force garbage collection
         gc.collect()
-        
+
         # Clear internal caches if they've grown large
         if len(self.snapshots) > 100:
             self.snapshots = self.snapshots[-50:]  # Keep recent snapshots
@@ -308,58 +322,65 @@ class MemoryOptimizer:
     def _log_memory_results(self, metrics: MemoryOptimizationMetrics):
         """Log memory optimization results"""
         if metrics.memory_saved_mb > 0:
-            print_success(f"Memory optimized for {metrics.operation_name}: "
-                         f"saved {metrics.memory_saved_mb:.1f}MB "
-                         f"({metrics.start_memory_mb:.1f}MB → {metrics.end_memory_mb:.1f}MB)")
+            print_success(
+                f"Memory optimized for {metrics.operation_name}: "
+                f"saved {metrics.memory_saved_mb:.1f}MB "
+                f"({metrics.start_memory_mb:.1f}MB → {metrics.end_memory_mb:.1f}MB)"
+            )
         elif metrics.end_memory_mb <= metrics.start_memory_mb * 1.1:  # Within 10%
-            console.log(f"[green]Memory stable for {metrics.operation_name}: "
-                       f"{metrics.start_memory_mb:.1f}MB → {metrics.end_memory_mb:.1f}MB[/green]")
+            console.log(
+                f"[green]Memory stable for {metrics.operation_name}: "
+                f"{metrics.start_memory_mb:.1f}MB → {metrics.end_memory_mb:.1f}MB[/green]"
+            )
         else:
-            console.log(f"[yellow]Memory increased for {metrics.operation_name}: "
-                       f"{metrics.start_memory_mb:.1f}MB → {metrics.end_memory_mb:.1f}MB[/yellow]")
+            console.log(
+                f"[yellow]Memory increased for {metrics.operation_name}: "
+                f"{metrics.start_memory_mb:.1f}MB → {metrics.end_memory_mb:.1f}MB[/yellow]"
+            )
 
     def create_memory_efficient_iterator(self, data: List[Any], batch_size: int = 100) -> Iterator[List[Any]]:
         """
         Create memory-efficient iterator for large datasets
-        
+
         Args:
             data: Large dataset to process
             batch_size: Size of each batch to yield
-            
+
         Yields:
             Batched data chunks for memory-efficient processing
         """
         for i in range(0, len(data), batch_size):
-            batch = data[i:i + batch_size]
+            batch = data[i : i + batch_size]
             yield batch
-            
+
             # Trigger GC every 10 batches to prevent memory buildup
             if i > 0 and (i // batch_size) % 10 == 0:
                 gc.collect()
 
-    def optimize_large_dict_processing(self, large_dict: Dict[str, Any], 
-                                     chunk_size: int = 1000) -> Iterator[Dict[str, Any]]:
+    def optimize_large_dict_processing(
+        self, large_dict: Dict[str, Any], chunk_size: int = 1000
+    ) -> Iterator[Dict[str, Any]]:
         """
         Memory-efficient large dictionary processing
-        
+
         Args:
             large_dict: Large dictionary to process
             chunk_size: Number of items to process per chunk
-            
+
         Yields:
             Dictionary chunks for processing
         """
         items = list(large_dict.items())
-        
+
         for i in range(0, len(items), chunk_size):
-            chunk_items = items[i:i + chunk_size]
+            chunk_items = items[i : i + chunk_size]
             chunk_dict = dict(chunk_items)
-            
+
             yield chunk_dict
-            
+
             # Clean up temporary variables
             del chunk_items, chunk_dict
-            
+
             # Periodic GC
             if i > 0 and (i // chunk_size) % 5 == 0:
                 gc.collect()
@@ -372,13 +393,13 @@ class MemoryOptimizer:
     def get_memory_usage_report(self) -> Dict[str, Any]:
         """Get comprehensive memory usage report"""
         current_snapshot = self._take_memory_snapshot("report_generation")
-        
+
         # Calculate statistics from recent snapshots
         recent_snapshots = self.snapshots[-20:] if self.snapshots else [current_snapshot]
-        
+
         avg_memory = sum(s.memory_mb for s in recent_snapshots) / len(recent_snapshots)
         peak_memory = max(s.peak_memory_mb for s in recent_snapshots)
-        
+
         return {
             "current_memory_mb": current_snapshot.memory_mb,
             "average_memory_mb": avg_memory,
@@ -390,7 +411,7 @@ class MemoryOptimizer:
             "gc_collections": current_snapshot.gc_collections,
             "tracked_objects": len(self._tracked_objects),
             "memory_status": self._get_memory_status(current_snapshot),
-            "optimization_recommendations": self._get_optimization_recommendations(current_snapshot)
+            "optimization_recommendations": self._get_optimization_recommendations(current_snapshot),
         }
 
     def _get_memory_status(self, snapshot: MemorySnapshot) -> str:
@@ -407,21 +428,21 @@ class MemoryOptimizer:
     def _get_optimization_recommendations(self, snapshot: MemorySnapshot) -> List[str]:
         """Generate memory optimization recommendations"""
         recommendations = []
-        
+
         if snapshot.memory_mb > self.critical_threshold_mb:
             recommendations.append("Immediate garbage collection required")
             recommendations.append("Consider batch processing for large operations")
-            
+
         if snapshot.memory_mb > self.warning_threshold_mb:
             recommendations.append("Monitor memory usage closely")
             recommendations.append("Implement streaming processing for large datasets")
-            
+
         if snapshot.active_objects > 100000:
             recommendations.append("High object count detected - review object lifecycle")
-            
+
         if len(self._tracked_objects) > 0:
             recommendations.append(f"{len(self._tracked_objects)} objects being tracked for leaks")
-            
+
         return recommendations
 
     def create_memory_summary_table(self) -> None:
@@ -429,9 +450,9 @@ class MemoryOptimizer:
         if not self.metrics:
             console.print("[yellow]No memory optimization metrics available[/yellow]")
             return
-            
+
         print_header("Memory Optimization Summary", "SRE Memory Management")
-        
+
         # Create metrics table
         table = create_table(
             title="Memory Optimization Results",
@@ -442,38 +463,38 @@ class MemoryOptimizer:
                 {"name": "End (MB)", "style": "white", "justify": "right"},
                 {"name": "Saved (MB)", "style": "green", "justify": "right"},
                 {"name": "Optimizations", "style": "dim", "justify": "left", "max_width": 25},
-                {"name": "Status", "style": "white", "justify": "center"}
-            ]
+                {"name": "Status", "style": "white", "justify": "center"},
+            ],
         )
-        
+
         for metrics in self.metrics:
-            status_icon = STATUS_INDICATORS['success'] if metrics.success else STATUS_INDICATORS['error']
-            status_color = 'green' if metrics.success else 'red'
-            
+            status_icon = STATUS_INDICATORS["success"] if metrics.success else STATUS_INDICATORS["error"]
+            status_color = "green" if metrics.success else "red"
+
             saved_text = f"+{metrics.memory_saved_mb:.1f}" if metrics.memory_saved_mb > 0 else "0.0"
-            
+
             table.add_row(
                 metrics.operation_name,
                 f"{metrics.start_memory_mb:.1f}",
                 f"{metrics.peak_memory_mb:.1f}",
                 f"{metrics.end_memory_mb:.1f}",
                 saved_text,
-                ", ".join(metrics.optimization_techniques_applied[:2]) + ("..." if len(metrics.optimization_techniques_applied) > 2 else ""),
-                f"[{status_color}]{status_icon}[/]"
+                ", ".join(metrics.optimization_techniques_applied[:2])
+                + ("..." if len(metrics.optimization_techniques_applied) > 2 else ""),
+                f"[{status_color}]{status_icon}[/]",
             )
-        
+
         console.print(table)
-        
+
         # Current memory status
         report = self.get_memory_usage_report()
-        status_color = {
-            'good': 'green',
-            'moderate': 'yellow', 
-            'warning': 'yellow',
-            'critical': 'red'
-        }.get(report['memory_status'], 'white')
-        
-        console.print(f"\n[{status_color}]Current Memory: {report['current_memory_mb']:.1f}MB ({report['memory_status'].upper()})[/{status_color}]")
+        status_color = {"good": "green", "moderate": "yellow", "warning": "yellow", "critical": "red"}.get(
+            report["memory_status"], "white"
+        )
+
+        console.print(
+            f"\n[{status_color}]Current Memory: {report['current_memory_mb']:.1f}MB ({report['memory_status'].upper()})[/{status_color}]"
+        )
 
     def clear_optimization_data(self):
         """Clear all optimization tracking data"""
@@ -488,14 +509,12 @@ class MemoryOptimizer:
 _memory_optimizer: Optional[MemoryOptimizer] = None
 
 
-def get_memory_optimizer(warning_threshold_mb: float = 1024,
-                        critical_threshold_mb: float = 2048) -> MemoryOptimizer:
+def get_memory_optimizer(warning_threshold_mb: float = 1024, critical_threshold_mb: float = 2048) -> MemoryOptimizer:
     """Get or create global memory optimizer instance"""
     global _memory_optimizer
     if _memory_optimizer is None:
         _memory_optimizer = MemoryOptimizer(
-            warning_threshold_mb=warning_threshold_mb,
-            critical_threshold_mb=critical_threshold_mb
+            warning_threshold_mb=warning_threshold_mb, critical_threshold_mb=critical_threshold_mb
         )
     return _memory_optimizer
 
@@ -511,14 +530,17 @@ def create_memory_report():
 # Memory optimization decorators
 def memory_optimized(operation_name: str = None, enable_monitoring: bool = True):
     """Decorator for memory-optimized function execution"""
+
     def decorator(func: Callable) -> Callable:
         def wrapper(*args, **kwargs):
             optimizer = get_memory_optimizer()
             op_name = operation_name or f"{func.__module__}.{func.__name__}"
-            
+
             with optimizer.optimize_memory_usage(op_name, enable_monitoring):
                 return func(*args, **kwargs)
+
         return wrapper
+
     return decorator
 
 
@@ -527,7 +549,7 @@ __all__ = [
     "MemoryOptimizer",
     "MemorySnapshot",
     "MemoryOptimizationMetrics",
-    "get_memory_optimizer", 
+    "get_memory_optimizer",
     "create_memory_report",
-    "memory_optimized"
-]
+    "memory_optimized",
+]