Update src/utils/performance.py

src/utils/performance.py

@@ -1,90 +1,125 @@
-# src/utils/performance.py
-import time
-import psutil
-import threading
-from dataclasses import dataclass
-from typing import Dict, Any, List
-from src.utils.logger import get_logger
-
-logger = get_logger(__name__)
-
-@dataclass
-class PerformanceMetrics:
-    start_time: float
-    end_time: float
-    duration: float
-    cpu_usage: List[float]
-    memory_usage: List[float]
-    peak_memory: float
-    operation_name: str
-
-class PerformanceMonitor:
-    def __init__(self, operation_name: str):
-        self.operation_name = operation_name
-        self.start_time = None
-        self.end_time = None
-        self.cpu_usage = []
-        self.memory_usage = []
-        self.monitoring = False
-        self.monitor_thread = None
-    
-    def __enter__(self):
-        self.start_monitoring()
-        return self
-    
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.stop_monitoring()
-    
-    def start_monitoring(self):
-        """Start performance monitoring"""
-        self.start_time = time.time()
-        self.monitoring = True
-        
-        self.monitor_thread = threading.Thread(target=self._monitor_resources)
-        self.monitor_thread.daemon = True
-        self.monitor_thread.start()
-        
-        logger.info(f"Started monitoring: {self.operation_name}")
-    
-    def stop_monitoring(self):
-        """Stop performance monitoring and return metrics"""
-        self.end_time = time.time()
-        self.monitoring = False
-        
-        if self.monitor_thread:
-            self.monitor_thread.join(timeout=1.0)
-        
-        metrics = PerformanceMetrics(
-            start_time=self.start_time,
-            end_time=self.end_time,
-            duration=self.end_time - self.start_time,
-            cpu_usage=self.cpu_usage,
-            memory_usage=self.memory_usage,
-            peak_memory=max(self.memory_usage) if self.memory_usage else 0,
-            operation_name=self.operation_name
-        )
-        
-        self._log_metrics(metrics)
-        return metrics
-    
-    def _monitor_resources(self):
-        """Monitor CPU and memory usage in background"""
-        while self.monitoring:
-            try:
-                cpu_percent = psutil.cpu_percent(interval=0.1)
-                memory_info = psutil.virtual_memory()
-                
-                self.cpu_usage.append(cpu_percent)
-                self.memory_usage.append(memory_info.used / 1024 / 1024)  # Convert to MB
-                
-                time.sleep(0.5)  # Monitor every 0.5 seconds
-            except Exception as e:
-                logger.error(f"Error monitoring resources: {str(e)}")
-                break
-    
-    def _log_metrics(self, metrics: PerformanceMetrics):
-        """Log performance metrics"""
-        logger.info(f"Performance Report for: {metrics.operation_name}")
-        logger.info(f"  Duration: {metrics.duration:.2f} seconds")
-        logger.info(f"  Peak Memory: {metrics.peak_memory:.2f} MB")
-        logger.info(f"  Average CPU: {sum(metrics.cpu_usage)/len(metrics.cpu_usage):.1f}%" if metrics.cpu_usage else "N/A")
+import time
+import psutil
+import threading
+import os
+from dataclasses import dataclass
+from typing import Dict, Any, List
+from src.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+@dataclass
+class PerformanceMetrics:
+    start_time: float
+    end_time: float
+    duration: float
+    cpu_usage: List[float]
+    memory_usage: List[float]
+    peak_memory: float
+    operation_name: str
+
+class PerformanceMonitor:
+    def __init__(self, operation_name: str):
+        self.operation_name = operation_name
+        self.start_time = None
+        self.end_time = None
+        self.cpu_usage = []
+        self.memory_usage = []
+        self.monitoring = False
+        self.monitor_thread = None
+        # Get current process for accurate memory tracking
+        self.process = psutil.Process(os.getpid())
+        self.baseline_memory = self.process.memory_info().rss / 1024 / 1024  # MB
+    
+    def __enter__(self):
+        """Context manager entry - start monitoring"""
+        self.start_monitoring()
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - stop monitoring"""
+        return self.stop_monitoring()
+    
+    def start_monitoring(self):
+        """Start monitoring system resources"""
+        self.start_time = time.time()
+        self.monitoring = True
+        self.cpu_usage = []
+        self.memory_usage = []
+        
+        # Start monitoring thread
+        self.monitor_thread = threading.Thread(target=self._monitor_resources)
+        self.monitor_thread.daemon = True
+        self.monitor_thread.start()
+        
+        logger.info(f"Started monitoring: {self.operation_name}")
+    
+    def stop_monitoring(self):
+        """Stop resource monitoring and return metrics"""
+        self.end_time = time.time()
+        self.monitoring = False
+        
+        if self.monitor_thread:
+            self.monitor_thread.join()
+        
+        # Calculate memory increase from baseline
+        if self.memory_usage:
+            current_memory = self.process.memory_info().rss / 1024 / 1024
+            memory_increase = max(self.memory_usage) - self.baseline_memory
+            peak_memory = max(memory_increase, 0)  # Ensure non-negative
+        else:
+            peak_memory = 0
+        
+        metrics = PerformanceMetrics(
+            start_time=self.start_time,
+            end_time=self.end_time,
+            duration=self.end_time - self.start_time,
+            cpu_usage=self.cpu_usage,
+            memory_usage=self.memory_usage,
+            peak_memory=peak_memory,  # Memory increase from baseline
+            operation_name=self.operation_name
+        )
+        
+        self._log_metrics(metrics)
+        return metrics
+    
+    def _monitor_resources(self):
+        """Monitor CPU and memory usage in background"""
+        while self.monitoring:
+            try:
+                # Get process-specific CPU usage
+                cpu_percent = self.process.cpu_percent()
+                
+                # Get process-specific memory usage (RSS - Resident Set Size)
+                memory_info = self.process.memory_info()
+                memory_mb = memory_info.rss / 1024 / 1024  # Convert to MB
+                
+                self.cpu_usage.append(cpu_percent)
+                self.memory_usage.append(memory_mb)
+                
+                time.sleep(0.1)  # Monitor every 0.1 seconds for more precision
+            except Exception as e:
+                logger.error(f"Error monitoring resources: {str(e)}")
+                break
+    
+    def _log_metrics(self, metrics: PerformanceMetrics):
+        """Log performance metrics"""
+        logger.info(f"Performance Report for: {metrics.operation_name}")
+        logger.info(f"  Duration: {metrics.duration:.2f} seconds")
+        logger.info(f"  Peak Memory: {metrics.peak_memory:.2f} MB")
+        if metrics.cpu_usage:
+            avg_cpu = sum(metrics.cpu_usage) / len(metrics.cpu_usage)
+            logger.info(f"  Average CPU: {avg_cpu:.1f}%")
+    
+    def get_current_memory_usage(self):
+        """Get current process memory usage in MB"""
+        return self.process.memory_info().rss / 1024 / 1024
+
+def get_system_info():
+    """Get system information for performance context"""
+    process = psutil.Process()
+    return {
+        "cpu_count": psutil.cpu_count(),
+        "memory_total_gb": psutil.virtual_memory().total / 1024 / 1024 / 1024,
+        "process_memory_mb": process.memory_info().rss / 1024 / 1024
+    }