scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

Merge 6498d2f93b into 67c7b97fd2

Browse source
This commit is contained in:
Emerson Gomes 2026-08-27 15:42:57 -05:00 • committed by GitHub
commit a380107d12
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 345 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/hooks/model_max_budget_limiter.py
View file

@ -1,3 +1,4 @@
import asyncio
import json
import time
from collections.abc import Iterable, Mapping
@ -260,6 +261,9 @@ 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._redis_increment_flush_lock = asyncio.Lock()
self._detached_increment_operations = None
self.deployment_budget_config = None
async def is_key_within_model_budget(

114
litellm/router_strategy/budget_limiter.py
View file

@ -26,7 +26,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
from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperation
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
@ -100,6 +100,9 @@ class RouterBudgetLimiting(CustomLogger):
):
self.dual_cache = dual_cache
self.redis_increment_operation_queue: list[RedisPipelineIncrementOperation] = []
self._redis_increment_queue_lock = asyncio.Lock()
self._redis_increment_flush_lock = asyncio.Lock()
self._detached_increment_operations: tuple[RedisPipelineIncrementOperation, ...] | None = None
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 +396,78 @@ 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 _detach_queued_increment_operations(self) -> tuple[RedisPipelineIncrementOperation, ...]:
async with self._get_redis_increment_queue_lock():
if self._detached_increment_operations is not None:
return self._detached_increment_operations
increment_operations_to_flush: Final = tuple(self.redis_increment_operation_queue)
self.redis_increment_operation_queue = [] # mutable-ok: emptied queue must stay appendable
self._detached_increment_operations = increment_operations_to_flush
return increment_operations_to_flush
async def _clear_detached_increment_operations(self) -> None:
async with self._get_redis_increment_queue_lock():
self._detached_increment_operations = None
async def _requeue_detached_increment_operations(self) -> None:
async with self._get_redis_increment_queue_lock():
detached_increment_operations: Final = self._detached_increment_operations
if detached_increment_operations is None:
return
self.redis_increment_operation_queue = (
list( # mutable-ok: restored flush batch must stay appendable
detached_increment_operations
)
+ self.redis_increment_operation_queue
)
self._detached_increment_operations = None
async def _finish_increment_pipeline_after_cancellation(
self,
pipeline_task: asyncio.Task[object],
) -> None:
try:
await pipeline_task
except Exception:
await self._requeue_detached_increment_operations()
verbose_router_logger.exception("Error pushing queued Redis increment operations to Redis")
return
await self._clear_detached_increment_operations()
async def _flush_queued_increment_operations(self, redis_cache: RedisCache) -> bool:
increment_operations_to_flush: Final = await self._detach_queued_increment_operations()
if len(increment_operations_to_flush) == 0:
return True
verbose_router_logger.debug(
"Pushing Redis Increment Pipeline for queue: %s",
increment_operations_to_flush,
)
increment_list: Final = list( # mutable-ok: Redis pipeline contract requires a list
increment_operations_to_flush
)
pipeline_task: Final = asyncio.create_task(
redis_cache.async_increment_pipeline(
increment_list=increment_list,
)
)
try:
await asyncio.shield(pipeline_task)
except Exception:
await asyncio.shield(self._requeue_detached_increment_operations())
verbose_router_logger.exception("Error pushing queued Redis increment operations to Redis")
return False
except asyncio.CancelledError:
await asyncio.shield(self._finish_increment_pipeline_after_cancellation(pipeline_task))
raise
await self._clear_detached_increment_operations()
return True
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
"""Original method now uses helper functions"""
@ -519,33 +593,25 @@ class RouterBudgetLimiting(CustomLogger):
DEFAULT_REDIS_SYNC_INTERVAL
) # Still wait DEFAULT_REDIS_SYNC_INTERVAL seconds on error before retrying
async def _push_in_memory_increments_to_redis(self):
async def _push_in_memory_increments_to_redis(self) -> bool:
"""
How this works:
- async_log_success_event collects all provider spend increments in `redis_increment_operation_queue`
- This function pushes all increments to Redis in a batched pipeline to optimize performance
Only runs if Redis is initialized
Only runs if Redis is initialized. Returns False when the detached batch could not be
written, so callers must not treat Redis as up to date.
"""
try:
if not self.dual_cache.redis_cache:
return # Redis is not initialized
redis_cache: Final = self.dual_cache.redis_cache
if redis_cache is None:
return True
verbose_router_logger.debug(
"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,
)
)
self.redis_increment_operation_queue = []
except Exception as e:
verbose_router_logger.error("Error syncing in-memory cache with Redis: %s", e)
async with self._redis_increment_flush_lock:
try:
return await self._flush_queued_increment_operations(redis_cache)
except asyncio.CancelledError:
await asyncio.shield(self._requeue_detached_increment_operations())
raise
async def _sync_in_memory_spend_with_redis(self):
"""
@ -566,7 +632,9 @@ class RouterBudgetLimiting(CustomLogger):
return
# 1. Push all provider spend increments to Redis
await self._push_in_memory_increments_to_redis()
flush_succeeded: Final = await self._push_in_memory_increments_to_redis()
if not flush_succeeded:
return
# 2. Fetch all current provider spend from Redis to update in-memory cache
cache_keys: Final = []

18
tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py
View file

@ -16,6 +16,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import (
resolve_model_budget,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import BudgetConfig as GenericBudgetInfo
@ -484,6 +485,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,

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

@ -0,0 +1,232 @@
import asyncio
from types import SimpleNamespace
import pytest
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import BudgetConfig
_SPEND_KEY = "provider_spend:openai:1d"
def _increment(increment_value: float) -> RedisPipelineIncrementOperation:
return RedisPipelineIncrementOperation(key=_SPEND_KEY, increment_value=increment_value, ttl=86400)
class _MockRedisCache:
def __init__(
self,
initial_values: dict[str, float],
pipeline_started: asyncio.Event | None = None,
allow_pipeline_to_complete: asyncio.Event | None = 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, float | None]:
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)
def _new_router_budget_limiter(
*,
redis_cache: object,
in_memory_cache: object | None = None,
redis_increment_operation_queue: list[RedisPipelineIncrementOperation] | None = None,
provider_budget_config: dict[str, BudgetConfig] | None = None,
) -> RouterBudgetLimiting:
budget_limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting)
budget_limiter.dual_cache = SimpleNamespace(
redis_cache=redis_cache,
in_memory_cache=in_memory_cache if in_memory_cache is not None else SimpleNamespace(),
)
budget_limiter.provider_budget_config = provider_budget_config
budget_limiter.deployment_budget_config = None
budget_limiter.tag_budget_config = None
budget_limiter.redis_increment_operation_queue = (
list(redis_increment_operation_queue) if redis_increment_operation_queue is not None else []
)
budget_limiter._redis_increment_queue_lock = asyncio.Lock()
budget_limiter._redis_increment_flush_lock = asyncio.Lock()
budget_limiter._detached_increment_operations = None
return budget_limiter
@pytest.mark.asyncio
async def test_should_await_redis_pipeline_before_sync_reads() -> None:
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 = _new_router_budget_limiter(
redis_cache=redis_cache,
in_memory_cache=in_memory_cache,
redis_increment_operation_queue=[_increment(60.0)],
provider_budget_config={"openai": BudgetConfig(time_period="1d", budget_limit=500.0)},
)
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:
redis_cache = _MockRedisCache(initial_values={}, should_fail_pipeline=True)
budget_limiter = _new_router_budget_limiter(
redis_cache=redis_cache,
redis_increment_operation_queue=[_increment(10.0)],
)
flush_succeeded = await budget_limiter._push_in_memory_increments_to_redis()
assert flush_succeeded is False
assert budget_limiter.redis_increment_operation_queue == [_increment(10.0)]
assert budget_limiter._detached_increment_operations is None
@pytest.mark.asyncio
async def test_should_keep_new_increments_when_pipeline_flush_fails() -> None:
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 = _new_router_budget_limiter(
redis_cache=redis_cache,
in_memory_cache=in_memory_cache,
redis_increment_operation_queue=[_increment(10.0)],
)
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 == [_increment(10.0), _increment(20.0)]
@pytest.mark.asyncio
async def test_should_keep_in_memory_spend_when_redis_pipeline_fails() -> None:
redis_cache = _MockRedisCache(initial_values={_SPEND_KEY: 100.0}, should_fail_pipeline=True)
in_memory_cache = _MockInMemoryCache(initial_values={_SPEND_KEY: 160.0})
budget_limiter = _new_router_budget_limiter(
redis_cache=redis_cache,
in_memory_cache=in_memory_cache,
redis_increment_operation_queue=[_increment(60.0)],
provider_budget_config={"openai": BudgetConfig(time_period="1d", budget_limit=500.0)},
)
await budget_limiter._sync_in_memory_spend_with_redis()
assert in_memory_cache.values[_SPEND_KEY] == 160.0
assert redis_cache.values[_SPEND_KEY] == 100.0
assert budget_limiter.redis_increment_operation_queue == [_increment(60.0)]
assert "batch_get" not in redis_cache.events
@pytest.mark.asyncio
async def test_should_keep_increments_when_flush_is_cancelled_after_success() -> None:
pipeline_started = asyncio.Event()
allow_pipeline_to_complete = asyncio.Event()
redis_cache = _MockRedisCache(
initial_values={_SPEND_KEY: 0.0},
pipeline_started=pipeline_started,
allow_pipeline_to_complete=allow_pipeline_to_complete,
)
budget_limiter = _new_router_budget_limiter(
redis_cache=redis_cache,
redis_increment_operation_queue=[_increment(10.0)],
)
push_task = asyncio.create_task(budget_limiter._push_in_memory_increments_to_redis())
await asyncio.wait_for(pipeline_started.wait(), timeout=1)
push_task.cancel()
allow_pipeline_to_complete.set()
with pytest.raises(asyncio.CancelledError):
await push_task
assert redis_cache.values[_SPEND_KEY] == 10.0
assert budget_limiter.redis_increment_operation_queue == []
assert budget_limiter._detached_increment_operations is None
@pytest.mark.asyncio
async def test_should_requeue_increments_when_flush_is_cancelled_and_redis_fails() -> None:
pipeline_started = asyncio.Event()
allow_pipeline_to_complete = asyncio.Event()
redis_cache = _MockRedisCache(
initial_values={_SPEND_KEY: 0.0},
pipeline_started=pipeline_started,
allow_pipeline_to_complete=allow_pipeline_to_complete,
should_fail_pipeline=True,
)
budget_limiter = _new_router_budget_limiter(
redis_cache=redis_cache,
redis_increment_operation_queue=[_increment(10.0)],
)
push_task = asyncio.create_task(budget_limiter._push_in_memory_increments_to_redis())
await asyncio.wait_for(pipeline_started.wait(), timeout=1)
push_task.cancel()
allow_pipeline_to_complete.set()
with pytest.raises(asyncio.CancelledError):
await push_task
assert redis_cache.values[_SPEND_KEY] == 0.0
assert budget_limiter.redis_increment_operation_queue == [_increment(10.0)]
assert budget_limiter._detached_increment_operations is None