/mlflow/gateway/budget_tracker/__init__.py

"""Budget tracker for AI Gateway cost management.

Provides an abstract BudgetTracker interface and window computation helpers.
The concrete InMemoryBudgetTracker lives in ``budget_tracker.in_memory``.
"""

from __future__ import annotations

import threading
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone

from mlflow.entities.gateway_budget_policy import (
    BudgetDuration,
    BudgetDurationUnit,
    BudgetTargetScope,
    GatewayBudgetPolicy,
)
from mlflow.environment_variables import (
    MLFLOW_GATEWAY_BUDGET_REDIS_URL,
    MLFLOW_GATEWAY_BUDGET_REFRESH_INTERVAL,
)
from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME

_EPOCH = datetime.fromtimestamp(0, tz=timezone.utc)
# Sunday-aligned epoch for WEEKS windows (Dec 28, 1969 is the Sunday before Jan 1, 1970)
_EPOCH_SUNDAY = _EPOCH - timedelta(days=4)

# Module-level singleton
_budget_tracker: BudgetTracker | None = None
_tracker_lock = threading.Lock()


def get_budget_tracker() -> BudgetTracker:
    """Get or create the module-level BudgetTracker singleton."""
    global _budget_tracker
    if _budget_tracker is None:
        with _tracker_lock:
            if _budget_tracker is None:
                if redis_url := MLFLOW_GATEWAY_BUDGET_REDIS_URL.get():
                    from mlflow.gateway.budget_tracker.redis import RedisBudgetTracker

                    _budget_tracker = RedisBudgetTracker(_redis_url=redis_url)
                else:
                    from mlflow.gateway.budget_tracker.in_memory import InMemoryBudgetTracker

                    _budget_tracker = InMemoryBudgetTracker()
    return _budget_tracker


@dataclass
class BudgetWindow:
    """Tracks spend within a single fixed time window for one policy."""

    policy: GatewayBudgetPolicy
    window_start: datetime
    window_end: datetime
    cumulative_spend: float = 0.0
    exceeded: bool = False


class BudgetTracker(ABC):
    """Abstract base class for budget trackers.

    Defines the interface for tracking cumulative cost per budget policy
    within fixed time windows. Concrete implementations may store state
    in memory, Redis, or other backends.
    """

    _last_refresh_time: float = float("-inf")

    def needs_refresh(self) -> bool:
        """Check whether policies should be re-fetched from the database."""
        return (
            time.monotonic() - self._last_refresh_time
        ) >= MLFLOW_GATEWAY_BUDGET_REFRESH_INTERVAL.get()

    def mark_refreshed(self) -> None:
        """Mark the tracker as just refreshed."""
        self._last_refresh_time = time.monotonic()

    def invalidate(self) -> None:
        """Reset the refresh timer so the next needs_refresh() call returns True."""
        self._last_refresh_time = float("-inf")

    @abstractmethod
    def refresh_policies(self, policies: list[GatewayBudgetPolicy]) -> list[BudgetWindow]:
        """Load or refresh policies from the database.

        Preserves accumulated cost for unchanged windows. Removes windows
        for policies that no longer exist.

        Returns:
            Windows that should be synced against authoritative trace data.
        """

    @abstractmethod
    def record_cost(
        self,
        cost_usd: float,
        workspace: str | None = None,
    ) -> list[BudgetWindow]:
        """Record a cost against all applicable policies.

        Args:
            cost_usd: The cost in USD to record.
            workspace: The workspace the request was made from (None for default).

        Returns:
            List of windows that were newly exceeded (limit exceeded for the first
            time in this window). Used to trigger webhook alerts.
        """

    @abstractmethod
    def should_reject_request(
        self,
        workspace: str | None = None,
    ) -> tuple[bool, BudgetWindow | None]:
        """Check if any REJECT-capable policy is exceeded.

        Args:
            workspace: The workspace to check against.

        Returns:
            Tuple of (exceeded, window). If exceeded is True, window is the
            first exceeded window found.
        """

    @abstractmethod
    def backfill_spend(self, spend_by_policy: dict[str, float]) -> None:
        """Set cumulative spend on windows from historical data.

        Used to seed newly created windows with spend from trace data
        so that budget tracking survives server restarts.

        Args:
            spend_by_policy: Dict mapping budget_policy_id to historical spend amount.
        """

    @abstractmethod
    def get_all_windows(self) -> list[BudgetWindow]:
        """Get the current window info for all tracked policies."""

    @abstractmethod
    def _get_window_info(self, budget_policy_id: str) -> BudgetWindow | None:
        """Get the current window info for a policy (for payload construction)."""


