From 7ba47a5b6e0a1c31f80b8ebdec8bce890aee859a Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 12:12:15 -0700 Subject: [PATCH 1/5] fix(budgets): page end-user cache invalidation after a budget reset The budget-tier reset read every customer linked to an expiring tier into one result set before the write, then invalidated their caches one key at a time. Both of those scale with the customer count, so a large enough deployment can OOM the proxy pod on the read, and the tail of the population sits on a stale spend counter while the per-key invalidations drain PR #40639 moved the reset write itself to a link-based UPDATE, so that pre-commit read no longer feeds the write. It only fed cache invalidation and the service-logging counts, which means it can move after the commit. This replaces it with a keyset walk over litellm_endusertable ordered by user_id, taking RESET_BUDGET_JOB_BATCH_SIZE rows per page, the same shape _reset_windows_for_source already uses, with no per-run page cap for the same reason that walk has none: the cursor cannot survive the run, so a cap would restart at the first customer on every tick and never reach the tail Each page's counter and cache keys now go out as one batched delete through a new DualCache.async_delete_cache_keys, which drops the in-memory entries and chunks the Redis DELETE at DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE num_endusers_found and num_endusers_updated now report the customers whose caches were invalidated after the commit rather than the rows read before it, so both read 0 when the cascade write fails --- litellm/caching/dual_cache.py | 17 ++ .../proxy/common_utils/reset_budget_job.py | 155 +++++++++++------- .../test_proxy_budget_reset.py | 22 ++- tests/test_litellm/caching/test_dual_cache.py | 32 ++++ .../common_utils/test_reset_budget_job.py | 109 ++++++++++-- 5 files changed, 262 insertions(+), 73 deletions(-) diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 81e2af45686..f98e4cca5d1 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -521,6 +521,23 @@ class DualCache(BaseCache): if self.redis_cache is not None: await self.redis_cache.async_delete_cache(key) + async def async_delete_cache_keys(self, keys: Sequence[str]) -> None: + """Batch twin of ``async_delete_cache``: one Redis round trip per chunk + instead of one per key. + + Chunked because Redis takes the whole list as a single DELETE command, + and a caller holding a population-sized list would otherwise build one + command out of it. + """ + if not keys: + return + for key in keys: + self.in_memory_cache.delete_cache(key) + if self.redis_cache is None: + return + for start in range(0, len(keys), DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE): + await self.redis_cache.delete_cache_keys(keys[start : start + DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE]) + async def async_get_ttl(self, key: str) -> int | None: """ Get the remaining TTL of a key in in-memory cache or redis diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index acb51e73daf..d9f7ab37eaa 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -26,7 +26,6 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import ( DB_RETRY_SAFE_ERROR_TYPES, LiteLLM_BudgetTableFull, - LiteLLM_EndUserTable, Litellm_EntityType, LiteLLM_TeamTable, LiteLLM_UserTable, @@ -193,13 +192,6 @@ def _enduser_cache_keys(row: _EndUserRow) -> tuple[str, ...]: return (end_user_cache_key(row.user_id),) -def _enduser_carried_spend(row: _EndUserRow, caps: Mapping[str, float]) -> float: - if not caps: - return 0.0 - effective_budget_id: Final[str | None] = row.budget_id or litellm.max_end_user_budget_id - return _carried_spend(row.spend, caps.get(effective_budget_id) if effective_budget_id is not None else None) - - def _budget_link_where( budget_ids: Sequence[str], extra: Mapping[str, object] = MappingProxyType({}), @@ -207,6 +199,21 @@ def _budget_link_where( return {"budget_id": {"in": list(budget_ids)}, **extra} +def _enduser_invalidation_where(budget_ids: Sequence[str]) -> dict[str, object]: + """Customers whose cached spend a committed reset of these tiers invalidated. + + Mirrors ``_queue_enduser_resets``: the link, plus the NULL-budget_id rows + that ride the default tier when that tier is one of the expiring ones. The + write's ``spend > 0`` filter has no twin here because the commit already + zeroed those rows, so post-commit it would match nobody. + """ + linked: Final = _budget_link_where(budget_ids) + default_budget_id: Final = litellm.max_end_user_budget_id + if default_budget_id is None or default_budget_id not in budget_ids: + return linked + return {"OR": [linked, {"budget_id": None}]} # mutable-ok: prisma where filter must be a dict + + def _queue_budget_linked_resets( writes: LinkedSpendResetWrites, cascade: "_BudgetCascade", @@ -265,7 +272,6 @@ class _BudgetCascade: budgets: tuple[LiteLLM_BudgetTableFull, ...] = () budget_ids: tuple[str, ...] = () budget_resets: tuple[tuple[str, datetime], ...] = () - endusers: tuple[_EndUserRow, ...] = () counter_resets: tuple[tuple[str, float], ...] = () cache_keys: tuple[str, ...] = () rollover_caps: Mapping[str, float] = field(default_factory=lambda: MappingProxyType({})) @@ -275,6 +281,7 @@ class _BudgetCascade: class _BudgetCascadeCommitted: cascade: _BudgetCascade advanced: int + endusers_invalidated: int = 0 @dataclass(frozen=True, slots=True) @@ -416,10 +423,10 @@ _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = ( ) -def _budget_cascade_event_metadata(cascade: _BudgetCascade) -> dict[str, object]: +def _budget_cascade_event_metadata(cascade: _BudgetCascade, endusers_invalidated: int = 0) -> dict[str, object]: return { "num_budgets_found": len(cascade.budgets), - "num_endusers_found": len(cascade.endusers), + "num_endusers_found": endusers_invalidated, } @@ -593,6 +600,32 @@ class ResetBudgetJob: e, ) + @staticmethod + async def _invalidate_caches(counter_keys: Sequence[str], cache_keys: Sequence[str]) -> None: + """Batch twin of ``_invalidate_spend_counter`` and + ``_invalidate_user_api_key_cache_entry``, carrying the same + after-the-commit requirement as both. + + One round trip per chunk rather than one per key: a tier's dependent + population is unbounded, and awaiting each key in turn makes the last + dependent wait out every dependent ahead of it. + """ + if not counter_keys and not cache_keys: + return + try: + from litellm.proxy.proxy_server import spend_counter_cache, user_api_key_cache + + await spend_counter_cache.async_delete_cache_keys(counter_keys) + await user_api_key_cache.async_delete_cache_keys(cache_keys) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to invalidate %d spend counters and %d user_api_key_cache entries: %s. " + "Budgets may be over-enforced until the counters expire.", + len(counter_keys), + len(cache_keys), + e, + ) + async def _fetch_linked_rows( self, table: SpendLinkedTable[_RowT], @@ -612,18 +645,54 @@ class ResetBudgetJob: verbose_proxy_logger.warning("Failed to fetch %s for counter invalidation: %s", log_subject, e) return () - async def _collect_endusers_to_reset(self, budget_ids: Sequence[str]) -> tuple[_EndUserRow, ...]: - linked: Final[Sequence[_EndUserRow] | None] = await self._with_db_retry( - lambda: self.prisma_client.get_data( - table_name="enduser", - query_type="find_all", - budget_id_list=list(budget_ids), - ), - reason="reset_budget_read_endusers_failure", - ) - if litellm.max_end_user_budget_id is None or litellm.max_end_user_budget_id not in budget_ids: - return tuple(linked or ()) - return (*(linked or ()), *await self._get_endusers_with_no_budget_id()) + async def _invalidate_enduser_caches(self, budget_ids: Sequence[str]) -> int: + """Drop the cached spend of every customer the committed tier reset zeroed. + + Walked a page at a time with a keyset cursor, for the same reason + ``_reset_windows_for_source`` is: the customers sharing one tier are + unbounded, so reading them into one result set puts a + customer-count-sized list in the proxy's heap on every tick, and a + deployment large enough turns that into an OOM rather than a slow tick. + + No per-run page cap, also for that walk's reason: the position cannot + survive the run, so a cap would restart at the first customer every tick + and never reach the tail. The cursor strictly advances, so this + terminates on its own. + """ + if not budget_ids: + return 0 + where: Final = _enduser_invalidation_where(budget_ids) + cursor = "" + invalidated = 0 + while True: + rows = await self._fetch_enduser_page(where=where, cursor=cursor) + if not rows: + return invalidated + await self._invalidate_caches( + counter_keys=tuple(_enduser_counter_key(row) for row in rows), + cache_keys=tuple(key for row in rows for key in _enduser_cache_keys(row)), + ) + invalidated += len(rows) + if len(rows) < RESET_BUDGET_JOB_BATCH_SIZE: + return invalidated + cursor = rows[-1].user_id + + async def _fetch_enduser_page(self, where: Mapping[str, object], cursor: str) -> tuple[_EndUserRow, ...]: + """One keyset page of customers, ordered by primary key so the cursor never repeats a row.""" + try: + return tuple( + await self._with_db_retry( + lambda: EndUserRepository(self.prisma_client).table.find_many( + where={**where, "user_id": {"gt": cursor}}, # mutable-ok: prisma where filter must be a dict + order={"user_id": "asc"}, # mutable-ok: prisma order filter must be a dict + take=RESET_BUDGET_JOB_BATCH_SIZE, + ), + reason="reset_budget_read_endusers_failure", + ) + ) + except Exception as e: + verbose_proxy_logger.warning("Failed to fetch end users for cache invalidation: %s", e) + return () async def _collect_budget_cascade(self, budgets_to_reset: Sequence[LiteLLM_BudgetTableFull]) -> _BudgetCascade: """Resolve every row the expiring budget tiers gate, before any write. @@ -670,7 +739,6 @@ class ResetBudgetJob: if _rollover_enabled() else {} # mutable-ok: empty sentinel immediately frozen by MappingProxyType ) - endusers: Final[tuple[_EndUserRow, ...]] = await self._collect_endusers_to_reset(budget_ids) return _BudgetCascade( budgets=tuple(budgets_to_reset), budget_ids=budget_ids, @@ -682,7 +750,6 @@ class ResetBudgetJob: for b in budgets_to_reset if b.budget_id is not None and b.budget_duration is not None ), - endusers=endusers, counter_resets=( *( (_team_membership_counter_key(row), _row_carried_spend(row, rollover_caps)) @@ -695,7 +762,6 @@ class ResetBudgetJob: (_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in model_access_groups ), - *((_enduser_counter_key(row), _enduser_carried_spend(row, rollover_caps)) for row in endusers), ), rollover_caps=rollover_caps, cache_keys=( @@ -704,7 +770,6 @@ class ResetBudgetJob: *(key for row in orgs for key in _org_cache_keys(row)), *(key for row in tags for key in _tag_cache_keys(row)), *(key for row in model_access_groups for key in _model_access_group_cache_keys(row)), - *(key for row in endusers for key in _enduser_cache_keys(row)), ), ) @@ -736,10 +801,10 @@ class ResetBudgetJob: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) async def _invalidate_budget_cascade_caches(self, cascade: _BudgetCascade) -> None: - for counter_key, _ in cascade.counter_resets: - await self._invalidate_spend_counter(counter_key) - for cache_key in cascade.cache_keys: - await self._invalidate_user_api_key_cache_entry(cache_key) + await self._invalidate_caches( + counter_keys=tuple(counter_key for counter_key, _ in cascade.counter_resets), + cache_keys=cascade.cache_keys, + ) async def _reset_expired_budget_cascade(self) -> _BudgetCascadeCommitted | _BudgetCascadeFailed: now: Final = datetime.now(timezone.utc) @@ -769,6 +834,7 @@ class ResetBudgetJob: (reset_at for _, reset_at in cascade.budget_resets), cutoff=datetime.now(timezone.utc), ), + endusers_invalidated=await self._invalidate_enduser_caches(cascade.budget_ids), ) async def reset_budget_for_litellm_budget_table(self) -> None: @@ -788,7 +854,7 @@ class ResetBudgetJob: end_time: Final = time.time() match outcome: - case _BudgetCascadeCommitted(cascade=cascade, advanced=advanced): + case _BudgetCascadeCommitted(cascade=cascade, advanced=advanced, endusers_invalidated=endusers_invalidated): asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_success_hook( service=ServiceTypes.RESET_BUDGET_JOB, @@ -797,8 +863,8 @@ class ResetBudgetJob: start_time=start_time, end_time=end_time, event_metadata={ - **_budget_cascade_event_metadata(cascade), - "num_endusers_updated": len(cascade.endusers), + **_budget_cascade_event_metadata(cascade, endusers_invalidated), + "num_endusers_updated": endusers_invalidated, "num_endusers_failed": 0, }, ) @@ -827,27 +893,6 @@ class ResetBudgetJob: case _: assert_never(outcome) - async def _get_endusers_with_no_budget_id( - self, - ) -> list[LiteLLM_EndUserTable]: - """ - Fetch end users that have no explicit budget_id set (NULL) and have - accumulated spend > 0. These are implicitly-created end users that - rely on the default budget (litellm.max_end_user_budget_id) applied - in-memory during auth checks. - """ - table: Final = EndUserRepository(self.prisma_client).table - rows: Final = await self._with_db_retry( - lambda: table.find_many( - where={ - "budget_id": None, - "spend": {"gt": 0}, - }, - ), - reason="reset_budget_read_endusers_without_budget_id_failure", - ) - return [LiteLLM_EndUserTable.model_validate(row.model_dump()) for row in rows] - async def _write_key_reset_updates(self, updated_keys: Sequence[_RowReset[LiteLLM_VerificationToken]]) -> None: """ Write per-row {spend, budget_reset_at} updates for keys. diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index fe3c38a771f..32bcee7cb2a 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -102,21 +102,24 @@ def _wire_batcher_for_test(prisma_client, fail_commit=False): return batch_calls -def _wire_cascade_reads_for_test(prisma_client): +def _wire_cascade_reads_for_test(prisma_client, endusers=()): """ The budget tier's cascade reads the rows it is about to zero, so their spend counters can be invalidated after the commit. Give each of those tables an awaitable find_many so the reads resolve instead of falling into the job's warn-and-continue path. + + End users are read by the post-commit invalidation walk rather than by + ``get_data``, so callers that care about customers pass them here. """ for table in ( "litellm_teammembership", "litellm_verificationtoken", "litellm_organizationtable", "litellm_tagtable", - "litellm_endusertable", ): getattr(prisma_client.db, table).find_many = AsyncMock(return_value=[]) + prisma_client.db.litellm_endusertable.find_many = AsyncMock(return_value=list(endusers)) @pytest.mark.asyncio @@ -556,7 +559,7 @@ async def test_reset_budget_continues_other_categories_on_failure(): **{u["user_id"]: u["spend"] for u in [user2]}, **{t["team_id"]: t["spend"] for t in [team1, team2]}, } - _wire_cascade_reads_for_test(prisma_client) + _wire_cascade_reads_for_test(prisma_client, endusers=[enduser1]) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -607,7 +610,10 @@ async def test_reset_budget_continues_other_categories_on_failure(): called_tables = { call.kwargs.get("table_name") for call in prisma_client.get_data.await_args_list } - assert called_tables == {"key", "user", "team", "budget", "enduser"} + assert called_tables == {"key", "user", "team", "budget"} + # Customers are not part of that set: the cascade zeroes them by budget link + # and reads them only afterwards, to invalidate their cached spend. + prisma_client.db.litellm_endusertable.find_many.assert_awaited() # Every category writes through the batch path now, so update_data is unused. prisma_client.update_data.assert_not_awaited() @@ -1029,7 +1035,7 @@ async def test_service_logger_endusers_success(): prisma_client.get_data = AsyncMock(side_effect=fake_get_data) prisma_client.update_data = AsyncMock() batch_calls = _wire_batcher_for_test(prisma_client) - _wire_cascade_reads_for_test(prisma_client) + _wire_cascade_reads_for_test(prisma_client, endusers=endusers) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -1094,7 +1100,7 @@ async def test_service_logger_endusers_failure(): prisma_client.get_data = AsyncMock(side_effect=fake_get_data) prisma_client.update_data = AsyncMock() _wire_batcher_for_test(prisma_client, fail_commit=True) - _wire_cascade_reads_for_test(prisma_client) + _wire_cascade_reads_for_test(prisma_client, endusers=endusers) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -1121,7 +1127,9 @@ async def test_service_logger_endusers_failure(): ) = proxy_logging_obj.service_logging_obj.async_service_failure_hook.call_args event_metadata = kwargs.get("event_metadata", {}) assert event_metadata.get("num_budgets_found") == len(budgets) - assert event_metadata.get("num_endusers_found") == len(endusers) + # Customers are read by the post-commit invalidation walk, which a failed + # commit never reaches, so a failure reports none touched. + assert event_metadata.get("num_endusers_found") == 0 assert "endusers_found" not in event_metadata assert "budgets_found" not in event_metadata proxy_logging_obj.service_logging_obj.async_service_success_hook.assert_not_called() diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 95395878c25..5f59de9cca5 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache, _redis_circuit_breaker_guard, _redis_circuit_breaker_guard_sync @@ -759,3 +760,34 @@ async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplo " (199 more Redis timeouts since the previous Redis timeout line were logged at DEBUG)", ) ] + + +@pytest.mark.asyncio +async def test_async_delete_cache_keys_drops_memory_and_chunks_redis(): + """Batch delete clears both layers, and chunks Redis so one caller's large + key list cannot become a single oversized DELETE command.""" + redis_cache = MagicMock(spec=RedisCache) + redis_cache.delete_cache_keys = AsyncMock() + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + keys = [f"key-{i}" for i in range(DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE + 7)] + for key in keys: + dual_cache.in_memory_cache.set_cache(key=key, value=1) + + await dual_cache.async_delete_cache_keys(keys) + + assert all(dual_cache.in_memory_cache.get_cache(key=key) is None for key in keys) + sent = [call.args[0] for call in redis_cache.delete_cache_keys.await_args_list] + assert [len(chunk) for chunk in sent] == [DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE, 7] + assert [key for chunk in sent for key in chunk] == keys + + +@pytest.mark.asyncio +async def test_async_delete_cache_keys_on_empty_list_touches_no_backend(): + """An empty page must not reach Redis: DELETE with no arguments is an error.""" + redis_cache = MagicMock(spec=RedisCache) + redis_cache.delete_cache_keys = AsyncMock() + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + + await dual_cache.async_delete_cache_keys([]) + + redis_cache.delete_cache_keys.assert_not_awaited() diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 943a6c905c0..0f39af3dee3 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -4,7 +4,7 @@ import sys import types from datetime import datetime, timedelta, timezone from datetime import time as dt_time -from typing import Any, Dict, Final, List +from typing import Any, Dict, Final, List, Optional from unittest.mock import AsyncMock, MagicMock import httpx @@ -16,6 +16,7 @@ from litellm.proxy._types import LiteLLM_VerificationToken from litellm.proxy.common_utils import reset_budget_job as reset_budget_job_module from litellm.constants import ( PROXY_BUDGET_RESCHEDULER_MIN_TIME, + RESET_BUDGET_JOB_BATCH_SIZE, RESET_BUDGET_JOB_LOCK_TTL_SECONDS, RESET_BUDGET_JOB_NAME, ) @@ -35,9 +36,24 @@ class MockTable: def set_find_many_results(self, results: List[Any]): self._find_many_results = results - async def find_many(self, where: Dict[str, Any]) -> List[Any]: - self.find_many_calls.append({"where": where}) - return self._find_many_results + async def find_many( + self, + where: Dict[str, Any], + order: Optional[Dict[str, str]] = None, + take: Optional[int] = None, + ) -> List[Any]: + """Replays canned rows, honouring the keyset cursor + ``take`` a paged + caller relies on: without that a paged walk never advances and the + test would hang instead of failing.""" + paging = {k: v for k, v in (("order", order), ("take", take)) if v is not None} + self.find_many_calls.append({"where": where, **paging}) + rows = list(self._find_many_results) + for field, condition in where.items(): + if isinstance(condition, dict) and "gt" in condition and field != "spend": + rows = [row for row in rows if getattr(row, field, "") > condition["gt"]] + for field, direction in (order or {}).items(): + rows.sort(key=lambda row: getattr(row, field, ""), reverse=direction == "desc") + return rows[:take] if take is not None else rows async def update_many(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: self.update_many_calls.append({"where": where, "data": data}) @@ -784,10 +800,16 @@ def test_reset_budget_resets_endusers_with_null_budget_id(reset_budget_job, mock }, ] - # Verify find_many was called to fetch NULL-budget-id end users + # The post-commit invalidation walk covers both branches, so implicitly + # created customers on the default tier get their cached spend dropped too, + # and it is paged rather than reading the whole customer population. find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls assert len(find_many_calls) == 1 - assert find_many_calls[0]["where"] == {"budget_id": None, "spend": {"gt": 0}} + assert find_many_calls[0]["where"]["OR"] == [ + {"budget_id": {"in": [default_budget_id]}}, + {"budget_id": None}, + ] + assert find_many_calls[0]["take"] == RESET_BUDGET_JOB_BATCH_SIZE litellm.max_end_user_budget_id = None @@ -818,9 +840,12 @@ def test_reset_budget_skips_null_budget_id_endusers_when_default_not_configured( asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - # Should NOT have queried for NULL-budget-id end users + # The invalidation walk must not reach for NULL-budget-id customers: they + # ride a default tier that is not expiring, so their spend stays put. find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls - assert len(find_many_calls) == 0 + assert [call["where"] for call in find_many_calls] == [ + {"budget_id": {"in": ["some-budget"]}, "user_id": {"gt": ""}} + ] litellm.max_end_user_budget_id = None @@ -855,9 +880,12 @@ def test_reset_budget_skips_null_budget_id_endusers_when_default_not_in_reset_li asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - # Should NOT have queried for NULL-budget-id end users + # The invalidation walk must not reach for NULL-budget-id customers: they + # ride a default tier that is not expiring, so their spend stays put. find_many_calls = mock_prisma_client.db.litellm_endusertable.find_many_calls - assert len(find_many_calls) == 0 + assert [call["where"] for call in find_many_calls] == [ + {"budget_id": {"in": ["other-budget"]}, "user_id": {"gt": ""}} + ] litellm.max_end_user_budget_id = None @@ -1235,6 +1263,21 @@ def _make_counter_invalidation_job(monkeypatch): user_api_key_cache = MagicMock() user_api_key_cache.async_delete_cache = AsyncMock() + # Batch deletes fan out to the same per-key calls the real DualCache makes, + # so an assertion reads "this key was invalidated" whether the caller went + # one key at a time or a page at a time. + async def _delete_counter_keys(keys): + for key in keys: + spend_counter_cache.in_memory_cache.delete_cache(key=key) + await spend_counter_cache.redis_cache.async_delete_cache(key=key) + + async def _delete_management_keys(keys): + for key in keys: + await user_api_key_cache.async_delete_cache(key=key) + + spend_counter_cache.async_delete_cache_keys = AsyncMock(side_effect=_delete_counter_keys) + user_api_key_cache.async_delete_cache_keys = AsyncMock(side_effect=_delete_management_keys) + fake_module = types.ModuleType("litellm.proxy.proxy_server") fake_module.spend_counter_cache = spend_counter_cache fake_module.user_api_key_cache = user_api_key_cache @@ -1569,7 +1612,7 @@ def test_budget_table_reset_invalidates_enduser_counter_and_cache(reset_budget_j "user_id": "customer-42", }, ) - mock_prisma_client.data["enduser"] = [test_enduser] + mock_prisma_client.db.litellm_endusertable.set_find_many_results([test_enduser]) asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) @@ -1579,6 +1622,50 @@ def test_budget_table_reset_invalidates_enduser_counter_and_cache(reset_budget_j assert "end_user_id:customer-42" in deleted +def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma_client, monkeypatch): + """The post-commit invalidation walk stays bounded in memory and in round trips. + + Reading every customer on an expiring tier into one result set puts a + customer-count-sized list in the proxy's heap on every tick, which is an OOM + on a large enough deployment rather than a slow tick. Awaiting one cache call + per customer makes the last customer wait out every customer ahead of it. + Both regress silently, so pin the page size, the strictly advancing cursor, + and one batched call per page. + """ + counter_cache: Final = _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")] + population: Final = RESET_BUDGET_JOB_BATCH_SIZE * 2 + 3 + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [ + type("EndUser", (), {"user_id": f"cust-{i:06d}", "spend": 5.0, "budget_id": "budget-1"}) + for i in range(population) + ] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + reads: Final = mock_prisma_client.db.litellm_endusertable.find_many_calls + assert [read["take"] for read in reads] == [RESET_BUDGET_JOB_BATCH_SIZE] * 3 + assert [read["where"]["user_id"]["gt"] for read in reads] == [ + "", + f"cust-{RESET_BUDGET_JOB_BATCH_SIZE - 1:06d}", + f"cust-{RESET_BUDGET_JOB_BATCH_SIZE * 2 - 1:06d}", + ] + + assert counter_cache.async_delete_cache_keys.await_count == 3 + assert counter_cache.user_api_key_cache.async_delete_cache_keys.await_count == 3 + counter_cache.async_delete_cache.assert_not_called() + + invalidated: Final = { + key for call in counter_cache.async_delete_cache_keys.await_args_list for key in call.args[0] + } + assert invalidated == {f"spend:end_user:cust-{i:06d}" for i in range(population)} + evicted: Final = { + key for call in counter_cache.user_api_key_cache.async_delete_cache_keys.await_args_list for key in call.args[0] + } + assert evicted == {f"end_user_id:cust-{i:06d}" for i in range(population)} + + def test_budget_table_reset_commits_even_when_cache_eviction_fails(reset_budget_job, mock_prisma_client, monkeypatch): """Eviction runs after the commit, so a broken cache cannot undo the write.""" From cdf0142f4a09b150c4efda2a5dd5b91d7f1d5b88 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 16:14:08 -0700 Subject: [PATCH 2/5] fix(proxy): isolate each cache and each page in the budget reset invalidation Greptile review follow-ups on the paged end-user cache invalidation. UserApiKeyCache keeps hashed token keys in a second in-memory partition, and routes delete_cache / async_delete_cache there. It inherited the new batch delete unchanged, so a budget cascade cleared the main partition and left the key object sitting on its pre-reset spend. Override it the way async_set_cache_pipeline already partitions its entries. The spend counters and the management cache shared one exception handler, so a Redis failure on the counters returned before the management cache was touched at all. Each cache gets its own await and its own handler now. A failed page read returned the same empty tuple that ends the walk normally, so a truncated pass was reported as a complete one. The window is advanced by then and no later tick comes back for the customers past that page, so the walk now says it was cut short and the service log carries it. --- .../proxy/common_utils/reset_budget_job.py | 107 ++++++++++++------ .../proxy/common_utils/user_api_key_cache.py | 8 ++ .../common_utils/test_reset_budget_job.py | 64 +++++++++++ .../common_utils/test_user_api_key_cache.py | 25 ++++ 4 files changed, 171 insertions(+), 33 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index d9f7ab37eaa..dd16a642342 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -277,11 +277,23 @@ class _BudgetCascade: rollover_caps: Mapping[str, float] = field(default_factory=lambda: MappingProxyType({})) +@dataclass(frozen=True, slots=True) +class _EndUserInvalidation: + """How far the post-commit customer walk got, and whether a failed page read + cut it short of the tail.""" + + invalidated: int = 0 + truncated: bool = False + + +_NO_ENDUSERS_INVALIDATED: Final = _EndUserInvalidation() + + @dataclass(frozen=True, slots=True) class _BudgetCascadeCommitted: cascade: _BudgetCascade advanced: int - endusers_invalidated: int = 0 + endusers: _EndUserInvalidation @dataclass(frozen=True, slots=True) @@ -292,6 +304,10 @@ class _BudgetCascadeFailed: _EMPTY_CASCADE: Final = _BudgetCascade() +#: Which of the proxy's two caches a batch of keys belongs to. ``spend_counter_cache`` +#: holds the live running spend; ``user_api_key_cache`` holds the cached management rows. +_InvalidatedCache = Literal["spend counter", "user_api_key_cache"] + @dataclass(frozen=True, slots=True) class _ChunkOutcome: @@ -423,10 +439,12 @@ _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = ( ) -def _budget_cascade_event_metadata(cascade: _BudgetCascade, endusers_invalidated: int = 0) -> dict[str, object]: +def _budget_cascade_event_metadata( + cascade: _BudgetCascade, endusers: _EndUserInvalidation = _NO_ENDUSERS_INVALIDATED +) -> dict[str, object]: return { "num_budgets_found": len(cascade.budgets), - "num_endusers_found": endusers_invalidated, + "num_endusers_found": endusers.invalidated, } @@ -610,19 +628,30 @@ class ResetBudgetJob: population is unbounded, and awaiting each key in turn makes the last dependent wait out every dependent ahead of it. """ - if not counter_keys and not cache_keys: + await ResetBudgetJob._invalidate_cache("spend counter", counter_keys) + await ResetBudgetJob._invalidate_cache("user_api_key_cache", cache_keys) + + @staticmethod + async def _invalidate_cache(cache: _InvalidatedCache, keys: Sequence[str]) -> None: + """One cache's share of a batch, awaited separately from the other's so a + failure against either still leaves the other one invalidated.""" + if not keys: return try: from litellm.proxy.proxy_server import spend_counter_cache, user_api_key_cache - await spend_counter_cache.async_delete_cache_keys(counter_keys) - await user_api_key_cache.async_delete_cache_keys(cache_keys) + match cache: + case "spend counter": + await spend_counter_cache.async_delete_cache_keys(keys) + case "user_api_key_cache": + await user_api_key_cache.async_delete_cache_keys(keys) + case _: + assert_never(cache) except Exception as e: verbose_proxy_logger.warning( - "Failed to invalidate %d spend counters and %d user_api_key_cache entries: %s. " - "Budgets may be over-enforced until the counters expire.", - len(counter_keys), - len(cache_keys), + "Failed to invalidate %d %s entries: %s. Budgets may be over-enforced until they expire.", + len(keys), + cache, e, ) @@ -645,7 +674,7 @@ class ResetBudgetJob: verbose_proxy_logger.warning("Failed to fetch %s for counter invalidation: %s", log_subject, e) return () - async def _invalidate_enduser_caches(self, budget_ids: Sequence[str]) -> int: + async def _invalidate_enduser_caches(self, budget_ids: Sequence[str]) -> _EndUserInvalidation: """Drop the cached spend of every customer the committed tier reset zeroed. Walked a page at a time with a keyset cursor, for the same reason @@ -658,41 +687,52 @@ class ResetBudgetJob: survive the run, so a cap would restart at the first customer every tick and never reach the tail. The cursor strictly advances, so this terminates on its own. + + A page that fails to read stops the walk short of the tail. The window is + already advanced by then, so no later tick comes back for the customers + past it, which is why the walk reports that it was cut short instead of + passing the part it managed off as the whole. """ if not budget_ids: - return 0 + return _NO_ENDUSERS_INVALIDATED where: Final = _enduser_invalidation_where(budget_ids) cursor = "" invalidated = 0 while True: - rows = await self._fetch_enduser_page(where=where, cursor=cursor) + try: + rows = await self._fetch_enduser_page(where=where, cursor=cursor) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to fetch end users for cache invalidation after %s customers (cursor %r): %s. " + "The customers past that page keep their cached spend until it expires.", + invalidated, + cursor, + e, + ) + return _EndUserInvalidation(invalidated=invalidated, truncated=True) if not rows: - return invalidated + return _EndUserInvalidation(invalidated=invalidated) await self._invalidate_caches( counter_keys=tuple(_enduser_counter_key(row) for row in rows), cache_keys=tuple(key for row in rows for key in _enduser_cache_keys(row)), ) invalidated += len(rows) if len(rows) < RESET_BUDGET_JOB_BATCH_SIZE: - return invalidated + return _EndUserInvalidation(invalidated=invalidated) cursor = rows[-1].user_id async def _fetch_enduser_page(self, where: Mapping[str, object], cursor: str) -> tuple[_EndUserRow, ...]: """One keyset page of customers, ordered by primary key so the cursor never repeats a row.""" - try: - return tuple( - await self._with_db_retry( - lambda: EndUserRepository(self.prisma_client).table.find_many( - where={**where, "user_id": {"gt": cursor}}, # mutable-ok: prisma where filter must be a dict - order={"user_id": "asc"}, # mutable-ok: prisma order filter must be a dict - take=RESET_BUDGET_JOB_BATCH_SIZE, - ), - reason="reset_budget_read_endusers_failure", - ) + return tuple( + await self._with_db_retry( + lambda: EndUserRepository(self.prisma_client).table.find_many( + where={**where, "user_id": {"gt": cursor}}, # mutable-ok: prisma where filter must be a dict + order={"user_id": "asc"}, # mutable-ok: prisma order filter must be a dict + take=RESET_BUDGET_JOB_BATCH_SIZE, + ), + reason="reset_budget_read_endusers_failure", ) - except Exception as e: - verbose_proxy_logger.warning("Failed to fetch end users for cache invalidation: %s", e) - return () + ) async def _collect_budget_cascade(self, budgets_to_reset: Sequence[LiteLLM_BudgetTableFull]) -> _BudgetCascade: """Resolve every row the expiring budget tiers gate, before any write. @@ -834,7 +874,7 @@ class ResetBudgetJob: (reset_at for _, reset_at in cascade.budget_resets), cutoff=datetime.now(timezone.utc), ), - endusers_invalidated=await self._invalidate_enduser_caches(cascade.budget_ids), + endusers=await self._invalidate_enduser_caches(cascade.budget_ids), ) async def reset_budget_for_litellm_budget_table(self) -> None: @@ -854,7 +894,7 @@ class ResetBudgetJob: end_time: Final = time.time() match outcome: - case _BudgetCascadeCommitted(cascade=cascade, advanced=advanced, endusers_invalidated=endusers_invalidated): + case _BudgetCascadeCommitted() as committed: asyncio.create_task( self.proxy_logging_obj.service_logging_obj.async_service_success_hook( service=ServiceTypes.RESET_BUDGET_JOB, @@ -863,13 +903,14 @@ class ResetBudgetJob: start_time=start_time, end_time=end_time, event_metadata={ - **_budget_cascade_event_metadata(cascade, endusers_invalidated), - "num_endusers_updated": endusers_invalidated, + **_budget_cascade_event_metadata(committed.cascade, committed.endusers), + "num_endusers_updated": committed.endusers.invalidated, "num_endusers_failed": 0, + "enduser_invalidation_truncated": committed.endusers.truncated, }, ) ) - return _ChunkOutcome(fetched=len(cascade.budgets), advanced=advanced) + return _ChunkOutcome(fetched=len(committed.cascade.budgets), advanced=committed.advanced) case _BudgetCascadeFailed(cascade=cascade, error=error): verbose_proxy_logger.exception( "Failed to reset the budget table cascade (team member, enduser, org, tag and model access " diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 1c7a379897f..4b2f12dcc27 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -221,6 +221,14 @@ class UserApiKeyCache(DualCache): return await super().async_delete_cache(key) + async def async_delete_cache_keys(self, keys: Sequence[str]) -> None: + key_object_keys: Final = tuple(key for key in keys if is_user_key_cache_key(key)) + other_keys: Final = tuple(key for key in keys if not is_user_key_cache_key(key)) + if key_object_keys: + await self.key_object_cache.async_delete_cache_keys(key_object_keys) + if other_keys: + await super().async_delete_cache_keys(other_keys) + def flush_cache(self) -> None: super().flush_cache() self.key_object_cache.in_memory_cache.flush_cache() diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 0f39af3dee3..0606723f6dd 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -1667,6 +1667,70 @@ def test_enduser_invalidation_is_paged_and_batched(reset_budget_job, mock_prisma +def test_enduser_invalidation_reports_a_page_read_failure_instead_of_a_clean_finish( + mock_prisma_client, monkeypatch +): + """A page that fails to read is not the end of the customer list. + + The tier's window is already advanced by the time this walk runs, so no later + tick comes back for the customers past the page that failed: their cached + spend goes on rejecting requests until it expires. Returning the same empty + page normal end-of-data returns hid that behind a report of a clean pass. + """ + _make_counter_invalidation_job(monkeypatch) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")] + endusers: Final = mock_prisma_client.db.litellm_endusertable + endusers.set_find_many_results( + [ + type("EndUser", (), {"user_id": f"cust-{i:06d}", "spend": 5.0, "budget_id": "budget-1"}) + for i in range(RESET_BUDGET_JOB_BATCH_SIZE + 3) + ] + ) + read_page: Final = endusers.find_many + + async def fail_after_the_first_page(**kwargs): + if endusers.find_many_calls: + raise RuntimeError("connection reset while paging customers") + return await read_page(**kwargs) + + endusers.find_many = fail_after_the_first_page + logging_obj: Final = RecordingProxyLogging() + job: Final = ResetBudgetJob(proxy_logging_obj=logging_obj, prisma_client=mock_prisma_client) + + _run_and_drain_hooks(job.reset_budget_for_litellm_budget_table) + + metadata: Final = logging_obj.service_logging_obj.success_calls[0]["event_metadata"] + assert metadata["enduser_invalidation_truncated"] is True + assert metadata["num_endusers_updated"] == RESET_BUDGET_JOB_BATCH_SIZE + + +def test_a_failed_counter_batch_still_evicts_the_management_cache( + reset_budget_job, mock_prisma_client, monkeypatch +): + """The spend counters and the management cache are invalidated independently. + + Sharing one handler meant a Redis failure on the counters returned before the + management cache was touched at all. The commit has already zeroed those rows + by then, so the cached objects keep authorizing against their pre-reset spend + until they expire. + """ + counter_cache: Final = _make_counter_invalidation_job(monkeypatch) + counter_cache.async_delete_cache_keys = AsyncMock(side_effect=RuntimeError("redis unavailable")) + mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-1")] + mock_prisma_client.db.litellm_endusertable.set_find_many_results( + [type("EndUser", (), {"user_id": "customer-42", "spend": 5.0, "budget_id": "budget-1"})] + ) + + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + evicted: Final = { + key + for call in counter_cache.user_api_key_cache.async_delete_cache_keys.await_args_list + for key in call.args[0] + } + assert "end_user_id:customer-42" in evicted + + def test_budget_table_reset_commits_even_when_cache_eviction_fails(reset_budget_job, mock_prisma_client, monkeypatch): """Eviction runs after the commit, so a broken cache cannot undo the write.""" counter_cache = _make_counter_invalidation_job(monkeypatch) diff --git a/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py index 2d5d76ed542..262cb91d670 100644 --- a/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py +++ b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py @@ -82,6 +82,10 @@ class FakeRedisCache(RedisCache): async def async_delete_cache(self, key: str): # type: ignore[override] self._store.pop(key, None) + async def delete_cache_keys(self, keys): # type: ignore[override] + for key in keys: + self._store.pop(key, None) + def _make_key_obj(token: str = "tok") -> UserAPIKeyAuth: # Minimal object (UserAPIKeyAuth inherits token from base view). @@ -331,6 +335,27 @@ class TestUserKeyObjectPartition: assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None assert await redis.async_get_cache(HASHED_TOKEN) is None + @pytest.mark.asyncio + async def test_batch_delete_routes_each_key_to_its_partition(self): + """A batch delete has to clear the same partition the single delete does. + + ``DualCache``'s batch delete only knows about the main in-memory cache, so + inheriting it unchanged leaves a key object sitting in ``key_object_cache`` + with its pre-reset spend, and the next request is authorized against that + stale copy until the local entry expires. + """ + redis = FakeRedisCache() + cache = UserApiKeyCache(redis_cache=redis) + await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + await cache.async_set_cache(end_user_cache_key("u1"), {"user_id": "u1"}) + + await cache.async_delete_cache_keys([HASHED_TOKEN, end_user_cache_key("u1")]) + + assert await cache.async_get_cache(HASHED_TOKEN, model_type=UserAPIKeyAuth) is None + assert await cache.async_get_cache(end_user_cache_key("u1")) is None + assert await redis.async_get_cache(HASHED_TOKEN) is None + assert await redis.async_get_cache(end_user_cache_key("u1")) is None + @pytest.mark.asyncio async def test_pipeline_write_routes_each_entry_to_its_partition(self): cache = UserApiKeyCache(in_memory_cache=InMemoryCache(max_size_in_memory=2)) From 14fbd623d7de87fc01131a969d8321b87501cd98 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 16:28:16 -0700 Subject: [PATCH 3/5] fix(proxy): clear both cache partitions and carry the walk position as a value UserApiKeyCache's batch delete ran the two partitions in sequence, so a Redis failure on the hashed token partition returned before the ordinary management keys were touched. Both partitions are attempted now and the first failure is re-raised for the caller to report. The customer walk kept its position in two locals it reassigned each page. It now mirrors the window walk in the same file: a page helper returns where the walk goes next, and the driver rebinds one value. --- .../proxy/common_utils/reset_budget_job.py | 70 +++++++++++-------- .../proxy/common_utils/user_api_key_cache.py | 21 ++++-- .../common_utils/test_user_api_key_cache.py | 28 ++++++++ 3 files changed, 84 insertions(+), 35 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index dd16a642342..ddabf91bff6 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -278,22 +278,24 @@ class _BudgetCascade: @dataclass(frozen=True, slots=True) -class _EndUserInvalidation: - """How far the post-commit customer walk got, and whether a failed page read - cut it short of the tail.""" +class _EndUserWalk: + """Where the post-commit customer walk stands: the keyset cursor its next + page resumes from, None once there is no next page, how many customers it + has reached, and whether a failed page read cut it short of the tail.""" + cursor: str | None = "" invalidated: int = 0 truncated: bool = False -_NO_ENDUSERS_INVALIDATED: Final = _EndUserInvalidation() +_ENDUSER_WALK_DONE: Final = _EndUserWalk(cursor=None) @dataclass(frozen=True, slots=True) class _BudgetCascadeCommitted: cascade: _BudgetCascade advanced: int - endusers: _EndUserInvalidation + endusers: _EndUserWalk @dataclass(frozen=True, slots=True) @@ -440,7 +442,7 @@ _WINDOW_SOURCES: Final[tuple[_WindowSource, ...]] = ( def _budget_cascade_event_metadata( - cascade: _BudgetCascade, endusers: _EndUserInvalidation = _NO_ENDUSERS_INVALIDATED + cascade: _BudgetCascade, endusers: _EndUserWalk = _ENDUSER_WALK_DONE ) -> dict[str, object]: return { "num_budgets_found": len(cascade.budgets), @@ -674,7 +676,7 @@ class ResetBudgetJob: verbose_proxy_logger.warning("Failed to fetch %s for counter invalidation: %s", log_subject, e) return () - async def _invalidate_enduser_caches(self, budget_ids: Sequence[str]) -> _EndUserInvalidation: + async def _invalidate_enduser_caches(self, budget_ids: Sequence[str]) -> _EndUserWalk: """Drop the cached spend of every customer the committed tier reset zeroed. Walked a page at a time with a keyset cursor, for the same reason @@ -694,32 +696,38 @@ class ResetBudgetJob: passing the part it managed off as the whole. """ if not budget_ids: - return _NO_ENDUSERS_INVALIDATED + return _ENDUSER_WALK_DONE where: Final = _enduser_invalidation_where(budget_ids) - cursor = "" - invalidated = 0 - while True: - try: - rows = await self._fetch_enduser_page(where=where, cursor=cursor) - except Exception as e: - verbose_proxy_logger.warning( - "Failed to fetch end users for cache invalidation after %s customers (cursor %r): %s. " - "The customers past that page keep their cached spend until it expires.", - invalidated, - cursor, - e, - ) - return _EndUserInvalidation(invalidated=invalidated, truncated=True) - if not rows: - return _EndUserInvalidation(invalidated=invalidated) - await self._invalidate_caches( - counter_keys=tuple(_enduser_counter_key(row) for row in rows), - cache_keys=tuple(key for row in rows for key in _enduser_cache_keys(row)), + walk = _EndUserWalk() + while walk.cursor is not None: + walk = await self._invalidate_enduser_page(where=where, cursor=walk.cursor, reached=walk.invalidated) + return walk + + async def _invalidate_enduser_page( + self, where: Mapping[str, object], cursor: str, reached: int + ) -> _EndUserWalk: + """Invalidate one page of customers and say where the walk goes next.""" + try: + rows: Final = await self._fetch_enduser_page(where=where, cursor=cursor) + except Exception as e: + verbose_proxy_logger.warning( + "Failed to fetch end users for cache invalidation after %s customers (cursor %r): %s. " + "The customers past that page keep their cached spend until it expires.", + reached, + cursor, + e, ) - invalidated += len(rows) - if len(rows) < RESET_BUDGET_JOB_BATCH_SIZE: - return _EndUserInvalidation(invalidated=invalidated) - cursor = rows[-1].user_id + return _EndUserWalk(cursor=None, invalidated=reached, truncated=True) + if not rows: + return _EndUserWalk(cursor=None, invalidated=reached) + await self._invalidate_caches( + counter_keys=tuple(_enduser_counter_key(row) for row in rows), + cache_keys=tuple(key for row in rows for key in _enduser_cache_keys(row)), + ) + walked: Final = reached + len(rows) + if len(rows) < RESET_BUDGET_JOB_BATCH_SIZE: + return _EndUserWalk(cursor=None, invalidated=walked) + return _EndUserWalk(cursor=rows[-1].user_id, invalidated=walked) async def _fetch_enduser_page(self, where: Mapping[str, object], cursor: str) -> tuple[_EndUserRow, ...]: """One keyset page of customers, ordered by primary key so the cursor never repeats a row.""" diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 4b2f12dcc27..5a8e3a9482d 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import re from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, TypeVar, cast, overload @@ -222,12 +223,24 @@ class UserApiKeyCache(DualCache): await super().async_delete_cache(key) async def async_delete_cache_keys(self, keys: Sequence[str]) -> None: + """Batch twin of ``async_delete_cache``, partitioned the way + ``async_set_cache_pipeline`` partitions its writes. + + Both partitions are cleared even when one of them raises: a caller + batching these has already committed the rows they cache, so a partition + left holding pre-reset spend goes on being authorized against until the + entry expires. The first failure is re-raised for the caller to report. + """ key_object_keys: Final = tuple(key for key in keys if is_user_key_cache_key(key)) other_keys: Final = tuple(key for key in keys if not is_user_key_cache_key(key)) - if key_object_keys: - await self.key_object_cache.async_delete_cache_keys(key_object_keys) - if other_keys: - await super().async_delete_cache_keys(other_keys) + outcomes: Final = await asyncio.gather( + self.key_object_cache.async_delete_cache_keys(key_object_keys), + super().async_delete_cache_keys(other_keys), + return_exceptions=True, + ) + failed: Final = tuple(outcome for outcome in outcomes if isinstance(outcome, BaseException)) + if failed: + raise failed[0] def flush_cache(self) -> None: super().flush_cache() diff --git a/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py index 262cb91d670..f24175a1922 100644 --- a/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py +++ b/tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py @@ -87,6 +87,15 @@ class FakeRedisCache(RedisCache): self._store.pop(key, None) +class PartitionFailingRedisCache(FakeRedisCache): + """Fails the batch delete for the key-object partition and no other.""" + + async def delete_cache_keys(self, keys): # type: ignore[override] + if any(is_user_key_cache_key(key) for key in keys): + raise ConnectionError("redis unavailable") + await super().delete_cache_keys(keys) + + def _make_key_obj(token: str = "tok") -> UserAPIKeyAuth: # Minimal object (UserAPIKeyAuth inherits token from base view). return UserAPIKeyAuth(token=token) @@ -356,6 +365,25 @@ class TestUserKeyObjectPartition: assert await redis.async_get_cache(HASHED_TOKEN) is None assert await redis.async_get_cache(end_user_cache_key("u1")) is None + @pytest.mark.asyncio + async def test_batch_delete_clears_the_other_partition_when_one_fails(self): + """One partition failing must not cost the other its deletions. + + A caller batching these has already committed the rows they cache, so a + partition that is skipped keeps authorizing against pre-reset spend until + the entry expires. The failure is still raised for the caller to report. + """ + redis = PartitionFailingRedisCache() + cache = UserApiKeyCache(redis_cache=redis) + await cache.async_set_cache(HASHED_TOKEN, _make_key_obj(HASHED_TOKEN), model_type=UserAPIKeyAuth) + await cache.async_set_cache(end_user_cache_key("u1"), {"user_id": "u1"}) + + with pytest.raises(ConnectionError): + await cache.async_delete_cache_keys([HASHED_TOKEN, end_user_cache_key("u1")]) + + assert await cache.async_get_cache(end_user_cache_key("u1")) is None + assert await redis.async_get_cache(end_user_cache_key("u1")) is None + @pytest.mark.asyncio async def test_pipeline_write_routes_each_entry_to_its_partition(self): cache = UserApiKeyCache(in_memory_cache=InMemoryCache(max_size_in_memory=2)) From 0d8b46b88ccd48556111423ba8cdc440815acecb Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 16:41:13 -0700 Subject: [PATCH 4/5] refactor(proxy): trim the invalidation docstrings and inject the page read failure Cuts the new docstrings back to the parts a reader cannot get from the code, and fixes a stale reference: the walk this one is modelled on is _reset_windows_for, not _reset_windows_for_source. The truncation test reached in and replaced MockTable.find_many. The mock takes a scheduled read failure instead, the way it already takes canned rows. --- litellm/caching/dual_cache.py | 9 +--- .../proxy/common_utils/reset_budget_job.py | 44 +++++-------------- .../proxy/common_utils/user_api_key_cache.py | 10 ++--- .../common_utils/test_reset_budget_job.py | 17 +++---- 4 files changed, 26 insertions(+), 54 deletions(-) diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index f98e4cca5d1..66be77dbb40 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -522,13 +522,8 @@ class DualCache(BaseCache): await self.redis_cache.async_delete_cache(key) async def async_delete_cache_keys(self, keys: Sequence[str]) -> None: - """Batch twin of ``async_delete_cache``: one Redis round trip per chunk - instead of one per key. - - Chunked because Redis takes the whole list as a single DELETE command, - and a caller holding a population-sized list would otherwise build one - command out of it. - """ + """Batch twin of ``async_delete_cache``, chunked because Redis takes the + whole list as one DELETE command.""" if not keys: return for key in keys: diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index ddabf91bff6..2260d890e46 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -202,10 +202,8 @@ def _budget_link_where( def _enduser_invalidation_where(budget_ids: Sequence[str]) -> dict[str, object]: """Customers whose cached spend a committed reset of these tiers invalidated. - Mirrors ``_queue_enduser_resets``: the link, plus the NULL-budget_id rows - that ride the default tier when that tier is one of the expiring ones. The - write's ``spend > 0`` filter has no twin here because the commit already - zeroed those rows, so post-commit it would match nobody. + Mirrors ``_queue_enduser_resets`` without its ``spend > 0`` filter, which + post-commit would match nobody. """ linked: Final = _budget_link_where(budget_ids) default_budget_id: Final = litellm.max_end_user_budget_id @@ -279,9 +277,8 @@ class _BudgetCascade: @dataclass(frozen=True, slots=True) class _EndUserWalk: - """Where the post-commit customer walk stands: the keyset cursor its next - page resumes from, None once there is no next page, how many customers it - has reached, and whether a failed page read cut it short of the tail.""" + """Where the customer walk stands. ``cursor`` is None once it is done, and + ``truncated`` says a failed page read cut it short of the tail.""" cursor: str | None = "" invalidated: int = 0 @@ -306,8 +303,6 @@ class _BudgetCascadeFailed: _EMPTY_CASCADE: Final = _BudgetCascade() -#: Which of the proxy's two caches a batch of keys belongs to. ``spend_counter_cache`` -#: holds the live running spend; ``user_api_key_cache`` holds the cached management rows. _InvalidatedCache = Literal["spend counter", "user_api_key_cache"] @@ -623,20 +618,15 @@ class ResetBudgetJob: @staticmethod async def _invalidate_caches(counter_keys: Sequence[str], cache_keys: Sequence[str]) -> None: """Batch twin of ``_invalidate_spend_counter`` and - ``_invalidate_user_api_key_cache_entry``, carrying the same - after-the-commit requirement as both. - - One round trip per chunk rather than one per key: a tier's dependent - population is unbounded, and awaiting each key in turn makes the last - dependent wait out every dependent ahead of it. - """ + ``_invalidate_user_api_key_cache_entry``, after the commit like both: + one round trip per chunk where a tier's dependents are unbounded.""" await ResetBudgetJob._invalidate_cache("spend counter", counter_keys) await ResetBudgetJob._invalidate_cache("user_api_key_cache", cache_keys) @staticmethod async def _invalidate_cache(cache: _InvalidatedCache, keys: Sequence[str]) -> None: - """One cache's share of a batch, awaited separately from the other's so a - failure against either still leaves the other one invalidated.""" + """One cache's share of a batch, awaited separately so either failing + still leaves the other invalidated.""" if not keys: return try: @@ -679,21 +669,9 @@ class ResetBudgetJob: async def _invalidate_enduser_caches(self, budget_ids: Sequence[str]) -> _EndUserWalk: """Drop the cached spend of every customer the committed tier reset zeroed. - Walked a page at a time with a keyset cursor, for the same reason - ``_reset_windows_for_source`` is: the customers sharing one tier are - unbounded, so reading them into one result set puts a - customer-count-sized list in the proxy's heap on every tick, and a - deployment large enough turns that into an OOM rather than a slow tick. - - No per-run page cap, also for that walk's reason: the position cannot - survive the run, so a cap would restart at the first customer every tick - and never reach the tail. The cursor strictly advances, so this - terminates on its own. - - A page that fails to read stops the walk short of the tail. The window is - already advanced by then, so no later tick comes back for the customers - past it, which is why the walk reports that it was cut short instead of - passing the part it managed off as the whole. + Paged like ``_reset_windows_for``, and capless for its reason too: the + customers on one tier are unbounded, and a cap cannot keep its position + across pod elections, so it would restart at the first customer forever. """ if not budget_ids: return _ENDUSER_WALK_DONE diff --git a/litellm/proxy/common_utils/user_api_key_cache.py b/litellm/proxy/common_utils/user_api_key_cache.py index 5a8e3a9482d..89ff113c6d3 100644 --- a/litellm/proxy/common_utils/user_api_key_cache.py +++ b/litellm/proxy/common_utils/user_api_key_cache.py @@ -223,13 +223,11 @@ class UserApiKeyCache(DualCache): await super().async_delete_cache(key) async def async_delete_cache_keys(self, keys: Sequence[str]) -> None: - """Batch twin of ``async_delete_cache``, partitioned the way - ``async_set_cache_pipeline`` partitions its writes. + """Batch twin of ``async_delete_cache``, partitioned like + ``async_set_cache_pipeline``. - Both partitions are cleared even when one of them raises: a caller - batching these has already committed the rows they cache, so a partition - left holding pre-reset spend goes on being authorized against until the - entry expires. The first failure is re-raised for the caller to report. + Both partitions are cleared even when one raises, because a caller + batching these has already committed the rows they cache. """ key_object_keys: Final = tuple(key for key in keys if is_user_key_cache_key(key)) other_keys: Final = tuple(key for key in keys if not is_user_key_cache_key(key)) diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 0606723f6dd..5bc3c549098 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -32,10 +32,16 @@ class MockTable: self.find_many_calls: List[Dict[str, Any]] = [] self.update_many_calls: List[Dict[str, Any]] = [] self._find_many_results: List[Any] = [] + self._find_many_error: Optional[tuple[int, Exception]] = None def set_find_many_results(self, results: List[Any]): self._find_many_results = results + def set_find_many_error(self, after_reads: int, error: Exception): + """Fail every read past the first ``after_reads``, the way a connection + dropping partway through a paged walk does.""" + self._find_many_error = (after_reads, error) + async def find_many( self, where: Dict[str, Any], @@ -45,6 +51,8 @@ class MockTable: """Replays canned rows, honouring the keyset cursor + ``take`` a paged caller relies on: without that a paged walk never advances and the test would hang instead of failing.""" + if self._find_many_error is not None and len(self.find_many_calls) >= self._find_many_error[0]: + raise self._find_many_error[1] paging = {k: v for k, v in (("order", order), ("take", take)) if v is not None} self.find_many_calls.append({"where": where, **paging}) rows = list(self._find_many_results) @@ -1686,14 +1694,7 @@ def test_enduser_invalidation_reports_a_page_read_failure_instead_of_a_clean_fin for i in range(RESET_BUDGET_JOB_BATCH_SIZE + 3) ] ) - read_page: Final = endusers.find_many - - async def fail_after_the_first_page(**kwargs): - if endusers.find_many_calls: - raise RuntimeError("connection reset while paging customers") - return await read_page(**kwargs) - - endusers.find_many = fail_after_the_first_page + endusers.set_find_many_error(1, RuntimeError("connection reset while paging customers")) logging_obj: Final = RecordingProxyLogging() job: Final = ResetBudgetJob(proxy_logging_obj=logging_obj, prisma_client=mock_prisma_client) From fd411373fcdc012668003fe6b1228828cce32338 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 16 Sep 2026 16:44:39 -0700 Subject: [PATCH 5/5] style(proxy): collapse the enduser page signature onto one line --- litellm/proxy/common_utils/reset_budget_job.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 2260d890e46..1299a4df243 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -681,9 +681,7 @@ class ResetBudgetJob: walk = await self._invalidate_enduser_page(where=where, cursor=walk.cursor, reached=walk.invalidated) return walk - async def _invalidate_enduser_page( - self, where: Mapping[str, object], cursor: str, reached: int - ) -> _EndUserWalk: + async def _invalidate_enduser_page(self, where: Mapping[str, object], cursor: str, reached: int) -> _EndUserWalk: """Invalidate one page of customers and say where the walk goes next.""" try: rows: Final = await self._fetch_enduser_page(where=where, cursor=cursor)