backend/scheduler/resource_allocator.py

"""
Resource Allocator for Adaptive GPU Allocation

Handles intelligent resource allocation:
- GPU allocation decisions
- Node pool selection
- Spot vs on-demand selection
- Batch size optimization
- Inference mode selection

Supports adaptive resource allocation based on cluster load.
"""

import uuid
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional

from pydantic import BaseModel


class InferenceMode(str, Enum):
    """Inference modes with different resource requirements."""
    LIGHTWEIGHT = "lightweight"  # Fast, less accurate
    STANDARD = "standard"        # Balanced
    FULL = "full"                 # Complete, more accurate


class NodePoolType(str, Enum):
    """Node pool types for different workloads."""
    CPU = "cpu"
    GPU_STANDARD = "gpu_standard"
    GPU_HIGH_MEMORY = "gpu_high_memory"
    GPU_AMPERE = "gpu_ampere"
    GPU_HOPPER = "gpu_hopper"


class AllocationDecision(BaseModel):
    """Resource allocation decision for a job."""
    job_id: uuid.UUID
    
    # Allocation details
    gpu_count: int = 1
    gpu_type: str = "v100"
    node_pool: str = "gpu_standard"
    inference_mode: str = "standard"
    batch_size: int = 4
    
    # Spot vs on-demand
    use_spot: bool = True
    
    # Optimization hints
    enable_batching: bool = False
    reduce_mutation_depth: bool = False
    use_lightweight_hallucination: bool = False
    
    # Metadata
    allocated_at: datetime = None
    
    def __init__(self, **data):
        if "allocated_at" not in data or data["allocated_at"] is None:
            data["allocated_at"] = datetime.utcnow()
        super().__init__(**data)


