backend/scheduler/priority_engine.py

"""
Priority Engine for Intelligent Job Scheduling

Computes priority scores for jobs based on multiple factors:
- Tenant plan weight (Enterprise > Pro > Basic > Free)
- SLA deadline urgency
- Cost sensitivity (inverse of budget buffer)
- Job waiting time (urgency factor)

Score = w_p * P + w_s * SLA + w_c * CostSensitivity + w_u * Urgency
"""

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

from pydantic import BaseModel

# Plan type weights as defined in Day 4 requirements
PLAN_WEIGHTS = {
    "free": 0.2,
    "basic": 0.4,
    "pro": 0.5,
    "enterprise": 1.0,
}


class PriorityQueueType(str, Enum):
    """Priority queue types for job classification."""
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"


class JobPriorityScore(BaseModel):
    """Complete priority score breakdown for a job."""
    job_id: uuid.UUID
    tenant_id: uuid.UUID
    
    # Component scores (all normalized to [0, 1])
    plan_score: float = 0.0  # Based on tenant plan
    sla_score: float = 0.0  # Based on deadline urgency
    cost_sensitivity_score: float = 0.0  # Based on budget buffer
    urgency_score: float = 0.0  # Based on wait time
    
    # Weighted total
    total_score: float = 0.0
    
    # Assigned queue
    queue_type: PriorityQueueType = PriorityQueueType.MEDIUM
    
    # Metadata
    calculated_at: datetime = None
    estimated_gpu_hours: float = 0.0
    estimated_cost: float = 0.0
    
    def __init__(self, **data):
        if "calculated_at" not in data or data["calculated_at"] is None:
            data["calculated_at"] = datetime.utcnow()
        super().__init__(**data)


