resource_manager.py

"""
AgentFile Advanced Resource Manager
Goes beyond normal MoE resource management
Provides intelligent memory, compute, and quality optimization
"""

import torch
import torch.nn as nn
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
from enum import Enum
import psutil
import threading
import time
from queue import Queue
import numpy as np
import logging

logger = logging.getLogger(__name__)

class ResourcePriority(Enum):
    """Resource allocation priorities"""
    CRITICAL = 1
    HIGH = 2
    NORMAL = 3
    LOW = 4
    BACKGROUND = 5

@dataclass
class ResourceBudget:
    """Resource budget configuration"""
    memory_gb: float = 8.0
    compute_units: float = 100.0
    priority: ResourcePriority = ResourcePriority.NORMAL
    max_experts: int = 4
    quality_threshold: float = 0.8

class IntelligentResourceManager:
    """Intelligent Resource Manager - Goes beyond normal MoE"""
    
    def __init__(self, budget: ResourceBudget):
        self.budget = budget
        self.memory_usage = {}
        self.compute_usage = {}
        self.quality_metrics = {}
        
        # Predictive models
        self.memory_predictor = self._create_memory_predictor()
        self.compute_predictor = self._create_compute_predictor()
        
        # Resource pools
        self.memory_pool = Queue()
        self.compute_pool = Queue()
        
        # Monitoring
        self.monitor_thread = None
        self.running = False
        
        # Statistics
        self.stats = {
            'total_memory_allocated': 0,
            'total_compute_allocated': 0,
            'avg_quality': 0.0,
            'optimization_count': 0
        }
    
    def _create_memory_predictor(self):
        """Create a simple memory usage predictor"""
        
        class MemoryPredictor:
            def __init__(self):
                self.history = []
                
            def predict(self, model_size: int, batch_size: int) -> float:
                base_memory = model_size * 2  # 2x for forward/backward
                batch_factor = batch_size * 0.1  # 10% per batch element
                return base_memory * (1 + batch_factor)
        
        return MemoryPredictor()
    
    def _create_compute_predictor(self):
        """Create a compute usage predictor"""
        
        class ComputePredictor:
            def __init__(self):
                self.flops_per_parameter = 6  # Approximate FLOPs per parameter
                
            def predict(self, model_size: int, sequence_length: int) -> float:
                flops = model_size * self.flops_per_parameter * sequence_length
                return flops / 1e9  # Convert to GFLOPs
        
        return ComputePredictor()
    
    def allocate_resources(
        self, 
        model_size: int, 
        batch_size: int, 
        sequence_length: int
    ) -> Dict:
        """Intelligently allocate resources"""
        
        # Predict requirements
        memory_needed = self.memory_predictor.predict(model_size, batch_size)
        compute_needed = self.compute_predictor.predict(model_size, sequence_length)
        
        # Check availability
        available_memory = self._get_available_memory()
        available_compute = self._get_available_compute()
        
        # Allocate based on availability and priority
        allocation = {
            'memory': min(memory_needed, available_memory * 0.8),
            'compute': min(compute_needed, available_compute * 0.8),
            'batch_size': batch_size,
            'sequence_length': sequence_length,
            'optimization_level': self._calculate_optimization_level(
                memory_needed, available_memory
            )
        }
        
        # Update usage tracking
        self.memory_usage[id(allocation)] = allocation['memory']
        self.compute_usage[id(allocation)] = allocation['compute']
        
        return allocation
    
    def optimize_allocation(
        self, 
        current_allocation: Dict, 
        quality_feedback: float
    ) -> Dict:
        """Optimize allocation based on quality feedback"""
        
        optimized = current_allocation.copy()
        
        # If quality is low, increase resources
        if quality_feedback < self.budget.quality_threshold:
            optimized['memory'] = min(
                optimized['memory'] * 1.2,
                self.budget.memory_gb
            )
            
            optimized['compute'] = min(
                optimized['compute'] * 1.15,
                self.budget.compute_units
            )
            
            # Reduce batch size if memory constrained
            if optimized['memory'] >= self.budget.memory_gb * 0.9:
                optimized['batch_size'] = max(1, optimized['batch_size'] // 2)
        
        # If quality is high, we can be more efficient
        elif quality_feedback > 0.95:
            optimized['memory'] = optimized['memory'] * 0.9
            optimized['batch_size'] = int(optimized['batch_size'] * 1.1)
        
        self.stats['optimization_count'] += 1
        
        return optimized
    
    def get_optimal_expert_count(
        self, 
        input_complexity: float, 
        available_resources: Dict
    ) -> int:
        """Determine optimal number of experts based on input and resources"""
        
        base_experts = int(input_complexity * self.budget.max_experts)
        
        memory_factor = available_resources.get('memory', 0) / self.budget.memory_gb
        compute_factor = available_resources.get('compute', 0) / self.budget.compute_units
        
        resource_factor = (memory_factor + compute_factor) / 2
        
        optimal_experts = int(base_experts * resource_factor)
        
        return max(1, min(optimal_experts, self.budget.max_experts))
    
    def _get_available_memory(self) -> float:
        """Get available memory in GB"""
        
        try:
            if torch.cuda.is_available():
                return torch.cuda.get_device_properties(0).total_mem / 1e9
            else:
                return psutil.virtual_memory().available / 1e9
        except:
            return 8.0
    
    def _get_available_compute(self) -> float:
        """Get available compute units"""
        
        try:
            return psutil.cpu_percent() / 100
        except:
            return 100.0
    
    def _calculate_optimization_level(self, needed: float, available: float) -> str:
        """Calculate optimization level needed"""
        
        ratio = needed / available if available > 0 else 1.0
        
        if ratio < 0.5:
            return "none"
        elif ratio < 0.7:
            return "light"
        elif ratio < 0.9:
            return "moderate"
        else:
            return "aggressive"
    
    def start_monitoring(self):
        """Start resource monitoring thread"""
        
        self.running = True
        self.monitor_thread = threading.Thread(target=self._monitor_loop)
        self.monitor_thread.daemon = True
        self.monitor_thread.start()
        logger.info("Resource monitoring started")
    
    def stop_monitoring(self):
        """Stop resource monitoring"""
        
        self.running = False
        if self.monitor_thread:
            self.monitor_thread.join()
        logger.info("Resource monitoring stopped")
    
    def _monitor_loop(self):
        """Main monitoring loop"""
        
        while self.running:
            self._update_stats()
            self._check_resource_leaks()
            self._auto_optimize()
            time.sleep(1.0)
    
    def _update_stats(self):
        """Update resource statistics"""
        
        total_memory = sum(self.memory_usage.values())
        total_compute = sum(self.compute_usage.values())
        
        self.stats['total_memory_allocated'] = total_memory
        self.stats['total_compute_allocated'] = total_compute
        
        if self.quality_metrics:
            self.stats['avg_quality'] = np.mean(list(self.quality_metrics.values()))
    
    def _check_resource_leaks(self):
        """Check for and clean up resource leaks"""
        
        for alloc_id, memory in list(self.memory_usage.items()):
            # In real implementation, you'd track timestamps
            pass
    
    def _auto_optimize(self):
        """Automatically optimize resource allocation"""
        
        if self.stats['total_memory_allocated'] > self.budget.memory_gb * 0.9:
            self._reduce_memory_usage()
        
        if self.stats['total_compute_allocated'] < self.budget.compute_units * 0.5:
            self._increase_compute_usage()
    
    def _reduce_memory_usage(self):
        """Reduce memory usage"""
        
        for alloc_id in list(self.memory_usage.keys()):
            # In real implementation, you'd modify the actual allocations
            pass
    
    def _increase_compute_usage(self):
        """Increase compute usage for better throughput"""
        
        pass
    
    def get_status(self) -> Dict:
        """Get current resource status"""
        
        return {
            'memory_usage': self.memory_usage,
            'compute_usage': self.compute_usage,
            'quality_metrics': self.quality_metrics,
            'stats': self.stats,
            'budget': {
                'memory_gb': self.budget.memory_gb,
                'compute_units': self.budget.compute_units,
                'max_experts': self.budget.max_experts
            }
        }

class AdaptiveBatchScheduler:
    """Adaptive Batch Scheduler - Dynamically adjusts batch sizes"""
    
    def __init__(self, resource_manager: IntelligentResourceManager):
        self.resource_manager = resource_manager
        self.batch_history = []
        self.quality_history = []
        
    def get_optimal_batch_size(
        self, 
        current_batch_size: int, 
        quality_feedback: float
    ) -> int:
        """Get optimal batch size based on feedback"""
        
        self.batch_history.append(current_batch_size)
        self.quality_history.append(quality_feedback)
        
        if len(self.batch_history) > 100:
            self.batch_history = self.batch_history[-100:]
            self.quality_history = self.quality_history[-100:]
        
        if len(self.quality_history) > 10:
            recent_quality = np.mean(self.quality_history[-10:])
            older_quality = np.mean(self.quality_history[-20:-10]) if len(self.quality_history) > 20 else recent_quality
            quality_trend = recent_quality - older_quality
        else:
            quality_trend = 0
        
        if quality_trend > 0.05:
            new_batch_size = int(current_batch_size * 1.1)
        elif quality_trend < -0.05:
            new_batch_size = int(current_batch_size * 0.9)
        else:
            new_batch_size = current_batch_size
        
        new_batch_size = max(1, min(new_batch_size, 32))
        
        return new_batch_size

class QualityAwareRouter:
    """Quality-Aware Router - Routes based on quality requirements"""
    
    def __init__(self, num_experts: int, quality_threshold: float = 0.8):
        self.num_experts = num_experts
        self.quality_threshold = quality_threshold
        
        # Expert quality scores
        self.expert_quality_scores = [0.5] * num_experts
        
        # Input complexity analyzer
        self.complexity_analyzer = self._create_complexity_analyzer()
        
    def _create_complexity_analyzer(self):
        """Create input complexity analyzer"""
        
        class ComplexityAnalyzer:
            def analyze(self, input_ids: torch.Tensor) -> float:
                unique_tokens = len(torch.unique(input_ids))
                total_tokens = input_ids.numel()
                return unique_tokens / total_tokens
        
        return ComplexityAnalyzer()
    
    def route(
        self, 
        input_ids: torch.Tensor, 
        available_experts: List[int]
    ) -> List[Tuple[int, float]]:
        """Route input to experts based on quality requirements"""
        
        complexity = self.complexity_analyzer.analyze(input_ids)
        
        selected_experts = []
        
        for expert_id in available_experts:
            quality_score = self.expert_quality_scores[expert_id]
            
            if quality_score >= self.quality_threshold:
                weight = self._calculate_expert_weight(expert_id, complexity)
                selected_experts.append((expert_id, weight))
        
        selected_experts.sort(key=lambda x: x[1], reverse=True)
        top_k = min(4, len(selected_experts))
        
        return selected_experts[:top_k]
    
    def _calculate_expert_weight(self, expert_id: int, complexity: float) -> float:
        """Calculate expert weight based on complexity"""
        
        specializations = {
            0: 0.3,  # Simple tasks
            1: 0.5,  # Medium tasks
            2: 0.7,  # Complex tasks
            3: 0.9   # Very complex tasks
        }
        
        expert_specialization = specializations.get(expert_id, 0.5)
        
        weight = 1.0 - abs(complexity - expert_specialization)
        weight *= self.expert_quality_scores[expert_id]
        
        return weight
    
    def update_quality_scores(self, expert_id: int, quality: float):
        """Update expert quality scores based on feedback"""
        
        alpha = 0.1
        self.expert_quality_scores[expert_id] = (
            alpha * quality + 
            (1 - alpha) * self.expert_quality_scores[expert_id]
        )

class DynamicExpertPool:
    """Dynamic Expert Pool - Manages experts dynamically"""
    
    def __init__(self, max_experts: int = 4):
        self.max_experts = max_experts
        self.loaded_experts = {}
        self.expert_usage = {}
        
    def load_expert(self, expert_id: int, model_path: str):
        """Load an expert model"""
        
        if len(self.loaded_experts) >= self.max_experts:
            self._unload_least_used()
        
        from transformers import AutoModelForCausalLM
        model = AutoModelForCausalLM.from_pretrained(model_path)
        
        self.loaded_experts[expert_id] = model
        self.expert_usage[expert_id] = 0
        
        logger.info(f"Loaded expert {expert_id}")
    
    def unload_expert(self, expert_id: int):
        """Unload an expert model"""
        
        if expert_id in self.loaded_experts:
            del self.loaded_experts[expert_id]
            del self.expert_usage[expert_id]
            
            import gc
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            logger.info(f"Unloaded expert {expert_id}")
    
    def _unload_least_used(self):
        """Unload the least used expert"""
        
        if not self.expert_usage:
            return
        
        least_used_id = min(self.expert_usage, key=self.expert_usage.get)
        self.unload_expert(least_used_id)
    
    def get_expert(self, expert_id: int):
        """Get an expert model"""
        
        if expert_id in self.loaded_experts:
            self.expert_usage[expert_id] += 1
            return self.loaded_experts[expert_id]
        return None
    
    def get_loaded_experts(self) -> List[int]:
        """Get list of loaded expert IDs"""
        
        return list(self.loaded_experts.keys())
    
    def optimize_memory(self):
        """Optimize memory usage"""
        
        current_time = time.time()
        for expert_id in list(self.expert_usage.keys()):
            # In real implementation, you'd track last usage time
            pass

def create_resource_manager(
    memory_budget: float = 8.0,
    compute_budget: float = 100.0,
    max_experts: int = 4,
    quality_threshold: float = 0.8
) -> IntelligentResourceManager:
    """Convenience function to create a resource manager"""
    
    budget = ResourceBudget(
        memory_gb=memory_budget,
        compute_units=compute_budget,
        max_experts=max_experts,
        quality_threshold=quality_threshold
    )
    
    return IntelligentResourceManager(budget)