def _compute_window_start(
    duration: BudgetDuration,
    now: datetime,
) -> datetime:
    """Compute the start of the current fixed window for a given policy.

    Windows are aligned to:
    - MINUTES: aligned to epoch minutes
    - HOURS: aligned to epoch hours (e.g., duration.value=2 → 0:00, 2:00, 4:00, …)
    - DAYS: aligned to epoch days (e.g., duration.value=7 → weekly from epoch)
    - WEEKS: aligned to Sunday-based weeks (e.g., duration.value=1 → every Sunday from epoch)
    - MONTHS: aligned to first of months (e.g., duration.value=3 → Jan 1, Apr 1, Jul 1, …)
    """
    if duration.value <= 0:
        raise ValueError(f"duration.value must be positive, got {duration.value}")

    if duration.unit == BudgetDurationUnit.MINUTES:
        epoch = _EPOCH
        minutes_since_epoch = (now - epoch).total_seconds() / 60
        window_index = int(minutes_since_epoch) // duration.value
        window_start_minutes = window_index * duration.value
        return epoch + timedelta(minutes=window_start_minutes)

    elif duration.unit == BudgetDurationUnit.HOURS:
        epoch = _EPOCH
        hours_since_epoch = (now - epoch).total_seconds() / 3600
        window_index = int(hours_since_epoch) // duration.value
        window_start_hours = window_index * duration.value
        return epoch + timedelta(hours=window_start_hours)

    elif duration.unit == BudgetDurationUnit.DAYS:
        epoch = _EPOCH
        days_since_epoch = (now - epoch).days
        window_index = days_since_epoch // duration.value
        window_start_days = window_index * duration.value
        return epoch + timedelta(days=window_start_days)

    elif duration.unit == BudgetDurationUnit.WEEKS:
        days_since_sunday_epoch = (now - _EPOCH_SUNDAY).days
        window_index = days_since_sunday_epoch // (7 * duration.value)
        window_start_days = window_index * (7 * duration.value)
        return _EPOCH_SUNDAY + timedelta(days=window_start_days)

    elif duration.unit == BudgetDurationUnit.MONTHS:
        year = now.year
        month = now.month
        total_months = (year - 1970) * 12 + (month - 1)
        window_index = total_months // duration.value
        window_start_months = window_index * duration.value
        start_year = 1970 + window_start_months // 12
        start_month = (window_start_months % 12) + 1
        return datetime(start_year, start_month, 1, tzinfo=timezone.utc)

    raise ValueError(f"Unknown duration type: {duration.unit}")


def _compute_window_end(
    duration: BudgetDuration,
    window_start: datetime,
) -> datetime:
    """Compute the end of the current fixed window."""
    if duration.unit == BudgetDurationUnit.MINUTES:
        return window_start + timedelta(minutes=duration.value)
    elif duration.unit == BudgetDurationUnit.HOURS:
        return window_start + timedelta(hours=duration.value)
    elif duration.unit == BudgetDurationUnit.DAYS:
        return window_start + timedelta(days=duration.value)
    elif duration.unit == BudgetDurationUnit.WEEKS:
        return window_start + timedelta(weeks=duration.value)
    elif duration.unit == BudgetDurationUnit.MONTHS:
        year = window_start.year
        month = window_start.month + duration.value
        while month > 12:
            month -= 12
            year += 1
        return datetime(year, month, 1, tzinfo=timezone.utc)

    raise ValueError(f"Unknown duration type: {duration.unit}")


def _policy_applies(policy: GatewayBudgetPolicy, workspace: str | None) -> bool:
    """Check if a policy applies to a given workspace.

    GLOBAL policies apply to all workspaces. WORKSPACE policies only apply
    when the request workspace matches the policy's workspace.
    """
    if policy.target_scope == BudgetTargetScope.GLOBAL:
        return True
    effective_workspace = workspace or DEFAULT_WORKSPACE_NAME
    return policy.workspace == effective_workspace