diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index d4c47914394..4eccde5742b 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -20,6 +20,8 @@ from contextvars import ContextVar from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast +from pydantic import TypeAdapter + import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import ( @@ -80,11 +82,22 @@ class _AsyncRedisCommands(Protocol): def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ... + def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ... + _BREAKER_GUARD_FRAME_NAMES: Final = frozenset( {"", "wrapper", "_run_under_circuit_breaker", "_run_under_circuit_breaker_sync"} ) +_INCREMENT_WITH_FLOOR_LUA: Final = ( + "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) " + "if count < 0 then redis.call('SET', KEYS[1], 0) count = 0 end " + "redis.call('EXPIRE', KEYS[1], ARGV[2]) " + "return count" +) + +_LUA_COUNT: Final = TypeAdapter(int) + def _get_call_stack_info(num_frames: int = 2) -> str: """ @@ -680,7 +693,7 @@ class RedisCache(BaseCache): # NON blocking - notify users Redis is throwing an exception print_verbose(f"litellm.caching.caching: set() - Got exception from REDIS : {e}") - def increment_cache(self, key, value: int, ttl: float | None = None, refresh_ttl: bool = False, **kwargs) -> int: + def increment_cache(self, key, value: int, ttl: float | None = None, **kwargs) -> int: _redis_client: Final = self.redis_client start_time = time.time() set_ttl: Final = self.get_ttl(ttl=ttl) @@ -701,7 +714,7 @@ class RedisCache(BaseCache): if set_ttl is not None: # check if key already has ttl, if not -> set ttl start_time = time.time() - current_ttl: Final = -1 if refresh_ttl else _redis_client.ttl(key) + current_ttl: Final = _redis_client.ttl(key) end_time = time.time() _duration = end_time - start_time self.service_logger_obj.service_success_hook( @@ -736,6 +749,20 @@ class RedisCache(BaseCache): ) raise e + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + """Add ``value`` to ``key``, clamp the result at zero, and refresh the TTL, in one Lua call. + + A counter whose key expired while a request was still in flight would otherwise be + recreated negative by that request's decrement. Clamping inside the same call is what + keeps it safe: a separate corrective write could land after another pod's increment and + erase it. Returns the resulting count. + """ + namespaced_key: Final = self.check_and_fix_namespace(key=key) + count: Final[object] = self.redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval + _INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl + ) + return _LUA_COUNT.validate_python(count) + @_redis_circuit_breaker_guard async def async_scan_iter(self, pattern: str, count: int = 100) -> list: start_time: Final = time.time() @@ -1241,6 +1268,14 @@ class RedisCache(BaseCache): result = result.decode() return float(result) + @_redis_circuit_breaker_guard + async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: + """Async twin of ``increment_with_floor``, sharing its Lua script and its guarantees.""" + _redis_client: Final = self._async_commands() + namespaced_key: Final = self.check_and_fix_namespace(key=key) + count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl) + return _LUA_COUNT.validate_python(count) + async def flush_cache_buffer(self): print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}") await self.async_set_cache_pipeline(self.redis_batch_writing_buffer) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 00d27b5f8d9..771d2bb4328 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -188,9 +188,7 @@ class LeastBusyLoggingHandler(CustomLogger): self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) if redis_cache is None: return - shared: Final = redis_cache.increment_cache(key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True) - if shared < 0: - redis_cache.set_cache(key, 0, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) @@ -208,10 +206,6 @@ class LeastBusyLoggingHandler(CustomLogger): await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) if redis_cache is None: return - shared: Final = await redis_cache.async_increment( - key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True - ) - if shared < 0: - await redis_cache.async_set_cache(key, 0, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/test_litellm/router_strategy/test_least_busy.py index 55702c6fe73..eb9591811ac 100644 --- a/tests/test_litellm/router_strategy/test_least_busy.py +++ b/tests/test_litellm/router_strategy/test_least_busy.py @@ -40,17 +40,16 @@ class SharedRedisCounters: async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: self.set_cache(key, value) - def increment_cache(self, key: str, value: int, ttl: float | None = None, refresh_ttl: bool = False) -> int: + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: current: Final = self.count(key) or 0 assert isinstance(current, int) - incremented: Final = current + value + incremented: Final = max(0, current + value) self.encoded[key] = json.dumps(incremented) - if ttl is not None and (refresh_ttl or key not in self.ttls): - self.ttls[key] = ttl + self.ttls[key] = ttl return incremented - async def async_increment(self, key: str, value: float, ttl: int | None = None, refresh_ttl: bool = False) -> float: - return self.increment_cache(key, int(value), ttl, refresh_ttl) + async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: + return self.increment_with_floor(key, value, ttl) def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: return {key: self.count(key) for key in key_list} @@ -102,12 +101,14 @@ def test_sync_pick_reads_the_shared_counts() -> None: def test_redis_counts_keep_a_refreshed_ttl() -> None: shared: Final = SharedRedisCounters() worker: Final = _worker(shared) + key: Final = f"{GROUP}_request_count:dep-a" + shared.ttls[key] = 5 worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) worker.log_success_event(_call_kwargs("dep-a"), None, None, None) - assert shared.count(f"{GROUP}_request_count:dep-a") == 0 - assert shared.ttls == {f"{GROUP}_request_count:dep-a": IN_FLIGHT_COUNT_TTL_SECONDS} + assert shared.count(key) == 0 + assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS} @pytest.mark.asyncio @@ -129,7 +130,7 @@ class UnavailableRedis(SharedRedisCounters): def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: raise ConnectionError("redis is down") - def increment_cache(self, key: str, value: int, ttl: float | None = None, refresh_ttl: bool = False) -> int: + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: raise ConnectionError("redis is down") @@ -159,6 +160,7 @@ def test_a_shared_counter_that_expired_mid_request_cannot_go_negative() -> None: worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + assert shared.count(f"{GROUP}_request_count:dep-a") == 1 assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B