backend/scheduler/usage_tracker.py

"""
Usage Tracker for Tenant Budget Management

Tracks tenant resource usage:
- Total GPU hours consumed
- Total cost incurred
- Budget limit enforcement
- Rolling billing period tracking

Supports budget enforcement policies:
- REJECT: Reject new jobs when budget exceeded
- DOWNGRADE: Move to lower priority queue
- THROTTLE: Force throughput mode
"""

import uuid
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum
from typing import Any, Dict, List, Optional

from pydantic import BaseModel


class BudgetEnforcementPolicy(str, Enum):
    """Budget enforcement policies."""
    REJECT = "reject"  # Reject new jobs
    DOWNGRADE = "downgrade"  # Move to lower priority
    THROTTLE = "throttle"  # Force throughput mode


class UsageRecord(BaseModel):
    """Record of resource usage for a job."""
    job_id: uuid.UUID
    tenant_id: uuid.UUID
    
    # Usage metrics
    gpu_hours: float = 0.0
    cost: float = 0.0
    samples_processed: int = 0
    
    # Timestamps
    started_at: datetime
    completed_at: Optional[datetime] = None


class TenantUsage(BaseModel):
    """Aggregated usage for a tenant."""
    tenant_id: uuid.UUID
    
    # Current period usage
    gpu_hours_current: float = 0.0
    cost_current: float = 0.0
    jobs_completed: int = 0
    samples_processed: int = 0
    
    # Budget info
    budget_limit: Optional[float] = None
    budget_used_percent: float = 0.0
    
    # Period info
    period_start: datetime
    period_end: datetime
    
    # Enforcement
    enforcement_policy: str = "reject"


