diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c1e73607b72..359b7d23f08 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -3421,28 +3421,35 @@ async def delete_cache_key_objects( proxy_logging_obj: ProxyLogging | None, ) -> None: """ - Evict a batch of key objects, for callers that delete keys in bulk. Auth resolves a cached key + Evict a batch of deleted key objects here and on every other worker. Auth resolves a cached key object without re-reading its team, so a key left cached after its row is gone keeps buying access until its TTL expires. - Evicting locally only reaches this worker; `_delete_cache_key_object` itself broadcasts each - token, so a deleted key left in a peer worker's in-memory cache still authenticates there until - its TTL expires. - Best-effort per key: the rows are already deleted by the time this runs, so an unreachable cache backend must not abort the caller partway through its own cascade. - Every token is dropped from this worker's memory first, then from the auth cache and the - internal usage cache in one batch each, so a deleted key stops authenticating here the instant - its rows are gone instead of refilling from Redis on the next miss. Every other worker then - gets the eviction in one pipelined broadcast: a publish per key would pass the publisher's - pending cap on a bulk delete, and the dropped keys would keep authenticating there until their - TTL. + Every token is dropped from this worker's memory first, then from the auth cache's Redis layer, + so a deleted key stops authenticating here the instant its rows are gone instead of refilling + from Redis on the next miss. Every other worker then gets the eviction in one pipelined + broadcast: a publish per key would pass the publisher's pending cap on a bulk delete, and the + dropped keys would keep authenticating there until their TTL. The internal usage cache is not + an auth refill source, so its batch delete runs beside that chain and a hung Redis costs one + bounded wait there instead of another in line. """ evict_local(cache_keys=hashed_tokens, user_api_key_cache=user_api_key_cache) + usage_eviction: Final = ( + () + if proxy_logging_obj is None + else (evict_shared(cache_keys=hashed_tokens, cache=proxy_logging_obj.internal_usage_cache.dual_cache),) + ) + _ = await asyncio.gather( + _evict_auth_cache_then_broadcast(hashed_tokens=hashed_tokens, user_api_key_cache=user_api_key_cache), + *usage_eviction, + ) + + +async def _evict_auth_cache_then_broadcast(hashed_tokens: Sequence[str], user_api_key_cache: UserApiKeyCache) -> None: await evict_shared(cache_keys=hashed_tokens, cache=user_api_key_cache) - if proxy_logging_obj is not None: - await evict_shared(cache_keys=hashed_tokens, cache=proxy_logging_obj.internal_usage_cache.dual_cache) await broadcast_evictions(cache_keys=hashed_tokens) diff --git a/tests/unit/proxy/auth/test_auth_checks.py b/tests/unit/proxy/auth/test_auth_checks.py index e1afa81d372..29a036a4142 100644 --- a/tests/unit/proxy/auth/test_auth_checks.py +++ b/tests/unit/proxy/auth/test_auth_checks.py @@ -46,7 +46,7 @@ from litellm.proxy.auth.auth_checks import ( _team_soft_budget_check, common_checks, ) -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import InternalUsageCache, ProxyLogging from litellm.proxy.utils import CallInfo from litellm.types.passthrough_endpoints.pass_through_endpoints import LITELLM_PASS_THROUGH_ENDPOINT_MARKER @@ -2380,3 +2380,45 @@ async def test_bulk_eviction_clears_the_redis_copies_before_waiting_on_a_wedged_ assert refilled == [] assert redis.entries == {} + + +class _WedgedDeleteRedisCache(_AuthRedisCache): + def __init__(self) -> None: + super().__init__({}) + self.release = asyncio.Event() + + async def delete_cache_keys(self, keys: Sequence[str]) -> None: + await self.release.wait() + + +@pytest.mark.asyncio +async def test_bulk_eviction_broadcasts_while_the_usage_cache_redis_is_wedged( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(pubsub_module, "_in_flight_publishes", asyncio.Semaphore(16)) + monkeypatch.setattr(pubsub_module, "_pending_publishes", set()) + monkeypatch.setattr(pubsub_module, "_BULK_EVICTION_WAIT_SECONDS", 0.5) + client: Final = _SlowPublishRedisClient() + usage_redis: Final = _WedgedDeleteRedisCache() + proxy_logging_obj: Final = ProxyLogging(user_api_key_cache=None) + proxy_logging_obj.internal_usage_cache = InternalUsageCache( + dual_cache=DualCache(redis_cache=usage_redis) # pyright: ignore[reportArgumentType] # a dict-backed fake of the Redis layer + ) + hashed_tokens: Final = ("token-a", "token-b") + + with patch( + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache", + return_value=_PubSubRedisCache(client), + ): + eviction: Final = asyncio.create_task( + delete_cache_key_objects( + hashed_tokens=hashed_tokens, user_api_key_cache=UserApiKeyCache(), proxy_logging_obj=proxy_logging_obj + ) + ) + await asyncio.sleep(0.1) + broadcast_while_wedged: Final = {json.loads(message)["cache_key"] for message in client.published} + usage_redis.release.set() + await eviction + await _drain_publishes() + + assert broadcast_while_wedged == set(hashed_tokens)