Spaces:
Sleeping
Sleeping
Aditya-Jadhav150
feat: AEGIS-AI Detection Engine Overhaul v2 — Multi-model ensemble, forensic artifact localization, audio detection, decision fusion, rich forensic reporting (Code Only)
ea107b5 | """ | |
| Performance Monitoring and Instrumentation for AEGIS-AI. | |
| Tracks timing and resource usage for various pipeline stages. | |
| """ | |
| import os | |
| import time | |
| import logging | |
| from typing import Dict, Optional, Any | |
| from contextlib import contextmanager | |
| try: | |
| import psutil | |
| HAS_PSUTIL = True | |
| except ImportError: | |
| HAS_PSUTIL = False | |
| logger = logging.getLogger(__name__) | |
| class PerformanceMonitor: | |
| """ | |
| Monitors timing and memory usage across different stages of execution. | |
| """ | |
| def __init__(self): | |
| """Initialize the performance monitor.""" | |
| self.timings: Dict[str, float] = {} | |
| self._start_times: Dict[str, float] = {} | |
| # Valid stages defined in implementation plan | |
| self.VALID_STAGES = { | |
| 'upload_validation', | |
| 'face_detection', | |
| 'model_inference', | |
| 'heatmap_generation', | |
| 'audio_processing', | |
| 'metadata_analysis', | |
| 'fusion_engine', | |
| 'total_analysis', | |
| 'model_load_time' | |
| } | |
| def start(self, stage: str) -> None: | |
| """ | |
| Start timing a specific stage. | |
| Args: | |
| stage (str): Name of the stage to start timing. | |
| """ | |
| if stage not in self.VALID_STAGES: | |
| logger.warning(f"Unknown performance stage started: {stage}") | |
| self._start_times[stage] = time.perf_counter() | |
| logger.debug(f"Started timing stage: {stage}") | |
| def stop(self, stage: str) -> None: | |
| """ | |
| Stop timing a specific stage and record the duration. | |
| Args: | |
| stage (str): Name of the stage to stop timing. | |
| """ | |
| if stage not in self._start_times: | |
| logger.warning(f"Attempted to stop unstarted stage: {stage}") | |
| return | |
| duration = time.perf_counter() - self._start_times.pop(stage) | |
| duration_ms = duration * 1000.0 | |
| # Accumulate time if stage is run multiple times | |
| if stage in self.timings: | |
| self.timings[stage] += duration_ms | |
| else: | |
| self.timings[stage] = duration_ms | |
| logger.debug(f"Completed stage {stage} in {duration_ms:.2f}ms") | |
| def get_peak_ram_mb(self) -> float: | |
| """ | |
| Get the current RAM usage of the process in MB. | |
| Note: This is current RAM, which serves as a proxy for peak if called at end. | |
| Returns: | |
| float: RAM usage in MB. | |
| """ | |
| if not HAS_PSUTIL: | |
| return 0.0 | |
| process = psutil.Process(os.getpid()) | |
| return process.memory_info().rss / (1024 * 1024) | |
| def get_report(self) -> Dict[str, Any]: | |
| """ | |
| Get a full report of all timings and peak RAM. | |
| Returns: | |
| Dict[str, Any]: Dictionary containing timings (ms) and memory stats. | |
| """ | |
| report = {k: round(v, 2) for k, v in self.timings.items()} | |
| report['peak_ram_mb'] = round(self.get_peak_ram_mb(), 2) | |
| return report | |
| def track(self, stage: str): | |
| """ | |
| Context manager for timing a block of code. | |
| Args: | |
| stage (str): Name of the stage to track. | |
| Yields: | |
| None | |
| """ | |
| self.start(stage) | |
| try: | |
| yield | |
| finally: | |
| self.stop(stage) | |