diff --git a/litellm/repositories/usage_alert_health.py b/litellm/repositories/usage_alert_health.py index 646aa0421c0..f1ee7643672 100644 --- a/litellm/repositories/usage_alert_health.py +++ b/litellm/repositories/usage_alert_health.py @@ -1,5 +1,6 @@ """Sample successful aggregate flushes without adding request-time telemetry.""" +import asyncio from datetime import datetime, timedelta, timezone from typing import ( TYPE_CHECKING, @@ -11,6 +12,7 @@ from typing import ( from pydantic import BaseModel, ConfigDict, TypeAdapter +from litellm.repositories.analytics_budget import AnalyticsDatabase, BoundedAnalyticsDatabase, analytics_deadline from litellm.repositories.usage_alert_metrics import MetricsDatabase if TYPE_CHECKING: @@ -34,13 +36,18 @@ class UsageAlertHealthRepository: @classmethod def from_prisma(cls, prisma: "PrismaClient") -> "UsageAlertHealthRepository": - database: Final = cast(MetricsDatabase, prisma.writer_db) # cast-ok: generated Prisma delegate - return cls(database) + database: Final = cast(AnalyticsDatabase, prisma.writer_db) # cast-ok: generated Prisma transaction delegate + return cls(BoundedAnalyticsDatabase(database, background=True)) async def record(self, source: FlushSource, now: datetime) -> bool: hour: Final = now.astimezone(timezone.utc).replace(minute=0, second=0, microsecond=0, tzinfo=None) rows: Final = await self._database.query_raw( - """INSERT INTO "LiteLLM_UsageAlertFlush" (source,hour,observed_at) + """WITH removed AS ( + DELETE FROM "LiteLLM_UsageAlertFlush" WHERE (source,hour) IN ( + SELECT source,hour FROM "LiteLLM_UsageAlertFlush" + WHERE hour<$2::timestamp-INTERVAL '35 days' ORDER BY hour,source LIMIT 10 + ) + ) INSERT INTO "LiteLLM_UsageAlertFlush" (source,hour,observed_at) SELECT $1,$2::timestamp,$3::timestamp WHERE NOT EXISTS ( SELECT 1 FROM "LiteLLM_SpendLogDelivery" WHERE state='pending' AND created_at<$2::timestamp LIMIT 1 ) ON CONFLICT (source,hour) DO UPDATE SET observed_at=EXCLUDED.observed_at RETURNING source""", @@ -48,9 +55,6 @@ class UsageAlertHealthRepository: hour, now.astimezone(timezone.utc).replace(tzinfo=None), ) - await self._database.query_raw( - 'DELETE FROM "LiteLLM_UsageAlertFlush" WHERE hour<$1::timestamp RETURNING source', hour - timedelta(days=35) - ) return bool(rows) async def coverage(self, source: FlushSource, start: datetime) -> FlushCoverage: @@ -70,11 +74,11 @@ class UsageAlertHealthRepository: class UsageAlertFlushRecorder: - """At most one heartbeat attempt per source per worker per hour after success.""" + """Retry optional freshness once a minute, and stop after an hourly success.""" def __init__(self) -> None: - self._user_hour: datetime | None = None - self._team_hour: datetime | None = None + self._user_after: datetime | None = None + self._team_after: datetime | None = None async def record(self, prisma_client: "PrismaClient", source: FlushSource) -> None: from litellm._logging import verbose_proxy_logger @@ -84,17 +88,21 @@ class UsageAlertFlushRecorder: return now: Final = datetime.now(timezone.utc) hour: Final = now.replace(minute=0, second=0, microsecond=0) - prior: Final = self._user_hour if source == "user" else self._team_hour - if prior == hour: + prior: Final = self._user_after if source == "user" else self._team_after + if prior is not None and now < prior: return + self._defer(source, now + timedelta(minutes=1)) try: - database: Final = cast(MetricsDatabase, prisma_client.writer_db) # cast-ok: generated Prisma delegate - if not await UsageAlertHealthRepository(database).record(source, now): - return + with analytics_deadline(asyncio.get_running_loop().time() + 1): + if not await UsageAlertHealthRepository.from_prisma(prisma_client).record(source, now): + return except Exception: # noqa: BLE001 # observability must not fail or requeue a committed spend flush verbose_proxy_logger.warning("Could not record usage-alert aggregate freshness") return + self._defer(source, hour + timedelta(hours=1)) + + def _defer(self, source: FlushSource, until: datetime) -> None: if source == "user": - self._user_hour = hour + self._user_after = until else: - self._team_hour = hour + self._team_after = until diff --git a/tests/integration/database/test_usage_alert_evaluation.py b/tests/integration/database/test_usage_alert_evaluation.py index 4b926174e14..ddb3baf1b8b 100644 --- a/tests/integration/database/test_usage_alert_evaluation.py +++ b/tests/integration/database/test_usage_alert_evaluation.py @@ -1,10 +1,11 @@ +import asyncio from datetime import date, datetime, timedelta, timezone from typing import Final import pytest from integration._support.actors import ManagementDatabase -from litellm.repositories.usage_alert_health import UsageAlertHealthRepository +from litellm.repositories.usage_alert_health import UsageAlertFlushRecorder, UsageAlertHealthRepository from litellm.repositories.usage_alerts import UsageAlertsRepository from litellm.types.proxy.usage_alerts import UsageAlertRuleInput @@ -53,3 +54,17 @@ async def test_freshness_timestamps_use_the_real_prisma_parameter_types(db: Mana start.replace(tzinfo=None), (start + timedelta(days=1, hours=1)).replace(tzinfo=None), ) + + +@pytest.mark.asyncio(loop_scope="session") +async def test_freshness_lock_cannot_stall_committed_spend(db: ManagementDatabase) -> None: + recorder: Final = UsageAlertFlushRecorder() + async with db.prisma.writer_db.tx(timeout=timedelta(seconds=10)) as transaction: + await transaction.execute_raw('LOCK TABLE "LiteLLM_UsageAlertFlush" IN ACCESS EXCLUSIVE MODE') + await asyncio.wait_for(recorder.record(db.prisma, "user"), timeout=2) + await asyncio.wait_for(recorder.record(db.prisma, "user"), timeout=0.1) + await UsageAlertFlushRecorder().record(db.prisma, "user") + rows: Final = await db.prisma.writer_db.query_raw( + "SELECT source FROM \"LiteLLM_UsageAlertFlush\" WHERE source='user' AND hour=date_trunc('hour',CURRENT_TIMESTAMP AT TIME ZONE 'UTC')" + ) + assert rows == [{"source": "user"}]