refactor(alerts): build the key index in a migration instead of after startup

litellm-proxy-extras/litellm_proxy_extras/migrations/20261109090100_daily_user_spend_api_key_date_index/migration.sql

@@ -0,0 +1,2 @@
+CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_DailyUserSpend_api_key_date_idx"
+ON "LiteLLM_DailyUserSpend" ("api_key", "date");

litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py

@@ -59,7 +59,6 @@ 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

litellm/models/performance_alert.py

@@ -81,8 +81,6 @@ class ResourceUnreadableError(Exception):
 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: ...
 

litellm/proxy/spend_tracking/performance_alerts.py

@@ -12,10 +12,8 @@ from litellm.models.performance_alert import (
     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
@@ -75,14 +73,8 @@ async def run_performance_alerts(
         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)
-                await store.save(claim.owner, state)
-                continue
             state = await _step(store, claim.owner, state)
             await sleep(PAUSE_SECONDS)
         if state.completed:

litellm/repositories/performance_alert_repository.py

@@ -26,7 +26,6 @@ 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_TIMED_OUT: Final = "`57014`"
 _STATEMENT_TIMEOUT: Final = "2s"
 _TRANSACTION_TIMEOUT: Final = timedelta(seconds=5)
@@ -160,7 +159,6 @@ SET day = $2, scope = $3, resume_after = $4, wrapped = $5, checked = $6, skipped
 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"""
 
 
 def is_statement_timeout(error: Exception) -> bool:
@@ -200,14 +198,6 @@ class PerformanceAlertRepository:
         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."""
         from prisma.errors import RawQueryError

tests/integration/database/test_performance_alerts.py

@@ -10,7 +10,6 @@ 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
@@ -145,7 +144,6 @@ async def fixture() -> AsyncIterator[Fixture]:
 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)

tests/unit/proxy/spend_tracking/test_performance_alerts.py

@@ -25,7 +25,6 @@ class Store:
     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)
@@ -37,10 +36,7 @@ class Store:
     async def accounting_pending(self, day: date) -> bool:
         return self.pending
 
-    async def key_index_ready(self) -> bool:
+    def _keep(self, state: SweepState) -> None:
-        return self.index_ready
-
-    async def save(self, owner: str, state: SweepState) -> None:
         self.saved.append(state)
         self.state = state
 
@@ -57,13 +53,13 @@ class Store:
         if self.unreadable.intersection(ids):
             raise ResourceUnreadableError()
         advanced: Final = advance(state, ids, limit)
-        await self.save(owner, advanced)
+        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))
-        await self.save(owner, advanced)
+        self._keep(advanced)
         return advanced
 
 
@@ -129,14 +125,6 @@ async def test_an_unreadable_resource_is_retried_alone_then_skipped() -> None:
     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]