From 9f08db91f9fb2a25826cba0ccffbdb6a0e4d4bdb Mon Sep 17 00:00:00 2001 From: Michael Riad Zaky Date: Wed, 29 Apr 2026 14:07:05 -0700 Subject: [PATCH] Refresh Redis TTL on counter writes and skip stale in-memory on Redis miss --- litellm/caching/dual_cache.py | 11 ++- litellm/caching/redis_cache.py | 10 ++- litellm/proxy/db/spend_counter_reseed.py | 14 +-- litellm/proxy/proxy_server.py | 21 +++-- .../test_litellm/caching/test_redis_cache.py | 44 ++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 87 +++++++++++++++++++ 6 files changed, 165 insertions(+), 22 deletions(-) diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 34ae3638a5..2159f04296 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -392,15 +392,13 @@ class DualCache(BaseCache): value: float, parent_otel_span: Optional[Span] = None, local_only: bool = False, + refresh_ttl: bool = False, **kwargs, ) -> Optional[float]: """ - Key - the key in cache - - Value - float - the value you want to increment by - - Returns - the incremented value, or None if no cache backend is - available (in_memory_cache is None and Redis failed/is absent). + Increment counter in both caches. refresh_ttl bumps the Redis TTL + on every write (counter-style). Default preserves window-style + semantics (TTL set once on first write). """ result: Optional[float] = None try: @@ -415,6 +413,7 @@ class DualCache(BaseCache): value, parent_otel_span=parent_otel_span, ttl=kwargs.get("ttl", None), + refresh_ttl=refresh_ttl, ) return result diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 84a2887f52..deee4f6ea4 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -824,6 +824,7 @@ class RedisCache(BaseCache): value: float, ttl: Optional[int] = None, parent_otel_span: Optional[Span] = None, + refresh_ttl: bool = False, ) -> float: from redis.asyncio import Redis @@ -834,11 +835,12 @@ class RedisCache(BaseCache): try: result = await _redis_client.incrbyfloat(name=key, amount=value) if _used_ttl is not None: - # check if key already has ttl, if not -> set ttl - current_ttl = await _redis_client.ttl(key) - if current_ttl == -1: - # Key has no expiration + if refresh_ttl: await _redis_client.expire(key, _used_ttl) + else: + current_ttl = await _redis_client.ttl(key) + if current_ttl == -1: + await _redis_client.expire(key, _used_ttl) ## LOGGING ## end_time = time.time() diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index bf60a087c6..a979471dc8 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -129,7 +129,9 @@ class SpendCounterReseed: """ lock = await SpendCounterReseed._get_lock(counter_key) async with lock: - # Re-check after acquiring the lock - another waiter may have warmed it. + # Re-check after acquiring the lock. Skip in-memory on a clean + # Redis miss - in-memory is per-pod-stale. + redis_clean_miss = False if spend_counter_cache.redis_cache is not None: try: val = await spend_counter_cache.redis_cache.async_get_cache( @@ -137,11 +139,13 @@ class SpendCounterReseed: ) if val is not None: return float(val) + redis_clean_miss = True except Exception: pass - val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - if val is not None: - return float(val) + if not redis_clean_miss: + val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + if val is not None: + return float(val) db_spend = await SpendCounterReseed.from_db(prisma_client, counter_key) if db_spend is None: @@ -149,7 +153,7 @@ class SpendCounterReseed: # Warm even when 0 so subsequent reads hit cache, not DB. try: await spend_counter_cache.async_increment_cache( - key=counter_key, value=db_spend + key=counter_key, value=db_spend, refresh_ttl=True ) except Exception: verbose_proxy_logger.exception( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6f5ab1afb6..6cba6a3e96 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1798,12 +1798,16 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float: 3. Reseed from authoritative DB spend (counter expired, cross-pod stale) 4. Caller-supplied fallback (DB unavailable, cold start) """ - # 1. Try Redis first (cross-pod authoritative) + # 1. Redis first (cross-pod authoritative). On clean miss, skip + # in-memory: per-pod in-memory only has this pod's writes, so it + # would mask cross-pod increments. + redis_clean_miss = False if spend_counter_cache.redis_cache is not None: try: val = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) if val is not None: return float(val) + redis_clean_miss = True except Exception as e: verbose_proxy_logger.debug( "get_current_spend: Redis read failed for %s, falling back to in-memory: %s", @@ -1811,10 +1815,11 @@ async def get_current_spend(counter_key: str, fallback_spend: float) -> float: e, ) - # 2. Fall back to in-memory counter (single-instance or Redis failure) - val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - if val is not None: - return float(val) + # 2. In-memory only when Redis is unreachable. + if not redis_clean_miss: + val = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + if val is not None: + return float(val) # 3. Reseed from DB - fallback_spend lags cross-pod, would allow bypass. db_spend = await SpendCounterReseed.coalesced( @@ -1976,10 +1981,12 @@ async def _init_and_increment_spend_counter( base_spend = getattr(source, "spend", 0.0) or 0.0 if base_spend > 0: await spend_counter_cache.async_increment_cache( - key=counter_key, value=base_spend + key=counter_key, value=base_spend, refresh_ttl=True ) - await spend_counter_cache.async_increment_cache(key=counter_key, value=increment) + await spend_counter_cache.async_increment_cache( + key=counter_key, value=increment, refresh_ttl=True + ) async def update_cache( # noqa: PLR0915 diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index b39eb42821..78192400fb 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -50,6 +50,50 @@ async def test_redis_cache_async_increment(namespace, monkeypatch, redis_no_ping ) +@pytest.mark.asyncio +async def test_redis_cache_async_increment_refresh_ttl_true_bumps_existing_ttl( + monkeypatch, redis_no_ping +): + """With refresh_ttl=True, every increment should call expire() to bump + the TTL, even when the key already has a TTL (counter-style use).""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + mock_redis_instance = AsyncMock() + mock_redis_instance.__aenter__.return_value = mock_redis_instance + mock_redis_instance.__aexit__.return_value = None + mock_redis_instance.ttl.return_value = 42 # key already has ~42s left + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_increment( + key="spend:team_member:u:t", value=0.05, refresh_ttl=True + ) + + mock_redis_instance.expire.assert_awaited_once_with("spend:team_member:u:t", 60) + + +@pytest.mark.asyncio +async def test_redis_cache_async_increment_default_does_not_bump_existing_ttl( + monkeypatch, redis_no_ping +): + """Default (refresh_ttl=False) preserves window-style semantics: TTL is + set only on first creation, never refreshed (used by rate-limit windows).""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + mock_redis_instance = AsyncMock() + mock_redis_instance.__aenter__.return_value = mock_redis_instance + mock_redis_instance.__aexit__.return_value = None + mock_redis_instance.ttl.return_value = 42 # key already has ~42s left + + with patch.object( + redis_cache, "init_async_client", return_value=mock_redis_instance + ): + await redis_cache.async_increment(key="rate_limit:window", value=1) + + mock_redis_instance.expire.assert_not_awaited() + + @pytest.mark.asyncio async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping): monkeypatch.setenv("REDIS_HOST", "my-fake-host") diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index e0b6d229e2..37e5300565 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -5750,3 +5750,90 @@ class TestLazyFeatureMiddleware: assert attempts == [ "called" ], f"failing register_fn should be invoked once, not on every request; got {attempts}" + + +@pytest.mark.asyncio +async def test_get_current_spend_redis_clean_miss_skips_stale_in_memory(): + """When Redis is reachable and cleanly returns None (TTL expired, + counter genuinely absent), the read must reseed from DB - NOT fall + through to per-pod in-memory which only contains this pod's writes. + + Pre-fix in multi-pod deployments, in-memory contained a stale local + subset (e.g. $30) while DB had the true cross-pod total ($500). The + fall-through returned $30, enforcement passed, bypass. + """ + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import get_current_spend + + counter_cache = DualCache() + counter_key = "spend:team_member:user-1:team-1" + + # Per-pod stale in-memory: only this pod's writes, not cross-pod truth. + counter_cache.in_memory_cache.set_cache(key=counter_key, value=30.0) + + # Redis cleanly returns None (key expired or never written on this pod). + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(return_value=None) + fake_redis.async_increment = AsyncMock(return_value=500.0) + counter_cache.redis_cache = fake_redis + + # DB has the authoritative cross-pod spend. + db_row = MagicMock() + db_row.spend = 500.0 + fake_prisma = MagicMock() + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=db_row) + + import litellm.proxy.proxy_server as ps + + orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client + ps.spend_counter_cache = counter_cache + ps.prisma_client = fake_prisma + try: + spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) + assert spend == 500.0, ( + f"expected DB-authoritative 500.0 on clean Redis miss, got {spend} " + f"(stale per-pod in-memory $30 would have caused multi-pod bypass)" + ) + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma + + +@pytest.mark.asyncio +async def test_get_current_spend_redis_error_falls_back_to_in_memory(): + """When Redis raises, the read should still degrade to in-memory rather + than going straight to DB - in-memory is at least same-pod-fresh and + cheaper than a DB query during a Redis outage.""" + from litellm.caching.dual_cache import DualCache + from litellm.proxy.proxy_server import get_current_spend + + counter_cache = DualCache() + counter_key = "spend:team_member:user-1:team-1" + + counter_cache.in_memory_cache.set_cache(key=counter_key, value=42.0) + + fake_redis = AsyncMock() + fake_redis.async_get_cache = AsyncMock(side_effect=ConnectionError("redis down")) + counter_cache.redis_cache = fake_redis + + fake_prisma = MagicMock() + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( + return_value=MagicMock(spend=999.0) + ) + + import litellm.proxy.proxy_server as ps + + orig_counter, orig_prisma = ps.spend_counter_cache, ps.prisma_client + ps.spend_counter_cache = counter_cache + ps.prisma_client = fake_prisma + try: + spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) + assert spend == 42.0, ( + f"expected in-memory fallback 42.0 on Redis error, got {spend} " + f"(should not have hit DB when Redis errored)" + ) + # DB query should NOT have fired - in-memory short-circuits. + fake_prisma.db.litellm_teammembership.find_unique.assert_not_awaited() + finally: + ps.spend_counter_cache = orig_counter + ps.prisma_client = orig_prisma