from __future__ import annotations
from dataclasses import dataclass
from enum import Enum
from mlflow.entities._mlflow_object import _MlflowObject
from mlflow.protos.service_pb2 import BudgetAction as ProtoBudgetAction
from mlflow.protos.service_pb2 import BudgetDuration as ProtoBudgetDuration
from mlflow.protos.service_pb2 import BudgetDurationUnit as ProtoBudgetDurationUnit
from mlflow.protos.service_pb2 import BudgetTargetScope as ProtoBudgetTargetScope
from mlflow.protos.service_pb2 import BudgetUnit as ProtoBudgetUnit
from mlflow.protos.service_pb2 import GatewayBudgetPolicy as ProtoGatewayBudgetPolicy
from mlflow.utils.workspace_utils import resolve_entity_workspace_name
[docs]class BudgetDurationUnit(str, Enum):
"""Duration unit for budget policy fixed windows."""
MINUTES = "MINUTES"
HOURS = "HOURS"
DAYS = "DAYS"
WEEKS = "WEEKS"
MONTHS = "MONTHS"
[docs] @classmethod
def from_proto(cls, proto: ProtoBudgetDurationUnit) -> BudgetDurationUnit | None:
try:
return cls(ProtoBudgetDurationUnit.Name(proto))
except ValueError:
return None
[docs] def to_proto(self) -> ProtoBudgetDurationUnit:
return ProtoBudgetDurationUnit.Value(self.value)
[docs]class BudgetTargetScope(str, Enum):
"""Target scope for a budget policy."""
GLOBAL = "GLOBAL"
WORKSPACE = "WORKSPACE"
[docs] @classmethod
def from_proto(cls, proto: ProtoBudgetTargetScope) -> BudgetTargetScope | None:
try:
return cls(ProtoBudgetTargetScope.Name(proto))
except ValueError:
return None
[docs] def to_proto(self) -> ProtoBudgetTargetScope:
return ProtoBudgetTargetScope.Value(self.value)
[docs]class BudgetAction(str, Enum):
"""Action to take when a budget is exceeded."""
ALERT = "ALERT"
REJECT = "REJECT"
[docs] @classmethod
def from_proto(cls, proto: ProtoBudgetAction) -> BudgetAction | None:
try:
return cls(ProtoBudgetAction.Name(proto))
except ValueError:
return None
[docs] def to_proto(self) -> ProtoBudgetAction:
return ProtoBudgetAction.Value(self.value)
[docs]class BudgetUnit(str, Enum):
"""Budget measurement unit."""
USD = "USD"
[docs] @classmethod
def from_proto(cls, proto: ProtoBudgetUnit) -> BudgetUnit | None:
try:
return cls(ProtoBudgetUnit.Name(proto))
except ValueError:
return None
[docs] def to_proto(self) -> ProtoBudgetUnit:
return ProtoBudgetUnit.Value(self.value)
[docs]@dataclass
class BudgetDuration:
"""Fixed window duration: a (unit, value) pair defining the length of a budget window."""
unit: BudgetDurationUnit
value: int
def __post_init__(self):
if isinstance(self.unit, str):
self.unit = BudgetDurationUnit(self.unit)
[docs] def to_proto(self) -> ProtoBudgetDuration:
proto = ProtoBudgetDuration()
proto.unit = self.unit.to_proto()
proto.value = self.value
return proto
[docs] @classmethod
def from_proto(cls, proto: ProtoBudgetDuration) -> BudgetDuration:
return cls(
unit=BudgetDurationUnit.from_proto(proto.unit),
value=proto.value,
)
[docs]@dataclass
class GatewayBudgetPolicy(_MlflowObject):
"""
Represents a budget policy for the AI Gateway.
Budget policies set limits with fixed time windows,
supporting global or per-workspace scoping.
Args:
budget_policy_id: Unique identifier for this budget policy.
budget_unit: Budget measurement unit (e.g. USD).
budget_amount: Budget limit amount.
duration: Fixed time window (unit + length pair).
target_scope: Scope of the budget (GLOBAL or WORKSPACE).
budget_action: Action when budget is exceeded (ALERT, REJECT).
created_at: Timestamp (milliseconds) when the policy was created.
last_updated_at: Timestamp (milliseconds) when the policy was last updated.
created_by: User ID who created the policy.
last_updated_by: User ID who last updated the policy.
workspace: Workspace that owns the policy.
"""
budget_policy_id: str
budget_unit: BudgetUnit
budget_amount: float
duration: BudgetDuration
target_scope: BudgetTargetScope
budget_action: BudgetAction
created_at: int
last_updated_at: int
created_by: str | None = None
last_updated_by: str | None = None
workspace: str | None = None
def __post_init__(self):
self.workspace = resolve_entity_workspace_name(self.workspace)
if isinstance(self.budget_unit, str):
self.budget_unit = BudgetUnit(self.budget_unit)
if isinstance(self.target_scope, str):
self.target_scope = BudgetTargetScope(self.target_scope)
if isinstance(self.budget_action, str):
self.budget_action = BudgetAction(self.budget_action)
[docs] def to_proto(self):
proto = ProtoGatewayBudgetPolicy()
proto.budget_policy_id = self.budget_policy_id
proto.budget_unit = self.budget_unit.to_proto()
proto.budget_amount = self.budget_amount
proto.duration.CopyFrom(self.duration.to_proto())
proto.target_scope = self.target_scope.to_proto()
proto.budget_action = self.budget_action.to_proto()
proto.created_by = self.created_by or ""
proto.created_at = self.created_at
proto.last_updated_by = self.last_updated_by or ""
proto.last_updated_at = self.last_updated_at
return proto
[docs] @classmethod
def from_proto(cls, proto):
return cls(
budget_policy_id=proto.budget_policy_id,
budget_unit=BudgetUnit.from_proto(proto.budget_unit),
budget_amount=proto.budget_amount,
duration=BudgetDuration.from_proto(proto.duration),
target_scope=BudgetTargetScope.from_proto(proto.target_scope),
budget_action=BudgetAction.from_proto(proto.budget_action),
created_by=proto.created_by or None,
created_at=proto.created_at,
last_updated_by=proto.last_updated_by or None,
last_updated_at=proto.last_updated_at,
)