mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
feat(alerts): detect daily changes in usage for gateway, models, teams, users and keys
This commit is contained in:
moe-berri
parent
59f5bc42d4
commit
84bcf274f5
18 changed files with 1262 additions and 7 deletions
|
|
@ -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;
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
|
|
|||
211
litellm/models/performance_alert.py
Normal file
211
litellm/models/performance_alert.py
Normal file
|
|
@ -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
|
||||
)
|
||||
|
|
@ -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=(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
99
litellm/proxy/spend_tracking/performance_alerts.py
Normal file
99
litellm/proxy/spend_tracking/performance_alerts.py
Normal file
|
|
@ -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__)
|
||||
|
|
@ -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={
|
||||
|
|
|
|||
253
litellm/repositories/performance_alert_repository.py
Normal file
253
litellm/repositories/performance_alert_repository.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
199
tests/integration/database/test_performance_alerts.py
Normal file
199
tests/integration/database/test_performance_alerts.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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),
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
152
tests/unit/models/test_performance_alert.py
Normal file
152
tests/unit/models/test_performance_alert.py
Normal file
|
|
@ -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))
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
156
tests/unit/proxy/spend_tracking/test_performance_alerts.py
Normal file
156
tests/unit/proxy/spend_tracking/test_performance_alerts.py
Normal file
|
|
@ -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 == []
|
||||
Loading…
Add table
Reference in a new issue