From 2fc520329ffcbf6ec98b86d8f1c21e0da314e337 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 19:26:28 -0700 Subject: [PATCH] fix(router): keep the budget push off the request callback path The provider budget push runs inside the request success callback, so awaiting the Redis pipeline there made every request wait for the round trip. Hand it back to a task whose failure is logged through the breaker aware logger, so an open breaker stays a debug line and a real Redis error is one error line instead of an unretrieved task traceback --- litellm/router_strategy/budget_limiter.py | 11 +++- .../test_budget_limiter_hotpath.py | 56 +++++++++++++++++++ 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index acdc57d706d..3e094df7ac8 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -28,7 +28,7 @@ from typing import Any, Final import litellm from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisPipelineIncrementOperation, log_redis_failure +from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperation, log_redis_failure from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, @@ -93,6 +93,13 @@ class _LiteLLMParamsDictView: return dict(self._params) +async def _push_increments_to_redis(redis_cache: RedisCache, queued: list[RedisPipelineIncrementOperation]) -> None: + try: + await redis_cache.async_increment_pipeline(increment_list=queued) + except Exception as e: + log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) + + class RouterBudgetLimiting(CustomLogger): def __init__( self, @@ -540,7 +547,7 @@ class RouterBudgetLimiting(CustomLogger): queued: Final = self.redis_increment_operation_queue self.redis_increment_operation_queue = [] if queued: - await self.dual_cache.redis_cache.async_increment_pipeline(increment_list=queued) + asyncio.create_task(_push_increments_to_redis(self.dual_cache.redis_cache, queued)) except Exception as e: log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) diff --git a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py index 128b2dbd842..4cc8fe78811 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py @@ -340,3 +340,59 @@ async def test_sync_refused_by_the_open_circuit_breaker_is_quiet_and_leaks_no_ta unretrieved.assert_not_called() assert limiter.redis_increment_operation_queue == [] assert redis_cache.async_increment_pipeline.await_count == 1 + + +async def _limiter_with_redis(redis_cache: MagicMock) -> RouterBudgetLimiting: + limiter = RouterBudgetLimiting( + dual_cache=DualCache(redis_cache=redis_cache), + provider_budget_config={"openai": BudgetConfig(max_budget=1.0, budget_duration="1d")}, + ) + await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) + limiter.redis_increment_operation_queue = [{"key": "provider_spend:openai:1d", "increment_value": 0.5, "ttl": 60}] + return limiter + + +@pytest.mark.asyncio +async def test_push_returns_before_redis_answers(disable_budget_sync): + """The push runs inside the request success callback, so it must hand the Redis round trip to a task instead of waiting on it.""" + redis_answered = asyncio.Event() + + async def wait_for_redis(**_: object) -> None: + await redis_answered.wait() + + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_increment_pipeline = AsyncMock(side_effect=wait_for_redis) + limiter = await _limiter_with_redis(redis_cache) + + await asyncio.wait_for(limiter._push_in_memory_increments_to_redis(), timeout=1) + await asyncio.sleep(0) + + assert not redis_answered.is_set() + assert redis_cache.async_increment_pipeline.await_count == 1 + assert limiter.redis_increment_operation_queue == [] + redis_answered.set() + await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) + + +@pytest.mark.asyncio +async def test_push_task_failure_is_logged_once_and_not_leaked(disable_budget_sync, caplog): + """A real Redis failure on the background push must surface as one error line, never as an unretrieved task exception.""" + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_increment_pipeline = AsyncMock(side_effect=ConnectionError("Error 61 connecting to 127.0.0.1:6379")) + limiter = await _limiter_with_redis(redis_cache) + loop = asyncio.get_running_loop() + unretrieved = MagicMock() + loop.set_exception_handler(unretrieved) + + try: + with caplog.at_level(logging.ERROR): + await limiter._push_in_memory_increments_to_redis() + await asyncio.sleep(0) + gc.collect() + finally: + loop.set_exception_handler(None) + + assert [record.getMessage() for record in caplog.records] == [ + "Error syncing in-memory cache with Redis: Error 61 connecting to 127.0.0.1:6379" + ] + unretrieved.assert_not_called()