mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(proxy): broadcast cache invalidation for end-users, access groups, revoked keys, and config params
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo
parent
330a09235d
commit
9123692a50
10 changed files with 405 additions and 31 deletions
|
|
@ -1149,7 +1149,7 @@ async def get_end_user_object(
|
|||
if end_user_id is None:
|
||||
return None
|
||||
|
||||
_key: Final = f"end_user_id:{end_user_id}"
|
||||
_key: Final = _end_user_cache_key(end_user_id)
|
||||
|
||||
# Check cache first
|
||||
cached_user_obj: Final = await user_api_key_cache.async_get_cache(
|
||||
|
|
@ -1191,7 +1191,7 @@ async def get_end_user_object(
|
|||
|
||||
# Save to cache
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=f"end_user_id:{end_user_id}",
|
||||
key=_key,
|
||||
value=_response,
|
||||
model_type=LiteLLM_EndUserTable,
|
||||
)
|
||||
|
|
@ -1206,6 +1206,39 @@ _END_USER_VALIDATION_NEGATIVE_TTL: Final = 60
|
|||
_END_USER_VALIDATION_POSITIVE_TTL: Final = 300
|
||||
|
||||
|
||||
def _end_user_cache_key(end_user_id: str) -> str:
|
||||
return f"end_user_id:{end_user_id}"
|
||||
|
||||
|
||||
def _end_user_validation_cache_key(end_user_id: str) -> str:
|
||||
return f"end_user_validation:{end_user_id}"
|
||||
|
||||
|
||||
async def delete_cached_end_user_object(
|
||||
end_user_id: str,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
) -> None:
|
||||
"""
|
||||
Every endpoint that mutates litellm_endusertable must call this: the auth path
|
||||
reads the end-user object and the id-validation verdict cache-first, so without
|
||||
invalidation a blocked, rebudgeted, or deleted customer keeps being served until
|
||||
the TTL expires. Best-effort on both steps, and broadcast so workers that did not
|
||||
handle the mutation drop their in-memory copy too.
|
||||
"""
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation
|
||||
|
||||
for cache_key in (_end_user_cache_key(end_user_id), _end_user_validation_cache_key(end_user_id)):
|
||||
try:
|
||||
await user_api_key_cache.async_delete_cache(key=cache_key)
|
||||
except Exception as e: # noqa: BLE001 # best-effort eviction: a cache backend error must not fail the mutation
|
||||
verbose_proxy_logger.warning(
|
||||
"Failed to evict cached end-user entry %s; a stale end-user may be served until its TTL expires: %s",
|
||||
cache_key,
|
||||
e,
|
||||
)
|
||||
await publish_auth_cache_invalidation(cache_key=cache_key)
|
||||
|
||||
|
||||
async def resolve_and_validate_end_user_id(
|
||||
raw_end_user_id: str | None,
|
||||
prisma_client: PrismaClient | None,
|
||||
|
|
@ -1241,7 +1274,7 @@ async def resolve_and_validate_end_user_id(
|
|||
if prisma_client is None:
|
||||
return raw_end_user_id
|
||||
|
||||
cache_key: Final = f"end_user_validation:{raw_end_user_id}"
|
||||
cache_key: Final = _end_user_validation_cache_key(raw_end_user_id)
|
||||
cached: Final = await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
if cached == "valid":
|
||||
return raw_end_user_id
|
||||
|
|
@ -1904,7 +1937,19 @@ async def _delete_cache_key_object(
|
|||
hashed_token: str,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging | None,
|
||||
broadcast_invalidation: bool = True,
|
||||
):
|
||||
"""
|
||||
``broadcast_invalidation`` is what makes a revoked or regenerated key stop
|
||||
working fleet-wide right away; the local delete plus the Redis delete leave
|
||||
every other worker serving its own in-memory copy until the TTL expires.
|
||||
|
||||
Callers that evict a key each worker can already invalidate on its own (an
|
||||
expired key, whose expiry every worker reads off the cached object) pass
|
||||
False so a client replaying that key can't drive a broadcast per request.
|
||||
"""
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation
|
||||
|
||||
key: Final = hashed_token
|
||||
|
||||
user_api_key_cache.delete_cache(key=key)
|
||||
|
|
@ -1913,6 +1958,9 @@ async def _delete_cache_key_object(
|
|||
if proxy_logging_obj is not None:
|
||||
await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key)
|
||||
|
||||
if broadcast_invalidation:
|
||||
await publish_auth_cache_invalidation(cache_key=key)
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
async def _get_team_db_check(team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None):
|
||||
|
|
@ -2105,6 +2153,8 @@ async def _delete_cache_access_object(
|
|||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging | None = None,
|
||||
):
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation
|
||||
|
||||
key: Final = f"access_group_id:{access_group_id}"
|
||||
|
||||
user_api_key_cache.delete_cache(key=key)
|
||||
|
|
@ -2113,6 +2163,8 @@ async def _delete_cache_access_object(
|
|||
if proxy_logging_obj is not None:
|
||||
await proxy_logging_obj.internal_usage_cache.dual_cache.async_delete_cache(key=key)
|
||||
|
||||
await publish_auth_cache_invalidation(cache_key=key)
|
||||
|
||||
|
||||
@log_db_metrics
|
||||
async def get_access_object(
|
||||
|
|
|
|||
|
|
@ -1573,6 +1573,7 @@ async def _user_api_key_auth_builder(
|
|||
hashed_token=hash_token(api_key),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
broadcast_invalidation=False,
|
||||
)
|
||||
raise ProxyException(
|
||||
message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}",
|
||||
|
|
|
|||
|
|
@ -1,7 +1,15 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from typing import ( # noqa: TID251 # cast narrows a json str to the Literal
|
||||
TYPE_CHECKING,
|
||||
Final,
|
||||
Literal,
|
||||
TypeAlias,
|
||||
cast,
|
||||
get_args,
|
||||
)
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy.common_utils.config_sync_pubsub import (
|
||||
|
|
@ -11,8 +19,13 @@ from litellm.proxy.common_utils.config_sync_pubsub import (
|
|||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
InvalidationTarget: TypeAlias = Literal["user_api_key", "config_param"]
|
||||
|
||||
DEFAULT_INVALIDATION_TARGET: Final[InvalidationTarget] = "user_api_key"
|
||||
_INVALIDATION_TARGETS: Final[frozenset[str]] = frozenset(get_args(InvalidationTarget))
|
||||
|
||||
AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation"
|
||||
_POLL_TIMEOUT_SECONDS: Final = 1.0
|
||||
|
|
@ -29,32 +42,41 @@ def auth_cache_invalidation_channel(redis_cache: "RedisCache") -> str:
|
|||
@dataclass(frozen=True, slots=True)
|
||||
class _CacheInvalidationMessage:
|
||||
cache_key: str
|
||||
cache: InvalidationTarget
|
||||
|
||||
|
||||
def _cache_invalidation_message_json(cache_key: str) -> str:
|
||||
return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key)))
|
||||
def _cache_invalidation_message_json(cache_key: str, target: InvalidationTarget) -> str:
|
||||
return json.dumps(asdict(_CacheInvalidationMessage(cache_key=cache_key, cache=target)))
|
||||
|
||||
|
||||
def _cache_key_from_message_data(data: object) -> str | None:
|
||||
if isinstance(data, bytes):
|
||||
data = data.decode("utf-8", errors="replace")
|
||||
if not isinstance(data, str):
|
||||
def _invalidation_from_message_data(data: object) -> _CacheInvalidationMessage | None:
|
||||
decoded: Final = data.decode("utf-8", errors="replace") if isinstance(data, bytes) else data
|
||||
if not isinstance(decoded, str):
|
||||
return None
|
||||
try:
|
||||
parsed: Final = json.loads(data)
|
||||
parsed: Final = json.loads(decoded)
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
if not isinstance(parsed, dict):
|
||||
return None
|
||||
cache_key: Final = parsed.get("cache_key")
|
||||
return cache_key if isinstance(cache_key, str) else None
|
||||
if not isinstance(cache_key, str):
|
||||
return None
|
||||
raw_target: Final = parsed.get("cache", DEFAULT_INVALIDATION_TARGET)
|
||||
if raw_target not in _INVALIDATION_TARGETS:
|
||||
return None
|
||||
target: Final = cast(InvalidationTarget, raw_target) # cast-ok: membership checked against the Literal's args
|
||||
return _CacheInvalidationMessage(cache_key=cache_key, cache=target)
|
||||
|
||||
|
||||
async def publish_auth_cache_invalidation(cache_key: str) -> None:
|
||||
async def publish_auth_cache_invalidation(
|
||||
cache_key: str,
|
||||
target: InvalidationTarget = DEFAULT_INVALIDATION_TARGET,
|
||||
) -> None:
|
||||
"""
|
||||
Best-effort broadcast so every worker drops its local in-memory copy of a
|
||||
mutated management object; without this, only the handling worker and Redis
|
||||
are evicted and other workers keep serving the stale object until its TTL.
|
||||
mutated object; without this, only the handling worker and Redis are
|
||||
evicted and other workers keep serving the stale object until its TTL.
|
||||
"""
|
||||
redis_cache: Final = coordination_redis_cache()
|
||||
if redis_cache is None:
|
||||
|
|
@ -67,21 +89,24 @@ async def publish_auth_cache_invalidation(cache_key: str) -> None:
|
|||
cache_key,
|
||||
)
|
||||
return
|
||||
await client.publish(auth_cache_invalidation_channel(redis_cache), _cache_invalidation_message_json(cache_key))
|
||||
await client.publish(
|
||||
auth_cache_invalidation_channel(redis_cache),
|
||||
_cache_invalidation_message_json(cache_key=cache_key, target=target),
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors
|
||||
verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e)
|
||||
|
||||
|
||||
class AuthCacheInvalidationSubscriber:
|
||||
__slots__ = ("_redis_cache", "_task", "_user_api_key_cache")
|
||||
__slots__ = ("_in_memory_caches", "_redis_cache", "_task")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
redis_cache: "RedisCache",
|
||||
user_api_key_cache: "UserApiKeyCache",
|
||||
in_memory_caches: Mapping[InvalidationTarget, "InMemoryCache"],
|
||||
) -> None:
|
||||
self._redis_cache = redis_cache
|
||||
self._user_api_key_cache = user_api_key_cache
|
||||
self._in_memory_caches = in_memory_caches
|
||||
self._task: asyncio.Task[None] | None = None
|
||||
|
||||
def start(self) -> None:
|
||||
|
|
@ -138,12 +163,12 @@ class AuthCacheInvalidationSubscriber:
|
|||
|
||||
def _apply_message(self, message: object) -> None:
|
||||
data: Final = message.get("data") if isinstance(message, dict) else None
|
||||
cache_key: Final = _cache_key_from_message_data(data)
|
||||
if cache_key is None:
|
||||
invalidation: Final = _invalidation_from_message_data(data)
|
||||
if invalidation is None:
|
||||
return
|
||||
in_memory_cache: Final = self._user_api_key_cache.in_memory_cache
|
||||
in_memory_cache: Final = self._in_memory_caches.get(invalidation.cache)
|
||||
if in_memory_cache is not None:
|
||||
in_memory_cache.delete_cache(cache_key)
|
||||
in_memory_cache.delete_cache(invalidation.cache_key)
|
||||
|
||||
@staticmethod
|
||||
async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None:
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ All /customer management endpoints
|
|||
"""
|
||||
|
||||
#### END-USER/CUSTOMER MANAGEMENT ####
|
||||
from collections.abc import Iterable
|
||||
from datetime import datetime, timedelta
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -21,6 +22,7 @@ import litellm
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.auth_checks import delete_cached_end_user_object
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
|
||||
from litellm.proxy.management_helpers.object_permission_utils import (
|
||||
|
|
@ -43,6 +45,13 @@ from litellm.types.proxy.management_endpoints.customer_endpoints import (
|
|||
router: Final = APIRouter()
|
||||
|
||||
|
||||
async def _invalidate_cached_end_users(end_user_ids: Iterable[str]) -> None:
|
||||
from litellm.proxy.proxy_server import user_api_key_cache
|
||||
|
||||
for end_user_id in end_user_ids:
|
||||
await delete_cached_end_user_object(end_user_id=end_user_id, user_api_key_cache=user_api_key_cache)
|
||||
|
||||
|
||||
def _to_customer_response(record: BaseModel) -> CustomerResponse:
|
||||
"""Validate a raw end-user DB row into the typed customer response.
|
||||
|
||||
|
|
@ -96,6 +105,7 @@ async def block_user(data: BlockUsers):
|
|||
},
|
||||
)
|
||||
records.append(record)
|
||||
await _invalidate_cached_end_users(data.user_ids)
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
|
|
@ -389,6 +399,8 @@ async def new_end_user(
|
|||
include={"litellm_budget_table": True, "object_permission": True},
|
||||
)
|
||||
|
||||
await _invalidate_cached_end_users((data.user_id,))
|
||||
|
||||
return _to_customer_response(end_user_record)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
|
|
@ -632,6 +644,8 @@ async def update_end_user(
|
|||
raise ValueError(f"Failed updating customer data. User ID does not exist passed user_id={data.user_id}")
|
||||
verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response)
|
||||
|
||||
await _invalidate_cached_end_users((data.user_id,))
|
||||
|
||||
return _to_customer_response(response)
|
||||
else:
|
||||
raise ValueError(f"user_id is required, passed user_id = {data.user_id}")
|
||||
|
|
@ -705,6 +719,7 @@ async def delete_end_user(
|
|||
where={"user_id": {"in": data.user_ids}}
|
||||
)
|
||||
verbose_proxy_logger.debug("received response from updating prisma client. response=%s", response)
|
||||
await _invalidate_cached_end_users(data.user_ids)
|
||||
return DeleteCustomersResponse(
|
||||
deleted_customers=response,
|
||||
message="Successfully deleted customers with ids: " + str(data.user_ids),
|
||||
|
|
|
|||
|
|
@ -6411,11 +6411,18 @@ class ProxyConfig:
|
|||
redis_cache: RedisCache | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
) -> None:
|
||||
from litellm.proxy.utils import litellm_config_cache
|
||||
|
||||
if redis_cache is None or self.auth_cache_invalidation_subscriber is not None:
|
||||
return
|
||||
subscriber: Final = AuthCacheInvalidationSubscriber(
|
||||
redis_cache=redis_cache,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
in_memory_caches=MappingProxyType(
|
||||
{
|
||||
"user_api_key": user_api_key_cache.in_memory_cache,
|
||||
"config_param": litellm_config_cache.in_memory_cache,
|
||||
}
|
||||
),
|
||||
)
|
||||
self.auth_cache_invalidation_subscriber = subscriber
|
||||
subscriber.start()
|
||||
|
|
|
|||
|
|
@ -101,6 +101,7 @@ from litellm.proxy._types import (
|
|||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation
|
||||
from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.db.create_views import (
|
||||
|
|
@ -2964,8 +2965,14 @@ async def evict_config_param(param_name: str) -> None:
|
|||
|
||||
|
||||
async def invalidate_config_param(param_name: str) -> None:
|
||||
"""Evict from both cache layers; call after every LiteLLM_Config write."""
|
||||
"""Evict from both cache layers; call after every LiteLLM_Config write.
|
||||
|
||||
``evict_config_param`` only clears this worker's in-memory layer and the
|
||||
shared Redis entry, so the broadcast is what stops the other workers from
|
||||
serving their own in-memory copy of the pre-write value until it expires.
|
||||
"""
|
||||
await evict_config_param(param_name)
|
||||
await publish_auth_cache_invalidation(cache_key=_config_cache_key(param_name), target="config_param")
|
||||
await publish_config_param_change(param_name)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -5615,3 +5615,111 @@ def test_is_user_proxy_admin_rejects_view_only_admin():
|
|||
assert _is_user_proxy_admin(user_obj=viewer) is False
|
||||
assert _is_user_proxy_admin(user_obj=admin) is True
|
||||
assert _is_user_proxy_admin(user_obj=None) is False
|
||||
|
||||
class _RecordedInvalidations:
|
||||
def __init__(self) -> None:
|
||||
self.published: list = []
|
||||
|
||||
async def publish(self, cache_key: str, target: str = "user_api_key") -> None:
|
||||
self.published.append((cache_key, target))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def recorded_invalidations():
|
||||
recorder = _RecordedInvalidations()
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation",
|
||||
recorder.publish,
|
||||
):
|
||||
yield recorder
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_cache_key_object_broadcasts_revocation(recorded_invalidations):
|
||||
"""
|
||||
A revoked or regenerated key stays usable on every other worker until its
|
||||
in-memory entry expires unless the eviction is broadcast.
|
||||
"""
|
||||
from litellm.proxy.auth.auth_checks import _delete_cache_key_object
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
cache.in_memory_cache.set_cache("hashed-token", {"key_name": "sk-1"})
|
||||
|
||||
await _delete_cache_key_object(
|
||||
hashed_token="hashed-token",
|
||||
user_api_key_cache=cache,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
|
||||
assert cache.in_memory_cache.get_cache("hashed-token") is None
|
||||
assert recorded_invalidations.published == [("hashed-token", "user_api_key")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_cache_key_object_skips_broadcast_when_opted_out(recorded_invalidations):
|
||||
"""The expired-key auth path evicts on every worker on its own, so it must not
|
||||
let a client replaying that key drive one broadcast per request."""
|
||||
from litellm.proxy.auth.auth_checks import _delete_cache_key_object
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
await _delete_cache_key_object(
|
||||
hashed_token="hashed-token",
|
||||
user_api_key_cache=UserApiKeyCache(),
|
||||
proxy_logging_obj=None,
|
||||
broadcast_invalidation=False,
|
||||
)
|
||||
|
||||
assert recorded_invalidations.published == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_cache_access_object_broadcasts(recorded_invalidations):
|
||||
from litellm.proxy.auth.auth_checks import _delete_cache_access_object
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
cache.in_memory_cache.set_cache("access_group_id:ag-1", {"models": []})
|
||||
|
||||
await _delete_cache_access_object(access_group_id="ag-1", user_api_key_cache=cache)
|
||||
|
||||
assert cache.in_memory_cache.get_cache("access_group_id:ag-1") is None
|
||||
assert recorded_invalidations.published == [("access_group_id:ag-1", "user_api_key")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_cached_end_user_object_evicts_and_broadcasts_both_entries(recorded_invalidations):
|
||||
"""
|
||||
The auth path caches the end-user row and, separately, the verdict on whether
|
||||
the id resolves at all. Both survive a /customer mutation without this.
|
||||
"""
|
||||
from litellm.proxy.auth.auth_checks import delete_cached_end_user_object
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
cache.in_memory_cache.set_cache("end_user_id:eu-1", {"user_id": "eu-1", "blocked": False})
|
||||
cache.in_memory_cache.set_cache("end_user_validation:eu-1", "valid")
|
||||
|
||||
await delete_cached_end_user_object(end_user_id="eu-1", user_api_key_cache=cache)
|
||||
|
||||
assert cache.in_memory_cache.get_cache("end_user_id:eu-1") is None
|
||||
assert cache.in_memory_cache.get_cache("end_user_validation:eu-1") is None
|
||||
assert recorded_invalidations.published == [
|
||||
("end_user_id:eu-1", "user_api_key"),
|
||||
("end_user_validation:eu-1", "user_api_key"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_cached_end_user_object_still_broadcasts_when_eviction_fails(recorded_invalidations):
|
||||
from litellm.proxy.auth.auth_checks import delete_cached_end_user_object
|
||||
|
||||
cache = MagicMock()
|
||||
cache.async_delete_cache = AsyncMock(side_effect=ConnectionError("redis down"))
|
||||
|
||||
await delete_cached_end_user_object(end_user_id="eu-1", user_api_key_cache=cache)
|
||||
|
||||
assert recorded_invalidations.published == [
|
||||
("end_user_id:eu-1", "user_api_key"),
|
||||
("end_user_validation:eu-1", "user_api_key"),
|
||||
]
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from unittest.mock import patch
|
|||
import pytest
|
||||
from redis.asyncio import Redis
|
||||
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
|
||||
AUTH_CACHE_INVALIDATION_CHANNEL,
|
||||
AuthCacheInvalidationSubscriber,
|
||||
|
|
@ -69,8 +70,9 @@ class _FakeRedisCache:
|
|||
return self._client
|
||||
|
||||
|
||||
def _invalidation_message(cache_key: str) -> dict:
|
||||
return {"type": "message", "data": json.dumps({"cache_key": cache_key}).encode()}
|
||||
def _invalidation_message(cache_key: str, cache: Optional[str] = None) -> dict:
|
||||
payload = {"cache_key": cache_key} if cache is None else {"cache_key": cache_key, "cache": cache}
|
||||
return {"type": "message", "data": json.dumps(payload).encode()}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -82,7 +84,12 @@ async def test_publish_sends_cache_key_json_on_channel() -> None:
|
|||
):
|
||||
await publish_auth_cache_invalidation(cache_key="project_id:p-1")
|
||||
|
||||
assert client.published == [(AUTH_CACHE_INVALIDATION_CHANNEL, json.dumps({"cache_key": "project_id:p-1"}))]
|
||||
assert client.published == [
|
||||
(
|
||||
AUTH_CACHE_INVALIDATION_CHANNEL,
|
||||
json.dumps({"cache_key": "project_id:p-1", "cache": "user_api_key"}),
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -129,7 +136,7 @@ async def test_subscriber_deletes_local_cache_entry_on_message() -> None:
|
|||
pubsub = _QueuePubSub(initial_messages=[_invalidation_message("project_id:p-1")])
|
||||
subscriber = AuthCacheInvalidationSubscriber(
|
||||
redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[pubsub])),
|
||||
user_api_key_cache=cache,
|
||||
in_memory_caches={"user_api_key": cache.in_memory_cache},
|
||||
)
|
||||
subscriber.start()
|
||||
try:
|
||||
|
|
@ -151,11 +158,64 @@ async def test_subscriber_ignores_malformed_messages() -> None:
|
|||
|
||||
subscriber = AuthCacheInvalidationSubscriber(
|
||||
redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])),
|
||||
user_api_key_cache=cache,
|
||||
in_memory_caches={"user_api_key": cache.in_memory_cache},
|
||||
)
|
||||
subscriber._apply_message({"type": "message", "data": b"not json"})
|
||||
subscriber._apply_message({"type": "message", "data": json.dumps({"other": "x"}).encode()})
|
||||
subscriber._apply_message({"type": "message", "data": json.dumps({"cache_key": "project_id:p-1", "cache": "nope"}).encode()})
|
||||
subscriber._apply_message("raw string")
|
||||
subscriber._apply_message(None)
|
||||
|
||||
assert cache.in_memory_cache.get_cache("project_id:p-1") is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_targets_the_named_cache() -> None:
|
||||
client = _RecordingRedisClient()
|
||||
with patch(
|
||||
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
|
||||
return_value=_FakeRedisCache(client=client),
|
||||
):
|
||||
await publish_auth_cache_invalidation(cache_key="litellm_config:param:general_settings", target="config_param")
|
||||
|
||||
assert json.loads(client.published[0][1]) == {
|
||||
"cache_key": "litellm_config:param:general_settings",
|
||||
"cache": "config_param",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subscriber_routes_message_to_the_targeted_cache_only() -> None:
|
||||
"""
|
||||
A config-param broadcast must land on the config cache. Routing by target keeps
|
||||
one worker's config write from evicting an unrelated, identically keyed auth entry.
|
||||
"""
|
||||
auth_cache = InMemoryCache()
|
||||
config_cache = InMemoryCache()
|
||||
shared_key = "litellm_config:param:general_settings"
|
||||
auth_cache.set_cache(shared_key, {"from": "auth"})
|
||||
config_cache.set_cache(shared_key, {"from": "config"})
|
||||
|
||||
subscriber = AuthCacheInvalidationSubscriber(
|
||||
redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])),
|
||||
in_memory_caches={"user_api_key": auth_cache, "config_param": config_cache},
|
||||
)
|
||||
subscriber._apply_message(_invalidation_message(shared_key, cache="config_param"))
|
||||
|
||||
assert config_cache.get_cache(shared_key) is None
|
||||
assert auth_cache.get_cache(shared_key) is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subscriber_defaults_untargeted_message_to_the_auth_cache() -> None:
|
||||
"""Messages published by a pre-upgrade worker carry no target during a rolling deploy."""
|
||||
auth_cache = InMemoryCache()
|
||||
auth_cache.set_cache("project_id:p-1", {"models": []})
|
||||
|
||||
subscriber = AuthCacheInvalidationSubscriber(
|
||||
redis_cache=_FakeRedisCache(client=_ScriptedPubSubRedisClient(pubsubs=[_QueuePubSub()])),
|
||||
in_memory_caches={"user_api_key": auth_cache},
|
||||
)
|
||||
subscriber._apply_message(_invalidation_message("project_id:p-1"))
|
||||
|
||||
assert auth_cache.get_cache("project_id:p-1") is None
|
||||
|
|
|
|||
|
|
@ -781,3 +781,82 @@ def test_char_delete_body(mock_prisma_client, mock_user_api_key_auth):
|
|||
"deleted_customers": 2,
|
||||
"message": "Successfully deleted customers with ids: ['c1', 'c2']",
|
||||
}
|
||||
|
||||
@pytest.fixture
|
||||
def seeded_end_user_cache():
|
||||
"""A cache already holding the auth-path entries for customer c1."""
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
cache.in_memory_cache.set_cache("end_user_id:c1", {"user_id": "c1", "blocked": False})
|
||||
cache.in_memory_cache.set_cache("end_user_validation:c1", "valid")
|
||||
with patch("litellm.proxy.proxy_server.user_api_key_cache", cache):
|
||||
yield cache
|
||||
|
||||
|
||||
def _cached_end_user_entries(cache):
|
||||
return (
|
||||
cache.in_memory_cache.get_cache("end_user_id:c1"),
|
||||
cache.in_memory_cache.get_cache("end_user_validation:c1"),
|
||||
)
|
||||
|
||||
|
||||
def test_block_customer_evicts_cached_end_user(mock_prisma_client, mock_user_api_key_auth, seeded_end_user_cache):
|
||||
"""
|
||||
Without eviction the auth path keeps serving the pre-block end-user object,
|
||||
so a blocked customer's requests still pass until the cache entry expires.
|
||||
"""
|
||||
mock_prisma_client.db.litellm_endusertable.upsert = AsyncMock(
|
||||
return_value=LiteLLM_EndUserTable(user_id="c1", blocked=True)
|
||||
)
|
||||
|
||||
response = client.post("/customer/block", json={"user_ids": ["c1"]}, headers={"Authorization": "Bearer k"})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert _cached_end_user_entries(seeded_end_user_cache) == (None, None)
|
||||
|
||||
|
||||
def test_update_customer_evicts_cached_end_user(mock_prisma_client, mock_user_api_key_auth, seeded_end_user_cache):
|
||||
mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock(
|
||||
return_value=_row({"user_id": "c1", "blocked": False})
|
||||
)
|
||||
mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=_row(_FULL_DB_ROW))
|
||||
|
||||
response = client.post(
|
||||
"/customer/update",
|
||||
json={"user_id": "c1", "alias": "Acme"},
|
||||
headers={"Authorization": "Bearer k"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert _cached_end_user_entries(seeded_end_user_cache) == (None, None)
|
||||
|
||||
|
||||
def test_delete_customer_evicts_cached_end_user(mock_prisma_client, mock_user_api_key_auth, seeded_end_user_cache):
|
||||
mock_prisma_client.db.litellm_endusertable.find_many = AsyncMock(
|
||||
return_value=[LiteLLM_EndUserTable(user_id="c1", blocked=False)]
|
||||
)
|
||||
mock_prisma_client.db.litellm_endusertable.delete_many = AsyncMock(return_value=1)
|
||||
|
||||
response = client.post("/customer/delete", json={"user_ids": ["c1"]}, headers={"Authorization": "Bearer k"})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert _cached_end_user_entries(seeded_end_user_cache) == (None, None)
|
||||
|
||||
|
||||
def test_new_customer_evicts_cached_negative_validation(mock_prisma_client, mock_user_api_key_auth):
|
||||
"""
|
||||
The id-validation cache stores 'invalid' verdicts too, so creating a customer
|
||||
for an id that was just rejected must drop that verdict.
|
||||
"""
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
cache = UserApiKeyCache()
|
||||
cache.in_memory_cache.set_cache("end_user_validation:c1", "invalid")
|
||||
mock_prisma_client.db.litellm_endusertable.create = AsyncMock(return_value=_row(_FULL_DB_ROW))
|
||||
|
||||
with patch("litellm.proxy.proxy_server.user_api_key_cache", cache):
|
||||
response = client.post("/customer/new", json={"user_id": "c1"}, headers={"Authorization": "Bearer k"})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert cache.in_memory_cache.get_cache("end_user_validation:c1") is None
|
||||
|
|
|
|||
|
|
@ -265,3 +265,23 @@ async def test_prefetch_config_params_swallows_db_error_without_caching(
|
|||
prisma.db.litellm_config.find_many = AsyncMock(side_effect=RuntimeError("boom"))
|
||||
await prefetch_config_params(prisma, ["a", "b"])
|
||||
assert _swap_config_cache._store == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalidate_config_param_broadcasts_the_cache_key(
|
||||
_swap_config_cache: Any,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""
|
||||
Evicting only this worker's copy leaves every other worker serving the
|
||||
pre-write param value for the rest of the 60s TTL, so the write has to be
|
||||
broadcast at the config cache.
|
||||
"""
|
||||
published: List[Any] = []
|
||||
|
||||
async def _record(cache_key: str, target: str = "user_api_key") -> None:
|
||||
published.append((cache_key, target))
|
||||
|
||||
monkeypatch.setattr(utils_mod, "publish_auth_cache_invalidation", _record)
|
||||
await invalidate_config_param("p5")
|
||||
|
||||
assert published == [("litellm_config:param:p5", "config_param")]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue