""" Model Availability Tracker Implements circuit breaker pattern for AI models """ import time import logging from typing import Dict, Optional from dataclasses import dataclass from datetime import datetime, timedelta logger = logging.getLogger(__name__) @dataclass class ModelStatus: """Status information for a model""" name: str consecutive_failures: int = 0 last_failure_time: Optional[float] = None is_available: bool = True total_failures: int = 0 total_successes: int = 0 last_success_time: Optional[float] = None class ModelAvailabilityTracker: """ Tracks model availability and implements circuit breaker pattern. - Marks model unavailable after 3 consecutive failures - Skips unavailable models for 5 minutes before retrying - Provides status reporting for health checks """ # Configuration MAX_CONSECUTIVE_FAILURES = 3 COOLDOWN_SECONDS = 300 # 5 minutes def __init__(self): """Initialize tracker with empty status for all models""" self._models: Dict[str, ModelStatus] = { 'qwen2vl': ModelStatus(name='Qwen2-VL'), 'nemotron': ModelStatus(name='Nemotron'), 'ollama': ModelStatus(name='Ollama') } def record_success(self, model_name: str) -> None: """ Record a successful model execution. Resets consecutive failures and marks model as available. Args: model_name: Name of the model ('qwen2vl', 'nemotron', 'ollama') """ if model_name not in self._models: logger.warning(f"Unknown model: {model_name}") return status = self._models[model_name] status.consecutive_failures = 0 status.is_available = True status.total_successes += 1 status.last_success_time = time.time() logger.debug(f"[{status.name}] Success recorded (total: {status.total_successes})") def record_failure(self, model_name: str, error: Optional[Exception] = None) -> None: """ Record a model failure. Increments consecutive failures and marks unavailable after threshold. Args: model_name: Name of the model ('qwen2vl', 'nemotron', 'ollama') error: Optional exception that caused the failure """ if model_name not in self._models: logger.warning(f"Unknown model: {model_name}") return status = self._models[model_name] status.consecutive_failures += 1 status.total_failures += 1 status.last_failure_time = time.time() # Mark unavailable after threshold if status.consecutive_failures >= self.MAX_CONSECUTIVE_FAILURES: status.is_available = False logger.warning( f"[{status.name}] Marked UNAVAILABLE after {status.consecutive_failures} " f"consecutive failures. Will retry in {self.COOLDOWN_SECONDS}s" ) else: logger.warning( f"[{status.name}] Failure {status.consecutive_failures}/{self.MAX_CONSECUTIVE_FAILURES}" + (f": {error}" if error else "") ) def is_available(self, model_name: str) -> bool: """ Check if a model is available for use. If model was marked unavailable, checks if cooldown period has elapsed. If cooldown elapsed, marks model as available for retry. Args: model_name: Name of the model ('qwen2vl', 'nemotron', 'ollama') Returns: True if model is available, False if in cooldown period """ if model_name not in self._models: logger.warning(f"Unknown model: {model_name}") return False status = self._models[model_name] # If already available, return True if status.is_available: return True # Check if cooldown period has elapsed if status.last_failure_time is None: return True time_since_failure = time.time() - status.last_failure_time if time_since_failure >= self.COOLDOWN_SECONDS: # Cooldown elapsed, allow retry logger.info( f"[{status.name}] Cooldown period elapsed ({time_since_failure:.0f}s). " "Allowing retry..." ) status.is_available = True status.consecutive_failures = 0 # Reset for fresh start return True # Still in cooldown remaining = self.COOLDOWN_SECONDS - time_since_failure logger.debug( f"[{status.name}] Still unavailable. Retry in {remaining:.0f}s" ) return False def get_status(self, model_name: str) -> Optional[Dict]: """ Get detailed status for a specific model. Args: model_name: Name of the model ('qwen2vl', 'nemotron', 'ollama') Returns: Dictionary with status information or None if model unknown """ if model_name not in self._models: return None status = self._models[model_name] # Calculate time since last failure time_since_failure = None cooldown_remaining = None if status.last_failure_time: time_since_failure = time.time() - status.last_failure_time if not status.is_available: cooldown_remaining = max(0, self.COOLDOWN_SECONDS - time_since_failure) # Calculate time since last success time_since_success = None if status.last_success_time: time_since_success = time.time() - status.last_success_time return { 'name': status.name, 'available': status.is_available, 'consecutive_failures': status.consecutive_failures, 'total_failures': status.total_failures, 'total_successes': status.total_successes, 'last_failure_time': datetime.fromtimestamp(status.last_failure_time).isoformat() if status.last_failure_time else None, 'last_success_time': datetime.fromtimestamp(status.last_success_time).isoformat() if status.last_success_time else None, 'time_since_failure_seconds': round(time_since_failure, 1) if time_since_failure else None, 'time_since_success_seconds': round(time_since_success, 1) if time_since_success else None, 'cooldown_remaining_seconds': round(cooldown_remaining, 1) if cooldown_remaining else None } def get_all_status(self) -> Dict[str, Dict]: """ Get status for all tracked models. Returns: Dictionary mapping model names to their status information """ return { model_name: self.get_status(model_name) for model_name in self._models.keys() } def reset(self, model_name: Optional[str] = None) -> None: """ Reset tracking for a model or all models. Args: model_name: Name of model to reset, or None to reset all """ if model_name: if model_name in self._models: self._models[model_name] = ModelStatus(name=self._models[model_name].name) logger.info(f"[{model_name}] Status reset") else: for name in self._models.keys(): self._models[name] = ModelStatus(name=self._models[name].name) logger.info("All model statuses reset") # Global singleton instance _tracker = ModelAvailabilityTracker() def get_tracker() -> ModelAvailabilityTracker: """Get the global model availability tracker instance""" return _tracker