mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
feat(alerts): detect daily changes in recorded usage (#50666)
* feat(alerts): detect daily changes in usage for gateway, models, teams, users and keys
* fix(alerts): drop unused suppressions
* fix(alerts): satisfy strict lint for model usage timing
* fix(alerts): skip only resources that time out or hold invalid data
* fix(alerts): log an unfinished daily sweep at info level
* fix(alerts): keep the deployment-wide run every five minutes
* fix(alerts): record alert and sweep timestamps in UTC
* fix(alerts): repair detector tests and tracing
* refactor(alerts): build the key index in a migration instead of after startup
* refactor(alerts): build the key index after startup again
A concurrent index build inside the boot migration can exceed the 600 second migrate deploy deadline on a large LiteLLM_DailyUserSpend and stop the proxy from starting. This reverts 393db6e4e0.
* fix(alerts): keep keys pending, bound the accounting wait, resume unfinished scopes
Keys waiting for their index no longer mark the day complete. Only deliveries created during the evaluated day hold back its sweep. A day that did not finish restarts the next day from where it stopped instead of a hash boundary that never rotated non-hex IDs.
* refactor(alerts): drop the pending delivery wait, alerts are best effort on flushed usage
* test(alerts): cover detector persistence and recovery
* fix(alerts): recognize Prisma statement timeouts
This commit is contained in:
moe-berri
GitHub
parent
26971a8f1f
commit
701639ea22
23 changed files with 1488 additions and 7 deletions
|
|
@ -0,0 +1,21 @@
|
|||
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,
|
||||
"start_scope" VARCHAR(16),
|
||||
"start_after" VARCHAR(1024) NOT NULL DEFAULT '',
|
||||
"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;
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -897,6 +897,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])
|
||||
|
|
@ -1383,6 +1384,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
|
||||
|
||||
|
|
@ -2224,3 +2227,19 @@ 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)
|
||||
start_scope String? @client.VarChar(16)
|
||||
start_after String @default("") @client.VarChar(1024)
|
||||
completed_at DateTime?
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
}
|
||||
|
|
|
|||
|
|
@ -231,6 +231,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",
|
||||
|
|
@ -248,6 +250,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)),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -261,6 +261,7 @@ _PRISMA_VIEWS: Final[frozenset[str]] = frozenset(
|
|||
_PRISMA_MODELS: Final[frozenset[str]] = frozenset(
|
||||
(
|
||||
"LiteLLM_Alert",
|
||||
"LiteLLM_PerformanceAlertState",
|
||||
"LiteLLM_BudgetTable",
|
||||
"LiteLLM_CredentialsTable",
|
||||
"LiteLLM_ProxyModelTable",
|
||||
|
|
|
|||
217
litellm/models/performance_alert.py
Normal file
217
litellm/models/performance_alert.py
Normal file
|
|
@ -0,0 +1,217 @@
|
|||
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. Scopes are walked in ID order. When the previous day stopped partway, that scope
|
||||
starts where it stopped and wraps around, so an unfinished sweep does not leave out the same resources every 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
|
||||
start_scope: PerformanceScope | None = None
|
||||
start_after: str = Field(default="", max_length=1024)
|
||||
|
||||
@property
|
||||
def start(self) -> str:
|
||||
return self.start_after if self.scope == self.start_scope else ""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Claim:
|
||||
owner: str
|
||||
state: SweepState | None
|
||||
|
||||
|
||||
class SweepOwnershipLostError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class ResourceUnreadableError(Exception):
|
||||
"""A page that timed out or held invalid data. Retried alone, then stepped over.
|
||||
Any other database error ends the run so the next run resumes from the same place."""
|
||||
|
||||
|
||||
class PerformanceAlertStore(Protocol):
|
||||
async def claim(self, interval: timedelta) -> Claim | None: ...
|
||||
async def key_index_ready(self) -> bool: ...
|
||||
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 new_sweep(day: date, previous: SweepState | None = None) -> SweepState:
|
||||
if previous is None or previous.completed or previous.scope == "gateway":
|
||||
return SweepState(day=day)
|
||||
return SweepState(day=day, start_scope=previous.scope, start_after=previous.cursor)
|
||||
|
||||
|
||||
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})
|
||||
moved: Final = state.model_copy(update={"scope": SCOPES[position + 1], "wrapped": False})
|
||||
return moved.model_copy(update={"cursor": moved.start})
|
||||
|
||||
|
||||
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 or not state.start:
|
||||
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
|
||||
)
|
||||
|
|
@ -2967,6 +2967,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=(
|
||||
|
|
|
|||
|
|
@ -69,6 +69,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:
|
||||
|
|
@ -108,6 +109,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"]
|
||||
return ModelUsageTransaction(
|
||||
key=ModelUsageKey(
|
||||
date=date,
|
||||
|
|
@ -119,8 +122,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,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -157,6 +161,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,
|
||||
|
|
@ -174,6 +179,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
|
||||
|
|
|
|||
|
|
@ -11919,6 +11919,26 @@ class ProxyStartupEvent:
|
|||
read_general_settings=_current_general_settings,
|
||||
)
|
||||
|
||||
from litellm.proxy.spend_tracking.performance_alerts import RUN_INTERVAL, run_performance_alerts
|
||||
from litellm.repositories.performance_alert_repository import PerformanceAlertRepository
|
||||
|
||||
performance_alert_repository: Final = PerformanceAlertRepository(prisma_client)
|
||||
|
||||
async def run_performance_alerts_job() -> None:
|
||||
if _current_general_settings().get("performance_alerts_enabled") is True:
|
||||
await run_performance_alerts(performance_alert_repository)
|
||||
|
||||
scheduler.add_job( # pyright: ignore[reportUnknownMemberType] # APScheduler does not provide typed job methods
|
||||
run_performance_alerts_job,
|
||||
"interval",
|
||||
seconds=RUN_INTERVAL.total_seconds(),
|
||||
id="performance_alerts",
|
||||
replace_existing=True,
|
||||
max_instances=1,
|
||||
coalesce=True,
|
||||
misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME,
|
||||
)
|
||||
|
||||
### PTU DAILY ROLLUP ###
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import (
|
||||
is_ptu_cost_attribution_enabled,
|
||||
|
|
|
|||
|
|
@ -897,6 +897,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])
|
||||
|
|
@ -1383,6 +1384,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
|
||||
|
||||
|
|
@ -2224,3 +2227,19 @@ 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)
|
||||
start_scope String? @client.VarChar(16)
|
||||
start_after String @default("") @client.VarChar(1024)
|
||||
completed_at DateTime?
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
}
|
||||
|
|
|
|||
89
litellm/proxy/spend_tracking/performance_alerts.py
Normal file
89
litellm/proxy/spend_tracking/performance_alerts.py
Normal file
|
|
@ -0,0 +1,89 @@
|
|||
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,
|
||||
ResourceUnreadableError,
|
||||
SweepOwnershipLostError,
|
||||
SweepState,
|
||||
new_sweep,
|
||||
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."""
|
||||
try:
|
||||
return await store.evaluate_page(owner, state, PAGE_SIZE)
|
||||
except ResourceUnreadableError:
|
||||
pass
|
||||
try:
|
||||
return await store.evaluate_page(owner, state, 1)
|
||||
except ResourceUnreadableError:
|
||||
verbose_proxy_logger.warning(
|
||||
"Performance alerts skipped one %s it could not read within its limits", state.scope
|
||||
)
|
||||
return await store.skip_next(owner, state)
|
||||
|
||||
|
||||
def _start(claim: Claim, now: datetime) -> SweepState:
|
||||
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.info(
|
||||
"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,
|
||||
)
|
||||
return new_sweep(day, claim.state)
|
||||
|
||||
|
||||
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 = _start(claim, now())
|
||||
if 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 waiting for %s before checking keys", KEY_DATE_INDEX)
|
||||
break
|
||||
state = await _step(store, claim.owner, state)
|
||||
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__)
|
||||
|
|
@ -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 (
|
||||
TYPE_CHECKING,
|
||||
Final,
|
||||
Protocol,
|
||||
cast, # noqa: TID251 # Prisma exposes generated table actions dynamically
|
||||
)
|
||||
|
||||
from litellm.models.alert import Alert, AlertDetection, AlertPage, AlertStatus
|
||||
from litellm.repositories.prisma_protocols import PrismaRecord, TableActions
|
||||
|
|
@ -9,6 +15,34 @@ 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, created_at, updated_at)
|
||||
SELECT detection_key, type, resource_type, resource_id, window_start, window_end, evidence,
|
||||
now() AT TIME ZONE 'UTC', now() AT TIME ZONE 'UTC'
|
||||
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 +60,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={
|
||||
|
|
|
|||
263
litellm/repositories/performance_alert_repository.py
Normal file
263
litellm/repositories/performance_alert_repository.py
Normal file
|
|
@ -0,0 +1,263 @@
|
|||
import uuid
|
||||
from collections.abc import AsyncGenerator, Mapping, Sequence
|
||||
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from datetime import date, timedelta
|
||||
from itertools import chain
|
||||
from typing import TYPE_CHECKING, Final, Protocol, cast # noqa: TID251 # Prisma generates transaction types
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
from typing_extensions import LiteralString
|
||||
|
||||
from litellm.models.performance_alert import (
|
||||
BASELINE_DAYS,
|
||||
Claim,
|
||||
PerformanceScope,
|
||||
ResourceUnreadableError,
|
||||
ResourceWindow,
|
||||
SweepOwnershipLostError,
|
||||
SweepState,
|
||||
advance,
|
||||
detect,
|
||||
)
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
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
|
||||
start_scope: PerformanceScope | None
|
||||
start_after: str
|
||||
|
||||
|
||||
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() AT TIME ZONE 'UTC')
|
||||
ON CONFLICT (id) DO UPDATE SET owner = EXCLUDED.owner, next_run_at = EXCLUDED.next_run_at,
|
||||
updated_at = now() AT TIME ZONE 'UTC'
|
||||
WHERE "LiteLLM_PerformanceAlertState".next_run_at <= now() + interval '30 seconds'
|
||||
RETURNING day, scope, resume_after AS cursor, wrapped, checked, skipped, completed_at IS NOT NULL AS completed,
|
||||
start_scope, start_after"""
|
||||
_SAVE: Final = """UPDATE "LiteLLM_PerformanceAlertState"
|
||||
SET day = $2, scope = $3, resume_after = $4, wrapped = $5, checked = $6, skipped = $7,
|
||||
start_scope = $9, start_after = $10, completed_at = CASE WHEN $8 THEN now() AT TIME ZONE 'UTC' END, updated_at = now() AT TIME ZONE 'UTC'
|
||||
WHERE id = 1 AND owner = $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. A worker whose
|
||||
timer fires up to 30 seconds early still takes it, so the claim's own delay never pushes a run back a whole
|
||||
interval."""
|
||||
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 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 evaluate_page(self, owner: str, state: SweepState, limit: int) -> SweepState:
|
||||
"""Evaluate the next `limit` resources and save their alerts together with the new position."""
|
||||
from prisma.errors import RawQueryError
|
||||
|
||||
try:
|
||||
return await self._evaluate_page(owner, state, limit)
|
||||
except ValidationError as exc:
|
||||
raise ResourceUnreadableError() from exc
|
||||
except RawQueryError as exc:
|
||||
if PrismaDBExceptionHandler.is_statement_timeout_error(exc):
|
||||
raise ResourceUnreadableError() from exc
|
||||
raise
|
||||
|
||||
async def _evaluate_page(self, owner: str, state: SweepState, limit: int) -> SweepState:
|
||||
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, state.start)
|
||||
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,
|
||||
state.start_scope,
|
||||
state.start_after,
|
||||
)
|
||||
if updated != 1:
|
||||
raise SweepOwnershipLostError()
|
||||
|
|
@ -897,6 +897,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])
|
||||
|
|
@ -1383,6 +1384,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
|
||||
|
||||
|
|
@ -2224,3 +2227,19 @@ 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)
|
||||
start_scope String? @client.VarChar(16)
|
||||
start_after String @default("") @client.VarChar(1024)
|
||||
completed_at DateTime?
|
||||
updated_at DateTime @default(now()) @updatedAt
|
||||
}
|
||||
|
|
|
|||
217
tests/integration/database/test_performance_alerts.py
Normal file
217
tests/integration/database/test_performance_alerts.py
Normal file
|
|
@ -0,0 +1,217 @@
|
|||
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()
|
||||
|
||||
(age,) = read_rows(
|
||||
"SELECT max(abs(extract(epoch FROM created_at - (now() AT TIME ZONE 'UTC'))))::int AS seconds"
|
||||
' FROM "LiteLLM_Alert" WHERE resource_id LIKE %s',
|
||||
(f"{p}%",),
|
||||
)
|
||||
assert int(str(age["seconds"])) < 300
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("seconds_until_due", "claimed"), ((20, True), (60, False)))
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_timer_firing_just_before_the_run_is_due_still_takes_it(
|
||||
fixture: Fixture, seconds_until_due: int, claimed: bool
|
||||
) -> None:
|
||||
write_rows(
|
||||
'INSERT INTO "LiteLLM_PerformanceAlertState" (id, owner, next_run_at) VALUES (1, %s, now() + %s::interval)',
|
||||
("earlier-worker", f"{seconds_until_due} seconds"),
|
||||
)
|
||||
assert (await fixture.repository.claim(timedelta(minutes=5)) is not None) is claimed
|
||||
|
|
@ -277,6 +277,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),
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
166
tests/unit/models/test_performance_alert.py
Normal file
166
tests/unit/models/test_performance_alert.py
Normal file
|
|
@ -0,0 +1,166 @@
|
|||
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,
|
||||
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_beginning() -> None:
|
||||
assert advance(new_sweep(DAY), ("gateway",), 25) == SweepState(day=DAY, scope="model", checked=1)
|
||||
|
||||
|
||||
def test_a_full_page_continues_from_its_last_resource() -> None:
|
||||
state: Final = SweepState(day=DAY, scope="team")
|
||||
assert advance(state, ("t1", "t2"), 2) == state.model_copy(update={"cursor": "t2", "checked": 2})
|
||||
|
||||
|
||||
def test_a_scope_without_a_carried_start_moves_on_after_one_pass() -> None:
|
||||
finished: Final = advance(SweepState(day=DAY, scope="team", cursor="t9"), ("t10",), 25)
|
||||
assert (finished.scope, finished.cursor, finished.wrapped) == ("user", "", False)
|
||||
|
||||
|
||||
def test_an_unfinished_day_restarts_its_scope_where_it_stopped_then_wraps() -> None:
|
||||
previous: Final = SweepState(day=DAY - timedelta(days=1), scope="team", cursor="t5", checked=40)
|
||||
today: Final = new_sweep(DAY, previous)
|
||||
assert (today.scope, today.start_scope, today.start_after, today.checked) == ("gateway", "team", "t5", 0)
|
||||
at_team: Final = next_scope(next_scope(today))
|
||||
assert (at_team.scope, at_team.cursor, at_team.wrapped) == ("team", "t5", False)
|
||||
after_tail: Final = advance(at_team, ("t6", "t7"), 25)
|
||||
assert (after_tail.scope, after_tail.cursor, after_tail.wrapped) == ("team", "", True)
|
||||
after_head: Final = advance(after_tail, ("t1", "t5"), 25)
|
||||
assert (after_head.scope, after_head.cursor, after_head.wrapped) == ("user", "", False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"previous",
|
||||
(
|
||||
SweepState(day=DAY - timedelta(days=1), scope="team", cursor="t5", completed=True),
|
||||
SweepState(day=DAY - timedelta(days=1), scope="gateway"),
|
||||
),
|
||||
)
|
||||
def test_a_finished_or_barely_started_day_carries_nothing(previous: SweepState) -> None:
|
||||
assert new_sweep(DAY, previous) == SweepState(day=DAY)
|
||||
|
||||
|
||||
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})
|
||||
|
|
@ -51,11 +51,14 @@ def _payload(**overrides: Any) -> dict[str, Any]:
|
|||
"request_tags": "[]",
|
||||
"custom_llm_provider": "openai",
|
||||
"status": "success",
|
||||
"request_duration_ms": None,
|
||||
**overrides,
|
||||
}
|
||||
|
||||
|
||||
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 +67,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 +132,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 +158,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 +201,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
|
||||
|
||||
|
|
|
|||
|
|
@ -253,6 +253,7 @@ async def test_model_insights_reads_back_what_the_rollup_wrote() -> None:
|
|||
"request_tags": '["task:debugging"]',
|
||||
"custom_llm_provider": "openai",
|
||||
"status": "success",
|
||||
"request_duration_ms": None,
|
||||
}
|
||||
|
||||
transactions = (
|
||||
|
|
|
|||
170
tests/unit/proxy/spend_tracking/test_performance_alerts.py
Normal file
170
tests/unit/proxy/spend_tracking/test_performance_alerts.py
Normal file
|
|
@ -0,0 +1,170 @@
|
|||
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,
|
||||
ResourceUnreadableError,
|
||||
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
|
||||
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 key_index_ready(self) -> bool:
|
||||
return self.index_ready
|
||||
|
||||
def _keep(self, 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))
|
||||
after: Final = tuple(name for name in names if name > state.cursor)
|
||||
return tuple(name for name in after if not state.wrapped or name <= state.start)[: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):
|
||||
raise ResourceUnreadableError()
|
||||
advanced: Final = advance(state, ids, limit)
|
||||
self._keep(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))
|
||||
self._keep(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_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_resumes_its_scope_where_it_stopped_and_still_checks_everything() -> 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)
|
||||
users: Final = tuple(state for state in store.saved if state.scope == "user")
|
||||
assert (users[0].cursor, users[0].wrapped) == ("user-1", False)
|
||||
assert any(state.wrapped for state in users)
|
||||
assert store.state is not None and store.state.completed and store.state.checked == 1 + 3 * 4
|
||||
|
||||
|
||||
@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_without_finishing_the_day_then_run_once_it_exists() -> 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 not store.state.completed and store.state.scope == "key"
|
||||
store.index_ready = True
|
||||
await _run(store)
|
||||
assert store.state is not None and store.state.completed and store.state.checked == 1 + 3 * 4
|
||||
|
||||
|
||||
@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 == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_database_wide_failure_ends_the_run_without_skipping_anything() -> None:
|
||||
class Unavailable(Store):
|
||||
async def evaluate_page(self, owner: str, state: SweepState, limit: int) -> SweepState:
|
||||
from prisma.errors import RawQueryError
|
||||
|
||||
raise RawQueryError({"user_facing_error": {"message": "FATAL: sorry, too many clients already"}})
|
||||
|
||||
store: Final = Unavailable()
|
||||
await _run(store)
|
||||
assert store.saved == []
|
||||
|
|
@ -1664,6 +1664,7 @@ class MockPrisma:
|
|||
self.database_url = database_url
|
||||
self.proxy_logging_obj = proxy_logging_obj
|
||||
self.http_client = http_client
|
||||
self.writer_db = MagicMock()
|
||||
|
||||
async def connect(self):
|
||||
pass
|
||||
|
|
|
|||
159
tests/unit/repositories/test_performance_alert_repository.py
Normal file
159
tests/unit/repositories/test_performance_alert_repository.py
Normal file
|
|
@ -0,0 +1,159 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from datetime import date, timedelta
|
||||
from typing import Final
|
||||
from unittest.mock import ANY, AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.models.performance_alert import ResourceUnreadableError, SweepOwnershipLostError, SweepState
|
||||
from litellm.repositories.performance_alert_repository import PerformanceAlertRepository
|
||||
|
||||
|
||||
def _raw_query_error(message: str, code: str = "") -> Exception:
|
||||
from prisma.errors import RawQueryError
|
||||
|
||||
return RawQueryError({"user_facing_error": {"is_panic": False, "message": message, "meta": {"code": code}}})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("message", "code", "retryable"),
|
||||
(
|
||||
("ERROR: canceling statement due to statement timeout", "57014", True),
|
||||
("Query canceled", "57014", True),
|
||||
("ERROR: canceling statement due to statement timeout", "", True),
|
||||
("FATAL: sorry, too many clients already", "53300", False),
|
||||
("FATAL: terminating connection due to administrator command", "57P01", False),
|
||||
("ERROR: deadlock detected", "40P01", False),
|
||||
),
|
||||
)
|
||||
async def test_only_a_statement_timeout_marks_a_resource_unreadable(message: str, code: str, retryable: bool) -> None:
|
||||
error: Final = _raw_query_error(message, code)
|
||||
repository, _query, execute = _repository(error)
|
||||
with pytest.raises(ResourceUnreadableError if retryable else type(error)):
|
||||
await repository.evaluate_page("worker-a", SweepState(day=date(2026, 1, 8), scope="model"), 25)
|
||||
execute.assert_not_awaited()
|
||||
|
||||
|
||||
def _repository(
|
||||
rows: Sequence[Mapping[str, object]] | Exception,
|
||||
) -> tuple[PerformanceAlertRepository, AsyncMock, AsyncMock]:
|
||||
query: Final = AsyncMock(side_effect=((), rows))
|
||||
execute: Final = AsyncMock(return_value=1)
|
||||
transaction: Final = MagicMock(query_raw=query, execute_raw=execute)
|
||||
context: Final = MagicMock(__aenter__=AsyncMock(return_value=transaction), __aexit__=AsyncMock(return_value=False))
|
||||
client: Final = MagicMock(writer_db=MagicMock(tx=MagicMock(return_value=context)))
|
||||
return PerformanceAlertRepository(client), query, execute
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("day", (None, date(2026, 1, 8)))
|
||||
async def test_claim_returns_the_saved_position_or_a_fresh_sweep(day: date | None) -> None:
|
||||
repository, query, _execute = _repository(
|
||||
(
|
||||
{
|
||||
"day": day,
|
||||
"scope": "team",
|
||||
"cursor": "team-a",
|
||||
"wrapped": True,
|
||||
"checked": 12,
|
||||
"skipped": 2,
|
||||
"completed": False,
|
||||
"start_scope": "team",
|
||||
"start_after": "team-z",
|
||||
},
|
||||
)
|
||||
)
|
||||
claim: Final = await repository.claim(timedelta(minutes=5))
|
||||
assert claim is not None
|
||||
assert claim.state == (
|
||||
None
|
||||
if day is None
|
||||
else SweepState(
|
||||
day=day,
|
||||
scope="team",
|
||||
cursor="team-a",
|
||||
wrapped=True,
|
||||
checked=12,
|
||||
skipped=2,
|
||||
start_scope="team",
|
||||
start_after="team-z",
|
||||
)
|
||||
)
|
||||
query.assert_awaited_with(ANY, claim.owner, "300 seconds")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_claim_taken_by_another_worker_returns_no_work() -> None:
|
||||
repository, _query, execute = _repository(())
|
||||
assert await repository.claim(timedelta(minutes=5)) is None
|
||||
execute.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def spend_window() -> Mapping[str, object]:
|
||||
return {
|
||||
"resource_id": "model-a",
|
||||
"spend": 100.0,
|
||||
"base_spend": 70.0,
|
||||
"requests": 100,
|
||||
"base_requests": 700,
|
||||
"failed": 0,
|
||||
"base_failed": 0,
|
||||
"response_ms": 0,
|
||||
"timed": 0,
|
||||
"base_response_ms": 0,
|
||||
"base_timed": 0,
|
||||
"active_days": 7,
|
||||
"timed_days": 0,
|
||||
"exists": True,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_page_saves_the_detected_evidence_and_its_resume_position(spend_window: Mapping[str, object]) -> None:
|
||||
repository, _query, execute = _repository((spend_window,))
|
||||
state: Final = SweepState(day=date(2026, 1, 8), scope="model")
|
||||
result: Final = await repository.evaluate_page("worker-a", state, 1)
|
||||
assert result == state.model_copy(update={"cursor": "model-a", "checked": 1})
|
||||
payload: Final = TypeAdapter(str).validate_python(execute.await_args_list[0].args[1])
|
||||
records: Final = TypeAdapter(tuple[dict[str, object], ...]).validate_json(payload)
|
||||
assert len(records) == 1
|
||||
assert records[0]["resource_id"] == "model-a"
|
||||
assert records[0]["type"] == "spend"
|
||||
assert records[0]["evidence"] == {"value": 100.0, "baseline": 10.0, "unit": "usd", "sample_size": 100}
|
||||
execute.assert_awaited_with(ANY, "worker-a", "2026-01-08", "model", "model-a", False, 1, 0, False, None, "")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lost_ownership_aborts_the_page_transaction(spend_window: Mapping[str, object]) -> None:
|
||||
repository, _query, execute = _repository((spend_window,))
|
||||
execute.side_effect = (1, 0)
|
||||
with pytest.raises(SweepOwnershipLostError):
|
||||
await repository.evaluate_page("previous-worker", SweepState(day=date(2026, 1, 8), scope="model"), 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure", ("timeout", "invalid_data", "connection_lost"))
|
||||
async def test_unreadable_data_is_retryable_but_database_failures_propagate(failure: str) -> None:
|
||||
error: Final = _raw_query_error("query failed", "57014" if failure == "timeout" else "57P01")
|
||||
repository, _query, execute = _repository(({},) if failure == "invalid_data" else error)
|
||||
with pytest.raises(ResourceUnreadableError if failure != "connection_lost" else type(error)):
|
||||
await repository.evaluate_page("worker-a", SweepState(day=date(2026, 1, 8), scope="model"), 25)
|
||||
execute.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("wrapped", (False, True))
|
||||
async def test_skipping_an_unreadable_resource_saves_progress_without_counting_it_as_checked(wrapped: bool) -> None:
|
||||
repository, query, execute = _repository(({"resource_id": "team-b"},))
|
||||
state: Final = SweepState(
|
||||
day=date(2026, 1, 8), scope="team", cursor="team-a", wrapped=wrapped, start_scope="team", start_after="team-z"
|
||||
)
|
||||
result: Final = await repository.skip_next("worker-a", state)
|
||||
assert result == state.model_copy(update={"cursor": "team-b", "skipped": 1})
|
||||
query.assert_awaited_with(ANY, "2026-01-08", "2026-01-01", 1, "team-a", *(("team-z",) if wrapped else ()))
|
||||
execute.assert_awaited_once_with(
|
||||
ANY, "worker-a", "2026-01-08", "team", "team-b", wrapped, 0, 1, False, "team", "team-z"
|
||||
)
|
||||
6
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
6
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -31505,6 +31505,12 @@ export interface components {
|
|||
* @description If True (default), a locally-managed user's password must contain an uppercase letter.
|
||||
*/
|
||||
password_policy_require_uppercase?: boolean | null;
|
||||
/**
|
||||
* Performance Alerts Enabled
|
||||
* @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.
|
||||
* @default false
|
||||
*/
|
||||
performance_alerts_enabled: boolean;
|
||||
/**
|
||||
* Plugins
|
||||
* @description external services registered as embeddable UI plugins
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue