Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 31 additions & 4 deletions litellm/caching/dual_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -150,13 +150,23 @@ def _backfill_kwargs(self, kwargs: "dict[str, object]") -> "dict[str, object]":
def set_cache(self, key, value, local_only: bool = False, **kwargs):
# Update both Redis and in-memory cache
try:
# ``kwargs`` is this call's own dict, so injecting a ttl into it is safe. What must not
# happen is one tier's injected ttl reaching the other, so note whether the caller
# supplied one before either tier writes back.
ttl_given: Final = "ttl" in kwargs
if self.in_memory_cache is not None:
if "ttl" not in kwargs and self.default_in_memory_ttl is not None:
if not ttl_given and self.default_in_memory_ttl is not None:
kwargs["ttl"] = self.default_in_memory_ttl

self.in_memory_cache.set_cache(key, value, **kwargs)

if self.redis_cache is not None and local_only is False:
if not ttl_given:
redis_ttl = self.default_redis_ttl
if redis_ttl is None:
redis_ttl = self.default_in_memory_ttl
if redis_ttl is not None:
kwargs["ttl"] = redis_ttl
self.redis_cache.set_cache(key, value, **kwargs)
except Exception as e:
print_verbose(e)
Expand Down Expand Up @@ -505,12 +515,19 @@ def _log_shared_batch_get_failure(e: Exception) -> None:
async def async_set_cache(self, key, value, local_only: bool = False, **kwargs):
print_verbose(f"async set cache: cache key: {key}; local_only: {local_only}; value: {value}")
try:
ttl_given: Final = "ttl" in kwargs
if self.in_memory_cache is not None:
if "ttl" not in kwargs and self.default_in_memory_ttl is not None:
if not ttl_given and self.default_in_memory_ttl is not None:
kwargs["ttl"] = self.default_in_memory_ttl
await self.in_memory_cache.async_set_cache(key, value, **kwargs)

