mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(proxy): stop stale auth cache re-publish so key updates and deletes propagate across replicas (#33565)
With enable_redis_auth_cache and multiple replicas, /key/update and
/key/delete delete the Redis auth blob and the handling pod's in-memory
entry, but two read-path writers re-published the stale blob from any other
replica's per-pod memory back to Redis with a fresh 60s TTL on every
request: the post-auth re-cache in user_api_key_auth and the spend writeback
in update_cache. Replicas whose in-memory entries expired then re-primed
themselves from the poisoned Redis entry, so key limit and access changes
never took effect fleet-wide while traffic continued, and a deleted key kept
authenticating.
The auth object is now written only by the DB-load paths
(IdentityStore._resolve_key, get_key_object): the post-auth re-cache is
removed outright (even a local-only write could race an invalidation and
resurrect a revoked key on this worker) and spend tracking no longer writes
the auth object back at all; spend is tracked through the spend🔑*
counters. The remaining spend writebacks for user, team, end-user, and tag
objects become local-only so they cannot republish stale management objects
either, with one deliberate exception: the proxy-wide
{litellm_proxy_admin_name}:spend scalar keeps its shared Redis write because
the global max_budget check reads it between authoritative DB reloads, and
it carries no limits or permissions so sharing it cannot resurrect an
invalidated auth blob.
DualCache's redis-to-memory read backfill also ignored default_in_memory_ttl,
pinning backfilled entries for InMemoryCache's 600s default instead of the
configured 60s auth TTL; the backfill now injects the configured default like
every write path already does, so a replica primed from Redis converges
within the auth cache TTL as well.
Consolidates the sibling stale-auth-recache branch; the delete-propagation
case is the duplicate ticket LIT-4350.
Resolves LIT-4219
(cherry picked from commit ae8dc1f39f)
This commit is contained in:
parent
052b5a2169
commit
a68439ec0c
7 changed files with 348 additions and 34 deletions
|
|
@ -103,6 +103,18 @@ class DualCache(BaseCache):
|
|||
if default_redis_ttl is not None:
|
||||
self.default_redis_ttl = default_redis_ttl
|
||||
|
||||
def _backfill_kwargs(self, kwargs: "dict[str, object]") -> "dict[str, object]":
|
||||
"""
|
||||
Kwargs for writing a Redis read result into the in-memory tier.
|
||||
|
||||
Applies ``default_in_memory_ttl`` exactly like the write paths do;
|
||||
without it, backfilled entries fall to ``InMemoryCache``'s own default
|
||||
TTL and can outlive the TTL this cache was configured with.
|
||||
"""
|
||||
if "ttl" not in kwargs and self.default_in_memory_ttl is not None:
|
||||
return {**kwargs, "ttl": self.default_in_memory_ttl}
|
||||
return kwargs
|
||||
|
||||
def set_cache(self, key, value, local_only: bool = False, **kwargs):
|
||||
# Update both Redis and in-memory cache
|
||||
try:
|
||||
|
|
@ -160,7 +172,7 @@ class DualCache(BaseCache):
|
|||
|
||||
if redis_result is not None:
|
||||
# Update in-memory cache with the value from Redis
|
||||
self.in_memory_cache.set_cache(key, redis_result, **kwargs)
|
||||
self.in_memory_cache.set_cache(key, redis_result, **self._backfill_kwargs(kwargs))
|
||||
|
||||
result = redis_result
|
||||
|
||||
|
|
@ -226,7 +238,7 @@ class DualCache(BaseCache):
|
|||
|
||||
if redis_result is not None:
|
||||
# Update in-memory cache with the value from Redis
|
||||
await self.in_memory_cache.async_set_cache(key, redis_result, **kwargs)
|
||||
await self.in_memory_cache.async_set_cache(key, redis_result, **self._backfill_kwargs(kwargs))
|
||||
|
||||
result = redis_result
|
||||
|
||||
|
|
@ -318,7 +330,7 @@ class DualCache(BaseCache):
|
|||
result[key_to_index[key]] = value
|
||||
|
||||
if value is not None and self.in_memory_cache is not None:
|
||||
await self.in_memory_cache.async_set_cache(key, value, **kwargs)
|
||||
await self.in_memory_cache.async_set_cache(key, value, **self._backfill_kwargs(kwargs))
|
||||
|
||||
return result
|
||||
except Exception:
|
||||
|
|
|
|||
|
|
@ -1992,16 +1992,6 @@ async def _user_api_key_auth_builder(
|
|||
raise HTTPException(401, detail="Invalid API key, no token associated")
|
||||
api_key = valid_token.token
|
||||
|
||||
# Add hashed token to cache
|
||||
asyncio.create_task(
|
||||
_cache_key_object(
|
||||
hashed_token=api_key,
|
||||
user_api_key_obj=valid_token,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
)
|
||||
|
||||
valid_token_dict = valid_token.model_dump(exclude_none=True)
|
||||
valid_token_dict.pop("token", None)
|
||||
# budget_throttle_pct is excluded from model_dump (it must not leak
|
||||
|
|
|
|||
|
|
@ -2768,21 +2768,6 @@ async def update_cache(
|
|||
)
|
||||
# set cooldown on alert
|
||||
|
||||
if existing_spend_obj is not None and getattr(existing_spend_obj, "team_spend", None) is not None:
|
||||
existing_team_spend = existing_spend_obj.team_spend or 0
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
existing_spend_obj.team_spend = existing_team_spend + response_cost
|
||||
|
||||
if existing_spend_obj is not None and getattr(existing_spend_obj, "team_member_spend", None) is not None:
|
||||
existing_team_member_spend = existing_spend_obj.team_member_spend or 0
|
||||
# Calculate the new cost by adding the existing cost and response_cost
|
||||
existing_spend_obj.team_member_spend = existing_team_member_spend + response_cost
|
||||
|
||||
# Existing spend_obj is mutated; UserApiKeyCache.async_set_cache_pipeline turns
|
||||
# BaseModel values into dicts for Redis (same Codec path as async_set_cache).
|
||||
existing_spend_obj.spend = new_spend
|
||||
values_to_update_in_cache.append((hashed_token, existing_spend_obj))
|
||||
|
||||
### UPDATE USER SPEND ###
|
||||
async def _update_user_cache():
|
||||
## UPDATE CACHE FOR USER ID + GLOBAL PROXY
|
||||
|
|
@ -2986,13 +2971,27 @@ async def update_cache(
|
|||
if tags is not None:
|
||||
await _update_tag_cache()
|
||||
|
||||
asyncio.create_task(
|
||||
user_api_key_cache.async_set_cache_pipeline(
|
||||
cache_list=values_to_update_in_cache,
|
||||
ttl=get_management_object_ttl(user_api_key_cache),
|
||||
litellm_parent_otel_span=parent_otel_span,
|
||||
global_proxy_spend_key = "{}:spend".format(litellm_proxy_admin_name)
|
||||
local_object_updates = tuple((k, v) for k, v in values_to_update_in_cache if k != global_proxy_spend_key)
|
||||
shared_scalar_updates = tuple((k, v) for k, v in values_to_update_in_cache if k == global_proxy_spend_key)
|
||||
|
||||
if local_object_updates:
|
||||
asyncio.create_task(
|
||||
user_api_key_cache.async_set_cache_pipeline(
|
||||
cache_list=list(local_object_updates),
|
||||
ttl=get_management_object_ttl(user_api_key_cache),
|
||||
litellm_parent_otel_span=parent_otel_span,
|
||||
local_only=True,
|
||||
)
|
||||
)
|
||||
if shared_scalar_updates:
|
||||
asyncio.create_task(
|
||||
user_api_key_cache.async_set_cache_pipeline(
|
||||
cache_list=list(shared_scalar_updates),
|
||||
ttl=get_management_object_ttl(user_api_key_cache),
|
||||
litellm_parent_otel_span=parent_otel_span,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def run_ollama_serve():
|
||||
|
|
|
|||
|
|
@ -88,6 +88,60 @@ async def test_dual_cache_async_set_cache_injects_default_in_memory_ttl():
|
|||
assert expiry <= after + 60
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dual_cache_redis_backfill_injects_default_in_memory_ttl():
|
||||
"""
|
||||
A Redis-hit backfill into the in-memory tier must honor
|
||||
default_in_memory_ttl the same way the write paths do. Without it, the
|
||||
backfilled entry falls to InMemoryCache's own default_ttl (600s), so a
|
||||
replica that primed a management object (e.g. a virtual key's auth blob)
|
||||
from Redis keeps serving it for 10 minutes after the object was updated
|
||||
and invalidated, instead of re-reading within the configured TTL.
|
||||
"""
|
||||
in_memory_cache = InMemoryCache(default_ttl=600)
|
||||
redis_cache = MagicMock()
|
||||
redis_cache.async_get_cache = AsyncMock(return_value="redis_value")
|
||||
dual_cache = DualCache(
|
||||
in_memory_cache=in_memory_cache,
|
||||
redis_cache=redis_cache,
|
||||
default_in_memory_ttl=60,
|
||||
)
|
||||
|
||||
before = time.time()
|
||||
result = await dual_cache.async_get_cache(key="backfill_key")
|
||||
after = time.time()
|
||||
|
||||
assert result == "redis_value"
|
||||
expiry = in_memory_cache.ttl_dict["backfill_key"]
|
||||
assert expiry >= before + 60
|
||||
assert expiry <= after + 60
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dual_cache_batch_redis_backfill_injects_default_in_memory_ttl():
|
||||
"""async_batch_get_cache's Redis-to-memory backfill must honor
|
||||
default_in_memory_ttl, same as the single-key path."""
|
||||
in_memory_cache = InMemoryCache(default_ttl=600)
|
||||
mock_redis = MagicMock(spec=RedisCache)
|
||||
mock_redis.async_batch_get_cache = AsyncMock(
|
||||
return_value={"batch_backfill_key": "redis_value"}
|
||||
)
|
||||
dual_cache = DualCache(
|
||||
in_memory_cache=in_memory_cache,
|
||||
redis_cache=mock_redis,
|
||||
default_in_memory_ttl=60,
|
||||
)
|
||||
|
||||
before = time.time()
|
||||
result = await dual_cache.async_batch_get_cache(keys=["batch_backfill_key"])
|
||||
after = time.time()
|
||||
|
||||
assert result == ["redis_value"]
|
||||
expiry = in_memory_cache.ttl_dict["batch_backfill_key"]
|
||||
assert expiry >= before + 60
|
||||
assert expiry <= after + 60
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dual_cache_async_set_cache_respects_explicit_ttl():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -522,6 +522,49 @@ async def test_get_key_object_should_raise_if_reconnect_fails_on_db_connection_e
|
|||
assert mock_prisma_client.get_data.await_count == 1
|
||||
|
||||
|
||||
def _fake_redis_cache():
|
||||
fake_redis = MagicMock()
|
||||
fake_redis.async_get_cache = AsyncMock(return_value=None)
|
||||
fake_redis.async_set_cache = AsyncMock()
|
||||
fake_redis.async_set_cache_pipeline = AsyncMock()
|
||||
fake_redis.async_delete_cache = AsyncMock()
|
||||
return fake_redis
|
||||
|
||||
|
||||
class TestAuthCacheRedisWritePolicy:
|
||||
"""Redis auth-cache entries may only be written from fresh DB loads.
|
||||
|
||||
With ``enable_redis_auth_cache`` and multiple replicas, a pod that re-publishes
|
||||
a cache-derived key object to Redis can resurrect a stale auth blob after
|
||||
``/key/update`` or ``/key/delete`` already deleted it, so limit changes never
|
||||
propagate fleet-wide while traffic keeps refreshing the stale entry's TTL.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_key_object_db_load_publishes_to_redis(self):
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.get_data = AsyncMock(
|
||||
return_value=UserAPIKeyAuth(token="hashed-token-db")
|
||||
)
|
||||
|
||||
fake_redis = _fake_redis_cache()
|
||||
cache = UserApiKeyCache()
|
||||
cache.redis_cache = fake_redis
|
||||
|
||||
key_obj = await get_key_object(
|
||||
hashed_token="hashed-token-db",
|
||||
prisma_client=mock_prisma_client,
|
||||
user_api_key_cache=cache,
|
||||
)
|
||||
|
||||
assert key_obj.token == "hashed-token-db"
|
||||
fake_redis.async_set_cache.assert_awaited_once()
|
||||
assert (
|
||||
fake_redis.async_set_cache.await_args.kwargs.get("key")
|
||||
or fake_redis.async_set_cache.await_args.args[0]
|
||||
) == "hashed-token-db"
|
||||
|
||||
|
||||
def test_get_cli_jwt_auth_token_default_expiration(valid_sso_user_defined_values):
|
||||
"""Test generating CLI JWT token with default 24-hour expiration"""
|
||||
token = ExperimentalUIJWTToken.get_cli_jwt_auth_token(valid_sso_user_defined_values)
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
|
@ -4087,6 +4088,99 @@ async def test_auth_path_caches_team_object_under_canonical_team_id_key():
|
|||
assert cache.get_cache(key=None) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_does_not_rewrite_cached_key_object_back_into_cache():
|
||||
"""A cache-hit auth must not write the token back into the cache.
|
||||
|
||||
Re-writing on every auth let a replica holding a stale in-memory token
|
||||
republish it to shared Redis with a fresh TTL on each request, so
|
||||
/key/update and /key/delete never propagated across replicas or regional
|
||||
Redis while the key kept calling (stale auth re-cache feedback loop).
|
||||
Only the DB-load paths (IdentityStore._resolve_key / get_key_object) may
|
||||
populate the cache.
|
||||
"""
|
||||
from fastapi import Request
|
||||
from starlette.datastructures import URL
|
||||
|
||||
import litellm.proxy.proxy_server as _proxy_server_mod
|
||||
from litellm.proxy.auth.user_api_key_auth import _user_api_key_auth_builder
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.proxy_server import hash_token
|
||||
|
||||
api_key = "sk-lit-cached-key-no-rewrite"
|
||||
hashed_key = hash_token(api_key)
|
||||
|
||||
key_cache = UserApiKeyCache()
|
||||
stale_token = UserAPIKeyAuth(
|
||||
api_key=api_key,
|
||||
token=hashed_key,
|
||||
metadata={"model_rpm_limit": {"gpt-5.4-mini": 3}},
|
||||
last_refreshed_at=1000.0,
|
||||
)
|
||||
await key_cache.async_set_cache(
|
||||
key=hashed_key, value=stale_token, model_type=UserAPIKeyAuth
|
||||
)
|
||||
|
||||
fetch_from_db = AsyncMock(
|
||||
side_effect=AssertionError("cache-hit auth must not touch the DB")
|
||||
)
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.internal_usage_cache = MagicMock()
|
||||
proxy_logging_obj.internal_usage_cache.dual_cache = AsyncMock()
|
||||
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
|
||||
|
||||
attrs = {
|
||||
"prisma_client": MagicMock(),
|
||||
"user_api_key_cache": key_cache,
|
||||
"proxy_logging_obj": proxy_logging_obj,
|
||||
"master_key": "sk-test-master",
|
||||
"general_settings": {},
|
||||
"llm_model_list": [],
|
||||
"llm_router": None,
|
||||
"open_telemetry_logger": None,
|
||||
"model_max_budget_limiter": MagicMock(),
|
||||
"user_custom_auth": None,
|
||||
"jwt_handler": None,
|
||||
"litellm_proxy_admin_name": "admin",
|
||||
}
|
||||
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
|
||||
try:
|
||||
for k, v in attrs.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
request = Request(scope={"type": "http"})
|
||||
request._url = URL(url="/chat/completions")
|
||||
with patch(
|
||||
"litellm.proxy.auth.resolvers.store._fetch_key_object_from_db_with_reconnect",
|
||||
fetch_from_db,
|
||||
):
|
||||
result = await _user_api_key_auth_builder(
|
||||
request=request,
|
||||
api_key=f"Bearer {api_key}",
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
request_data={},
|
||||
)
|
||||
pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
|
||||
if pending:
|
||||
await asyncio.wait(pending, timeout=5)
|
||||
|
||||
assert result.token == hashed_key
|
||||
fetch_from_db.assert_not_called()
|
||||
|
||||
cached_after = await key_cache.async_get_cache(
|
||||
key=hashed_key, model_type=UserAPIKeyAuth
|
||||
)
|
||||
assert cached_after is not None
|
||||
assert cached_after.last_refreshed_at == 1000.0
|
||||
assert cached_after.metadata == {"model_rpm_limit": {"gpt-5.4-mini": 3}}
|
||||
finally:
|
||||
for k, v in originals.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
|
||||
|
||||
class TestCheckKeyModelBudgetWithFallback:
|
||||
"""`_check_key_model_budget_with_fallback` must reroute a request to the
|
||||
first configured `budget_fallbacks` entry still within its own budget,
|
||||
|
|
|
|||
|
|
@ -4490,6 +4490,128 @@ async def test_update_cache_pipeline_honors_user_api_key_cache_ttl():
|
|||
setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spend_tracking_never_writes_the_auth_object_back():
|
||||
"""Spend tracking must never write the auth object back into the cache.
|
||||
|
||||
Writing the mutated auth object back after every priced request let a
|
||||
stale copy be re-published with a fresh TTL: to shared Redis it defeated
|
||||
/key/update and /key/delete across replicas, and even a local-only write
|
||||
could race an invalidation and resurrect a revoked key on this worker.
|
||||
Spend is tracked through the spend:key:* counters, so the auth object is
|
||||
only ever written by the DB-load paths.
|
||||
"""
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
|
||||
original_cache = litellm.proxy.proxy_server.user_api_key_cache
|
||||
cache = UserApiKeyCache()
|
||||
setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
|
||||
try:
|
||||
hashed_token = "spend-tracking-no-writeback-token"
|
||||
await cache.async_set_cache(
|
||||
key=hashed_token,
|
||||
value=UserAPIKeyAuth(token=hashed_token, spend=1.0),
|
||||
model_type=UserAPIKeyAuth,
|
||||
)
|
||||
with (
|
||||
patch.object(
|
||||
cache, "async_set_cache_pipeline", new=AsyncMock()
|
||||
) as mock_pipeline,
|
||||
patch.object(cache, "async_set_cache", new=AsyncMock()) as mock_set,
|
||||
):
|
||||
await litellm.proxy.proxy_server.update_cache(
|
||||
token=hashed_token,
|
||||
user_id=None,
|
||||
end_user_id=None,
|
||||
team_id=None,
|
||||
response_cost=5.0,
|
||||
parent_otel_span=None,
|
||||
)
|
||||
pending = [
|
||||
t for t in asyncio.all_tasks() if t is not asyncio.current_task()
|
||||
]
|
||||
if pending:
|
||||
await asyncio.wait(pending, timeout=5)
|
||||
|
||||
key_pipeline_writes = [
|
||||
call
|
||||
for call in mock_pipeline.call_args_list
|
||||
if any(k == hashed_token for k, _ in call.kwargs["cache_list"])
|
||||
]
|
||||
assert key_pipeline_writes == []
|
||||
mock_set.assert_not_called()
|
||||
finally:
|
||||
setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_cache_global_proxy_spend_scalar_stays_shared():
|
||||
"""
|
||||
The proxy-wide spend estimate must keep flowing to Redis when the spend
|
||||
writeback goes per-pod: the global max_budget check reads the
|
||||
``{litellm_proxy_admin_name}:spend`` cache entry between authoritative DB
|
||||
reloads, so keeping it pod-local would let traffic spread across replicas
|
||||
exceed the proxy budget by roughly a factor of the replica count within a
|
||||
cache TTL. Sharing this scalar is safe because it carries no limits or
|
||||
permissions, so it cannot resurrect an invalidated auth blob.
|
||||
"""
|
||||
from litellm.caching.caching import DualCache
|
||||
|
||||
admin_name = litellm.proxy.proxy_server.litellm_proxy_admin_name
|
||||
global_key = "{}:spend".format(admin_name)
|
||||
|
||||
async def fake_get(key, **kwargs):
|
||||
if key == "user-lit":
|
||||
return {"user_id": "user-lit", "spend": 1.0}
|
||||
if key == global_key:
|
||||
return 10.0
|
||||
return None
|
||||
|
||||
original_cache = litellm.proxy.proxy_server.user_api_key_cache
|
||||
cache = DualCache(default_in_memory_ttl=300)
|
||||
setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache)
|
||||
try:
|
||||
with patch.object(
|
||||
cache, "async_get_cache", new=AsyncMock(side_effect=fake_get)
|
||||
):
|
||||
with patch.object(
|
||||
cache, "async_set_cache_pipeline", new=AsyncMock()
|
||||
) as mock_set_cache:
|
||||
await litellm.proxy.proxy_server.update_cache(
|
||||
token=None,
|
||||
user_id="user-lit",
|
||||
end_user_id=None,
|
||||
team_id=None,
|
||||
response_cost=5.0,
|
||||
parent_otel_span=None,
|
||||
)
|
||||
|
||||
pending = [
|
||||
t for t in asyncio.all_tasks() if t is not asyncio.current_task()
|
||||
]
|
||||
if pending:
|
||||
await asyncio.wait(pending, timeout=5)
|
||||
|
||||
calls = mock_set_cache.await_args_list
|
||||
local_keys = [
|
||||
k
|
||||
for c in calls
|
||||
if c.kwargs.get("local_only") is True
|
||||
for k, _ in c.kwargs["cache_list"]
|
||||
]
|
||||
shared_keys = [
|
||||
k
|
||||
for c in calls
|
||||
if c.kwargs.get("local_only") is not True
|
||||
for k, _ in c.kwargs["cache_list"]
|
||||
]
|
||||
assert "user-lit" in local_keys
|
||||
assert global_key not in local_keys
|
||||
assert shared_keys == [global_key]
|
||||
finally:
|
||||
setattr(litellm.proxy.proxy_server, "user_api_key_cache", original_cache)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_init_sso_settings_in_db():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue