scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

fix(proxy): isolate each cache and each page in the budget reset invalidation

Browse source 
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.
This commit is contained in:
ryan-crabbe-berri 2026-09-16 16:14:08 -07:00
parent 7ba47a5b6e
commit cdf0142f4a
4 changed files with 171 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

107
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -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 "

8
litellm/proxy/common_utils/user_api_key_cache.py
View file

@ -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()

64
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View file

@ -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)

25
tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py
View file

@ -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))