diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261109090100_daily_user_spend_api_key_date_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261109090100_daily_user_spend_api_key_date_index/migration.sql new file mode 100644 index 00000000000..5671aa395d1 --- /dev/null +++ b/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"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py index 91eddf46782..4c856c741d3 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/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 diff --git a/litellm/models/performance_alert.py b/litellm/models/performance_alert.py index 09a6210a5b0..7f772318bf3 100644 --- a/litellm/models/performance_alert.py +++ b/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: ... diff --git a/litellm/proxy/spend_tracking/performance_alerts.py b/litellm/proxy/spend_tracking/performance_alerts.py index 4902a22e05e..81347c8f78a 100644 --- a/litellm/proxy/spend_tracking/performance_alerts.py +++ b/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: diff --git a/litellm/repositories/performance_alert_repository.py b/litellm/repositories/performance_alert_repository.py index b312f6951de..241b7925162 100644 --- a/litellm/repositories/performance_alert_repository.py +++ b/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 diff --git a/tests/integration/database/test_performance_alerts.py b/tests/integration/database/test_performance_alerts.py index 735f135065c..44c3b6d524e 100644 --- a/tests/integration/database/test_performance_alerts.py +++ b/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) diff --git a/tests/unit/proxy/spend_tracking/test_performance_alerts.py b/tests/unit/proxy/spend_tracking/test_performance_alerts.py index edf3855007b..9885374a573 100644 --- a/tests/unit/proxy/spend_tracking/test_performance_alerts.py +++ b/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: - return self.index_ready - - async def save(self, owner: str, state: SweepState) -> None: + def _keep(self, 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]