From 32277795776508208b960fc060b958b4bd51ac15 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 30 Oct 2026 12:34:07 -0700 Subject: [PATCH] fix(proxy): seed a cold spend counter by adding when Redis refuses scripts The cold seed raises the counter with a Lua script. A Redis whose ACL denies scripting, or that renamed EVAL away, refused it, the reservation skipped and customer and tag budgets stopped being enforced. The seed now adds the cached spend there, as main does async_set_max also reports to the Redis service metrics like async_increment, so a failed seed still counts in litellm_redis_failed_requests --- litellm/caching/redis_cache.py | 45 ++++++++++++- litellm/integrations/otel/model/spans.py | 1 + litellm/proxy/proxy_server.py | 11 +++- tests/unit/caching/test_redis_cache.py | 40 ++++++++++++ .../proxy/proxy_server/test_spend_counters.py | 63 +++++++++++++++++++ 5 files changed, 155 insertions(+), 5 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index afe6f8a1eab..bd22c97b1f4 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -453,6 +453,18 @@ def is_redis_timeout_failure(exc: BaseException) -> bool: return any(isinstance(link, timeout_types) for link in _explicit_causes(exc)) +def is_redis_scripting_unavailable(exc: BaseException) -> bool: + """True when Redis refused to run a Lua script for this client: an ACL without + ``@scripting``, or ``EVAL`` renamed away.""" + try: + from redis.exceptions import NoPermissionError, ResponseError + except ImportError: + return False + if isinstance(exc, NoPermissionError): + return "permissions to run" in str(exc) + return isinstance(exc, ResponseError) and "unknown command" in str(exc).lower() + + def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseException) -> None: """Record a Redis failure that the calling method is about to swallow. @@ -1555,9 +1567,36 @@ class RedisCache(BaseCache): "return ARGV[1] end " "return cur" ) - result = cast( - "str | bytes | int | float | None", - await _redis_client.eval(lua, 1, key, str(value), str(int(_used_ttl or 0))), + start_time: Final = time.time() + try: + result = cast( + "str | bytes | int | float | None", + await _redis_client.eval(lua, 1, key, str(value), str(int(_used_ttl or 0))), + ) + except Exception as e: + failed_at: Final = time.time() + asyncio.create_task( + self.service_logger_obj.async_service_failure_hook( # pyright: ignore[reportUnknownMemberType] # bare dict in sig + service=ServiceTypes.REDIS, + duration=failed_at - start_time, + error=e, + call_type="async_set_max", + caller=_get_call_stack_info(), + start_time=start_time, + end_time=failed_at, + ) + ) + raise + end_time: Final = time.time() + asyncio.create_task( + self.service_logger_obj.async_service_success_hook( # pyright: ignore[reportUnknownMemberType] # bare dict in sig + service=ServiceTypes.REDIS, + duration=end_time - start_time, + call_type="async_set_max", + caller=_get_call_stack_info(), + start_time=start_time, + end_time=end_time, + ) ) if result is None: return None diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index d8be607150c..04387e14092 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -215,6 +215,7 @@ _SERVICE_VERB_BY_CALL_TYPE: Final[dict[str, str]] = { "async_increment": "incr", "async_increment_pipeline": "incr", "async_increment_if_present": "incr", + "async_set_max": "set", "delete_cache": "delete", "async_delete_cache": "delete", "async_rpush": "rpush", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 885eb3eb337..9d7e620e7c6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -313,7 +313,11 @@ from litellm.caching.redis_batch import ( active_post_call_redis_batch, drain_post_call_redis_batches, ) -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure +from litellm.caching.redis_cache import ( + RedisCircuitBreakerOpenError, + is_redis_scripting_unavailable, + is_redis_timeout_failure, +) from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.constants import ( _REALTIME_BODY_CACHE_SIZE, @@ -4169,7 +4173,10 @@ async def _lift_spend_counter_cache(counter_key: str, floor: float) -> None: if spend_counter_cache.redis_cache is not None: try: lifted: Final = await spend_counter_cache.redis_cache.async_set_max(key=counter_key, value=floor) - except Exception: + except Exception as e: + if is_redis_scripting_unavailable(e): + await increment_spend_counter_cache(counter_key=counter_key, increment=floor) + return await invalidate_spend_counter(counter_key=counter_key) raise if lifted is not None: diff --git a/tests/unit/caching/test_redis_cache.py b/tests/unit/caching/test_redis_cache.py index 2bdc7b9549b..f1f3f12de38 100644 --- a/tests/unit/caching/test_redis_cache.py +++ b/tests/unit/caching/test_redis_cache.py @@ -1697,3 +1697,43 @@ async def test_async_set_cache_pipeline_does_not_format_cached_values_for_loggin assert _FormatCountingStr.format_calls == 0 assert pipe.sets == [("k1", '"embedding"', timedelta(seconds=60))] + + +def _cache_reporting_to(service_logger: ServiceLogging) -> RedisCache: + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ): + return RedisCache(host="127.0.0.1", port=6379, service_logger_obj=service_logger) + + +@pytest.mark.asyncio +async def test_async_set_max_reports_a_failed_script_to_the_redis_failure_metric(redis_no_ping: None): + """The cold spend-counter seed runs through this script, so its Redis errors have to keep reaching + litellm_redis_failed_requests_total the way they did when the seed was an INCRBYFLOAT.""" + service_logger: Final = ServiceLogging(mock_testing=True) + cache: Final = _cache_reporting_to(service_logger) + client: Final = AsyncMock() + client.eval.side_effect = OSError("redis unavailable") + tasks_before: Final = asyncio.all_tasks() + + with patch.object(cache, "init_async_client", return_value=client): + with pytest.raises(OSError, match="redis unavailable"): + await cache.async_set_max(key="spend:end_user:customer-1", value=6.0) + await asyncio.wait_for(asyncio.gather(*(asyncio.all_tasks() - tasks_before)), timeout=5) + + assert (service_logger.mock_testing_async_success_hook, service_logger.mock_testing_async_failure_hook) == (0, 1) + + +@pytest.mark.asyncio +async def test_async_set_max_reports_a_completed_script_to_the_redis_latency_metric(redis_no_ping: None): + service_logger: Final = ServiceLogging(mock_testing=True) + cache: Final = _cache_reporting_to(service_logger) + client: Final = AsyncMock() + client.eval.return_value = b"6.0" + tasks_before: Final = asyncio.all_tasks() + + with patch.object(cache, "init_async_client", return_value=client): + assert await cache.async_set_max(key="spend:end_user:customer-1", value=6.0) == 6.0 + await asyncio.wait_for(asyncio.gather(*(asyncio.all_tasks() - tasks_before)), timeout=5) + + assert (service_logger.mock_testing_async_success_hook, service_logger.mock_testing_async_failure_hook) == (1, 0) diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index a684732327f..9f13b288dbb 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -33,6 +33,7 @@ from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import StatusCode +from redis.exceptions import NoPermissionError, ResponseError import litellm import litellm.proxy.proxy_server as ps @@ -1415,6 +1416,68 @@ async def test_a_redis_failure_while_lifting_a_cold_counter_drops_the_local_copy assert counters.in_memory_cache.get_cache(key=counter_key) is None +class _FakeRedisRefusingScripts(FakeAsyncRedis): + def __init__(self, server: FakeServer, refusal: ResponseError) -> None: + super().__init__(server=server) + self._refusal: Final = refusal + + async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> bytes | None: + raise self._refusal + + +def _cold_start_on(monkeypatch: pytest.MonkeyPatch, redis: FakeAsyncRedis, cached_spend: float) -> None: + objects: Final = MagicMock() + objects.async_get_cache = AsyncMock(return_value={"spend": cached_spend}) + monkeypatch.setattr( + ps, "spend_counter_cache", DualCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnClient(redis)) + ) + monkeypatch.setattr(ps, "user_api_key_cache", objects) + monkeypatch.setattr(ps, "prisma_client", None) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "refusal", + ( + NoPermissionError("User noscript has no permissions to run the 'eval' command"), + ResponseError("unknown command 'eval', with args beginning with: "), + ), + ids=("acl-without-scripting", "eval-renamed-away"), +) +async def test_a_cold_counter_is_still_seeded_on_a_redis_that_refuses_scripts( + monkeypatch: pytest.MonkeyPatch, refusal: ResponseError +) -> None: + redis: Final = _FakeRedisRefusingScripts(FakeServer(), refusal) + _cold_start_on(monkeypatch, redis, cached_spend=6.0) + + await ps.ensure_spend_counter_initialized(counter_key="spend:end_user:customer-1", source_cache_key="customer-1") + + assert await redis.get("spend:end_user:customer-1") == b"6" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "refusal", + ( + NoPermissionError("No permissions to access a key"), + ResponseError("WRONGTYPE Operation against a key holding the wrong kind of value"), + ), + ids=("key-outside-the-acl", "wrong-type"), +) +async def test_a_lift_refused_for_another_reason_is_raised_instead_of_seeded_by_adding( + monkeypatch: pytest.MonkeyPatch, refusal: ResponseError +) -> None: + redis: Final = _FakeRedisRefusingScripts(FakeServer(), refusal) + _cold_start_on(monkeypatch, redis, cached_spend=6.0) + + with pytest.raises(type(refusal)): + await ps.ensure_spend_counter_initialized( + counter_key="spend:end_user:customer-1", source_cache_key="customer-1" + ) + + assert await redis.get("spend:end_user:customer-1") is None + + # --------------------------------------------------------------------------- # _get_source_cache_base_spend # ---------------------------------------------------------------------------