class UsageTracker:
    """
    Tracks tenant resource usage and enforces budget limits.
    
    Features:
    - Per-tenant usage tracking
    - Rolling billing period (default: 30 days)
    - Budget limit enforcement
    - Usage reporting
    """
    
    # Default billing period
    DEFAULT_BILLING_PERIOD_DAYS = 30
    
    # Default budget enforcement
    DEFAULT_ENFORCEMENT_POLICY = BudgetEnforcementPolicy.REJECT
    
    def __init__(
        self,
        billing_period_days: int = DEFAULT_BILLING_PERIOD_DAYS,
        default_budget_limit: Optional[float] = None,
        default_enforcement: BudgetEnforcementPolicy = DEFAULT_ENFORCEMENT_POLICY,
    ):
        """
        Initialize usage tracker.
        
        Args:
            billing_period_days: Rolling billing period in days
            default_budget_limit: Default budget limit for new tenants
            default_enforcement: Default enforcement policy
        """
        self.billing_period_days = billing_period_days
        self.default_budget_limit = default_budget_limit
        self.default_enforcement = default_enforcement
        
        # In-memory storage (would use database in production)
        self._tenant_usage: Dict[uuid.UUID, TenantUsage] = {}
        self._job_records: Dict[uuid.UUID, UsageRecord] = {}
        
        # Tenant budget settings
        self._tenant_budgets: Dict[uuid.UUID, Dict[str, Any]] = {}
        
        # Initialize default budgets for known tenants
        self._initialize_default_budgets()
    
    def _initialize_default_budgets(self):
        """Initialize default budget configurations."""
        # Default budgets by plan type
        self._plan_defaults = {
            "free": {
                "budget_limit": 10.0,  # $10/month
                "enforcement": BudgetEnforcementPolicy.REJECT,
            },
            "basic": {
                "budget_limit": 100.0,  # $100/month
                "enforcement": BudgetEnforcementPolicy.DOWNGRADE,
            },
            "pro": {
                "budget_limit": 1000.0,  # $1000/month
                "enforcement": BudgetEnforcementPolicy.THROTTLE,
            },
            "enterprise": {
                "budget_limit": None,  # No limit
                "enforcement": BudgetEnforcementPolicy.THROTTLE,
            },
        }
    
    def register_tenant(
        self,
        tenant_id: uuid.UUID,
        plan_type: str = "free",
    ):
        """
        Register a new tenant with default budget.
        
        Args:
            tenant_id: Tenant identifier
            plan_type: Tenant plan type
        """
        now = datetime.utcnow()
        
        # Get plan defaults
        plan_defaults = self._plan_defaults.get(
            plan_type.lower(), self._plan_defaults["free"]
        )
        
        # Create tenant usage record
        self._tenant_usage[tenant_id] = TenantUsage(
            tenant_id=tenant_id,
            period_start=now,
            period_end=now + timedelta(days=self.billing_period_days),
            budget_limit=plan_defaults.get("budget_limit"),
            enforcement_policy=plan_defaults.get(
                "enforcement", BudgetEnforcementPolicy.REJECT
            ).value,
        )
        
        # Store budget settings
        self._tenant_budgets[tenant_id] = {
            "plan_type": plan_type,
            "budget_limit": plan_defaults.get("budget_limit"),
            "enforcement": plan_defaults.get(
                "enforcement", BudgetEnforcementPolicy.REJECT
            ),
        }
    
    def record_job_start(
        self,
        job_id: uuid.UUID,
        tenant_id: uuid.UUID,
    ):
        """
        Record the start of a job.
        
        Args:
            job_id: Job identifier
            tenant_id: Tenant identifier
        """
        # Ensure tenant is registered
        if tenant_id not in self._tenant_usage:
            self.register_tenant(tenant_id)
        
        # Create usage record
        self._job_records[job_id] = UsageRecord(
            job_id=job_id,
            tenant_id=tenant_id,
            started_at=datetime.utcnow(),
        )
    
    def record_job_completion(
        self,
        job_id: uuid.UUID,
        gpu_hours: float,
        cost: float,
        samples_processed: int,
    ):
        """
        Record job completion and update tenant usage.
        
        Args:
            job_id: Job identifier
            gpu_hours: GPU hours consumed
            cost: Total cost
            samples_processed: Number of samples processed
        """
        record = self._job_records.get(job_id)
        if record is None:
            return
        
        # Update record
        record.gpu_hours = gpu_hours
        record.cost = cost
        record.samples_processed = samples_processed
        record.completed_at = datetime.utcnow()
        
        # Update tenant usage
        tenant_id = record.tenant_id
        if tenant_id in self._tenant_usage:
            usage = self._tenant_usage[tenant_id]
            usage.gpu_hours_current += gpu_hours
            usage.cost_current += cost
            usage.jobs_completed += 1
            usage.samples_processed += samples_processed
            
            # Update budget used percentage
            if usage.budget_limit and usage.budget_limit > 0:
                usage.budget_used_percent = (
                    usage.cost_current / usage.budget_limit * 100
                )
    
    def check_budget(
        self,
        tenant_id: uuid.UUID,
        estimated_cost: float,
    ) -> tuple[bool, Optional[str]]:
        """
        Check if a job can be submitted within budget.
        
        Args:
            tenant_id: Tenant identifier
            estimated_cost: Estimated cost for the job
            
        Returns:
            Tuple of (allowed, enforcement_action)
        """
        # Ensure tenant is registered
        if tenant_id not in self._tenant_usage:
            self.register_tenant(tenant_id)
        
        usage = self._tenant_usage[tenant_id]
        
        # No budget limit - allow
        if usage.budget_limit is None:
            return True, None
        
        # Check if would exceed budget
        projected_total = usage.cost_current + estimated_cost
        
        if projected_total <= usage.budget_limit:
            return True, None
        
        # Budget exceeded - determine enforcement action
        policy = BudgetEnforcementPolicy(usage.enforcement_policy)
        
        if policy == BudgetEnforcementPolicy.REJECT:
            return False, "reject"
        elif policy == BudgetEnforcementPolicy.DOWNGRADE:
            return True, "downgrade"
        else:  # THROTTLE
            return True, "throttle"
    
    def get_tenant_usage(
        self,
        tenant_id: uuid.UUID,
    ) -> Optional[TenantUsage]:
        """
        Get current usage for a tenant.
        
        Args:
            tenant_id: Tenant identifier
            
        Returns:
            TenantUsage record or None
        """
        return self._tenant_usage.get(tenant_id)
    
    def get_all_tenant_usage(self) -> List[TenantUsage]:
        """Get usage for all tenants."""
        return list(self._tenant_usage.values())
    
    def update_budget(
        self,
        tenant_id: uuid.UUID,
        budget_limit: float,
        enforcement: BudgetEnforcementPolicy,
    ):
        """
        Update tenant budget settings.
        
        Args:
            tenant_id: Tenant identifier
            budget_limit: New budget limit
            enforcement: Enforcement policy
        """
        if tenant_id not in self._tenant_usage:
            self.register_tenant(tenant_id)
        
        usage = self._tenant_usage[tenant_id]
        usage.budget_limit = budget_limit
        usage.enforcement_policy = enforcement.value
        
        # Update percentage
        if budget_limit > 0:
            usage.budget_used_percent = (
                usage.cost_current / budget_limit * 100
            )
        
        # Update settings
        self._tenant_budgets[tenant_id] = {
            "budget_limit": budget_limit,
            "enforcement": enforcement,
        }
    
    def reset_usage(
        self,
        tenant_id: uuid.UUID,
    ):
        """
        Reset usage for a tenant (typically monthly).
        
        Args:
            tenant_id: Tenant identifier
        """
        if tenant_id not in self._tenant_usage:
            return
        
        now = datetime.utcnow()
        usage = self._tenant_usage[tenant_id]
        
        # Reset current period
        usage.gpu_hours_current = 0.0
        usage.cost_current = 0.0
        usage.jobs_completed = 0
        usage.samples_processed = 0
        usage.budget_used_percent = 0.0
        usage.period_start = now
        usage.period_end = now + timedelta(days=self.billing_period_days)
    
    def get_cost_metrics(
        self,
        tenant_id: uuid.UUID,
    ) -> Dict[str, Any]:
        """
        Get detailed cost metrics for a tenant.
        
        Args:
            tenant_id: Tenant identifier
            
        Returns:
            Dictionary with cost metrics
        """
        usage = self._tenant_usage.get(tenant_id)
        if usage is None:
            return {}
        
        # Calculate additional metrics
        avg_cost_per_job = (
            usage.cost_current / usage.jobs_completed
            if usage.jobs_completed > 0 else 0.0
        )
        
        avg_cost_per_sample = (
            usage.cost_current / usage.samples_processed
            if usage.samples_processed > 0 else 0.0
        )
        
        days_remaining = (
            usage.period_end - datetime.utcnow()
        ).total_seconds() / 86400
        
        projected_monthly_cost = (
            usage.cost_current / 
            max(1, (datetime.utcnow() - usage.period_start).total_seconds() / 86400)
            * 30
        )
        
        return {
            "tenant_id": str(tenant_id),
            "gpu_hours": usage.gpu_hours_current,
            "total_cost": usage.cost_current,
            "jobs_completed": usage.jobs_completed,
            "samples_processed": usage.samples_processed,
            "budget_limit": usage.budget_limit,
            "budget_used_percent": usage.budget_used_percent,
            "avg_cost_per_job": avg_cost_per_job,
            "avg_cost_per_sample": avg_cost_per_sample,
            "days_remaining": max(0, days_remaining),
            "projected_monthly_cost": projected_monthly_cost,
            "period_start": usage.period_start.isoformat(),
            "period_end": usage.period_end.isoformat(),
        }
    
    def get_efficiency_metric(
        self,
        tenant_id: uuid.UUID,
    ) -> float:
        """
        Calculate economic efficiency: Insights / GPUHours.
        
        Args:
            tenant_id: Tenant identifier
            
        Returns:
            Efficiency score (insights per GPU hour)
        """
        usage = self._tenant_usage.get(tenant_id)
        if usage is None or usage.gpu_hours_current <= 0:
            return 0.0
        
        return usage.samples_processed / usage.gpu_hours_current


# Global instance
_usage_tracker: Optional[UsageTracker] = None


def get_usage_tracker() -> UsageTracker:
    """Get or create the global UsageTracker instance."""
    global _usage_tracker
    if _usage_tracker is None:
        _usage_tracker = UsageTracker()
    return _usage_tracker


__all__ = [
    "UsageTracker",
    "UsageRecord",
    "TenantUsage",
    "BudgetEnforcementPolicy",
    "get_usage_tracker",
]