From 1291c2a14480a2bc21f91b08e100e4faab5efb09 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 26 Jul 2026 22:15:19 +0000 Subject: [PATCH] fix(router): make budget window resets atomic so concurrent spend is not lost Concurrent responses crossing the same budget window boundary all entered the reset path and each wrote only its own cost with a plain SET, so the shared counter ended up holding whichever write landed last. The window is now claimed atomically (Lua script when Redis is wired, an asyncio lock otherwise); the winner resets the counter, the losers increment on top of it. Queued Redis increments are also awaited instead of being dispatched with asyncio.create_task, and only the operations that were actually pushed are dropped from the queue, so a failed pipeline retries on the next sync. --- litellm/router_strategy/budget_limiter.py | 135 +++++++-- .../router_strategy/test_budget_limiter.py | 260 ++++++++++++++++++ 2 files changed, 370 insertions(+), 25 deletions(-) create mode 100644 tests/test_litellm/router_strategy/test_budget_limiter.py diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 067f38ab11c..17d3b00d43f 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -20,7 +20,7 @@ anthropic: import asyncio from datetime import datetime, timedelta, timezone -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Awaitable, Callable, Dict, List, Optional, Sequence, Tuple, Union import litellm from litellm._logging import verbose_router_logger @@ -43,6 +43,30 @@ from litellm.types.utils import GenericBudgetConfigType, StandardLoggingPayload DEFAULT_REDIS_SYNC_INTERVAL = 1 +# KEYS[1] = budget start time key, KEYS[2] = spend key +# ARGV[1] = current time, ARGV[2] = response cost, ARGV[3] = window length in seconds +# Returns {window start time, "1" if this caller opened the window else "0"} +BUDGET_WINDOW_CLAIM_SCRIPT = """ +local budget_start = redis.call('GET', KEYS[1]) +local current_time = tonumber(ARGV[1]) +local ttl = tonumber(ARGV[3]) + +if budget_start == false or (current_time - tonumber(budget_start)) > ttl then + redis.call('SET', KEYS[1], ARGV[1], 'EX', ttl) + redis.call('SET', KEYS[2], ARGV[2], 'EX', ttl) + return {ARGV[1], '1'} +end + +redis.call('INCRBYFLOAT', KEYS[2], ARGV[2]) +return {budget_start, '0'} +""" + + +def _decode_redis_value(value: Union[bytes, str, int, float]) -> str: + if isinstance(value, bytes): + return value.decode("utf-8") + return str(value) + class _LiteLLMParamsDictView: """ @@ -100,6 +124,8 @@ class RouterBudgetLimiting(CustomLogger): ): self.dual_cache = dual_cache self.redis_increment_operation_queue: List[RedisPipelineIncrementOperation] = [] + self._budget_window_lock = asyncio.Lock() + self._budget_window_claim_script: Callable[..., Awaitable[Any]] | None = None asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) self.provider_budget_config: Optional[GenericBudgetConfigType] = provider_budget_config self.deployment_budget_config: Optional[GenericBudgetConfigType] = None @@ -365,14 +391,81 @@ class RouterBudgetLimiting(CustomLogger): - The budget does not exist in cache, so we need to set it - The budget window has expired, so we need to reset everything - Does 2 things: - - stores key: `provider_spend:{provider}:1d`, value: response_cost - - stores key: `provider_budget_start_time:{provider}`, value: current_time. - This stores the start time of the new budget window + Concurrent requests crossing the same window boundary all reach this + path. Claiming the window is therefore atomic: exactly one caller resets + the spend counter to its own cost, and the callers that lose the race + increment on top of it instead of overwriting it. + + Returns the start time of the window this cost was recorded against. """ - await self.dual_cache.async_set_cache(key=spend_key, value=response_cost, ttl=ttl_seconds) - await self.dual_cache.async_set_cache(key=start_time_key, value=current_time, ttl=ttl_seconds) - return current_time + claimed_start = await self._claim_budget_window_in_redis( + spend_key=spend_key, + start_time_key=start_time_key, + current_time=current_time, + response_cost=response_cost, + ttl_seconds=ttl_seconds, + ) + if claimed_start is not None: + return claimed_start + + async with self._budget_window_lock: + budget_start = await self.dual_cache.async_get_cache(start_time_key) + if budget_start is not None and (current_time - float(budget_start)) <= ttl_seconds: + remaining_time = ttl_seconds - (current_time - float(budget_start)) + await self._increment_spend_in_current_window( + spend_key=spend_key, + response_cost=response_cost, + ttl=max(int(remaining_time), 1), + ) + return float(budget_start) + + await self.dual_cache.async_set_cache(key=spend_key, value=response_cost, ttl=ttl_seconds) + await self.dual_cache.async_set_cache(key=start_time_key, value=current_time, ttl=ttl_seconds) + return current_time + + async def _claim_budget_window_in_redis( + self, + spend_key: str, + start_time_key: str, + current_time: float, + response_cost: float, + ttl_seconds: int, + ) -> float | None: + """ + Claim the budget window through a Lua script so the reset is atomic + across proxy instances: only the caller that observes a missing or + expired start time resets the counter, everyone else racing across the + same boundary increments it. + + Returns the window start time, or None when Redis is unavailable, so the + caller can fall back to the local path. + """ + redis_cache = self.dual_cache.redis_cache + if redis_cache is None: + return None + + try: + if self._budget_window_claim_script is None: + self._budget_window_claim_script = redis_cache.async_register_script(BUDGET_WINDOW_CLAIM_SCRIPT) + result: Sequence[Union[bytes, str]] = await self._budget_window_claim_script( + keys=[start_time_key, spend_key], + args=[str(current_time), str(response_cost), ttl_seconds], + ) + budget_start = float(_decode_redis_value(result[0])) + window_claimed = _decode_redis_value(result[1]) == "1" + except Exception as e: + verbose_router_logger.error(f"Error claiming budget window for {spend_key} in Redis: {str(e)}") + return None + + if window_claimed: + await self.dual_cache.in_memory_cache.async_set_cache(key=spend_key, value=response_cost, ttl=ttl_seconds) + else: + remaining_time = ttl_seconds - (current_time - budget_start) + await self.dual_cache.in_memory_cache.async_increment( + key=spend_key, value=response_cost, ttl=max(int(remaining_time), 1) + ) + await self.dual_cache.in_memory_cache.async_set_cache(key=start_time_key, value=budget_start, ttl=ttl_seconds) + return budget_start async def _increment_spend_in_current_window(self, spend_key: str, response_cost: float, ttl: int): """ @@ -471,16 +564,7 @@ class RouterBudgetLimiting(CustomLogger): ttl_seconds=ttl_seconds, ) - if budget_start is None: - # First spend for this provider - budget_start = await self._handle_new_budget_window( - spend_key=spend_key, - start_time_key=start_time_key, - current_time=current_time, - response_cost=response_cost, - ttl_seconds=ttl_seconds, - ) - elif (current_time - budget_start) > ttl_seconds: + if (current_time - budget_start) > ttl_seconds: # Budget window expired - reset everything verbose_router_logger.debug("Budget window expired - resetting everything") budget_start = await self._handle_new_budget_window( @@ -535,14 +619,15 @@ class RouterBudgetLimiting(CustomLogger): "Pushing Redis Increment Pipeline for queue: %s", self.redis_increment_operation_queue, ) - 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, - ) - ) + pending_increments = list(self.redis_increment_operation_queue) + if len(pending_increments) == 0: + return - self.redis_increment_operation_queue = [] + await self.dual_cache.redis_cache.async_increment_pipeline( + increment_list=pending_increments, + ) + + self.redis_increment_operation_queue = self.redis_increment_operation_queue[len(pending_increments) :] except Exception as e: verbose_router_logger.error(f"Error syncing in-memory cache with Redis: {str(e)}") 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..0d6f6176deb --- /dev/null +++ b/tests/test_litellm/router_strategy/test_budget_limiter.py @@ -0,0 +1,260 @@ +import asyncio +from typing import Any, Dict, List, Optional, Sequence, Union + +import pytest + +from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisPipelineIncrementOperation +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.types.utils import BudgetConfig + + +@pytest.fixture +def disable_budget_sync(monkeypatch): + async def noop(*args, **kwargs): + return None + + monkeypatch.setattr( + "litellm.router_strategy.budget_limiter.RouterBudgetLimiting.periodic_sync_in_memory_spend_with_redis", + noop, + ) + + +class FakeAtomicRedisCache: + """ + Stands in for RedisCache, emulating the atomicity guarantee the budget + window Lua script relies on: the whole claim-or-increment runs without + yielding to other coroutines. + """ + + def __init__(self) -> None: + self.store: Dict[str, str] = {} + self.registered_scripts: List[str] = [] + self.script_calls: List[Dict[str, Any]] = [] + self.increment_pipeline_calls: List[List[RedisPipelineIncrementOperation]] = [] + + def async_register_script(self, script: str): + self.registered_scripts.append(script) + + async def run_script( + keys: Sequence[str], args: Sequence[Union[str, int, float]], client: Any = None + ) -> List[str]: + self.script_calls.append({"keys": list(keys), "args": list(args)}) + start_time_key, spend_key = keys + current_time, response_cost, ttl = str(args[0]), str(args[1]), float(args[2]) + + budget_start = self.store.get(start_time_key) + if budget_start is None or (float(current_time) - float(budget_start)) > ttl: + self.store[start_time_key] = current_time + self.store[spend_key] = response_cost + return [current_time, "1"] + + self.store[spend_key] = str(float(self.store.get(spend_key, "0")) + float(response_cost)) + return [budget_start, "0"] + + return run_script + + async def async_increment_pipeline(self, increment_list: List[RedisPipelineIncrementOperation], **kwargs): + await asyncio.sleep(0) + self.increment_pipeline_calls.append(list(increment_list)) + for op in increment_list: + self.store[op["key"]] = str(float(self.store.get(op["key"], "0")) + float(op["increment_value"])) + + +class YieldingDualCache(DualCache): + """ + DualCache whose reads and writes yield to the event loop, the way a real + network-backed cache does, so concurrent callers actually interleave. + """ + + async def async_get_cache(self, *args, **kwargs): + await asyncio.sleep(0) + return await super().async_get_cache(*args, **kwargs) + + async def async_set_cache(self, *args, **kwargs): + await asyncio.sleep(0) + return await super().async_set_cache(*args, **kwargs) + + +def _budget_limiter( + redis_cache: Optional[FakeAtomicRedisCache] = None, + dual_cache: Optional[DualCache] = None, +) -> RouterBudgetLimiting: + dual_cache = dual_cache or DualCache() + if redis_cache is not None: + dual_cache.redis_cache = redis_cache # type: ignore[assignment] + return RouterBudgetLimiting(dual_cache=dual_cache, provider_budget_config={}) + + +@pytest.mark.asyncio +async def test_concurrent_new_budget_windows_do_not_overwrite_each_other(disable_budget_sync): + """ + Two responses crossing the same window boundary both enter the reset path. + The window may only be reset once; the loser must add its cost on top. + """ + limiter = _budget_limiter() + spend_key = "provider_spend:synthetic:1h" + start_time_key = "provider_budget_start_time:synthetic" + + await asyncio.gather( + limiter._handle_new_budget_window( + spend_key=spend_key, + start_time_key=start_time_key, + current_time=1000.0, + response_cost=0.05, + ttl_seconds=3600, + ), + limiter._handle_new_budget_window( + spend_key=spend_key, + start_time_key=start_time_key, + current_time=1000.0, + response_cost=0.07, + ttl_seconds=3600, + ), + ) + + spend = await limiter.dual_cache.async_get_cache(spend_key) + assert float(spend) == pytest.approx(0.12) + assert float(await limiter.dual_cache.async_get_cache(start_time_key)) == 1000.0 + + +@pytest.mark.asyncio +async def test_concurrent_spend_increments_across_expired_window_keep_full_total(disable_budget_sync): + """ + Full path through `_increment_spend_for_key` with a window that expired long + ago: every concurrent cost must land in the freshly opened window. + """ + limiter = _budget_limiter(dual_cache=YieldingDualCache()) + budget_config = BudgetConfig(time_period="1h", budget_limit=100) + spend_key = "provider_spend:synthetic:1h" + start_time_key = "provider_budget_start_time:synthetic" + + await limiter.dual_cache.async_set_cache(key=start_time_key, value=1000.0, ttl=3600) + await limiter.dual_cache.async_set_cache(key=spend_key, value=42.0, ttl=3600) + + costs = [0.05, 0.07, 0.11] + await asyncio.gather( + *[ + limiter._increment_spend_for_key( + budget_config=budget_config, + spend_key=spend_key, + start_time_key=start_time_key, + response_cost=cost, + ) + for cost in costs + ] + ) + + spend = await limiter.dual_cache.async_get_cache(spend_key) + assert float(spend) == pytest.approx(sum(costs)) + + +@pytest.mark.asyncio +async def test_new_budget_window_claims_through_redis_when_available(disable_budget_sync): + """ + With Redis wired the reset must go through the atomic claim, and the local + cache must reflect whatever Redis decided. + """ + redis_cache = FakeAtomicRedisCache() + limiter = _budget_limiter(redis_cache=redis_cache) + spend_key = "provider_spend:synthetic:1h" + start_time_key = "provider_budget_start_time:synthetic" + + claimed_start = await limiter._handle_new_budget_window( + spend_key=spend_key, + start_time_key=start_time_key, + current_time=1000.0, + response_cost=0.05, + ttl_seconds=3600, + ) + losing_start = await limiter._handle_new_budget_window( + spend_key=spend_key, + start_time_key=start_time_key, + current_time=1000.0, + response_cost=0.07, + ttl_seconds=3600, + ) + + assert claimed_start == 1000.0 + assert losing_start == 1000.0 + assert float(redis_cache.store[spend_key]) == pytest.approx(0.12) + assert float(limiter.dual_cache.in_memory_cache.get_cache(spend_key)) == pytest.approx(0.12) + assert [call["keys"] for call in redis_cache.script_calls] == [[start_time_key, spend_key]] * 2 + + +@pytest.mark.asyncio +async def test_new_budget_window_falls_back_to_local_reset_when_redis_script_fails(disable_budget_sync): + class BrokenRedisCache(FakeAtomicRedisCache): + def async_register_script(self, script: str): + async def run_script(keys, args, client=None): + raise ConnectionError("redis is down") + + return run_script + + limiter = _budget_limiter(redis_cache=BrokenRedisCache()) + spend_key = "provider_spend:synthetic:1h" + + start_time = await limiter._handle_new_budget_window( + spend_key=spend_key, + start_time_key="provider_budget_start_time:synthetic", + current_time=1000.0, + response_cost=0.05, + ttl_seconds=3600, + ) + + assert start_time == 1000.0 + assert float(await limiter.dual_cache.async_get_cache(spend_key)) == pytest.approx(0.05) + + +@pytest.mark.asyncio +async def test_push_in_memory_increments_waits_for_redis_and_keeps_new_increments(disable_budget_sync): + """ + The queue may only be drained once Redis has actually acknowledged the + increments, and increments queued while that write is in flight must survive. + """ + redis_cache = FakeAtomicRedisCache() + limiter = _budget_limiter(redis_cache=redis_cache) + + pipeline_started = asyncio.Event() + release_pipeline = asyncio.Event() + original_increment_pipeline = redis_cache.async_increment_pipeline + + async def blocking_increment_pipeline(increment_list, **kwargs): + pipeline_started.set() + await release_pipeline.wait() + return await original_increment_pipeline(increment_list, **kwargs) + + redis_cache.async_increment_pipeline = blocking_increment_pipeline # type: ignore[method-assign] + + await limiter._increment_spend_in_current_window( + spend_key="provider_spend:synthetic:1h", response_cost=0.05, ttl=3600 + ) + push_task = asyncio.create_task(limiter._push_in_memory_increments_to_redis()) + + await pipeline_started.wait() + await limiter._increment_spend_in_current_window( + spend_key="provider_spend:synthetic:1h", response_cost=0.07, ttl=3600 + ) + release_pipeline.set() + await push_task + + assert [op["increment_value"] for op in redis_cache.increment_pipeline_calls[0]] == [0.05] + assert [op["increment_value"] for op in limiter.redis_increment_operation_queue] == [0.07] + + +@pytest.mark.asyncio +async def test_push_in_memory_increments_retains_queue_when_redis_write_fails(disable_budget_sync): + redis_cache = FakeAtomicRedisCache() + limiter = _budget_limiter(redis_cache=redis_cache) + + async def failing_increment_pipeline(increment_list, **kwargs): + raise ConnectionError("redis is down") + + redis_cache.async_increment_pipeline = failing_increment_pipeline # type: ignore[method-assign] + + await limiter._increment_spend_in_current_window( + spend_key="provider_spend:synthetic:1h", response_cost=0.05, ttl=3600 + ) + await limiter._push_in_memory_increments_to_redis() + + assert [op["increment_value"] for op in limiter.redis_increment_operation_queue] == [0.05]