From faa8bc2af05b5570763875f9720834adea52fc90 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 10 Nov 2026 05:53:39 +0800 Subject: [PATCH] fix(alerts): skip only resources that time out or hold invalid data --- litellm/models/performance_alert.py | 5 ++++ .../spend_tracking/performance_alerts.py | 12 +++------ .../performance_alert_repository.py | 20 +++++++++++++- .../spend_tracking/test_performance_alerts.py | 26 ++++++++++++++++--- .../test_performance_alert_repository.py | 25 ++++++++++++++++++ 5 files changed, 75 insertions(+), 13 deletions(-) create mode 100644 tests/unit/repositories/test_performance_alert_repository.py diff --git a/litellm/models/performance_alert.py b/litellm/models/performance_alert.py index 6e5a251240c..09a6210a5b0 100644 --- a/litellm/models/performance_alert.py +++ b/litellm/models/performance_alert.py @@ -73,6 +73,11 @@ 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 accounting_pending(self, day: date) -> bool: ... diff --git a/litellm/proxy/spend_tracking/performance_alerts.py b/litellm/proxy/spend_tracking/performance_alerts.py index f848298e377..41e366c2b5e 100644 --- a/litellm/proxy/spend_tracking/performance_alerts.py +++ b/litellm/proxy/spend_tracking/performance_alerts.py @@ -8,6 +8,7 @@ from litellm._logging import verbose_proxy_logger from litellm.models.performance_alert import ( Claim, PerformanceAlertStore, + ResourceUnreadableError, SweepOwnershipLostError, SweepState, new_sweep, @@ -25,20 +26,15 @@ 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): + except ResourceUnreadableError: pass try: return await store.evaluate_page(owner, state, 1) - except (PrismaError, ValidationError, TimeoutError) as exc: + except ResourceUnreadableError: verbose_proxy_logger.warning( - "Performance alerts skipped one %s it could not read within its limits (%s)", - state.scope, - type(exc).__name__, + "Performance alerts skipped one %s it could not read within its limits", state.scope ) return await store.skip_next(owner, state) diff --git a/litellm/repositories/performance_alert_repository.py b/litellm/repositories/performance_alert_repository.py index a49f1e64e70..125f5db8e74 100644 --- a/litellm/repositories/performance_alert_repository.py +++ b/litellm/repositories/performance_alert_repository.py @@ -6,13 +6,14 @@ 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 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, @@ -26,6 +27,7 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient KEY_DATE_INDEX: Final = "LiteLLM_DailyUserSpend_api_key_date_idx" +_STATEMENT_TIMED_OUT: Final = "`57014`" _STATEMENT_TIMEOUT: Final = "2s" _TRANSACTION_TIMEOUT: Final = timedelta(seconds=5) @@ -160,6 +162,10 @@ 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""" +def is_statement_timeout(error: Exception) -> bool: + return _STATEMENT_TIMED_OUT in str(error) + + class PerformanceAlertRepository: def __init__(self, prisma_client: "PrismaClient") -> None: self._database: Final = cast( # cast-ok: generated Prisma transactions expose the bounded SQL surface @@ -201,6 +207,18 @@ class PerformanceAlertRepository: 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 is_statement_timeout(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)) diff --git a/tests/unit/proxy/spend_tracking/test_performance_alerts.py b/tests/unit/proxy/spend_tracking/test_performance_alerts.py index ac284de0352..edf3855007b 100644 --- a/tests/unit/proxy/spend_tracking/test_performance_alerts.py +++ b/tests/unit/proxy/spend_tracking/test_performance_alerts.py @@ -4,7 +4,14 @@ from typing import Final import pytest -from litellm.models.performance_alert import Claim, SweepOwnershipLostError, SweepState, advance, new_sweep +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) @@ -48,9 +55,7 @@ class Store: 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") + raise ResourceUnreadableError() advanced: Final = advance(state, ids, limit) await self.save(owner, advanced) return advanced @@ -154,3 +159,16 @@ async def test_losing_the_claim_mid_run_stops_quietly() -> None: 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 == [] diff --git a/tests/unit/repositories/test_performance_alert_repository.py b/tests/unit/repositories/test_performance_alert_repository.py new file mode 100644 index 00000000000..a174be777e3 --- /dev/null +++ b/tests/unit/repositories/test_performance_alert_repository.py @@ -0,0 +1,25 @@ +import pytest + +from litellm.repositories.performance_alert_repository import is_statement_timeout + + +def _raw_query_error(message: str) -> Exception: + from prisma.errors import RawQueryError + + return RawQueryError({"user_facing_error": {"is_panic": False, "message": message, "meta": {}}}) + + +@pytest.mark.parametrize( + ("message", "expected"), + ( + ("Raw query failed. Code: `57014`. Message: `ERROR: canceling statement due to statement timeout`", True), + ("Error querying the database: FATAL: sorry, too many clients already", False), + ( + "Raw query failed. Code: `57P01`. Message: `FATAL: terminating connection due to administrator command`", + False, + ), + ("Raw query failed. Code: `40P01`. Message: `ERROR: deadlock detected`", False), + ), +) +def test_only_a_statement_timeout_marks_a_resource_unreadable(message: str, expected: bool) -> None: + assert is_statement_timeout(_raw_query_error(message)) is expected