class ResourceAllocator:
    """
    Intelligent resource allocator for evaluation jobs.
    
    Makes allocation decisions based on:
    - Job requirements (GPU count, memory)
    - Cluster state (available resources)
    - Cost optimization (spot vs on-demand)
    - Priority (high priority = better resources)
    """
    
    # Default configurations
    DEFAULT_BATCH_SIZE = 4
    DEFAULT_GPU_COUNT = 1
    
    # Node pool specifications
    NODE_POOL_SPECS = {
        NodePoolType.CPU: {
            "gpu_count": 0,
            "memory_gb": 32,
            "cost_per_hour": 0.50,
        },
        NodePoolType.GPU_STANDARD: {
            "gpu_count": 1,
            "gpu_type": "v100",
            "memory_gb": 60,
            "cost_per_hour": 2.48,
        },
        NodePoolType.GPU_HIGH_MEMORY: {
            "gpu_count": 1,
            "gpu_type": "v100",
            "memory_gb": 120,
            "cost_per_hour": 3.50,
        },
        NodePoolType.GPU_AMPERE: {
            "gpu_count": 1,
            "gpu_type": "a100",
            "memory_gb": 80,
            "cost_per_hour": 3.67,
        },
        NodePoolType.GPU_HOPPER: {
            "gpu_count": 1,
            "gpu_type": "h100",
            "memory_gb": 160,
            "cost_per_hour": 6.50,
        },
    }
    
    def __init__(
        self,
        default_node_pool: NodePoolType = NodePoolType.GPU_STANDARD,
        enable_spot_by_default: bool = True,
    ):
        """
        Initialize resource allocator.
        
        Args:
            default_node_pool: Default node pool type
            enable_spot_by_default: Use spot instances by default
        """
        self.default_node_pool = default_node_pool
        self.enable_spot_by_default = enable_spot_by_default
        
        # Cluster state (would be updated from monitoring)
        self._cluster_load: float = 0.0
        self._available_gpu_count: int = 0
    
    def allocate_resources(
        self,
        job_id: uuid.UUID,
        total_samples: int,
        priority_score: float = 0.5,
        required_gpu_memory_mb: int = 0,
        model_size: str = "7b",
        cluster_load: Optional[float] = None,
        available_gpus: Optional[int] = None,
    ) -> AllocationDecision:
        """
        Determine resource allocation for a job.
        
        Args:
            job_id: Unique job identifier
            total_samples: Number of samples to process
            priority_score: Job priority score (0-1)
            required_gpu_memory_mb: Required GPU memory in MB
            model_size: Model size (7b, 13b, 30b, 70b)
            cluster_load: Current cluster load (0-1)
            available_gpus: Number of available GPUs
            
        Returns:
            AllocationDecision with resource allocation details
        """
        # Use provided cluster state or defaults
        load = cluster_load if cluster_load is not None else self._cluster_load
        gpus = available_gpus if available_gpus is not None else self._available_gpu_count
        
        # Determine inference mode based on priority and load
        inference_mode = self._determine_inference_mode(priority_score, load)
        
        # Determine batch size
        batch_size = self._determine_batch_size(
            total_samples, load, inference_mode
        )
        
        # Determine GPU type and node pool
        gpu_type, node_pool = self._determine_gpu_and_pool(
            model_size, required_gpu_memory_mb
        )
        
        # Determine spot vs on-demand
        use_spot = self._determine_use_spot(priority_score, load)
        
        # Determine optimization flags
        enable_batching = self._should_enable_batching(load, total_samples)
        reduce_mutation_depth = self._should_reduce_mutation_depth(
            priority_score, load
        )
        use_lightweight_hallucination = self._should_use_lightweight_hallucination(
            load
        )
        
        return AllocationDecision(
            job_id=job_id,
            gpu_count=self.DEFAULT_GPU_COUNT,
            gpu_type=gpu_type,
            node_pool=node_pool.value,
            inference_mode=inference_mode.value,
            batch_size=batch_size,
            use_spot=use_spot,
            enable_batching=enable_batching,
            reduce_mutation_depth=reduce_mutation_depth,
            use_lightweight_hallucination=use_lightweight_hallucination,
        )
    
    def _determine_inference_mode(
        self,
        priority_score: float,
        cluster_load: float,
    ) -> InferenceMode:
        """Determine inference mode based on priority and load."""
        # High priority jobs get full mode
        if priority_score >= 0.7:
            return InferenceMode.FULL
        
        # Low load allows full mode
        if cluster_load < 0.5:
            return InferenceMode.FULL
        
        # Medium load - use standard
        if cluster_load < 0.8:
            return InferenceMode.STANDARD
        
        # High load - use lightweight
        return InferenceMode.LIGHTWEIGHT
    
    def _determine_batch_size(
        self,
        total_samples: int,
        cluster_load: float,
        inference_mode: InferenceMode,
    ) -> int:
        """Determine optimal batch size."""
        if inference_mode == InferenceMode.LIGHTWEIGHT:
            # Lightweight mode allows larger batches
            if cluster_load < 0.5:
                return min(16, max(4, total_samples // 10))
            return min(8, max(2, total_samples // 20))
        
        if inference_mode == InferenceMode.FULL:
            # Full mode requires smaller batches
            return min(4, max(1, total_samples // 50))
        
        # Standard mode
        return min(8, max(2, total_samples // 25))
    
    def _determine_gpu_and_pool(
        self,
        model_size: str,
        required_memory_mb: int,
    ) -> tuple[str, NodePoolType]:
        """Determine GPU type and node pool."""
        # Map model size to requirements
        model_requirements = {
            "7b": {"gpu_type": "v100", "memory_gb": 16},
            "13b": {"gpu_type": "v100", "memory_gb": 30},
            "30b": {"gpu_type": "a100", "memory_gb": 60},
            "70b": {"gpu_type": "a100", "memory_gb": 120},
        }
        
        req = model_requirements.get(model_size, model_requirements["7b"])
        
        # Check if more memory is required
        if required_memory_mb > req["memory_gb"] * 1024:
            return (req["gpu_type"], NodePoolType.GPU_HIGH_MEMORY)
        
        # Select node pool based on GPU type
        if req["gpu_type"] == "h100":
            return (req["gpu_type"], NodePoolType.GPU_HOPPER)
        elif req["gpu_type"] == "a100":
            return (req["gpu_type"], NodePoolType.GPU_AMPERE)
        else:
            return (req["gpu_type"], NodePoolType.GPU_STANDARD)
    
    def _determine_use_spot(
        self,
        priority_score: float,
        cluster_load: float,
    ) -> bool:
        """Determine whether to use spot instances."""
        # Don't use spot for critical jobs
        if priority_score >= 0.9:
            return False
        
        # Use spot by default if enabled
        if self.enable_spot_by_default:
            # But avoid spot during high load (may get preempted)
            if cluster_load > 0.9:
                return False
            return True
        
        return False
    
    def _should_enable_batching(
        self,
        cluster_load: float,
        total_samples: int,
    ) -> bool:
        """Determine if batching should be enabled."""
        # Enable for large jobs when load is moderate
        return cluster_load < 0.7 and total_samples > 50
    
    def _should_reduce_mutation_depth(
        self,
        priority_score: float,
        cluster_load: float,
    ) -> bool:
        """Determine if mutation depth should be reduced."""
        # Reduce for low priority jobs or high load
        return priority_score < 0.4 or cluster_load > 0.8
    
    def _should_use_lightweight_hallucination(
        self,
        cluster_load: float,
    ) -> bool:
        """Determine if lightweight hallucination detection should be used."""
        return cluster_load > 0.85
    
    def update_cluster_state(
        self,
        cluster_load: float,
        available_gpu_count: int,
    ):
        """
        Update cluster state for allocation decisions.
        
        Args:
            cluster_load: Current cluster load (0-1)
            available_gpu_count: Number of available GPUs
        """
        self._cluster_load = cluster_load
        self._available_gpu_count = available_gpu_count
    
    def get_allocation_cost_per_hour(
        self,
        allocation: AllocationDecision,
    ) -> float:
        """Calculate hourly cost for an allocation decision."""
        node_spec = self.NODE_POOL_SPECS.get(
            NodePoolType(allocation.node_pool),
            self.NODE_POOL_SPECS[NodePoolType.GPU_STANDARD]
        )
        
        base_cost = node_spec.get("cost_per_hour", 2.48)
        
        # Apply spot discount if using spot
        if allocation.use_spot:
            base_cost *= 0.35  # ~65% discount
        
        return base_cost


# Global instance
_resource_allocator: Optional[ResourceAllocator] = None


def get_resource_allocator() -> ResourceAllocator:
    """Get or create the global ResourceAllocator instance."""
    global _resource_allocator
    if _resource_allocator is None:
        _resource_allocator = ResourceAllocator()
    return _resource_allocator


__all__ = [
    "ResourceAllocator",
    "AllocationDecision",
    "InferenceMode",
    "NodePoolType",
    "get_resource_allocator",
]