diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261109090000_add_performance_alerts/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261109090000_add_performance_alerts/migration.sql new file mode 100644 index 00000000000..7df00c402bc --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261109090000_add_performance_alerts/migration.sql @@ -0,0 +1,19 @@ +SET LOCAL lock_timeout = '10s'; + +CREATE TABLE IF NOT EXISTS "LiteLLM_PerformanceAlertState" ( + "id" INTEGER PRIMARY KEY, + "owner" VARCHAR(32) NOT NULL, + "next_run_at" TIMESTAMPTZ(3) NOT NULL, + "day" VARCHAR(10), + "scope" VARCHAR(16), + "resume_after" VARCHAR(1024) NOT NULL DEFAULT '', + "wrapped" BOOLEAN NOT NULL DEFAULT false, + "checked" INTEGER NOT NULL DEFAULT 0, + "skipped" INTEGER NOT NULL DEFAULT 0, + "completed_at" TIMESTAMP(3), + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +ALTER TABLE "LiteLLM_DailyModelUsage" + ADD COLUMN IF NOT EXISTS "total_response_time_ms" BIGINT NOT NULL DEFAULT 0, + ADD COLUMN IF NOT EXISTS "timed_requests" BIGINT NOT NULL DEFAULT 0; diff --git a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py index 4c856c741d3..91eddf46782 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py @@ -59,6 +59,7 @@ class RequestLogIndex: REQUEST_LOG_INDEXES: Final = ( RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_api_key_startTime_idx", '("api_key", "startTime")'), RequestLogIndex("LiteLLM_SpendLogs", "LiteLLM_SpendLogs_litellm_call_id_idx", '("litellm_call_id")'), + RequestLogIndex("LiteLLM_DailyUserSpend", "LiteLLM_DailyUserSpend_api_key_date_idx", '("api_key", "date")'), ) _IDENTIFIER_MAX_BYTES: Final = 63 diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 1a6cdeae23f..d2dee203e7a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -894,6 +894,7 @@ model LiteLLM_DailyUserSpend { @@index([date]) @@index([user_id, date]) @@index([api_key]) + @@index([api_key, date]) @@index([model]) @@index([mcp_namespaced_tool_name]) @@index([endpoint]) @@ -1349,6 +1350,8 @@ model LiteLLM_DailyModelUsage { request_count BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt @@ -2185,3 +2188,17 @@ model LiteLLM_Alert { @@index([status, id]) } + +model LiteLLM_PerformanceAlertState { + id Int @id + owner String @client.VarChar(32) + next_run_at DateTime @client.Timestamptz(3) + day String? @client.VarChar(10) + scope String? @client.VarChar(16) + resume_after String @default("") @client.VarChar(1024) + wrapped Boolean @default(false) + checked Int @default(0) + skipped Int @default(0) + completed_at DateTime? + updated_at DateTime @default(now()) @updatedAt +} diff --git a/litellm/accounting/projection.py b/litellm/accounting/projection.py index dd52b49f101..6837bbc8354 100644 --- a/litellm/accounting/projection.py +++ b/litellm/accounting/projection.py @@ -230,6 +230,8 @@ def _model_contributions(row: Mapping[str, object], metadata: Mapping[str, objec model: Final = _text(row.get("model")) or "unknown" caller: Final = not origin success: Final = row.get("status") == "success" + duration: Final = row.get("request_duration_ms") + timed: Final = caller and success and isinstance(duration, int) and duration >= 0 return ( Contribution( "LiteLLM_DailyModelUsage", @@ -247,6 +249,8 @@ def _model_contributions(row: Mapping[str, object], metadata: Mapping[str, objec ("request_count", int(caller)), ("successful_requests", int(caller and success)), ("failed_requests", int(caller and not success)), + ("total_response_time_ms", _number(duration) if timed else 0), + ("timed_requests", int(timed)), ), ), ) diff --git a/litellm/models/performance_alert.py b/litellm/models/performance_alert.py new file mode 100644 index 00000000000..6e5a251240c --- /dev/null +++ b/litellm/models/performance_alert.py @@ -0,0 +1,211 @@ +import hashlib +from collections.abc import Sequence +from dataclasses import dataclass +from datetime import date, datetime, time, timedelta, timezone +from typing import Final, Literal, Protocol, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field + +from litellm.models.alert import MAX_ALERT_ID, AlertDetection, AlertType, MetricEvidence + +PerformanceScope: TypeAlias = Literal["gateway", "model", "team", "user", "key"] +SCOPES: Final[tuple[PerformanceScope, ...]] = ("gateway", "model", "team", "user", "key") +BASELINE_DAYS: Final = 7 + +MIN_ACTIVE_DAYS: Final = 4 +SPEND_MULTIPLIER: Final = 3 +SPEND_MIN_INCREASE_USD: Final = 50 +TRAFFIC_SPIKE_MULTIPLIER: Final = 10 +TRAFFIC_SPIKE_MIN_INCREASE: Final = 1000 +TRAFFIC_DROP_RATIO: Final = 0.2 +TRAFFIC_DROP_MIN_BASELINE: Final = 500 +LATENCY_MULTIPLIER: Final = 3 +LATENCY_MIN_INCREASE_MS: Final = 500 +LATENCY_MIN_TIMED_REQUESTS: Final = 100 +RELIABILITY_MULTIPLIER: Final = 3 +RELIABILITY_MIN_INCREASE: Final = 0.02 +RELIABILITY_MIN_REQUESTS: Final = 500 + + +class ResourceWindow(BaseModel): + """One resource's totals for the evaluated day, and summed over the seven days before it.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + resource_id: str = Field(min_length=1, max_length=1024) + spend: float = Field(allow_inf_nan=False) + base_spend: float = Field(allow_inf_nan=False) + requests: int = Field(ge=0, le=MAX_ALERT_ID) + base_requests: int = Field(ge=0, le=MAX_ALERT_ID) + failed: int = Field(ge=0, le=MAX_ALERT_ID) + base_failed: int = Field(ge=0, le=MAX_ALERT_ID) + response_ms: int = Field(ge=0, le=MAX_ALERT_ID) + timed: int = Field(ge=0, le=MAX_ALERT_ID) + base_response_ms: int = Field(ge=0, le=MAX_ALERT_ID) + base_timed: int = Field(ge=0, le=MAX_ALERT_ID) + active_days: int = Field(ge=0, le=BASELINE_DAYS) + timed_days: int = Field(ge=0, le=BASELINE_DAYS) + exists: bool + + +class SweepState(BaseModel): + """Where the daily sweep stands. Each scope is walked from a day-derived start to the end of its IDs, + then wraps to cover the IDs before that start, so an unfinished sweep leaves different resources out each day.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + + day: date + scope: PerformanceScope = "gateway" + cursor: str = Field(default="", max_length=1024) + wrapped: bool = False + checked: int = Field(default=0, ge=0) + skipped: int = Field(default=0, ge=0) + completed: bool = False + + +@dataclass(frozen=True, slots=True) +class Claim: + owner: str + state: SweepState | None + + +class SweepOwnershipLostError(Exception): + pass + + +class PerformanceAlertStore(Protocol): + async def claim(self, interval: timedelta) -> Claim | None: ... + async def accounting_pending(self, day: date) -> bool: ... + async def key_index_ready(self) -> bool: ... + async def save(self, owner: str, state: SweepState) -> None: ... + async def evaluate_page(self, owner: str, state: SweepState, limit: int) -> SweepState: ... + async def skip_next(self, owner: str, state: SweepState) -> SweepState: ... + + +def target_day(now: datetime) -> date: + """The last complete UTC day, switching to a new day at 01:00 UTC so late accounting has time to land.""" + return (now.astimezone(timezone.utc) - timedelta(hours=1)).date() - timedelta(days=1) + + +def rotation_start(day: date) -> str: + return hashlib.sha256(day.isoformat().encode()).hexdigest() + + +def new_sweep(day: date) -> SweepState: + return SweepState(day=day) + + +def next_scope(state: SweepState) -> SweepState: + position: Final = SCOPES.index(state.scope) + if position + 1 == len(SCOPES): + return state.model_copy(update={"cursor": "", "wrapped": False, "completed": True}) + following: Final = SCOPES[position + 1] + return state.model_copy(update={"scope": following, "cursor": rotation_start(state.day), "wrapped": False}) + + +def advance(state: SweepState, resource_ids: Sequence[str], limit: int, skipped: int = 0) -> SweepState: + counted: Final = state.model_copy( + update={"checked": state.checked + len(resource_ids) - skipped, "skipped": state.skipped + skipped} + ) + if state.scope == "gateway": + return next_scope(counted) + if len(resource_ids) == limit: + return counted.model_copy(update={"cursor": resource_ids[-1]}) + if state.wrapped: + return next_scope(counted) + return counted.model_copy(update={"cursor": "", "wrapped": True}) + + +def _ratio(numerator: float, denominator: float) -> float: + return numerator / denominator if denominator else 0.0 + + +def _alert( + scope: PerformanceScope, + window: ResourceWindow, + day: date, + kind: AlertType, + value: float, + baseline: float, + sample_size: int, +) -> AlertDetection: + start: Final = datetime.combine(day, time.min, timezone.utc) + unit: Final = {"spend": "usd", "traffic": "requests", "latency": "milliseconds", "reliability": "ratio"}[kind] + return AlertDetection.model_validate( + { + "type": kind, + "resource_type": scope, + "resource_id": "gateway" if scope == "gateway" else window.resource_id, + "window_start": start, + "window_end": start + timedelta(days=1), + "evidence": MetricEvidence.model_validate( + {"value": value, "baseline": baseline, "unit": unit, "sample_size": sample_size} + ), + } + ) + + +def _spend_spike(window: ResourceWindow) -> tuple[float, float] | None: + baseline: Final = window.base_spend / BASELINE_DAYS + triggered: Final = ( + window.active_days >= MIN_ACTIVE_DAYS + and baseline > 0 + and window.spend >= SPEND_MULTIPLIER * baseline + and window.spend - baseline >= SPEND_MIN_INCREASE_USD + ) + return (window.spend, baseline) if triggered else None + + +def _traffic_change(window: ResourceWindow) -> tuple[float, float] | None: + baseline: Final = window.base_requests / BASELINE_DAYS + spike: Final = ( + window.active_days >= MIN_ACTIVE_DAYS + and window.requests >= TRAFFIC_SPIKE_MULTIPLIER * baseline + and window.requests - baseline >= TRAFFIC_SPIKE_MIN_INCREASE + ) + drop: Final = ( + window.exists + and window.active_days == BASELINE_DAYS + and baseline >= TRAFFIC_DROP_MIN_BASELINE + and window.requests < TRAFFIC_DROP_RATIO * baseline + ) + return (window.requests, baseline) if spike or drop else None + + +def _slower_responses(window: ResourceWindow) -> tuple[float, float] | None: + current: Final = _ratio(window.response_ms, window.timed) + baseline: Final = _ratio(window.base_response_ms, window.base_timed) + triggered: Final = ( + window.timed >= LATENCY_MIN_TIMED_REQUESTS + and window.timed_days >= MIN_ACTIVE_DAYS + and baseline > 0 + and current >= LATENCY_MULTIPLIER * baseline + and current - baseline >= LATENCY_MIN_INCREASE_MS + ) + return (current, baseline) if triggered else None + + +def _more_failures(window: ResourceWindow) -> tuple[float, float] | None: + current: Final = min(1.0, _ratio(window.failed, window.requests)) + baseline: Final = min(1.0, _ratio(window.base_failed, window.base_requests)) + triggered: Final = ( + window.active_days >= MIN_ACTIVE_DAYS + and window.requests >= RELIABILITY_MIN_REQUESTS + and current >= RELIABILITY_MULTIPLIER * baseline + and current - baseline >= RELIABILITY_MIN_INCREASE + ) + return (current, baseline) if triggered else None + + +def detect(scope: PerformanceScope, window: ResourceWindow, day: date) -> tuple[AlertDetection, ...]: + checks: Final[tuple[tuple[AlertType, tuple[float, float] | None, int], ...]] = ( + ("spend", _spend_spike(window), window.requests), + ("traffic", _traffic_change(window), window.requests), + ("latency", _slower_responses(window), window.timed), + ("reliability", _more_failures(window), window.requests), + ) + return tuple( + _alert(scope, window, day, kind, found[0], found[1], sample_size) + for kind, found, sample_size in checks + if found is not None + ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e53b5d4ad52..1fe6ab32cb8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2928,6 +2928,13 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "every replica. On by default; set to tune the window, pin a job, or turn it off." ), ) + performance_alerts_enabled: bool = Field( + False, + description=( + "Check yesterday's usage for the gateway, models, teams, users and keys against the week before, " + "and save unusual spend, traffic, response time and failure rate changes as alerts. Off by default." + ), + ) spend_capture_rate_check: SpendCaptureRateCheckSettings | None = Field( None, description=( diff --git a/litellm/proxy/db/model_usage_rollup.py b/litellm/proxy/db/model_usage_rollup.py index f8cdabe9296..d0679df6822 100644 --- a/litellm/proxy/db/model_usage_rollup.py +++ b/litellm/proxy/db/model_usage_rollup.py @@ -62,6 +62,7 @@ class ModelUsageTransaction: completion_tokens: int successful: bool counts_as_request: bool = True + duration_ms: int | None = None def model_usage_task_type(request_tags: str) -> str: @@ -101,6 +102,8 @@ def build_model_usage_transaction(payload: SpendLogsPayload) -> ModelUsageTransa if date is None or (origin and origin != BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN): return None model: Final = payload["model"] or "unknown" + successful: Final = payload["status"] == "success" + duration: Final = payload["request_duration_ms"] if "request_duration_ms" in payload else None return ModelUsageTransaction( key=ModelUsageKey( date=date, @@ -112,8 +115,9 @@ def build_model_usage_transaction(payload: SpendLogsPayload) -> ModelUsageTransa spend=payload["spend"], prompt_tokens=payload["prompt_tokens"], completion_tokens=payload["completion_tokens"], - successful=payload["status"] == "success", + successful=successful, counts_as_request=not origin, + duration_ms=duration if not origin and successful and isinstance(duration, int) and duration >= 0 else None, ) @@ -150,6 +154,7 @@ async def flush_model_usage_transactions( requests = tuple(entry for entry in entries if entry.counts_as_request) successful = sum(1 for entry in requests if entry.successful) failed = len(requests) - successful + durations = tuple(entry.duration_ms for entry in requests if entry.duration_ms is not None) key_fields = { "date": key.date, "model_group": key.model_group, @@ -167,6 +172,8 @@ async def flush_model_usage_transactions( "request_count": len(requests), "successful_requests": successful, "failed_requests": failed, + "total_response_time_ms": sum(durations), + "timed_requests": len(durations), }, ) return diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6796fdbe895..4651b5f9f94 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11877,6 +11877,8 @@ class ProxyStartupEvent: read_general_settings=_current_general_settings, ) + cls._initialize_performance_alerts_job(scheduler=scheduler, prisma_client=prisma_client) + ### PTU DAILY ROLLUP ### from litellm.proxy.spend_tracking.ptu_feature_flag import ( is_ptu_cost_attribution_enabled, @@ -12253,6 +12255,28 @@ class ProxyStartupEvent: "LITELLM_EXPIRED_UI_SESSION_KEY_CLEANUP_ENABLED=true to enable)" ) + @classmethod + def _initialize_performance_alerts_job(cls, scheduler: AsyncIOScheduler, prisma_client: PrismaClient) -> None: + from litellm.proxy.spend_tracking.performance_alerts import RUN_INTERVAL, run_performance_alerts + from litellm.repositories.performance_alert_repository import PerformanceAlertRepository + + repository: Final = PerformanceAlertRepository(prisma_client) + + async def run() -> None: + if _current_general_settings().get("performance_alerts_enabled") is True: + await run_performance_alerts(repository) + + scheduler.add_job( # pyright: ignore[reportUnknownMemberType] # APScheduler does not provide typed job methods + run, + "interval", + seconds=RUN_INTERVAL.total_seconds(), + id="performance_alerts", + replace_existing=True, + max_instances=1, + coalesce=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + @classmethod def _initialize_spend_capture_rate_check_job( cls, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 1a6cdeae23f..d2dee203e7a 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -894,6 +894,7 @@ model LiteLLM_DailyUserSpend { @@index([date]) @@index([user_id, date]) @@index([api_key]) + @@index([api_key, date]) @@index([model]) @@index([mcp_namespaced_tool_name]) @@index([endpoint]) @@ -1349,6 +1350,8 @@ model LiteLLM_DailyModelUsage { request_count BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt @@ -2185,3 +2188,17 @@ model LiteLLM_Alert { @@index([status, id]) } + +model LiteLLM_PerformanceAlertState { + id Int @id + owner String @client.VarChar(32) + next_run_at DateTime @client.Timestamptz(3) + day String? @client.VarChar(10) + scope String? @client.VarChar(16) + resume_after String @default("") @client.VarChar(1024) + wrapped Boolean @default(false) + checked Int @default(0) + skipped Int @default(0) + completed_at DateTime? + updated_at DateTime @default(now()) @updatedAt +} diff --git a/litellm/proxy/spend_tracking/performance_alerts.py b/litellm/proxy/spend_tracking/performance_alerts.py new file mode 100644 index 00000000000..725d6139054 --- /dev/null +++ b/litellm/proxy/spend_tracking/performance_alerts.py @@ -0,0 +1,99 @@ +import asyncio +import time +from collections.abc import Awaitable, Callable +from datetime import datetime, timedelta, timezone +from typing import Final + +from litellm._logging import verbose_proxy_logger +from litellm.models.performance_alert import ( + Claim, + PerformanceAlertStore, + SweepOwnershipLostError, + SweepState, + new_sweep, + next_scope, + target_day, +) +from litellm.repositories.performance_alert_repository import KEY_DATE_INDEX + +RUN_INTERVAL: Final = timedelta(minutes=5) +RUN_BUDGET_SECONDS: Final = 10.0 +PAGE_SIZE: Final = 25 +PAUSE_SECONDS: Final = 0.1 + + +async def _step(store: PerformanceAlertStore, owner: str, state: SweepState) -> SweepState: + """One page, falling back to a single resource and then stepping over it, so one bad resource + cannot hold up the rest of the sweep.""" + from prisma.errors import PrismaError + from pydantic import ValidationError + + try: + return await store.evaluate_page(owner, state, PAGE_SIZE) + except (PrismaError, ValidationError, TimeoutError): + pass + try: + return await store.evaluate_page(owner, state, 1) + except (PrismaError, ValidationError, TimeoutError) as exc: + verbose_proxy_logger.warning( + "Performance alerts skipped one %s it could not read within its limits (%s)", + state.scope, + type(exc).__name__, + ) + return await store.skip_next(owner, state) + + +async def _start(store: PerformanceAlertStore, claim: Claim, now: datetime) -> SweepState | None: + day: Final = target_day(now) + if claim.state is not None and claim.state.day == day: + return claim.state + if claim.state is not None and not claim.state.completed: + verbose_proxy_logger.warning( + "Performance alerts did not finish %s: stopped in %s after %d resources, %d skipped", + claim.state.day, + claim.state.scope, + claim.state.checked, + claim.state.skipped, + ) + if await store.accounting_pending(day): + verbose_proxy_logger.info("Performance alerts waiting for %s usage to finish recording", day) + return None + return new_sweep(day) + + +async def run_performance_alerts( + store: PerformanceAlertStore, + interval: timedelta = RUN_INTERVAL, + now: Callable[[], datetime] = lambda: datetime.now(timezone.utc), + clock: Callable[[], float] = time.monotonic, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, +) -> None: + from prisma.errors import PrismaError + from pydantic import ValidationError + + try: + claim: Final = await store.claim(interval) + if claim is None: + return + started: Final = await _start(store, claim, now()) + if started is None or started.completed: + return + deadline: Final = clock() + RUN_BUDGET_SECONDS + keys_ready: Final = await store.key_index_ready() + state = started # rebind-ok: each page returns the next saved position + while not state.completed and clock() < deadline: + if state.scope == "key" and not keys_ready: + verbose_proxy_logger.info("Performance alerts skipping keys until %s is built", KEY_DATE_INDEX) + state = next_scope(state) # rebind-ok: keys wait for their index + await store.save(claim.owner, state) + continue + state = await _step(store, claim.owner, state) # rebind-ok: advance to the next saved position + await sleep(PAUSE_SECONDS) + if state.completed: + verbose_proxy_logger.info( + "Performance alerts checked %d resources for %s, %d skipped", state.checked, state.day, state.skipped + ) + except SweepOwnershipLostError: + verbose_proxy_logger.info("Performance alerts run ended, another worker took over") + except (PrismaError, ValidationError, TimeoutError) as exc: + verbose_proxy_logger.warning("Performance alerts run ended early (%s)", type(exc).__name__) diff --git a/litellm/repositories/alert_repository.py b/litellm/repositories/alert_repository.py index 1f02d3cdd81..9d9606b5d98 100644 --- a/litellm/repositories/alert_repository.py +++ b/litellm/repositories/alert_repository.py @@ -1,6 +1,12 @@ import hashlib import json -from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # Prisma exposes generated table actions dynamically +from collections.abc import Sequence +from typing import ( # noqa: TID251 # Prisma exposes generated table actions dynamically + TYPE_CHECKING, + Final, + Protocol, + cast, +) from litellm.models.alert import Alert, AlertDetection, AlertPage, AlertStatus from litellm.repositories.prisma_protocols import PrismaRecord, TableActions @@ -9,6 +15,33 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient +class AlertTransaction(Protocol): + async def execute_raw(self, query: str, *args: object) -> int: ... + + +def _detection_key(detection: AlertDetection) -> str: + return hashlib.sha256(json.dumps(detection.identity()).encode()).hexdigest() + + +async def record_alerts(transaction: AlertTransaction, detections: Sequence[AlertDetection]) -> None: + if not detections: + return + await transaction.execute_raw( + """INSERT INTO "LiteLLM_Alert" + (detection_key, type, resource_type, resource_id, window_start, window_end, evidence) + SELECT detection_key, type, resource_type, resource_id, window_start, window_end, evidence + FROM jsonb_to_recordset($1::jsonb) AS rows(detection_key text, type text, resource_type text, + resource_id text, window_start timestamp, window_end timestamp, evidence jsonb) + ON CONFLICT (detection_key) DO NOTHING""", + json.dumps( + tuple( + {**detection.model_dump(mode="json"), "detection_key": _detection_key(detection)} + for detection in detections + ) + ), + ) + + def _alert(row: PrismaRecord) -> Alert: values: Final = row.dict() return Alert.model_validate( @@ -26,7 +59,7 @@ class AlertRepository: from prisma import Json from prisma.errors import UniqueViolationError - key: Final = hashlib.sha256(json.dumps(detection.identity()).encode()).hexdigest() + key: Final = _detection_key(detection) try: row: Final = await self._table.create( data={ diff --git a/litellm/repositories/performance_alert_repository.py b/litellm/repositories/performance_alert_repository.py new file mode 100644 index 00000000000..a49f1e64e70 --- /dev/null +++ b/litellm/repositories/performance_alert_repository.py @@ -0,0 +1,253 @@ +import uuid +from collections.abc import AsyncGenerator, Mapping, Sequence +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from dataclasses import dataclass +from datetime import date, datetime, time, timedelta, timezone +from itertools import chain +from typing import TYPE_CHECKING, Final, Protocol, cast # noqa: TID251 # Prisma generates transaction types + +from pydantic import BaseModel, ConfigDict, TypeAdapter +from typing_extensions import LiteralString + +from litellm.models.performance_alert import ( + BASELINE_DAYS, + Claim, + PerformanceScope, + ResourceWindow, + SweepOwnershipLostError, + SweepState, + advance, + detect, + rotation_start, +) +from litellm.repositories.alert_repository import record_alerts + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +KEY_DATE_INDEX: Final = "LiteLLM_DailyUserSpend_api_key_date_idx" +_STATEMENT_TIMEOUT: Final = "2s" +_TRANSACTION_TIMEOUT: Final = timedelta(seconds=5) + + +class _Transaction(Protocol): + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ... + async def execute_raw(self, query: str, *args: object) -> int: ... + + +class _Database(Protocol): + def tx(self, *, timeout: timedelta) -> AbstractAsyncContextManager[_Transaction]: ... + + +class _StateRow(BaseModel): + model_config = ConfigDict(frozen=True) + + day: date | None + scope: PerformanceScope | None + cursor: str + wrapped: bool + checked: int + skipped: int + completed: bool + + +class _IdRow(BaseModel): + model_config = ConfigDict(frozen=True) + + resource_id: str + + +_STATE_ROWS: Final = TypeAdapter(tuple[_StateRow, ...]) +_WINDOWS: Final = TypeAdapter(tuple[ResourceWindow, ...]) +_ID_ROWS: Final = TypeAdapter(tuple[_IdRow, ...]) + + +@dataclass(frozen=True, slots=True) +class _Source: + table: LiteralString + resource_id: LiteralString + requests: LiteralString + filter: LiteralString + exists: LiteralString + + +_SOURCES: Final[Mapping[PerformanceScope, _Source]] = { + "gateway": _Source('"LiteLLM_DailyModelUsage"', "'gateway'", "request_count", "", "true"), + "model": _Source('"LiteLLM_DailyModelUsage"', "model_group", "request_count", "", "true"), + "team": _Source( + '"LiteLLM_DailyTeamSpend"', + "team_id", + "api_requests", + "", + 'EXISTS (SELECT 1 FROM "LiteLLM_TeamTable" t WHERE t.team_id = page.resource_id AND NOT t.blocked)', + ), + "user": _Source( + '"LiteLLM_DailyUserSpend"', + "user_id", + "api_requests", + "", + 'EXISTS (SELECT 1 FROM "LiteLLM_UserTable" u WHERE u.user_id = page.resource_id)', + ), + "key": _Source( + '"LiteLLM_DailyUserSpend"', + "api_key", + "api_requests", + " AND api_key ~ '^[0-9a-f]{64}$'", + 'EXISTS (SELECT 1 FROM "LiteLLM_VerificationToken" k WHERE k.token = page.resource_id' + " AND k.blocked IS NOT TRUE AND (k.expires IS NULL OR k.expires > now()))", + ), +} + + +def _bounds(source: _Source, wrapped: bool) -> LiteralString: + upper: Final = f" AND {source.resource_id} <= $5" if wrapped else "" + return ( + f" AND {source.resource_id} > $4{upper} AND {source.resource_id} <> ''" + f" AND octet_length({source.resource_id}) <= 1024{source.filter}" + ) + + +def _page_sql(scope: PerformanceScope, wrapped: bool) -> LiteralString: + source: Final = _SOURCES[scope] + where: Final = "" if scope == "gateway" else _bounds(source, wrapped) + group: Final = ( + " HAVING COUNT(*) > 0" + if scope == "gateway" + else f" GROUP BY {source.resource_id} ORDER BY {source.resource_id}" + ) + return f"""WITH page AS ( + SELECT {source.resource_id} AS resource_id, + COALESCE(SUM(spend) FILTER (WHERE date = $1), 0)::float8 AS spend, + COALESCE(SUM(spend) FILTER (WHERE date < $1), 0)::float8 AS base_spend, + COALESCE(SUM({source.requests}) FILTER (WHERE date = $1), 0)::bigint AS requests, + COALESCE(SUM({source.requests}) FILTER (WHERE date < $1), 0)::bigint AS base_requests, + COALESCE(SUM(failed_requests) FILTER (WHERE date = $1), 0)::bigint AS failed, + COALESCE(SUM(failed_requests) FILTER (WHERE date < $1), 0)::bigint AS base_failed, + COALESCE(SUM(total_response_time_ms) FILTER (WHERE date = $1), 0)::bigint AS response_ms, + COALESCE(SUM(timed_requests) FILTER (WHERE date = $1), 0)::bigint AS timed, + COALESCE(SUM(total_response_time_ms) FILTER (WHERE date < $1), 0)::bigint AS base_response_ms, + COALESCE(SUM(timed_requests) FILTER (WHERE date < $1), 0)::bigint AS base_timed, + COUNT(DISTINCT date) FILTER (WHERE date < $1 AND {source.requests} > 0) AS active_days, + COUNT(DISTINCT date) FILTER (WHERE date < $1 AND timed_requests > 0) AS timed_days + FROM {source.table} + WHERE date BETWEEN $2 AND $1{where}{group} + LIMIT $3::int +) +SELECT page.*, {source.exists} AS exists FROM page ORDER BY resource_id""" + + +def _next_id_sql(scope: PerformanceScope, wrapped: bool) -> LiteralString: + source: Final = _SOURCES[scope] + return f"""SELECT {source.resource_id} AS resource_id FROM {source.table} +WHERE date BETWEEN $2 AND $1{_bounds(source, wrapped)} +ORDER BY {source.resource_id} LIMIT $3::int""" + + +_LIMITS: Final = f"""SELECT set_config('statement_timeout', '{_STATEMENT_TIMEOUT}', true), + set_config('lock_timeout', '100ms', true), set_config('work_mem', '1MB', true), + set_config('max_parallel_workers_per_gather', '0', true)""" +_CLAIM: Final = """INSERT INTO "LiteLLM_PerformanceAlertState" (id, owner, next_run_at, updated_at) +VALUES (1, $1, now() + $2::interval, now()) +ON CONFLICT (id) DO UPDATE SET owner = EXCLUDED.owner, next_run_at = EXCLUDED.next_run_at, updated_at = now() +WHERE "LiteLLM_PerformanceAlertState".next_run_at <= now() +RETURNING day, scope, resume_after AS cursor, wrapped, checked, skipped, completed_at IS NOT NULL AS completed""" +_SAVE: Final = """UPDATE "LiteLLM_PerformanceAlertState" +SET day = $2, scope = $3, resume_after = $4, wrapped = $5, checked = $6, skipped = $7, + completed_at = CASE WHEN $8 THEN now() END, updated_at = now() +WHERE id = 1 AND owner = $1""" +_ACCOUNTING_PENDING: Final = """SELECT 1 AS resource_id FROM "LiteLLM_SpendLogDelivery" +WHERE state = 'pending' AND created_at < $1::timestamp LIMIT 1""" +_INDEX_READY: Final = """SELECT 1 AS resource_id FROM pg_index WHERE indexrelid = to_regclass($1) AND indisvalid""" + + +class PerformanceAlertRepository: + def __init__(self, prisma_client: "PrismaClient") -> None: + self._database: Final = cast( # cast-ok: generated Prisma transactions expose the bounded SQL surface + _Database, prisma_client.writer_db + ) + + @asynccontextmanager + async def _transaction(self) -> AsyncGenerator[_Transaction]: + async with self._database.tx(timeout=_TRANSACTION_TIMEOUT) as transaction: + await transaction.query_raw(_LIMITS) + yield transaction + + async def claim(self, interval: timedelta) -> Claim | None: + """Take this deployment's next run, at most one per `interval` across every worker and pod.""" + owner: Final = uuid.uuid4().hex + async with self._transaction() as transaction: + rows: Final = _STATE_ROWS.validate_python( + await transaction.query_raw(_CLAIM, owner, f"{int(interval.total_seconds())} seconds") + ) + if not rows: + return None + row: Final = rows[0] + if row.day is None or row.scope is None: + return Claim(owner, None) + return Claim(owner, SweepState.model_validate(row.model_dump())) + + async def accounting_pending(self, day: date) -> bool: + day_end: Final = datetime.combine(day + timedelta(days=1), time.min, timezone.utc) + async with self._transaction() as transaction: + return bool(await transaction.query_raw(_ACCOUNTING_PENDING, day_end.replace(tzinfo=None).isoformat())) + + async def key_index_ready(self) -> bool: + async with self._transaction() as transaction: + return bool(await transaction.query_raw(_INDEX_READY, f'"{KEY_DATE_INDEX}"')) + + async def save(self, owner: str, state: SweepState) -> None: + async with self._transaction() as transaction: + await self._save(transaction, owner, state) + + async def evaluate_page(self, owner: str, state: SweepState, limit: int) -> SweepState: + """Evaluate the next `limit` resources and save their alerts together with the new position.""" + async with self._transaction() as transaction: + windows: Final = _WINDOWS.validate_python( + await transaction.query_raw(_page_sql(state.scope, state.wrapped), *self._arguments(state, limit)) + ) + await record_alerts( + transaction, tuple(chain.from_iterable(detect(state.scope, window, state.day) for window in windows)) + ) + advanced: Final = advance(state, tuple(window.resource_id for window in windows), limit) + await self._save(transaction, owner, advanced) + return advanced + + async def skip_next(self, owner: str, state: SweepState) -> SweepState: + """Step over the next resource without evaluating it, for one that cannot be read within the limits.""" + async with self._transaction() as transaction: + rows: Final = ( + (_IdRow(resource_id="gateway"),) + if state.scope == "gateway" + else _ID_ROWS.validate_python( + await transaction.query_raw(_next_id_sql(state.scope, state.wrapped), *self._arguments(state, 1)) + ) + ) + advanced: Final = advance(state, tuple(row.resource_id for row in rows), 1, skipped=len(rows)) + await self._save(transaction, owner, advanced) + return advanced + + @staticmethod + def _arguments(state: SweepState, limit: int) -> tuple[str | int, ...]: + start: Final = state.day - timedelta(days=BASELINE_DAYS) + window: Final = (state.day.isoformat(), start.isoformat(), limit) + if state.scope == "gateway": + return window + if state.wrapped: + return (*window, state.cursor, rotation_start(state.day)) + return (*window, state.cursor) + + @staticmethod + async def _save(transaction: _Transaction, owner: str, state: SweepState) -> None: + updated: Final = await transaction.execute_raw( + _SAVE, + owner, + state.day.isoformat(), + state.scope, + state.cursor, + state.wrapped, + state.checked, + state.skipped, + state.completed, + ) + if updated != 1: + raise SweepOwnershipLostError() diff --git a/schema.prisma b/schema.prisma index 1a6cdeae23f..d2dee203e7a 100644 --- a/schema.prisma +++ b/schema.prisma @@ -894,6 +894,7 @@ model LiteLLM_DailyUserSpend { @@index([date]) @@index([user_id, date]) @@index([api_key]) + @@index([api_key, date]) @@index([model]) @@index([mcp_namespaced_tool_name]) @@index([endpoint]) @@ -1349,6 +1350,8 @@ model LiteLLM_DailyModelUsage { request_count BigInt @default(0) successful_requests BigInt @default(0) failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) created_at DateTime @default(now()) updated_at DateTime @updatedAt @@ -2185,3 +2188,17 @@ model LiteLLM_Alert { @@index([status, id]) } + +model LiteLLM_PerformanceAlertState { + id Int @id + owner String @client.VarChar(32) + next_run_at DateTime @client.Timestamptz(3) + day String? @client.VarChar(10) + scope String? @client.VarChar(16) + resume_after String @default("") @client.VarChar(1024) + wrapped Boolean @default(false) + checked Int @default(0) + skipped Int @default(0) + completed_at DateTime? + updated_at DateTime @default(now()) @updatedAt +} diff --git a/tests/integration/database/test_performance_alerts.py b/tests/integration/database/test_performance_alerts.py new file mode 100644 index 00000000000..0e5f8fbd31d --- /dev/null +++ b/tests/integration/database/test_performance_alerts.py @@ -0,0 +1,199 @@ +import asyncio +import hashlib +import os +import uuid +from collections.abc import AsyncIterator +from dataclasses import dataclass +from datetime import date, datetime, timedelta, timezone +from typing import Final + +import psycopg +import pytest +import pytest_asyncio +from litellm_proxy_extras.request_log_indexes import ensure_request_log_indexes +from psycopg import sql + +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.spend_tracking.performance_alerts import run_performance_alerts +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.alert_repository import AlertRepository +from litellm.repositories.performance_alert_repository import PerformanceAlertRepository +from tests.integration._support.database import read_rows, write_rows + +DAY: Final = date(2026, 1, 8) +NOW: Final = datetime(2026, 1, 9, 2, 0, tzinfo=timezone.utc) + + +@dataclass(frozen=True, slots=True) +class Day: + """One day of usage for a resource: spend, requests, failed requests, total response ms and timed requests.""" + + spend: float = 1.0 + requests: int = 1000 + failed: int = 10 + response_ms: int = 100_000 + timed: int = 1000 + + +STEADY: Final = Day() + + +_SPEND_TABLE_INSERT: Final = sql.SQL( + "INSERT INTO {table} (id, {column}, date, spend, api_requests, successful_requests, failed_requests," + " total_response_time_ms, timed_requests, api_key, model, custom_llm_provider, mcp_namespaced_tool_name," + " endpoint, updated_at) VALUES (gen_random_uuid()::text, %s, %s, %s, %s, %s, %s, %s, %s, 'k', 'm', 'p', ''," + " '', now())" +) +_KEY_INSERT: Final = ( + 'INSERT INTO "LiteLLM_DailyUserSpend" (id, api_key, date, spend, api_requests, successful_requests,' + " failed_requests, total_response_time_ms, timed_requests, user_id, model, custom_llm_provider," + " mcp_namespaced_tool_name, endpoint, updated_at) VALUES (gen_random_uuid()::text, %s, %s, %s, %s, %s, %s, %s," + " %s, '', 'm', 'p', '', '', now())" +) +_MODEL_INSERT: Final = ( + 'INSERT INTO "LiteLLM_DailyModelUsage" (model_group, date, spend, request_count, successful_requests,' + " failed_requests, total_response_time_ms, timed_requests, model, custom_llm_provider, task_type, updated_at)" + " VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 'm', 'p', 'uncategorized', now())" +) + + +def _seed(scope: str, resource: str, baseline: Day, evaluated: Day | None) -> None: + days: Final = tuple((DAY - timedelta(days=offset), baseline) for offset in range(1, 8)) + ( + ((DAY, evaluated),) if evaluated is not None else () + ) + rows: Final = tuple( + (resource, day.isoformat(), usage.spend, usage.requests, usage.requests - usage.failed, usage.failed) + + (usage.response_ms, usage.timed) + for day, usage in days + ) + query: Final = { + "team": _SPEND_TABLE_INSERT.format( + table=sql.Identifier("LiteLLM_DailyTeamSpend"), + column=sql.Identifier("team_id"), + ), + "user": _SPEND_TABLE_INSERT.format( + table=sql.Identifier("LiteLLM_DailyUserSpend"), + column=sql.Identifier("user_id"), + ), + "key": sql.SQL(_KEY_INSERT), + "model": sql.SQL(_MODEL_INSERT), + }[scope] + with psycopg.connect(os.environ["DATABASE_URL"]) as connection, connection.cursor() as cursor: + cursor.executemany(query, rows) + + +@dataclass(frozen=True, slots=True) +class Fixture: + prefix: str + repository: PerformanceAlertRepository + alerts: AlertRepository + + def key(self, name: str) -> str: + return hashlib.sha256(f"{self.prefix}-{name}".encode()).hexdigest() + + async def sweep(self) -> None: + for _ in range(20): + await run_performance_alerts(self.repository, interval=timedelta(0), now=lambda: NOW) + state = read_rows( + 'SELECT completed_at IS NOT NULL AS completed FROM "LiteLLM_PerformanceAlertState" WHERE id = 1', () + ) + if state and state[0]["completed"] is True: + return + pytest.fail("the sweep did not complete") + + def saved(self) -> dict[tuple[str, str], tuple[str, str]]: + rows: Final = read_rows( + 'SELECT resource_type, resource_id, type, status FROM "LiteLLM_Alert" WHERE resource_id LIKE %s' + " OR resource_id = %s", + (f"{self.prefix}%", self.key("slow")), + ) + return { + (str(row["resource_id"]), str(row["type"])): (str(row["resource_type"]), str(row["status"])) for row in rows + } + + +@pytest_asyncio.fixture(loop_scope="function") +async def fixture() -> AsyncIterator[Fixture]: + prefix: Final = f"perf-{uuid.uuid4().hex[:12]}" + client: Final = PrismaClient(os.environ["DATABASE_URL"], ProxyLogging(UserApiKeyCache())) + await client.connect() + write_rows('DELETE FROM "LiteLLM_PerformanceAlertState"', ()) + built: Final = Fixture(prefix, PerformanceAlertRepository(client), AlertRepository(client)) + try: + yield built + finally: + try: + pattern = f"{prefix}%" + write_rows('DELETE FROM "LiteLLM_DailyTeamSpend" WHERE team_id LIKE %s', (pattern,)) + write_rows( + 'DELETE FROM "LiteLLM_DailyUserSpend" WHERE user_id LIKE %s OR api_key = %s', + (pattern, built.key("slow")), + ) + write_rows('DELETE FROM "LiteLLM_DailyModelUsage" WHERE model_group LIKE %s', (pattern,)) + write_rows('DELETE FROM "LiteLLM_UserTable" WHERE user_id LIKE %s', (pattern,)) + write_rows( + 'DELETE FROM "LiteLLM_Alert" WHERE resource_id LIKE %s OR resource_id = %s' + " OR (resource_type = 'gateway' AND window_start = %s::timestamp)", + (pattern, built.key("slow"), DAY.isoformat()), + ) + write_rows('DELETE FROM "LiteLLM_PerformanceAlertState"', ()) + finally: + await client.disconnect() + + +@pytest.mark.asyncio +async def test_a_full_sweep_saves_one_alert_per_real_change_and_nothing_for_normal_or_new_usage( + fixture: Fixture, +) -> None: + assert ensure_request_log_indexes(os.environ["DATABASE_URL"], "public") + p: Final = fixture.prefix + _seed("team", f"{p}-team-spike", Day(spend=10.0), Day(spend=100.0)) + _seed("team", f"{p}-team-steady", STEADY, STEADY) + _seed("team", f"{p}-team-new", Day(requests=0, failed=0, timed=0, spend=0), Day(spend=900.0)) + write_rows('INSERT INTO "LiteLLM_UserTable" (user_id) VALUES (%s)', (f"{p}-user-stopped",)) + _seed("user", f"{p}-user-stopped", STEADY, None) + _seed("user", f"{p}-user-deleted", STEADY, None) + _seed("key", fixture.key("slow"), STEADY, Day(response_ms=1_000_000)) + _seed("model", f"{p}-model", STEADY, Day(failed=100)) + + await fixture.sweep() + + assert fixture.saved() == { + (f"{p}-team-spike", "spend"): ("team", "open"), + (f"{p}-user-stopped", "traffic"): ("user", "open"), + (fixture.key("slow"), "latency"): ("key", "open"), + (f"{p}-model", "reliability"): ("model", "open"), + } + + +@pytest.mark.asyncio +async def test_a_replayed_day_keeps_one_alert_and_its_dismissal(fixture: Fixture) -> None: + team: Final = f"{fixture.prefix}-team-spike" + _seed("team", team, Day(spend=10.0), Day(spend=100.0)) + await fixture.sweep() + (saved,) = read_rows('SELECT id FROM "LiteLLM_Alert" WHERE resource_id = %s', (team,)) + await fixture.alerts.set_status(int(str(saved["id"])), "dismissed", "admin") + + write_rows('DELETE FROM "LiteLLM_PerformanceAlertState"', ()) + await fixture.sweep() + + assert fixture.saved() == {(team, "spend"): ("team", "dismissed")} + + +@pytest.mark.asyncio +async def test_a_resource_that_cannot_be_read_is_skipped_without_blocking_the_rest(fixture: Fixture) -> None: + p: Final = fixture.prefix + _seed("team", f"{p}-team-a-broken", STEADY, Day(requests=-5, failed=0)) + _seed("team", f"{p}-team-b-spike", Day(spend=10.0), Day(spend=100.0)) + + await fixture.sweep() + + (state,) = read_rows('SELECT skipped FROM "LiteLLM_PerformanceAlertState" WHERE id = 1', ()) + assert int(str(state["skipped"])) >= 1 + assert fixture.saved() == {(f"{p}-team-b-spike", "spend"): ("team", "open")} + + +@pytest.mark.asyncio +async def test_only_one_of_many_concurrent_workers_gets_each_run(fixture: Fixture) -> None: + claims: Final = await asyncio.gather(*(fixture.repository.claim(timedelta(minutes=5)) for _ in range(8))) + assert sum(claim is not None for claim in claims) == 1 diff --git a/tests/unit/accounting/test_projection.py b/tests/unit/accounting/test_projection.py index c78d6bec757..d9ecbdb422c 100644 --- a/tests/unit/accounting/test_projection.py +++ b/tests/unit/accounting/test_projection.py @@ -260,6 +260,8 @@ async def test_entity_projection_matches_real_request_queues(monkeypatch: pytest "request_count": int(model.counts_as_request), "successful_requests": int(model.successful), "failed_requests": int(not model.successful), + "total_response_time_ms": model.duration_ms or 0, + "timed_requests": int(model.duration_ms is not None), } diff --git a/tests/unit/models/test_performance_alert.py b/tests/unit/models/test_performance_alert.py new file mode 100644 index 00000000000..8c146ecd52d --- /dev/null +++ b/tests/unit/models/test_performance_alert.py @@ -0,0 +1,152 @@ +from datetime import date, datetime, timedelta, timezone +from typing import Final + +import pytest + +from litellm.models.alert import AlertDetection, MetricEvidence +from litellm.models.performance_alert import ( + ResourceWindow, + SweepState, + advance, + detect, + new_sweep, + next_scope, + rotation_start, + target_day, +) + +DAY: Final = date(2026, 1, 8) +STEADY: Final = { + "resource_id": "team-a", + "spend": 10.0, + "base_spend": 70.0, + "requests": 1000, + "base_requests": 7000, + "failed": 10, + "base_failed": 70, + "response_ms": 100_000, + "timed": 1000, + "base_response_ms": 700_000, + "base_timed": 7000, + "active_days": 7, + "timed_days": 7, + "exists": True, +} + + +def _window(**changes: object) -> ResourceWindow: + return ResourceWindow.model_validate({**STEADY, **changes}) + + +def _types(**changes: object) -> tuple[str, ...]: + return tuple(alert.type for alert in detect("team", _window(**changes), DAY)) + + +def test_steady_usage_raises_nothing() -> None: + assert _types() == () + + +@pytest.mark.parametrize( + ("changes", "expected"), + ( + ({"spend": 60.0}, ("spend",)), + ({"spend": 59.9}, ()), + ({"base_spend": 280.0, "spend": 120.0}, ("spend",)), + ({"base_spend": 280.0, "spend": 119.9}, ()), + ({"spend": 5000.0, "active_days": 3}, ()), + ({"spend": 5000.0, "base_spend": 0.0}, ()), + ({"requests": 1100, "failed": 0, "response_ms": 110_000, "timed": 1100, "base_requests": 700}, ("traffic",)), + ({"requests": 1099, "failed": 0, "response_ms": 109_900, "timed": 1099, "base_requests": 700}, ()), + ({"requests": 50_000, "failed": 0, "active_days": 3}, ()), + ({"requests": 99, "failed": 0, "timed": 99, "response_ms": 9900, "base_requests": 3500}, ("traffic",)), + ({"requests": 100, "failed": 0, "timed": 100, "response_ms": 10_000, "base_requests": 3500}, ()), + ({"requests": 0, "failed": 0, "timed": 0, "response_ms": 0, "base_requests": 3500, "active_days": 6}, ()), + ({"requests": 0, "failed": 0, "timed": 0, "response_ms": 0, "base_requests": 3500, "exists": False}, ()), + ({"requests": 0, "failed": 0, "timed": 0, "response_ms": 0, "base_requests": 3493}, ()), + ({"response_ms": 600_000}, ("latency",)), + ({"response_ms": 599_000}, ()), + ({"response_ms": 60_000_000, "timed": 99}, ()), + ({"response_ms": 600_000, "timed_days": 3}, ()), + ({"failed": 40}, ("reliability",)), + ({"failed": 25}, ()), + ({"requests": 499, "failed": 400, "timed": 499, "response_ms": 49_900}, ()), + ({"base_failed": 0, "failed": 21}, ("reliability",)), + ), +) +def test_each_rule_fires_only_past_both_its_multiplier_and_its_floor( + changes: dict[str, object], expected: tuple[str, ...] +) -> None: + assert _types(**changes) == expected + + +def test_alert_saves_the_comparison_that_triggered_it() -> None: + alerts: Final = detect("team", _window(spend=60.0), DAY) + assert alerts == ( + AlertDetection.model_validate( + { + "type": "spend", + "resource_type": "team", + "resource_id": "team-a", + "window_start": "2026-01-08T00:00:00Z", + "window_end": "2026-01-09T00:00:00Z", + "evidence": {"value": 60.0, "baseline": 10.0, "unit": "usd", "sample_size": 1000}, + } + ), + ) + + +def test_latency_and_failure_rate_are_weighted_by_requests_not_averaged_per_day() -> None: + alerts: Final = detect("model", _window(response_ms=600_000, failed=40), DAY) + assert {alert.type: alert.evidence for alert in alerts} == { + "latency": MetricEvidence(value=600.0, baseline=100.0, unit="milliseconds", sample_size=1000), + "reliability": MetricEvidence(value=0.04, baseline=0.01, unit="ratio", sample_size=1000), + } + + +def test_gateway_alerts_use_the_gateway_resource_id() -> None: + alerts: Final = detect("gateway", _window(resource_id="anything", spend=60.0), DAY) + assert tuple((alert.resource_type, alert.resource_id) for alert in alerts) == (("gateway", "gateway"),) + + +@pytest.mark.parametrize( + ("now", "expected"), + ( + (datetime(2026, 1, 10, 0, 59, tzinfo=timezone.utc), date(2026, 1, 8)), + (datetime(2026, 1, 10, 1, 0, tzinfo=timezone.utc), date(2026, 1, 9)), + (datetime(2026, 1, 10, 9, 0, tzinfo=timezone(timedelta(hours=8))), date(2026, 1, 9)), + ), +) +def test_the_evaluated_day_changes_at_one_am_utc(now: datetime, expected: date) -> None: + assert target_day(now) == expected + + +def test_gateway_moves_to_models_starting_from_the_days_rotation_point() -> None: + assert advance(new_sweep(DAY), ("gateway",), 25) == SweepState( + day=DAY, scope="model", cursor=rotation_start(DAY), checked=1 + ) + + +def test_a_full_page_continues_from_its_last_resource() -> None: + state: Final = SweepState(day=DAY, scope="team", cursor=rotation_start(DAY)) + assert advance(state, ("t1", "t2"), 2) == state.model_copy(update={"cursor": "t2", "checked": 2}) + + +def test_a_short_page_wraps_to_the_ids_before_the_rotation_point_then_moves_on() -> None: + first_half: Final = advance(SweepState(day=DAY, scope="team", cursor="t9"), ("t10",), 25) + assert (first_half.scope, first_half.cursor, first_half.wrapped) == ("team", "", True) + second_half: Final = advance(first_half, ("a1",), 25) + assert (second_half.scope, second_half.cursor, second_half.wrapped) == ("user", rotation_start(DAY), False) + + +def test_the_sweep_completes_after_keys() -> None: + finished: Final = next_scope(SweepState(day=DAY, scope="key", cursor="k", wrapped=True, checked=9)) + assert finished.completed and finished.checked == 9 + + +def test_a_skipped_resource_counts_as_skipped_not_checked() -> None: + state: Final = SweepState(day=DAY, scope="user", cursor="u1", checked=4) + assert advance(state, ("u2",), 1, skipped=1) == state.model_copy(update={"cursor": "u2", "skipped": 1}) + + +def test_rotation_point_differs_between_days() -> None: + assert rotation_start(DAY) != rotation_start(DAY + timedelta(days=1)) diff --git a/tests/unit/proxy/db/test_model_usage_rollup.py b/tests/unit/proxy/db/test_model_usage_rollup.py index 1e79d06370b..6e60c575e34 100644 --- a/tests/unit/proxy/db/test_model_usage_rollup.py +++ b/tests/unit/proxy/db/test_model_usage_rollup.py @@ -55,7 +55,9 @@ def _payload(**overrides: Any) -> dict[str, Any]: } -def _transaction(model: str, spend: float, successful: bool = True) -> ModelUsageTransaction: +def _transaction( + model: str, spend: float, successful: bool = True, duration_ms: int | None = None +) -> ModelUsageTransaction: return ModelUsageTransaction( key=ModelUsageKey( date="2026-09-28", model_group=model, model=model, custom_llm_provider="openai", task_type="debugging" @@ -64,6 +66,7 @@ def _transaction(model: str, spend: float, successful: bool = True) -> ModelUsag prompt_tokens=10, completion_tokens=5, successful=successful, + duration_ms=duration_ms, ) @@ -128,10 +131,10 @@ async def test_flush_aggregates_each_rollup_row_into_one_upsert() -> None: await flush_model_usage_transactions( prisma_client=prisma, transactions=[ - _transaction("gpt-5", 0.5), + _transaction("gpt-5", 0.5, duration_ms=200), _transaction("claude", 1.0), _transaction("gpt-5", 0.25, successful=False), - _transaction("gpt-5", 0.25), + _transaction("gpt-5", 0.25, duration_ms=400), ], ) @@ -154,8 +157,39 @@ async def test_flush_aggregates_each_rollup_row_into_one_upsert() -> None: "request_count": {"increment": 3}, "successful_requests": {"increment": 2}, "failed_requests": {"increment": 1}, + "total_response_time_ms": {"increment": 600}, + "timed_requests": {"increment": 2}, } assert upserts["claude"]["create"]["request_count"] == 1 + assert (upserts["claude"]["create"]["total_response_time_ms"], upserts["claude"]["create"]["timed_requests"]) == ( + 0, + 0, + ) + + +@pytest.mark.parametrize( + ("overrides", "expected"), + ( + ({"request_duration_ms": 250}, 250), + ({"request_duration_ms": 0}, 0), + ({"request_duration_ms": 250, "status": "failure"}, None), + ({"request_duration_ms": None}, None), + ({"request_duration_ms": -1}, None), + ( + { + "request_duration_ms": 250, + "metadata": json.dumps({INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}), + }, + None, + ), + ), +) +def test_only_successful_caller_requests_with_a_duration_count_toward_response_time( + overrides: dict[str, Any], expected: int | None +) -> None: + transaction = build_model_usage_transaction(_payload(**overrides)) + assert transaction is not None + assert transaction.duration_ms == expected @pytest.mark.asyncio @@ -166,7 +200,10 @@ async def test_a_background_response_billing_row_adds_spend_but_no_request() -> prisma = _prisma(MagicMock(return_value=batcher)) create = build_model_usage_transaction(_payload(spend=0.0)) billing = build_model_usage_transaction( - _payload(spend=0.5, metadata=json.dumps({INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN})) + _payload( + spend=0.5, + metadata=json.dumps({INTERNAL_CALL_ORIGIN_METADATA_KEY: BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN}), + ) ) assert create is not None and billing is not None diff --git a/tests/unit/proxy/spend_tracking/test_performance_alerts.py b/tests/unit/proxy/spend_tracking/test_performance_alerts.py new file mode 100644 index 00000000000..ac284de0352 --- /dev/null +++ b/tests/unit/proxy/spend_tracking/test_performance_alerts.py @@ -0,0 +1,156 @@ +from dataclasses import dataclass, field +from datetime import date, datetime, timedelta, timezone +from typing import Final + +import pytest + +from litellm.models.performance_alert import Claim, SweepOwnershipLostError, SweepState, advance, new_sweep +from litellm.proxy.spend_tracking.performance_alerts import PAGE_SIZE, RUN_BUDGET_SECONDS, run_performance_alerts + +NOW: Final = datetime(2026, 1, 9, 2, 0, tzinfo=timezone.utc) +DAY: Final = date(2026, 1, 8) + + +@dataclass +class Store: + """Each scope holds `per_scope` resources. Resources named in `unreadable` fail any page that includes them.""" + + state: SweepState | None = None + claimable: bool = True + pending: bool = False + index_ready: bool = True + per_scope: int = 3 + unreadable: frozenset[str] = frozenset() + pages: list[tuple[str, int]] = field(default_factory=list) + saved: list[SweepState] = field(default_factory=list) + + async def claim(self, interval: timedelta) -> Claim | None: + return Claim("owner", self.state) if self.claimable else None + + async def accounting_pending(self, day: date) -> bool: + return self.pending + + async def key_index_ready(self) -> bool: + return self.index_ready + + async def save(self, owner: str, state: SweepState) -> None: + self.saved.append(state) + self.state = state + + def _next(self, state: SweepState, limit: int) -> tuple[str, ...]: + if state.scope == "gateway": + return ("gateway",) + names: Final = tuple(f"{state.scope}-{index}" for index in range(self.per_scope)) + remaining: Final = names[names.index(state.cursor) + 1 :] if state.cursor in names else names + return () if state.wrapped else remaining[:limit] + + async def evaluate_page(self, owner: str, state: SweepState, limit: int) -> SweepState: + self.pages.append((state.scope, limit)) + ids: Final = self._next(state, limit) + if self.unreadable.intersection(ids): + from prisma.errors import PrismaError + + raise PrismaError("canceling statement due to statement timeout") + advanced: Final = advance(state, ids, limit) + await self.save(owner, advanced) + return advanced + + async def skip_next(self, owner: str, state: SweepState) -> SweepState: + ids: Final = self._next(state, 1) + advanced: Final = advance(state, ids, 1, skipped=len(ids)) + await self.save(owner, advanced) + return advanced + + +async def _no_sleep(seconds: float) -> None: + return None + + +async def _run(store: Store, clock: list[float] | None = None) -> None: + ticks: Final = clock if clock is not None else [0.0] + await run_performance_alerts(store, now=lambda: NOW, clock=lambda: ticks[0], sleep=_no_sleep) + + +@pytest.mark.asyncio +async def test_a_fresh_sweep_checks_every_scope_and_completes() -> None: + store: Final = Store() + await _run(store) + assert store.state is not None and store.state.completed + assert (store.state.day, store.state.checked, store.state.skipped) == (DAY, 1 + 3 * 4, 0) + + +@pytest.mark.asyncio +async def test_no_work_happens_without_the_claim() -> None: + store: Final = Store(claimable=False) + await _run(store) + assert store.pages == [] and store.saved == [] + + +@pytest.mark.asyncio +async def test_a_new_day_waits_while_accounting_is_still_pending() -> None: + store: Final = Store(pending=True) + await _run(store) + assert store.pages == [] and store.saved == [] + + +@pytest.mark.asyncio +async def test_a_completed_day_does_no_more_work() -> None: + store: Final = Store(state=new_sweep(DAY).model_copy(update={"completed": True})) + await _run(store) + assert store.pages == [] + + +@pytest.mark.asyncio +async def test_an_unfinished_previous_day_is_replaced_by_the_new_day() -> None: + store: Final = Store(state=SweepState(day=DAY - timedelta(days=1), scope="user", cursor="user-1")) + await _run(store) + assert store.saved[0].day == DAY and store.pages[0] == ("gateway", PAGE_SIZE) + + +@pytest.mark.asyncio +async def test_the_run_resumes_the_saved_position_for_the_same_day() -> None: + store: Final = Store(state=SweepState(day=DAY, scope="key", cursor="key-1", checked=10)) + await _run(store) + assert store.pages[0] == ("key", PAGE_SIZE) + assert store.state is not None and store.state.checked == 11 + + +@pytest.mark.asyncio +async def test_an_unreadable_resource_is_retried_alone_then_skipped() -> None: + store: Final = Store(unreadable=frozenset({"team-1"})) + await _run(store) + assert ("team", 1) in store.pages + assert store.state is not None and store.state.completed + assert (store.state.checked, store.state.skipped) == (12, 1) + + +@pytest.mark.asyncio +async def test_keys_wait_for_their_index_while_other_scopes_run() -> None: + store: Final = Store(index_ready=False) + await _run(store) + assert all(scope != "key" for scope, _ in store.pages) + assert store.state is not None and store.state.completed and store.state.checked == 1 + 3 * 3 + + +@pytest.mark.asyncio +async def test_the_run_stops_at_its_time_budget_and_keeps_its_position() -> None: + clock: Final = [0.0] + store: Final = Store(per_scope=10_000) + + async def tick(seconds: float) -> None: + clock[0] += RUN_BUDGET_SECONDS * 0.4 + + await run_performance_alerts(store, now=lambda: NOW, clock=lambda: clock[0], sleep=tick) + assert len(store.pages) == 3 + assert store.state is not None and not store.state.completed + + +@pytest.mark.asyncio +async def test_losing_the_claim_mid_run_stops_quietly() -> None: + class TakenOver(Store): + async def evaluate_page(self, owner: str, state: SweepState, limit: int) -> SweepState: + raise SweepOwnershipLostError() + + store: Final = TakenOver() + await _run(store) + assert store.saved == []