if self.redis_cache is not None and local_only is False:
if not ttl_given:
redis_ttl = self.default_redis_ttl
if redis_ttl is None:
redis_ttl = self.default_in_memory_ttl
if redis_ttl is not None:
kwargs["ttl"] = redis_ttl
await self.redis_cache.async_set_cache(key, value, **kwargs)
except Exception as e:
log_redis_failure(
Expand All @@ -537,7 +554,10 @@ async def _set_on_batch(self, batch: RedisBatch, key: str, value: object, ttl: f
effective_ttl: Final = self.default_in_memory_ttl if ttl is None else ttl
if self.in_memory_cache is not None:
await self.in_memory_cache.async_set_cache(key, value, ttl=effective_ttl)
return batch.set(key, value, effective_ttl)
redis_ttl = self.default_redis_ttl if ttl is None else ttl
if redis_ttl is None:
redis_ttl = effective_ttl
return batch.set(key, value, redis_ttl)

# async_batch_set_cache
async def async_set_cache_pipeline(
Expand All @@ -548,12 +568,19 @@ async def async_set_cache_pipeline(
"""
print_verbose(f"async batch set cache: cache keys: {cache_list}; local_only: {local_only}")
try:
ttl_given: Final = "ttl" in kwargs
if self.in_memory_cache is not None:
if "ttl" not in kwargs and self.default_in_memory_ttl is not None:
if not ttl_given and self.default_in_memory_ttl is not None:
kwargs["ttl"] = self.default_in_memory_ttl
await self.in_memory_cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs)

if self.redis_cache is not None and local_only is False:
if not ttl_given:
redis_ttl = self.default_redis_ttl
if redis_ttl is None:
redis_ttl = self.default_in_memory_ttl
if redis_ttl is not None:
kwargs["ttl"] = redis_ttl
await self.redis_cache.async_set_cache_pipeline(
cache_list=cache_list, ttl=kwargs.pop("ttl", None), **kwargs
)
Expand Down
8 changes: 6 additions & 2 deletions litellm/proxy/common_utils/user_api_key_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,17 +69,21 @@ def __init__(
default_redis_ttl: float | None = None,
key_object_in_memory_cache: InMemoryCache | None = None,
) -> None:
# The auth cache contract is that Redis mirrors the in-memory TTL when no
# explicit Redis TTL is configured (see #43187). Compute the effective
# Redis TTL here so DualCache stays generic and does not special-case us.
effective_redis_ttl: Final = default_redis_ttl if default_redis_ttl is not None else default_in_memory_ttl
super().__init__(
in_memory_cache=in_memory_cache,
redis_cache=redis_cache,
default_in_memory_ttl=default_in_memory_ttl,
default_redis_ttl=default_redis_ttl,
default_redis_ttl=effective_redis_ttl,
)
self.key_object_cache: Final = DualCache(
in_memory_cache=key_object_in_memory_cache or InMemoryCache(),
redis_cache=redis_cache,
default_in_memory_ttl=default_in_memory_ttl,
default_redis_ttl=default_redis_ttl,
default_redis_ttl=effective_redis_ttl,
)

def in_memory_cache_for(self, key: str) -> InMemoryCache:
Expand Down
11 changes: 7 additions & 4 deletions litellm/proxy/proxy_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -4942,10 +4942,13 @@ def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache:
default_redis_ttl=CLI_SSO_SESSION_TTL_SECONDS,
)
if enable_redis_auth_cache is True:
user_api_key_cache.attach_redis_cache(
redis_cache,
default_redis_ttl=litellm.default_redis_ttl,
)
# The auth cache mirrors its own in-memory TTL on the Redis tier (see
# UserApiKeyCache.__init__ and #43187). Do not pass the global
# default_redis_ttl here: that would override the auth cache's shorter
# TTL whenever default_redis_ttl is set before Redis is attached.
# Operators who want a longer auth-cache TTL should set
# general_settings.user_api_key_cache_ttl.
user_api_key_cache.attach_redis_cache(redis_cache)
verbose_proxy_logger.info(
"enable_redis_auth_cache=True: attached Redis to "
"user_api_key_cache — virtual-key lookups are now "
Expand Down
79 changes: 79 additions & 0 deletions tests/unit/caching/test_dual_cache.py
Original file line number Diff line number Diff line change
Expand Up @@ -1022,3 +1022,82 @@ async def test_async_batch_reads_of_missing_keys_hit_redis_once_per_expiry_windo
dual_cache.last_redis_batch_access_time.update({key: time.time() - 61 for key in keys})
await dual_cache.async_batch_get_cache(keys)
assert redis_cache.async_batch_get_cache.await_count == 2

async def _write_without_a_ttl(dual_cache: DualCache, write_path: str) -> None:
"""One of the three write paths, called the way a caller that passes no ``ttl`` calls it."""
if write_path == "set_cache":
dual_cache.set_cache("ttl_key", "v")
elif write_path == "async_set_cache":
await dual_cache.async_set_cache("ttl_key", "v")
else:
await dual_cache.async_set_cache_pipeline([("ttl_key", "v")])


async def _write_with_a_ttl(dual_cache: DualCache, write_path: str, ttl: float) -> None:
"""The same three paths with an explicit ``ttl``, which must reach both tiers unchanged."""
if write_path == "set_cache":
dual_cache.set_cache("ttl_key", "v", ttl=ttl)
elif write_path == "async_set_cache":
await dual_cache.async_set_cache("ttl_key", "v", ttl=ttl)
else:
await dual_cache.async_set_cache_pipeline([("ttl_key", "v")], ttl=ttl)


def _ttl_recording_redis() -> MagicMock:
"""A Redis tier that records what it was written with, whichever of the two writes it uses."""
mock_redis = MagicMock()
mock_redis.async_set_cache = AsyncMock()
mock_redis.async_set_cache_pipeline = AsyncMock()
return mock_redis


def _written_redis_ttl(mock_redis: MagicMock, write_path: str) -> object:
"""The ttl the recording Redis tier was written with, on whichever of its writes that path uses."""
return getattr(mock_redis, write_path).call_args.kwargs["ttl"]


async def _tier_ttls(
write_path: str, ttl: float | None
) -> tuple[object, float]:
"""Write one key through one path and report the Redis ttl it was written with, plus the
in-memory expiry. ``ttl=None`` means the caller supplied no ``ttl`` at all."""
in_memory_cache = InMemoryCache(default_ttl=600)
mock_redis = _ttl_recording_redis()
dual_cache = DualCache(
in_memory_cache=in_memory_cache,
redis_cache=mock_redis,
default_in_memory_ttl=60,
default_redis_ttl=3600,
)
if ttl is None:
await _write_without_a_ttl(dual_cache, write_path)
else:
await _write_with_a_ttl(dual_cache, write_path, ttl)
return _written_redis_ttl(mock_redis, write_path), in_memory_cache.ttl_dict["ttl_key"]


@pytest.mark.asyncio
@pytest.mark.parametrize("write_path", ["set_cache", "async_set_cache", "async_set_cache_pipeline"])
async def test_dual_cache_writes_the_redis_tier_with_the_configured_default_redis_ttl(write_path: str):
"""
Regression for #43187: the Redis tier was written with default_in_memory_ttl, so a configured
default_redis_ttl never took effect. The in-memory tier still takes its own default.
"""
before = time.time()
written, expiry = await _tier_ttls(write_path, None)
after = time.time()

assert written == 3600
assert before + 60 <= expiry <= after + 60


@pytest.mark.asyncio
@pytest.mark.parametrize("write_path", ["set_cache", "async_set_cache", "async_set_cache_pipeline"])
async def test_dual_cache_gives_an_explicit_ttl_to_both_tiers_unchanged(write_path: str):
"""An explicit ``ttl`` still overrides both tier defaults rather than being replaced by them."""
before = time.time()
written, expiry = await _tier_ttls(write_path, 99.0)
after = time.time()

assert written == 99.0
assert before + 99 <= expiry <= after + 99
13 changes: 13 additions & 0 deletions tests/unit/caching/test_request_redis_batch_pre_call.py
Original file line number Diff line number Diff line change
Expand Up @@ -1024,6 +1024,19 @@ async def test_a_pipelined_management_write_without_a_ttl_expires_in_redis_like_
assert [(c[0], c[1], c[3]) for c in client.pipelines[0].commands] == [("SET", "team_id:t1", 5)]


@pytest.mark.asyncio
async def test_a_pipelined_write_without_a_ttl_uses_the_configured_default_redis_ttl():
"""Regression for #43187: the pipelined Redis SET was written with default_in_memory_ttl, so a
configured default_redis_ttl never reached it."""
client = FakeClient(_lua_ok_replies)
redis_cache = FakeRedisCache(client)
cache = DualCache(redis_cache=redis_cache, default_in_memory_ttl=5, default_redis_ttl=3600)
with request_redis_batch_scope() as request:
await cache.async_set_cache_pre_call("k1", {"v": 1}, None)
await request.flush_all()
assert [(c[0], c[1], c[3]) for c in client.pipelines[0].commands] == [("SET", "k1", 3600)]


@pytest.mark.asyncio
async def test_identity_prefetch_is_one_mget_after_which_hits_and_misses_alike_cost_no_read():
client = FakeClient(replies)
Expand Down
24 changes: 24 additions & 0 deletions tests/unit/proxy/test_redis_auth_cache_flag.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,3 +191,27 @@ def test_flag_absent_still_shares_the_model_budget_counters_over_redis(self):
ps._attach_redis_usage_cache(fake_redis, enable_redis_auth_cache=False)
assert limiter_cache.redis_cache is fake_redis
assert ps.user_api_key_cache.redis_cache is None

def test_auth_cache_redis_tier_keeps_in_memory_ttl_when_default_redis_ttl_is_set(monkeypatch):
"""
#43187: user_api_key_cache mirrors its in-memory TTL on Redis. A global
default_redis_ttl, from litellm_settings or cache_params, must not override
that when Redis is attached.
"""
from types import SimpleNamespace

from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache

monkeypatch.setattr(litellm, "default_redis_ttl", 3600)
auth_cache = UserApiKeyCache(default_in_memory_ttl=60)
with (
patch.object(ps, "user_api_key_cache", auth_cache),
patch.object(ps, "spend_counter_cache", DualCache()),
patch.object(ps, "cli_sso_session_cache", DualCache()),
patch.object(ps, "litellm_config_cache", SimpleNamespace(redis_cache=None)),
):
ps._attach_redis_usage_cache(_FakeRedisCache(), enable_redis_auth_cache=True)

assert auth_cache.redis_cache is not None
assert auth_cache.default_redis_ttl == 60
assert auth_cache.key_object_cache.default_redis_ttl == 60
Loading