Upload alert_system.py with huggingface_hub

alert_system.py

@@ -0,0 +1,785 @@
+#!/usr/bin/env python3
+"""
+Alert and Notification System
+
+Provides comprehensive alerting for GPU monitoring with multiple notification
+channels, threshold management, and alert history tracking.
+"""
+
+import time
+import json
+import logging
+import smtplib
+import subprocess
+import threading
+from typing import Dict, List, Optional, Callable, Any
+from dataclasses import dataclass, asdict
+from datetime import datetime, timedelta
+from email.mime.text import MIMEText
+from email.mime.multipart import MIMEMultipart
+from pathlib import Path
+
+from gpu_monitoring import GPUStatus, GPUDataManager
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AlertThreshold:
+    """Alert threshold configuration."""
+    metric: str
+    threshold: float
+    operator: str  # '>', '<', '>=', '<=', '=='
+    duration: int  # seconds to maintain threshold before alerting
+    enabled: bool = True
+    cooldown: int = 300  # seconds between same alerts
+
+
+@dataclass
+class Alert:
+    """Alert instance."""
+    id: str
+    timestamp: float
+    gpu_name: str
+    metric: str
+    value: float
+    threshold: float
+    message: str
+    severity: str  # 'info', 'warning', 'critical', 'emergency'
+    acknowledged: bool = False
+    resolved: bool = False
+    resolved_at: Optional[float] = None
+
+
+class NotificationChannel:
+    """Base class for notification channels."""
+    
+    def __init__(self, name: str, enabled: bool = True):
+        self.name = name
+        self.enabled = enabled
+    
+    def send(self, alert: Alert) -> bool:
+        """Send notification for alert."""
+        raise NotImplementedError
+
+
+class LogNotification(NotificationChannel):
+    """Log-based notification channel."""
+    
+    def send(self, alert: Alert) -> bool:
+        """Send alert to log."""
+        severity_colors = {
+            'info': '\033[94m',      # Blue
+            'warning': '\033[93m',   # Yellow
+            'critical': '\033[91m',  # Red
+            'emergency': '\033[41m'  # Red background
+        }
+        reset_color = '\033[0m'
+        
+        color = severity_colors.get(alert.severity, '')
+        
+        message = (
+            f"{color}[{alert.severity.upper()}] GPU Alert: {alert.gpu_name} - "
+            f"{alert.metric}: {alert.value} (threshold: {alert.threshold}){reset_color}"
+        )
+        
+        if alert.severity in ['critical', 'emergency']:
+            logger.critical(message)
+        elif alert.severity == 'warning':
+            logger.warning(message)
+        else:
+            logger.info(message)
+        
+        return True
+
+
+class DesktopNotification(NotificationChannel):
+    """Desktop notification channel using notify-send."""
+    
+    def send(self, alert: Alert) -> bool:
+        """Send desktop notification."""
+        try:
+            urgency = 'normal'
+            if alert.severity == 'critical':
+                urgency = 'critical'
+            elif alert.severity == 'warning':
+                urgency = 'normal'
+            else:
+                urgency = 'low'
+            
+            # Create notification command
+            cmd = [
+                'notify-send',
+                '--app-name=GPU Monitor',
+                f'--urgency={urgency}',
+                f'--icon=video-display',
+                f'GPU Alert - {alert.gpu_name}',
+                f'{alert.message}'
+            ]
+            
+            subprocess.run(cmd, check=True)
+            return True
+            
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Failed to send desktop notification: {e}")
+            return False
+        except FileNotFoundError:
+            logger.warning("notify-send not found, desktop notifications disabled")
+            return False
+
+
+class EmailNotification(NotificationChannel):
+    """Email notification channel."""
+    
+    def __init__(self, name: str, smtp_config: Dict[str, Any], recipients: List[str]):
+        super().__init__(name)
+        self.smtp_config = smtp_config
+        self.recipients = recipients
+    
+    def send(self, alert: Alert) -> bool:
+        """Send email notification."""
+        try:
+            # Create message
+            msg = MIMEMultipart()
+            msg['From'] = self.smtp_config['sender']
+            msg['To'] = ', '.join(self.recipients)
+            msg['Subject'] = f'GPU Alert: {alert.gpu_name} - {alert.severity.upper()}'
+            
+            # Create email body
+            body = f"""
+GPU Monitoring Alert
+
+GPU: {alert.gpu_name}
+Metric: {alert.metric}
+Current Value: {alert.value}
+Threshold: {alert.threshold}
+Severity: {alert.severity.upper()}
+Time: {datetime.fromtimestamp(alert.timestamp).strftime('%Y-%m-%d %H:%M:%S')}
+
+Message: {alert.message}
+
+This is an automated message from the GPU Monitoring System.
+"""
+            
+            msg.attach(MIMEText(body, 'plain'))
+            
+            # Send email
+            with smtplib.SMTP(self.smtp_config['server'], self.smtp_config['port']) as server:
+                if self.smtp_config.get('use_tls', True):
+                    server.starttls()
+                
+                if 'username' in self.smtp_config and 'password' in self.smtp_config:
+                    server.login(self.smtp_config['username'], self.smtp_config['password'])
+                
+                server.send_message(msg)
+            
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to send email notification: {e}")
+            return False
+
+
+class WebhookNotification(NotificationChannel):
+    """Webhook notification channel."""
+    
+    def __init__(self, name: str, webhook_url: str, headers: Optional[Dict[str, str]] = None):
+        super().__init__(name)
+        self.webhook_url = webhook_url
+        self.headers = headers or {}
+    
+    def send(self, alert: Alert) -> bool:
+        """Send webhook notification."""
+        try:
+            import requests
+            
+            payload = {
+                'alert_id': alert.id,
+                'timestamp': alert.timestamp,
+                'gpu_name': alert.gpu_name,
+                'metric': alert.metric,
+                'value': alert.value,
+                'threshold': alert.threshold,
+                'message': alert.message,
+                'severity': alert.severity,
+                'acknowledged': alert.acknowledged
+            }
+            
+            response = requests.post(
+                self.webhook_url,
+                json=payload,
+                headers=self.headers,
+                timeout=10
+            )
+            
+            return response.status_code == 200
+            
+        except Exception as e:
+            logger.error(f"Failed to send webhook notification: {e}")
+            return False
+
+
+class AlertManager:
+    """Main alert management system."""
+    
+    def __init__(self, config_file: str = "config/alerts.json", db_path: str = "data/gpu_monitoring.db"):
+        self.config_file = config_file
+        self.db_manager = GPUDataManager(db_path)
+        
+        # Alert state tracking
+        self.active_alerts = {}  # alert_id -> Alert
+        self.threshold_states = {}  # (gpu, metric) -> {'value': float, 'start_time': float}
+        self.last_alert_times = {}  # (gpu, metric) -> timestamp
+        
+        # Notification channels
+        self.channels = []
+        
+        # Configuration
+        self.thresholds = []
+        self.alert_settings = {}
+        
+        # Threading
+        self.running = False
+        self.thread = None
+        
+        # Load configuration
+        self.load_config()
+        self.setup_channels()
+    
+    def load_config(self):
+        """Load alert configuration from file."""
+        try:
+            if Path(self.config_file).exists():
+                with open(self.config_file, 'r') as f:
+                    config = json.load(f)
+                
+                # Load thresholds
+                self.thresholds = []
+                for threshold_data in config.get('thresholds', []):
+                    threshold = AlertThreshold(**threshold_data)
+                    self.thresholds.append(threshold)
+                
+                # Load settings
+                self.alert_settings = config.get('settings', {
+                    'check_interval': 5.0,
+                    'cleanup_interval': 3600.0,
+                    'max_alert_age': 86400.0  # 24 hours
+                })
+                
+                logger.info(f"Loaded {len(self.thresholds)} alert thresholds")
+            else:
+                self.create_default_config()
+                
+        except Exception as e:
+            logger.error(f"Error loading alert config: {e}")
+            self.create_default_config()
+    
+    def create_default_config(self):
+        """Create default alert configuration."""
+        default_config = {
+            'thresholds': [
+                {
+                    'metric': 'temperature',
+                    'threshold': 75.0,
+                    'operator': '>=',
+                    'duration': 10,
+                    'enabled': True,
+                    'cooldown': 300
+                },
+                {
+                    'metric': 'temperature',
+                    'threshold': 85.0,
+                    'operator': '>=',
+                    'duration': 5,
+                    'enabled': True,
+                    'cooldown': 600
+                },
+                {
+                    'metric': 'load',
+                    'threshold': 90.0,
+                    'operator': '>=',
+                    'duration': 30,
+                    'enabled': True,
+                    'cooldown': 600
+                },
+                {
+                    'metric': 'power_draw',
+                    'threshold': 200.0,
+                    'operator': '>=',
+                    'duration': 10,
+                    'enabled': True,
+                    'cooldown': 300
+                },
+                {
+                    'metric': 'fan_speed',
+                    'threshold': 95.0,
+                    'operator': '>=',
+                    'duration': 60,
+                    'enabled': True,
+                    'cooldown': 1800
+                }
+            ],
+            'settings': {
+                'check_interval': 5.0,
+                'cleanup_interval': 3600.0,
+                'max_alert_age': 86400.0
+            },
+            'notifications': {
+                'log': {'enabled': True},
+                'desktop': {'enabled': True},
+                'email': {
+                    'enabled': False,
+                    'smtp': {
+                        'server': 'smtp.gmail.com',
+                        'port': 587,
+                        'use_tls': True,
+                        'sender': 'your-email@gmail.com',
+                        'username': 'your-username',
+                        'password': 'your-app-password'
+                    },
+                    'recipients': ['admin@example.com']
+                },
+                'webhook': {
+                    'enabled': False,
+                    'url': 'https://your-webhook-url.com/alerts',
+                    'headers': {'Authorization': 'Bearer your-token'}
+                }
+            }
+        }
+        
+        # Save default config
+        Path(self.config_file).parent.mkdir(parents=True, exist_ok=True)
+        with open(self.config_file, 'w') as f:
+            json.dump(default_config, f, indent=2)
+        
+        logger.info("Created default alert configuration")
+    
+    def setup_channels(self):
+        """Setup notification channels."""
+        try:
+            if Path(self.config_file).exists():
+                with open(self.config_file, 'r') as f:
+                    config = json.load(f)
+                
+                notifications = config.get('notifications', {})
+                
+                # Log channel
+                if notifications.get('log', {}).get('enabled', True):
+                    self.channels.append(LogNotification('log'))
+                
+                # Desktop channel
+                if notifications.get('desktop', {}).get('enabled', True):
+                    self.channels.append(DesktopNotification('desktop'))
+                
+                # Email channel
+                email_config = notifications.get('email', {})
+                if email_config.get('enabled', False):
+                    smtp_config = email_config.get('smtp', {})
+                    recipients = email_config.get('recipients', [])
+                    if smtp_config and recipients:
+                        self.channels.append(EmailNotification(
+                            'email', smtp_config, recipients
+                        ))
+                
+                # Webhook channel
+                webhook_config = notifications.get('webhook', {})
+                if webhook_config.get('enabled', False):
+                    url = webhook_config.get('url')
+                    headers = webhook_config.get('headers', {})
+                    if url:
+                        self.channels.append(WebhookNotification('webhook', url, headers))
+                
+                logger.info(f"Setup {len(self.channels)} notification channels")
+                
+        except Exception as e:
+            logger.error(f"Error setting up notification channels: {e}")
+    
+    def add_threshold(self, threshold: AlertThreshold):
+        """Add a new alert threshold."""
+        self.thresholds.append(threshold)
+        self.save_config()
+    
+    def remove_threshold(self, metric: str, threshold_value: float):
+        """Remove an alert threshold."""
+        self.thresholds = [
+            t for t in self.thresholds 
+            if not (t.metric == metric and t.threshold == threshold_value)
+        ]
+        self.save_config()
+    
+    def save_config(self):
+        """Save current configuration to file."""
+        try:
+            config = {
+                'thresholds': [asdict(t) for t in self.thresholds],
+                'settings': self.alert_settings,
+                'notifications': {
+                    'log': {'enabled': any(isinstance(c, LogNotification) for c in self.channels)},
+                    'desktop': {'enabled': any(isinstance(c, DesktopNotification) for c in self.channels)},
+                }
+            }
+            
+            # Add email and webhook configs if channels exist
+            email_channel = next((c for c in self.channels if isinstance(c, EmailNotification)), None)
+            if email_channel:
+                config['notifications']['email'] = {
+                    'enabled': True,
+                    'smtp': email_channel.smtp_config,
+                    'recipients': email_channel.recipients
+                }
+            
+            webhook_channel = next((c for c in self.channels if isinstance(c, WebhookNotification)), None)
+            if webhook_channel:
+                config['notifications']['webhook'] = {
+                    'enabled': True,
+                    'url': webhook_channel.webhook_url,
+                    'headers': webhook_channel.headers
+                }
+            
+            with open(self.config_file, 'w') as f:
+                json.dump(config, f, indent=2)
+            
+        except Exception as e:
+            logger.error(f"Error saving alert config: {e}")
+    
+    def check_thresholds(self, gpu_name: str, status: GPUStatus):
+        """Check all thresholds against current status."""
+        current_time = time.time()
+        
+        # Check each threshold
+        for threshold in self.thresholds:
+            if not threshold.enabled:
+                continue
+            
+            # Get metric value
+            metric_value = self.get_metric_value(status, threshold.metric)
+            if metric_value is None:
+                continue
+            
+            # Check if threshold is exceeded
+            if self.check_operator(metric_value, threshold.operator, threshold.threshold):
+                # Start or continue threshold state
+                state_key = (gpu_name, threshold.metric, threshold.threshold)
+                
+                if state_key not in self.threshold_states:
+                    self.threshold_states[state_key] = {
+                        'value': metric_value,
+                        'start_time': current_time
+                    }
+                else:
+                    self.threshold_states[state_key]['value'] = metric_value
+                
+                # Check if duration threshold is met
+                state = self.threshold_states[state_key]
+                duration = current_time - state['start_time']
+                
+                if duration >= threshold.duration:
+                    # Check cooldown
+                    last_alert_key = state_key
+                    if last_alert_key in self.last_alert_times:
+                        time_since_last = current_time - self.last_alert_times[last_alert_key]
+                        if time_since_last < threshold.cooldown:
+                            continue
+                    
+                    # Create alert
+                    alert = self.create_alert(gpu_name, threshold, metric_value)
+                    self.trigger_alert(alert)
+                    
+                    # Update cooldown
+                    self.last_alert_times[last_alert_key] = current_time
+            
+            else:
+                # Reset threshold state if condition is no longer met
+                state_key = (gpu_name, threshold.metric, threshold.threshold)
+                if state_key in self.threshold_states:
+                    del self.threshold_states[state_key]
+    
+    def get_metric_value(self, status: GPUStatus, metric: str) -> Optional[float]:
+        """Get metric value from GPU status."""
+        if metric == 'temperature':
+            return status.temperature
+        elif metric == 'load':
+            return status.load
+        elif metric == 'power_draw':
+            return status.power_draw
+        elif metric == 'fan_speed':
+            return (status.fan_pwm / 255) * 100  # Convert PWM to percentage
+        elif metric == 'memory_usage':
+            if status.memory_total > 0:
+                return (status.memory_used / status.memory_total) * 100
+            return 0
+        elif metric == 'efficiency':
+            return status.efficiency
+        else:
+            return None
+    
+    def check_operator(self, value: float, operator: str, threshold: float) -> bool:
+        """Check if value meets threshold condition."""
+        if operator == '>':
+            return value > threshold
+        elif operator == '<':
+            return value < threshold
+        elif operator == '>=':
+            return value >= threshold
+        elif operator == '<=':
+            return value <= threshold
+        elif operator == '==':
+            return value == threshold
+        else:
+            return False
+    
+    def create_alert(self, gpu_name: str, threshold: AlertThreshold, value: float) -> Alert:
+        """Create an alert instance."""
+        alert_id = f"{gpu_name}_{threshold.metric}_{threshold.threshold}_{int(time.time())}"
+        
+        # Determine severity
+        severity = 'info'
+        if threshold.metric == 'temperature':
+            if threshold.threshold >= 85:
+                severity = 'emergency'
+            elif threshold.threshold >= 75:
+                severity = 'critical'
+            else:
+                severity = 'warning'
+        elif threshold.metric == 'load':
+            severity = 'warning'
+        elif threshold.metric == 'power_draw':
+            severity = 'critical'
+        elif threshold.metric == 'fan_speed':
+            severity = 'warning'
+        
+        message = f"{threshold.metric} ({value}) exceeded threshold ({threshold.threshold}) for {threshold.duration}s"
+        
+        alert = Alert(
+            id=alert_id,
+            timestamp=time.time(),
+            gpu_name=gpu_name,
+            metric=threshold.metric,
+            value=value,
+            threshold=threshold.threshold,
+            message=message,
+            severity=severity
+        )
+        
+        return alert
+    
+    def trigger_alert(self, alert: Alert):
+        """Trigger an alert and send notifications."""
+        # Store alert in database
+        self.db_manager.save_alert(alert)
+        
+        # Add to active alerts
+        self.active_alerts[alert.id] = alert
+        
+        # Send notifications
+        for channel in self.channels:
+            if channel.enabled:
+                try:
+                    success = channel.send(alert)
+                    if not success:
+                        logger.warning(f"Failed to send alert via {channel.name}")
+                except Exception as e:
+                    logger.error(f"Error sending alert via {channel.name}: {e}")
+        
+        logger.info(f"Alert triggered: {alert.message}")
+    
+    def resolve_alert(self, alert_id: str):
+        """Resolve an active alert."""
+        if alert_id in self.active_alerts:
+            alert = self.active_alerts[alert_id]
+            alert.resolved = True
+            alert.resolved_at = time.time()
+            
+            # Remove from active alerts
+            del self.active_alerts[alert_id]
+            
+            logger.info(f"Alert resolved: {alert.message}")
+    
+    def get_active_alerts(self) -> List[Alert]:
+        """Get list of active alerts."""
+        return list(self.active_alerts.values())
+    
+    def get_alert_history(self, hours: int = 24) -> List[Alert]:
+        """Get alert history from database."""
+        try:
+            cutoff_time = time.time() - (hours * 3600)
+            
+            with sqlite3.connect(self.db_manager.db_path) as conn:
+                cursor = conn.cursor()
+                
+                cursor.execute('''
+                    SELECT id, timestamp, gpu_name, metric, value, threshold, message, severity,
+                           acknowledged, resolved, resolved_at
+                    FROM alerts 
+                    WHERE timestamp >= ?
+                    ORDER BY timestamp DESC
+                ''', (cutoff_time,))
+                
+                rows = cursor.fetchall()
+                
+                if not rows:
+                    return []
+                
+                alerts = []
+                for row in rows:
+                    alert = Alert(
+                        id=row[0],
+                        timestamp=row[1],
+                        gpu_name=row[2],
+                        metric=row[3],
+                        value=row[4],
+                        threshold=row[5],
+                        message=row[6],
+                        severity=row[7],
+                        acknowledged=bool(row[8]),
+                        resolved=bool(row[9]),
+                        resolved_at=row[10]
+                    )
+                    alerts.append(alert)
+                
+                return alerts
+                
+        except Exception as e:
+            logger.error(f"Error getting alert history: {e}")
+            return []
+    
+    def cleanup_old_alerts(self):
+        """Clean up old alerts and threshold states."""
+        current_time = time.time()
+        max_age = self.alert_settings.get('max_alert_age', 86400.0)
+        
+        # Clean up old active alerts
+        old_alerts = []
+        for alert_id, alert in self.active_alerts.items():
+            if current_time - alert.timestamp > max_age:
+                old_alerts.append(alert_id)
+        
+        for alert_id in old_alerts:
+            del self.active_alerts[alert_id]
+        
+        # Clean up old threshold states
+        old_states = []
+        for state_key, state in self.threshold_states.items():
+            if current_time - state['start_time'] > max_age:
+                old_states.append(state_key)
+        
+        for state_key in old_states:
+            del self.threshold_states[state_key]
+        
+        # Clean up old last alert times
+        old_cooldowns = []
+        for key, timestamp in self.last_alert_times.items():
+            if current_time - timestamp > max_age:
+                old_cooldowns.append(key)
+        
+        for key in old_cooldowns:
+            del self.last_alert_times[key]
+    
+    def start(self):
+        """Start the alert manager."""
+        if self.running:
+            return
+        
+        self.running = True
+        self.thread = threading.Thread(target=self.run, daemon=True)
+        self.thread.start()
+        logger.info("Alert manager started")
+    
+    def stop(self):
+        """Stop the alert manager."""
+        self.running = False
+        if self.thread:
+            self.thread.join()
+        logger.info("Alert manager stopped")
+    
+    def run(self):
+        """Main alert manager loop."""
+        check_interval = self.alert_settings.get('check_interval', 5.0)
+        cleanup_interval = self.alert_settings.get('cleanup_interval', 3600.0)
+        last_cleanup = time.time()
+        
+        while self.running:
+            try:
+                # Get current GPU status
+                gpu_manager = GPUManager()
+                if gpu_manager.initialize():
+                    status_dict = gpu_manager.get_status()
+                    
+                    # Check thresholds for each GPU
+                    for gpu_name, status in status_dict.items():
+                        if status:
+                            self.check_thresholds(gpu_name, status)
+                
+                # Cleanup old data periodically
+                current_time = time.time()
+                if current_time - last_cleanup >= cleanup_interval:
+                    self.cleanup_old_alerts()
+                    last_cleanup = current_time
+                
+                time.sleep(check_interval)
+                
+            except Exception as e:
+                logger.error(f"Error in alert manager loop: {e}")
+                time.sleep(5)  # Wait before retrying
+
+
+class AlertAPI:
+    """API interface for alert management."""
+    
+    def __init__(self, alert_manager: AlertManager):
+        self.alert_manager = alert_manager
+    
+    def get_active_alerts(self) -> List[Dict[str, Any]]:
+        """Get active alerts."""
+        alerts = self.alert_manager.get_active_alerts()
+        return [asdict(alert) for alert in alerts]
+    
+    def get_alert_history(self, hours: int = 24) -> List[Dict[str, Any]]:
+        """Get alert history."""
+        alerts = self.alert_manager.get_alert_history(hours)
+        return [asdict(alert) for alert in alerts]
+    
+    def acknowledge_alert(self, alert_id: str) -> bool:
+        """Acknowledge an alert."""
+        if alert_id in self.alert_manager.active_alerts:
+            self.alert_manager.active_alerts[alert_id].acknowledged = True
+            return True
+        return False
+    
+    def resolve_alert(self, alert_id: str) -> bool:
+        """Resolve an alert."""
+        self.alert_manager.resolve_alert(alert_id)
+        return True
+    
+    def add_threshold(self, threshold_data: Dict[str, Any]) -> bool:
+        """Add a new threshold."""
+        try:
+            threshold = AlertThreshold(**threshold_data)
+            self.alert_manager.add_threshold(threshold)
+            return True
+        except Exception as e:
+            logger.error(f"Error adding threshold: {e}")
+            return False
+    
+    def remove_threshold(self, metric: str, threshold_value: float) -> bool:
+        """Remove a threshold."""
+        try:
+            self.alert_manager.remove_threshold(metric, threshold_value)
+            return True
+        except Exception as e:
+            logger.error(f"Error removing threshold: {e}")
+            return False
+
+
+if __name__ == "__main__":
+    # Test alert system
+    logging.basicConfig(level=logging.INFO)
+    
+    alert_manager = AlertManager()
+    alert_manager.start()
+    
+    try:
+        while True:
+            time.sleep(1)
+    except KeyboardInterrupt:
+        alert_manager.stop()