/mlflow/gateway/budget.py

"""Budget tracking and enforcement for the MLflow Gateway.

This module provides budget-related functions for recording costs, refreshing policies,
firing exceeded-budget webhooks, and creating on_complete callbacks for budget recording.
"""

import logging

from fastapi import HTTPException

import mlflow
from mlflow.entities import SpanType
from mlflow.entities.gateway_budget_policy import BudgetAction, BudgetTargetScope
from mlflow.entities.trace_location import MlflowExperimentLocation
from mlflow.entities.webhook import WebhookAction, WebhookEntity, WebhookEvent
from mlflow.gateway.budget_tracker import BudgetWindow, get_budget_tracker
from mlflow.gateway.tracing_utils import (
    _gateway_span_attributes,
    _gateway_span_name,
    _get_model_span_info,
)
from mlflow.store.tracking.gateway.entities import GatewayEndpointConfig
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
from mlflow.tracing.constant import CostKey, SpanAttributeKey
from mlflow.tracing.utils import calculate_cost_by_model_and_token_usage
from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME
from mlflow.webhooks.delivery import deliver_webhook
from mlflow.webhooks.types import BudgetPolicyExceededPayload

_logger = logging.getLogger(__name__)


def calculate_existing_cost_for_windows(
    store: SqlAlchemyStore, windows: list[BudgetWindow]
) -> dict[str, float]:
    """Calculate spend for budget windows from trace history.

    Queries historical trace cost data for each window so that budget tracking
    accounts for spend across all gateway workers and survives server restarts.

    Returns:
        Dict mapping budget_policy_id to historical spend amount.
    """
    result: dict[str, float] = {}
    if not windows:
        return result

    for window in windows:
        try:
            start_ms = int(window.window_start.timestamp() * 1000)
            end_ms = int(window.window_end.timestamp() * 1000)
            workspace = (
                window.policy.workspace
                if window.policy.target_scope == BudgetTargetScope.WORKSPACE
                else None
            )
            spend = store.sum_gateway_trace_cost(
                start_time_ms=start_ms,
                end_time_ms=end_ms,
                workspace=workspace,
            )
            if spend > 0:
                result[window.policy.budget_policy_id] = spend
        except Exception:
            _logger.debug(
                "Failed to calculate existing cost for policy %s",
                window.policy.budget_policy_id,
                exc_info=True,
            )
    return result


def maybe_refresh_budget_policies(store: SqlAlchemyStore) -> None:
    """Refresh budget policies from the database if stale."""
    tracker = get_budget_tracker()
    if tracker.needs_refresh():
        try:
            policies = store.list_budget_policies()
            windows = tracker.refresh_policies(policies)
            existing_spend = calculate_existing_cost_for_windows(store, windows)
            tracker.backfill_spend(existing_spend)
        except Exception:
            _logger.debug("Failed to refresh budget policies", exc_info=True)


def _compute_cost_from_child_spans(trace_id: str) -> float:
    """Sum total cost across child spans.

    Prefers ``LLM_COST`` if already set on the span (computed at span.end()),
    otherwise falls back to calculating from MODEL + CHAT_USAGE via LiteLLM.
    """
    total = 0.0
    for info in _get_model_span_info(trace_id):
        if llm_cost := info.attributes.get(SpanAttributeKey.LLM_COST):
            total += llm_cost.get(CostKey.TOTAL_COST, 0.0)
        else:
            model_name = info.attributes.get(SpanAttributeKey.MODEL)
            usage = info.attributes.get(SpanAttributeKey.CHAT_USAGE)
            if not usage:
                continue
            model_provider = info.attributes.get(SpanAttributeKey.MODEL_PROVIDER)
            if cost := calculate_cost_by_model_and_token_usage(model_name, usage, model_provider):
                total += cost.get(CostKey.TOTAL_COST, 0.0)
    return total


def fire_budget_exceeded_webhooks(
    newly_exceeded: list[BudgetWindow],
    workspace: str | None,
    registry_store,
) -> None:
    """Fire budget_policy.exceeded webhooks for newly-exceeded budget windows."""
    event = WebhookEvent(WebhookEntity.BUDGET_POLICY, WebhookAction.EXCEEDED)

    for window in newly_exceeded:
        policy = window.policy
        if policy.budget_action != BudgetAction.ALERT:
            continue

        payload = BudgetPolicyExceededPayload(
            budget_policy_id=policy.budget_policy_id,
            budget_unit=policy.budget_unit.value,
            budget_amount=policy.budget_amount,
            current_spend=window.cumulative_spend,
            duration_unit=policy.duration.unit.value,
            duration_value=policy.duration.value,
            target_scope=policy.target_scope.value,
            workspace=workspace or (policy.workspace or DEFAULT_WORKSPACE_NAME),
            window_start=int(window.window_start.timestamp() * 1000),
        )
        deliver_webhook(event=event, payload=payload, store=registry_store)


def _create_budget_error_trace(
    endpoint_config: GatewayEndpointConfig,
    exception: HTTPException,
) -> None:
    """Create an error trace for a budget limit rejection.

    Only creates a trace when usage tracking is enabled, matching the guard
    in ``maybe_traced_gateway_call``.
    """
    if not endpoint_config.usage_tracking:
        return
    try:
        with mlflow.start_span(
            name=_gateway_span_name(endpoint_config),
            span_type=SpanType.LLM,
            trace_destination=MlflowExperimentLocation(endpoint_config.experiment_id),
            attributes=_gateway_span_attributes(endpoint_config),
        ) as span:
            span.record_exception(exception)
    except Exception:
        _logger.debug("Failed to create budget error trace", exc_info=True)


def check_budget_limit(
    store: SqlAlchemyStore,
    endpoint_config: GatewayEndpointConfig,
    workspace: str | None = None,
) -> None:
    """Check if any REJECT-capable budget policy is exceeded.

    Raises HTTPException(429) with an error trace if the budget limit is exceeded.
    """
    maybe_refresh_budget_policies(store)
    tracker = get_budget_tracker()
    exceeded, window = tracker.should_reject_request(workspace=workspace)
    if exceeded:
        policy = window.policy
        unit = policy.duration.unit.value.lower()
        if policy.duration.value == 1:
            unit = unit.rstrip("s")
        reset_time = window.window_end.strftime("%Y-%m-%dT%H:%M:%SZ")
        detail = (
            f"Budget limit exceeded. "
            f"Limit: ${policy.budget_amount:.2f} USD per "
            f"{policy.duration.value} {unit}. "
            f"Budget resets at {reset_time}. "
            "Request rejected."
        )
        exc = HTTPException(status_code=429, detail=detail)
        _create_budget_error_trace(endpoint_config, exc)
        raise exc


def make_budget_on_complete(
    store: SqlAlchemyStore,
    workspace: str | None,
):
    """Create an on_complete callback that records budget cost from child span attributes."""
    from mlflow.server.handlers import _get_model_registry_store

    try:
        registry_store = _get_model_registry_store()
    except Exception:
        registry_store = None

    def on_complete():
        try:
            span = mlflow.get_current_active_span()
            if not span:
                return

            total_cost = _compute_cost_from_child_spans(span.trace_id)
            if total_cost <= 0:
                return

            maybe_refresh_budget_policies(store)
            tracker = get_budget_tracker()
            if newly_exceeded := tracker.record_cost(total_cost, workspace=workspace):
                if registry_store:
                    fire_budget_exceeded_webhooks(newly_exceeded, workspace, registry_store)
        except Exception:
            _logger.debug("Failed to record budget cost", exc_info=True)

    return on_complete