From 1e433da9698e3f650ae2e506e27b4962348d910a Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 17:17:48 -0700 Subject: [PATCH 1/6] test(proxy): cover overlapping cold starts and a lost counter on one and two proxies --- .../spend/test_end_user_budget_enforcement.py | 137 +++++++++++++++++- .../proxy/proxy_server/test_spend_counters.py | 72 ++++++--- 2 files changed, 185 insertions(+), 24 deletions(-) diff --git a/tests/integration/spend/test_end_user_budget_enforcement.py b/tests/integration/spend/test_end_user_budget_enforcement.py index 02c8aee8f5e..d9c31f9d15e 100644 --- a/tests/integration/spend/test_end_user_budget_enforcement.py +++ b/tests/integration/spend/test_end_user_budget_enforcement.py @@ -1,17 +1,20 @@ import json +import os import re import uuid from collections.abc import Callable, Mapping +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from typing import Final import httpx import pytest from integration._support.client import Gateway, Scenario, eventually, object_value, string_value -from integration._support.database import read_rows +from integration._support.database import read_rows, write_rows from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import JsonResponse from pydantic import BaseModel, JsonValue +from redis import Redis PROVIDER_AUTHORIZATION: Final = "Bearer integration-provider-key" CANNED_CONTENT: Final = "Hello! This is a mock response from the fake OpenAI endpoint." @@ -639,3 +642,135 @@ def test_streamed_call_is_refused_as_json_once_over_budget(gateway: Gateway) -> assert "data:" not in denied.text, denied.text _assert_budget_refusal(denied, end_user, 0.05) assert _observed_requests(upstream) == [] + + +_COLD_BURST: Final = 40 +_COLD_SEED_SPEND: Final = 6.0 +_COLD_BUDGET: Final = 10.0 +_CHAT_COST: Final = 20 * 0.001 + 20 * 0.002 + + +def _serves_model(gateway: Gateway, model: str) -> bool: + response: Final = gateway.request("GET", "/model/info") + return response.status_code == 200 and any( + object_value(entry)["model_name"] == model for entry in object_value(response.json())["data"] + ) + + +def _counter_cache() -> Redis: + return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) + + +def _counter(cache: Redis, end_user: str) -> float | None: + raw: Final = cache.get(f"spend:end_user:{end_user}") + return None if raw is None else float(raw) + + +def _lose_counter(cache: Redis, end_user: str) -> None: + cache.delete(f"spend:end_user:{end_user}") + + +def _set_end_user_row_spend(end_user: str, spend: float) -> None: + write_rows('UPDATE "LiteLLM_EndUserTable" SET spend = %s WHERE user_id = %s', (str(spend), end_user)) + + +def _await_row_spend(end_user: str, spend: float) -> None: + eventually( + lambda: _end_user_row(end_user), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) == pytest.approx(spend), + seconds=70, + ) + + +def _refused_spend(denied: httpx.Response, end_user: str) -> float: + error: Final = _denial_error(denied) + assert error["type"] == "budget_exceeded", denied.text + matched: Final = _REFUSAL_MESSAGE.match(string_value(error["message"])) + assert matched is not None, denied.text + assert matched.group("user") == end_user + return float(matched.group("spend")) + + +@pytest.mark.parametrize("spelling", ("user", "responses_user", "messages_metadata_user_id")) +def test_overlapping_cold_starts_count_the_end_user_spend_once(gateway: Gateway, peer: Gateway, spelling: str) -> None: + chosen: Final = SPELLINGS[spelling] + with ( + gateway.scenario() as scenario, + httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, + _counter_cache() as cache, + ): + deployment: Final = _deployment(scenario, chosen) + eventually(lambda: _serves_model(peer, deployment.model), bool, seconds=30) + key: Final = scenario.key(models=[deployment.model]) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + _new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + assert _customer_info(gateway, end_user).spend == _COLD_SEED_SPEND + assert _counter(cache, end_user) is None + upstream.delete("/__observations").raise_for_status() + + proxies: Final = (gateway, peer) + with ThreadPoolExecutor(max_workers=_COLD_BURST) as pool: + burst: Final = tuple( + pool.map( + lambda index: _send(proxies[index % 2], chosen, deployment, key, end_user, f"cold {index}"), + range(_COLD_BURST), + ) + ) + refused: Final = [response.text for response in burst if response.status_code != 200] + assert refused == [], refused + follow_up: Final = _send(gateway, chosen, deployment, key, end_user, "cold follow-up") + _assert_served(chosen, follow_up) + assert len(_observed_requests(upstream)) == _COLD_BURST + 1 + + expected_spend: Final = _COLD_SEED_SPEND + (_COLD_BURST + 1) * _CHAT_COST + _await_row_spend(end_user, expected_spend) + assert _customer_info(gateway, end_user).spend == pytest.approx(expected_spend) + counted: Final = eventually( + lambda: _counter(cache, end_user), + lambda value: value is not None and value == pytest.approx(expected_spend), + seconds=30, + return_last_on_timeout=True, + ) + assert counted == pytest.approx(expected_spend) + + +def test_a_lost_counter_keeps_spend_this_proxy_has_not_written_yet(gateway: Gateway) -> None: + with gateway.scenario() as scenario, _counter_cache() as cache: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.25) + key: Final = scenario.key(models=[model]) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + call_cost: Final = 20 * 0.001 + 20 * 0.25 + _new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + served: Final = _chat(gateway, model, key, end_user) + assert served.status_code == 200, served.text + _await_row_spend(end_user, _COLD_SEED_SPEND + call_cost) + _set_end_user_row_spend(end_user, _COLD_SEED_SPEND) + _lose_counter(cache, end_user) + + denied: Final = _chat(gateway, model, key, end_user) + + assert _refused_spend(denied, end_user) == pytest.approx(_COLD_SEED_SPEND + call_cost) + + +def test_a_lost_counter_keeps_spend_another_proxy_has_not_written_yet(gateway: Gateway, peer: Gateway) -> None: + with gateway.scenario() as scenario, _counter_cache() as cache: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.149) + eventually(lambda: _serves_model(peer, model), bool, seconds=30) + key: Final = scenario.key(models=[model]) + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + call_cost: Final = 20 * 0.001 + 20 * 0.149 + _new_customer(gateway, end_user, max_budget=_COLD_BUDGET, spend=_COLD_SEED_SPEND) + scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]}) + first: Final = _chat(peer, model, key, end_user) + assert first.status_code == 200, first.text + _await_row_spend(end_user, _COLD_SEED_SPEND + call_cost) + _set_end_user_row_spend(end_user, _COLD_SEED_SPEND) + _lose_counter(cache, end_user) + second: Final = _chat(peer, model, key, end_user) + assert second.status_code == 200, second.text + + denied: Final = _chat(gateway, model, key, end_user) + + assert _refused_spend(denied, end_user) == pytest.approx(_COLD_SEED_SPEND + 2 * call_cost) diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index d273ef8f42f..ff477221fa5 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -38,6 +38,7 @@ import litellm import litellm.proxy.proxy_server as ps from litellm._service_logger import ServiceLogging from litellm._internal_context import current_service_target +from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache, RedisCircuitBreaker from litellm.constants import GLOBAL_PROXY_SPEND_CACHE_KEY @@ -1329,32 +1330,57 @@ async def test_ensure_spend_counter_initialized_warm_skips_reseed_and_source( } -@pytest.mark.asyncio -async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache( - monkeypatch, -): - fake_cache = _make_spend_counter_cache(redis_get_value=None, redis_increment_value=7.0) - fake_user_cache = _make_user_api_key_cache(get_value={"spend": 7.0}) - monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) - monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) - monkeypatch.setattr(ps, "prisma_client", None) - monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) +class _SharedRedisCounters: + """The Redis counter commands a cold start uses, yielding first like a network call so overlapping starts + interleave the way they do across workers and pods.""" - await ps._ensure_spend_counter_initialized( - counter_key="spend:user:u", - source_cache_key="u", + def __init__(self) -> None: + self.values: Final[dict[str, float]] = {} # mutable-ok: stands in for the Redis keyspace + + async def async_get_cache(self, key: str, **_: object) -> float | None: + await asyncio.sleep(0) + return self.values.get(key) + + async def async_increment(self, key: str, value: float, **_: object) -> float: + await asyncio.sleep(0) + self.values[key] = self.values.get(key, 0.0) + value + return self.values[key] + + async def async_set_max(self, key: str, value: float, **_: object) -> float: + await asyncio.sleep(0) + self.values[key] = max(self.values.get(key, value), value) + return self.values[key] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("with_redis", (True, False)) +@pytest.mark.parametrize("counter_key", ("spend:end_user:customer-1", "spend:tag:tag-1")) +async def test_overlapping_cold_starts_count_the_cached_spend_once(monkeypatch, with_redis, counter_key): + cached_spend: Final = {"first-start": 2.5, "second-start": 6.0, "third-start": 6.0, "fourth-start": 2.5} + + async def _cached_object(key: str, **_: object) -> dict[str, float]: + await asyncio.sleep(0) + return {"spend": cached_spend[key]} + + redis: Final = _SharedRedisCounters() + counters: Final = DualCache( + in_memory_cache=InMemoryCache(), redis_cache=cast(RedisCache, redis) if with_redis else None + ) + objects: Final = MagicMock() + objects.async_get_cache = AsyncMock(side_effect=_cached_object) + monkeypatch.setattr(ps, "spend_counter_cache", counters) + monkeypatch.setattr(ps, "user_api_key_cache", objects) + monkeypatch.setattr(ps, "prisma_client", None) + + await asyncio.gather( + *( + ps.ensure_spend_counter_initialized(counter_key=counter_key, source_cache_key=start) + for start in cached_spend + ) ) - observed = { - "source_cache_called": fake_user_cache.async_get_cache.called, - "seed_increment_called": fake_cache.redis_cache.async_increment.called, - "warm_check_done": fake_cache.redis_cache.async_get_cache.called, - } - assert normalize(observed) == { - "source_cache_called": True, - "seed_increment_called": True, - "warm_check_done": True, - } + seeded: Final = redis.values.get(counter_key) if with_redis else counters.in_memory_cache.get_cache(key=counter_key) + assert seeded == 6.0 # --------------------------------------------------------------------------- From 6bdc5e1b01d40b57e6d593b4bbd0442e83d4a1cb Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 17:17:48 -0700 Subject: [PATCH 2/6] fix(proxy): lift a cold spend counter to the cached spend instead of adding it --- litellm/proxy/proxy_server.py | 46 +++++++++++++++++++++++------------ 1 file changed, 31 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 129cd7e1589..885eb3eb337 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3280,15 +3280,7 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None in-memory copy is guarded by a read-compare-write with no await in between, so it is atomic within the worker. """ - cached: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) - needs_update = True - if cached is not None: - try: - needs_update = float(cached) < db_spend - except (TypeError, ValueError): - needs_update = True - if needs_update: - spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend) + _lift_in_memory_spend_counter(counter_key=counter_key, floor=db_spend) if spend_counter_cache.redis_cache is not None: forget_spend_counter(counter_key) try: @@ -3304,6 +3296,18 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None record_spend_counter_value(counter_key, repaired) +def _lift_in_memory_spend_counter(counter_key: str, floor: float) -> None: + cached: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + needs_update = True + if cached is not None: + try: + needs_update = float(cached) < floor + except (TypeError, ValueError): + needs_update = True + if needs_update: + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=floor) + + async def reseed_spend_counter_from_db(counter_key: str) -> bool: """Recover a counter that the reservation reconcile found in an inconsistent state (missing, or where applying the reconcile delta would drive it @@ -3925,11 +3929,8 @@ async def _prepare_spend_counter_increment( 2. If not found, reseed from the DB via `SpendCounterReseed.coalesced`. Falls back to the cached object's `.spend` via user_api_key_cache only if prisma is unavailable, since that value can lag the flusher. - 3. Seed counter via async_increment_cache (not async_set_cache) to avoid a - check-then-set race: if two pods cold-start simultaneously, both may see - the counter as absent and seed it. Using increment means the worst case - is over-counting (conservative, blocks slightly early) rather than - under-counting (would allow overspend). + 3. Lift the counter to at least that spend rather than adding it, so + requests and pods that cold-start it together count it once. 4. Increment is returned for the caller to apply via pipeline """ with service_target(SPEND_COUNTERS_TARGET): @@ -4034,7 +4035,7 @@ async def ensure_spend_counter_initialized( # DB unavailable - fall back to in-process cache (may be stale). base_spend: Final = await _get_source_cache_base_spend(source_cache_key=source_cache_key) if base_spend > 0: - await increment_spend_counter_cache(counter_key=counter_key, increment=base_spend) + await _lift_spend_counter_cache(counter_key=counter_key, floor=base_spend) _ensure_spend_counter_initialized: Final = ensure_spend_counter_initialized @@ -4163,6 +4164,21 @@ async def increment_spend_counter_cache(counter_key: str, increment: float) -> f _increment_spend_counter_cache: Final = increment_spend_counter_cache +@with_service_target(SPEND_COUNTERS_TARGET) +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: + await invalidate_spend_counter(counter_key=counter_key) + raise + if lifted is not None: + _lift_in_memory_spend_counter(counter_key=counter_key, floor=lifted) + record_spend_counter_value(counter_key, lifted) + return + _lift_in_memory_spend_counter(counter_key=counter_key, floor=floor) + + @with_service_target(SPEND_COUNTERS_TARGET) async def invalidate_spend_counter(counter_key: str) -> None: forget_spend_counter(counter_key) From efefd3949c10aa261cee3ff1c502aea51a82b83f Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 17:58:32 -0700 Subject: [PATCH 3/6] test(proxy): cover a Redis failure while lifting a cold spend counter --- .../proxy/proxy_server/test_spend_counters.py | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index ff477221fa5..0dc9d0c4bff 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -1383,6 +1383,34 @@ async def test_overlapping_cold_starts_count_the_cached_spend_once(monkeypatch, assert seeded == 6.0 +class _RedisThatFailsTheLift: + async def async_get_cache(self, key: str, **_: object) -> None: + return None + + async def async_set_max(self, key: str, value: float, **_: object) -> float: + raise ConnectionError("redis went away") + + async def async_delete_cache(self, key: str, **_: object) -> None: + return None + + +@pytest.mark.asyncio +async def test_a_redis_failure_while_lifting_a_cold_counter_drops_the_local_copy_and_raises(monkeypatch): + counter_key: Final = "spend:end_user:customer-1" + counters: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=cast(RedisCache, _RedisThatFailsTheLift())) + counters.in_memory_cache.set_cache(key=counter_key, value=2.0) + objects: Final = MagicMock() + objects.async_get_cache = AsyncMock(return_value={"spend": 6.0}) + monkeypatch.setattr(ps, "spend_counter_cache", counters) + monkeypatch.setattr(ps, "user_api_key_cache", objects) + monkeypatch.setattr(ps, "prisma_client", None) + + with pytest.raises(ConnectionError): + await ps.ensure_spend_counter_initialized(counter_key=counter_key, source_cache_key="customer-1") + + assert counters.in_memory_cache.get_cache(key=counter_key) is None + + # --------------------------------------------------------------------------- # _get_source_cache_base_spend # --------------------------------------------------------------------------- From b8e12c91445293f1cc8cecced50c25352a1fe5f8 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 18:54:28 -0700 Subject: [PATCH 4/6] test(proxy): type the new spend counter test parameters --- tests/unit/proxy/proxy_server/test_spend_counters.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 0dc9d0c4bff..a684732327f 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -1355,7 +1355,9 @@ class _SharedRedisCounters: @pytest.mark.asyncio @pytest.mark.parametrize("with_redis", (True, False)) @pytest.mark.parametrize("counter_key", ("spend:end_user:customer-1", "spend:tag:tag-1")) -async def test_overlapping_cold_starts_count_the_cached_spend_once(monkeypatch, with_redis, counter_key): +async def test_overlapping_cold_starts_count_the_cached_spend_once( + monkeypatch: pytest.MonkeyPatch, with_redis: bool, counter_key: str +) -> None: cached_spend: Final = {"first-start": 2.5, "second-start": 6.0, "third-start": 6.0, "fourth-start": 2.5} async def _cached_object(key: str, **_: object) -> dict[str, float]: @@ -1395,7 +1397,9 @@ class _RedisThatFailsTheLift: @pytest.mark.asyncio -async def test_a_redis_failure_while_lifting_a_cold_counter_drops_the_local_copy_and_raises(monkeypatch): +async def test_a_redis_failure_while_lifting_a_cold_counter_drops_the_local_copy_and_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: counter_key: Final = "spend:end_user:customer-1" counters: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=cast(RedisCache, _RedisThatFailsTheLift())) counters.in_memory_cache.set_cache(key=counter_key, value=2.0) From 32277795776508208b960fc060b958b4bd51ac15 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 30 Oct 2026 12:34:07 -0700 Subject: [PATCH 5/6] 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 # --------------------------------------------------------------------------- From 1343f38da97b4f26d9dcfa3fabdac01d22ebaf5e Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 30 Oct 2026 13:12:51 -0700 Subject: [PATCH 6/6] refactor(caching): import redis exceptions directly in is_redis_scripting_unavailable Its only caller runs after a Redis call raised, so redis is always installed there, the same as the sibling helpers that import it lazily without a fallback --- litellm/caching/redis_cache.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index bd22c97b1f4..a707e31ef32 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -456,10 +456,8 @@ def is_redis_timeout_failure(exc: BaseException) -> bool: 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 + from redis.exceptions import NoPermissionError, ResponseError + if isinstance(exc, NoPermissionError): return "permissions to run" in str(exc) return isinstance(exc, ResponseError) and "unknown command" in str(exc).lower()