class PriorityEngine:
    """
    Multi-factor priority scoring engine for job scheduling.
    
    Computes priority scores based on:
    - Tenant plan tier (Enterprise gets scheduling preference)
    - SLA deadline urgency (closer deadlines = higher priority)
    - Cost sensitivity (lower budget buffer = higher priority)
    - Job waiting time (longer waits = higher priority)
    """
    
    # Default weights for priority components
    DEFAULT_WEIGHTS = {
        "plan": 0.30,  # w_p
        "sla": 0.25,   # w_s  
        "cost_sensitivity": 0.20,  # w_c
        "urgency": 0.25,  # w_u
    }
    
    # Queue score thresholds
    HIGH_QUEUE_THRESHOLD = 0.7
    MEDIUM_QUEUE_THRESHOLD = 0.4
    
    def __init__(
        self,
        weights: Optional[Dict[str, float]] = None,
        avg_inference_time_ms: float = 500.0,
    ):
        """
        Initialize priority engine.
        
        Args:
            weights: Custom weights for priority components.
                    If None, uses DEFAULT_WEIGHTS.
            avg_inference_time_ms: Average inference time in milliseconds.
                                   Used for GPU hour estimation.
        """
        self.weights = weights or self.DEFAULT_WEIGHTS
        self.avg_inference_time_ms = avg_inference_time_ms
        
        # Ensure weights sum to 1.0
        total = sum(self.weights.values())
        if total != 1.0:
            # Normalize weights
            self.weights = {k: v / total for k, v in self.weights.items()}
    
    def calculate_priority_score(
        self,
        job_id: uuid.UUID,
        tenant_id: uuid.UUID,
        plan_type: str,
        total_samples: int,
        deadline_timestamp: Optional[datetime] = None,
        budget_limit: Optional[float] = None,
        used_budget: Optional[float] = None,
        submitted_at: Optional[datetime] = None,
        current_time: Optional[datetime] = None,
    ) -> JobPriorityScore:
        """
        Calculate comprehensive priority score for a job.
        
        Args:
            job_id: Unique job identifier
            tenant_id: Tenant identifier
            plan_type: Tenant plan type (free, basic, pro, enterprise)
            total_samples: Number of samples to process
            deadline_timestamp: Optional deadline for SLA enforcement
            budget_limit: Optional budget limit for the tenant
            used_budget: Amount of budget already used
            submitted_at: When the job was submitted
            current_time: Current timestamp (defaults to now)
            
        Returns:
            JobPriorityScore with all component scores and total
        """
        now = current_time or datetime.utcnow()
        
        # 1. Calculate plan score (based on tenant tier)
        plan_score = self._calculate_plan_score(plan_type)
        
        # 2. Calculate SLA score (based on deadline urgency)
        sla_score = self._calculate_sla_score(deadline_timestamp, now)
        
        # 3. Calculate cost sensitivity (based on budget buffer)
        cost_sensitivity = self._calculate_cost_sensitivity(
            budget_limit, used_budget
        )
        
        # 4. Calculate urgency score (based on wait time)
        urgency_score = self._calculate_urgency_score(submitted_at, now)
        
        # Calculate weighted total
        total_score = (
            self.weights["plan"] * plan_score +
            self.weights["sla"] * sla_score +
            self.weights["cost_sensitivity"] * cost_sensitivity +
            self.weights["urgency"] * urgency_score
        )
        
        # Determine queue type
        queue_type = self._determine_queue_type(total_score)
        
        # Estimate GPU hours and cost
        estimated_gpu_hours = self._estimate_gpu_hours(total_samples)
        estimated_cost = self._estimate_job_cost(estimated_gpu_hours)
        
        return JobPriorityScore(
            job_id=job_id,
            tenant_id=tenant_id,
            plan_score=plan_score,
            sla_score=sla_score,
            cost_sensitivity_score=cost_sensitivity,
            urgency_score=urgency_score,
            total_score=total_score,
            queue_type=queue_type,
            calculated_at=now,
            estimated_gpu_hours=estimated_gpu_hours,
            estimated_cost=estimated_cost,
        )
    
    def _calculate_plan_score(self, plan_type: str) -> float:
        """
        Calculate score based on tenant plan type.
        
        Returns normalized score based on PLAN_WEIGHTS.
        """
        return PLAN_WEIGHTS.get(plan_type.lower(), PLAN_WEIGHTS["free"])
    
    def _calculate_sla_score(
        self,
        deadline_timestamp: Optional[datetime],
        current_time: datetime,
    ) -> float:
        """
        Calculate SLA urgency score.
        
        Jobs with imminent deadlines get higher priority.
        If no deadline, returns 0.5 (medium priority).
        """
        if deadline_timestamp is None:
            return 0.5  # No deadline = medium priority
        
        time_remaining = deadline_timestamp - current_time
        
        if time_remaining.total_seconds() <= 0:
            # Deadline passed - maximum urgency
            return 1.0
        
        # Score decreases as more time remains
        # Map to [0, 1] where 1 is urgent (very little time)
        max_urgency_hours = 1.0  # 1 hour = maximum urgency
        min_urgency_hours = 24.0  # 24 hours = minimum urgency
        
        hours_remaining = time_remaining.total_seconds() / 3600
        
        if hours_remaining <= max_urgency_hours:
            return 1.0
        elif hours_remaining >= min_urgency_hours:
            return 0.0
        else:
            # Linear interpolation
            return 1.0 - (hours_remaining - max_urgency_hours) / (min_urgency_hours - max_urgency_hours)
    
    def _calculate_cost_sensitivity(
        self,
        budget_limit: Optional[float],
        used_budget: Optional[float],
    ) -> float:
        """
        Calculate cost sensitivity based on budget buffer.
        
        Returns higher score when budget is nearly exhausted.
        """
        if budget_limit is None or used_budget is None:
            return 0.5  # No budget tracking = medium sensitivity
        
        if budget_limit <= 0:
            return 1.0  # No budget = maximum sensitivity
        
        # Calculate buffer remaining
        buffer_remaining = budget_limit - used_budget
        buffer_percent = buffer_remaining / budget_limit
        
        if buffer_percent <= 0:
            return 1.0  # Budget exhausted = maximum sensitivity
        elif buffer_percent >= 1.0:
            return 0.0  # Full budget = no sensitivity
        else:
            # Higher sensitivity as buffer decreases
            return 1.0 - buffer_percent
    
    def _calculate_urgency_score(
        self,
        submitted_at: Optional[datetime],
        current_time: datetime,
    ) -> float:
        """
        Calculate urgency based on job waiting time.
        
        Jobs waiting longer get higher priority to prevent starvation.
        """
        if submitted_at is None:
            return 0.5  # No submission time = medium urgency
        
        wait_time = current_time - submitted_at
        wait_hours = wait_time.total_seconds() / 3600
        
        # Maximum urgency after 4 hours of waiting
        # Minimum urgency (0) for jobs submitted < 5 minutes ago
        if wait_hours >= 4.0:
            return 1.0
        elif wait_hours <= 0.083:  # 5 minutes
            return 0.0
        else:
            # Linear interpolation
            return (wait_hours - 0.083) / (4.0 - 0.083)
    
    def _determine_queue_type(self, total_score: float) -> PriorityQueueType:
        """Determine which priority queue a job belongs to."""
        if total_score >= self.HIGH_QUEUE_THRESHOLD:
            return PriorityQueueType.HIGH
        elif total_score >= self.MEDIUM_QUEUE_THRESHOLD:
            return PriorityQueueType.MEDIUM
        else:
            return PriorityQueueType.LOW
    
    def _estimate_gpu_hours(self, total_samples: int) -> float:
        """
        Estimate GPU hours required for a job.
        
        GPUHours = (Samples * AvgInferenceTime) / 3600
        """
        # Average time per sample in seconds
        avg_time_per_sample_seconds = self.avg_inference_time_ms / 1000.0
        
        # Total time in seconds
        total_time_seconds = total_samples * avg_time_per_sample_seconds
        
        # Convert to hours
        gpu_hours = total_time_seconds / 3600.0
        
        return gpu_hours
    
    def _estimate_job_cost(self, estimated_gpu_hours: float) -> float:
        """
        Estimate cost based on GPU hours.
        
        Cost = GPUHours * CostPerGPUHour
        Default: $3.00 per GPU hour (AWS p3.2xlarge spot price)
        """
        cost_per_gpu_hour = 3.00  # Default cost
        return estimated_gpu_hours * cost_per_gpu_hour
    
    def should_promote_priority(
        self,
        current_score: JobPriorityScore,
        current_time: Optional[datetime] = None,
    ) -> bool:
        """
        Check if a job should be promoted to higher priority.
        
        Promotes if:
        - Deadline is approaching and job is not in high queue
        - Job has been waiting too long (starvation prevention)
        """
        now = current_time or datetime.utcnow()
        
        # Calculate time since submission
        if current_score.calculated_at:
            wait_time = now - current_score.calculated_at
            wait_hours = wait_time.total_seconds() / 3600
            
            # Promote if waiting > 2 hours in low queue
            if (current_score.queue_type == PriorityQueueType.LOW and 
                wait_hours > 2.0):
                return True
            
            # Promote if SLA is at risk
            if (current_score.sla_score > 0.8 and 
                current_score.queue_type != PriorityQueueType.HIGH):
                return True
        
        return False


# Global instance
_priority_engine: Optional[PriorityEngine] = None


def get_priority_engine() -> PriorityEngine:
    """Get or create the global PriorityEngine instance."""
    global _priority_engine
    if _priority_engine is None:
        _priority_engine = PriorityEngine()
    return _priority_engine


__all__ = [
    "PriorityEngine",
    "JobPriorityScore",
    "PriorityQueueType",
    "PLAN_WEIGHTS",
    "get_priority_engine",
]