From 2949c89bb9170fed3ca918eb0b11e46c39c64e26 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Thu, 9 Jul 2026 08:28:41 -0500 Subject: [PATCH] fix(router): await budget redis pipeline before sync reads --- .../proxy/hooks/model_max_budget_limiter.py | 2 + litellm/router_strategy/budget_limiter.py | 33 ++-- ...test_unit_test_max_model_budget_limiter.py | 18 ++ .../router_strategy/test_budget_limiter.py | 160 ++++++++++++++++++ 4 files changed, 202 insertions(+), 11 deletions(-) create mode 100644 tests/test_litellm/router_strategy/test_budget_limiter.py diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 215969ef899..d0a02153ff0 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1,3 +1,4 @@ +import asyncio import json from typing import Final @@ -28,6 +29,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): def __init__(self, dual_cache: DualCache): self.dual_cache = dual_cache self.redis_increment_operation_queue = [] + self._redis_increment_queue_lock = asyncio.Lock() self.deployment_budget_config = None async def is_key_within_model_budget( diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index d57d7da0410..c52040c8e3d 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -100,6 +100,7 @@ class RouterBudgetLimiting(CustomLogger): ): self.dual_cache = dual_cache self.redis_increment_operation_queue: list[RedisPipelineIncrementOperation] = [] + self._redis_increment_queue_lock = asyncio.Lock() asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) self.provider_budget_config: GenericBudgetConfigType | None = provider_budget_config self.deployment_budget_config: GenericBudgetConfigType | None = None @@ -393,7 +394,11 @@ class RouterBudgetLimiting(CustomLogger): increment_value=response_cost, ttl=ttl, ) - self.redis_increment_operation_queue.append(increment_op) + async with self._get_redis_increment_queue_lock(): + self.redis_increment_operation_queue.append(increment_op) + + def _get_redis_increment_queue_lock(self) -> asyncio.Lock: + return self._redis_increment_queue_lock async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Original method now uses helper functions""" @@ -527,25 +532,31 @@ class RouterBudgetLimiting(CustomLogger): Only runs if Redis is initialized """ + increment_operations_to_flush: list[RedisPipelineIncrementOperation] = [] try: if not self.dual_cache.redis_cache: return # Redis is not initialized + async with self._get_redis_increment_queue_lock(): + increment_operations_to_flush = self.redis_increment_operation_queue + self.redis_increment_operation_queue = [] + verbose_router_logger.debug( "Pushing Redis Increment Pipeline for queue: %s", - self.redis_increment_operation_queue, + increment_operations_to_flush, ) - if len(self.redis_increment_operation_queue) > 0: - asyncio.create_task( - self.dual_cache.redis_cache.async_increment_pipeline( - increment_list=self.redis_increment_operation_queue, - ) + if len(increment_operations_to_flush) > 0: + await self.dual_cache.redis_cache.async_increment_pipeline( + increment_list=increment_operations_to_flush, ) - self.redis_increment_operation_queue = [] - - except Exception as e: - verbose_router_logger.error("Error syncing in-memory cache with Redis: %s", e) + except Exception: + if len(increment_operations_to_flush) > 0: + async with self._get_redis_increment_queue_lock(): + self.redis_increment_operation_queue = ( + increment_operations_to_flush + self.redis_increment_operation_queue + ) + verbose_router_logger.exception("Error pushing queued Redis increment operations to Redis") async def _sync_in_memory_spend_with_redis(self): """ diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index 55459721906..18bfa58ca8d 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -14,6 +14,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( _PROXY_VirtualKeyModelMaxBudgetLimiter, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import BudgetConfig as GenericBudgetInfo @@ -452,6 +453,23 @@ async def test_async_log_success_event_pushes_redis_increments_when_redis_config mock_push.assert_awaited_once() +@pytest.mark.asyncio +async def test_model_budget_limiter_initializes_redis_increment_queue_lock(): + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + spend_key = "virtual_key_spend:test-key:gpt-4:1d" + + await limiter._increment_spend_in_current_window( + spend_key=spend_key, response_cost=0.01, ttl=86400 + ) + + assert limiter.redis_increment_operation_queue == [ + RedisPipelineIncrementOperation( + key=spend_key, increment_value=0.01, ttl=86400 + ) + ] + + @pytest.mark.asyncio async def test_get_fallback_model_within_budget_returns_none_without_fallbacks( budget_limiter, diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py new file mode 100644 index 00000000000..c5f19c260cc --- /dev/null +++ b/tests/test_litellm/router_strategy/test_budget_limiter.py @@ -0,0 +1,160 @@ +import asyncio +from types import SimpleNamespace +from typing import Optional + +import pytest + +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.utils import BudgetConfig + + +class _MockRedisCache: + def __init__( + self, + initial_values: dict[str, float], + pipeline_started: Optional[asyncio.Event] = None, + allow_pipeline_to_complete: Optional[asyncio.Event] = None, + should_fail_pipeline: bool = False, + ) -> None: + self.values = initial_values + self.events: list[str] = [] + self.pipeline_started = pipeline_started + self.allow_pipeline_to_complete = allow_pipeline_to_complete + self.should_fail_pipeline = should_fail_pipeline + + async def async_increment_pipeline( + self, increment_list: list[RedisPipelineIncrementOperation], **kwargs: object + ) -> None: + self.events.append("increment_pipeline:start") + if self.pipeline_started is not None: + self.pipeline_started.set() + if self.allow_pipeline_to_complete is not None: + await self.allow_pipeline_to_complete.wait() + if self.should_fail_pipeline: + raise RuntimeError("redis down") + for op in increment_list: + key = op["key"] + current = float(self.values.get(key, 0.0) or 0.0) + self.values[key] = current + float(op["increment_value"]) + self.events.append("increment_pipeline:done") + + async def async_batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, Optional[float]]: + self.events.append("batch_get") + return {key: self.values.get(key) for key in key_list} + + +class _MockInMemoryCache: + def __init__(self, initial_values: dict[str, float]) -> None: + self.values = initial_values + + async def async_increment(self, key: str, value: float, ttl: int, **kwargs: object) -> float: + current = float(self.values.get(key, 0.0) or 0.0) + self.values[key] = current + float(value) + return self.values[key] + + async def async_set_cache(self, key: str, value: float, **kwargs: object) -> None: + self.values[key] = float(value) + + +@pytest.mark.asyncio +async def test_should_await_redis_pipeline_before_sync_reads() -> None: + spend_key = "provider_spend:openai:1d" + pipeline_started = asyncio.Event() + allow_pipeline_to_complete = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={spend_key: 100.0}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + ) + in_memory_cache = _MockInMemoryCache(initial_values={spend_key: 160.0}) + + budget_limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting) + budget_limiter.dual_cache = SimpleNamespace( + redis_cache=redis_cache, + in_memory_cache=in_memory_cache, + ) + budget_limiter.provider_budget_config = {"openai": BudgetConfig(time_period="1d", budget_limit=500.0)} + budget_limiter.deployment_budget_config = None + budget_limiter.tag_budget_config = None + budget_limiter.redis_increment_operation_queue = [ + RedisPipelineIncrementOperation( + key=spend_key, + increment_value=60.0, + ttl=86400, + ) + ] + budget_limiter._redis_increment_queue_lock = asyncio.Lock() + + sync_task = asyncio.create_task(budget_limiter._sync_in_memory_spend_with_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + assert "batch_get" not in redis_cache.events + allow_pipeline_to_complete.set() + await sync_task + + assert redis_cache.values[spend_key] == 160.0 + assert in_memory_cache.values[spend_key] == 160.0 + assert budget_limiter.redis_increment_operation_queue == [] + assert redis_cache.events == [ + "increment_pipeline:start", + "increment_pipeline:done", + "batch_get", + ] + + +@pytest.mark.asyncio +async def test_should_requeue_increments_when_redis_pipeline_fails() -> None: + spend_key = "provider_spend:openai:1d" + redis_cache = _MockRedisCache( + initial_values={}, + should_fail_pipeline=True, + ) + budget_limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting) + budget_limiter.dual_cache = SimpleNamespace( + redis_cache=redis_cache, + in_memory_cache=SimpleNamespace(), + ) + budget_limiter.redis_increment_operation_queue = [ + RedisPipelineIncrementOperation(key=spend_key, increment_value=10.0, ttl=86400) + ] + budget_limiter._redis_increment_queue_lock = asyncio.Lock() + + await budget_limiter._push_in_memory_increments_to_redis() + + assert budget_limiter.redis_increment_operation_queue == [ + RedisPipelineIncrementOperation(key=spend_key, increment_value=10.0, ttl=86400) + ] + + +@pytest.mark.asyncio +async def test_should_keep_new_increments_when_pipeline_flush_fails() -> None: + spend_key = "provider_spend:openai:1d" + pipeline_started = asyncio.Event() + allow_pipeline_to_complete = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + should_fail_pipeline=True, + ) + in_memory_cache = _MockInMemoryCache(initial_values={spend_key: 0.0}) + budget_limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting) + budget_limiter.dual_cache = SimpleNamespace( + redis_cache=redis_cache, + in_memory_cache=in_memory_cache, + ) + budget_limiter.redis_increment_operation_queue = [ + RedisPipelineIncrementOperation(key=spend_key, increment_value=10.0, ttl=86400) + ] + budget_limiter._redis_increment_queue_lock = asyncio.Lock() + + push_task = asyncio.create_task(budget_limiter._push_in_memory_increments_to_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + await budget_limiter._increment_spend_in_current_window(spend_key=spend_key, response_cost=20.0, ttl=86400) + allow_pipeline_to_complete.set() + await push_task + + assert budget_limiter.redis_increment_operation_queue == [ + RedisPipelineIncrementOperation(key=spend_key, increment_value=10.0, ttl=86400), + RedisPipelineIncrementOperation(key=spend_key, increment_value=20.0, ttl=86400), + ]