Spaces:
Sleeping
Sleeping
| 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__) | |
| 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 | |
| } | |