scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

fix(router): make budget window resets atomic so concurrent spend is not lost

Browse source 
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.
This commit is contained in:
Devin AI 2026-07-26 22:15:19 +00:00
parent 24123269cc
commit 1291c2a144
2 changed files with 370 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

135
litellm/router_strategy/budget_limiter.py
View file

@ -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)}")

260
tests/test_litellm/router_strategy/test_budget_limiter.py Normal file
View file

@ -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]