scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 98

fix(router): await budget redis pipeline before sync reads (internal copy of #32618) (#43125)

Browse source 
* fix(router): await budget redis pipeline before sync reads

* refactor(router): remove superseded Redis flush helper

* fix(router): preserve concurrent spend during Redis synchronization

* fix(router): type per-key spend totals without loop Final bindings

* fix(router): log Redis failures before cancellable cleanup

* fix(router): finalize Redis batches before propagating cancellation

* test(router): reproduce cancellation while Redis cleanup is blocked

* test: align budget hotpath checks with awaited Redis flush

* fix(budgets): finish Redis flush after cancellation while queued

* test(budgets): consolidate Redis regressions in mapped tests

* fix(budgets): coalesce failed Redis increments by key

* test(budgets): assert spend behavior instead of batch state

* perf(router): sum pending budget spend by key once per sync

---------

Co-authored-by: Emerson Gomes <emerson.gomes@thalesgroup.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-24 22:27:13 -07:00 • committed by GitHub
parent 9915df6875
commit ccf866c801
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 592 additions and 93 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, Sequence
@ -297,6 +298,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(

213
litellm/router_strategy/budget_limiter.py
View file

@ -21,8 +21,10 @@ anthropic:
import asyncio
import builtins
import logging
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from itertools import groupby
from types import MappingProxyType
from typing import Any, Final
import litellm
@ -93,11 +95,12 @@ 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)
def _sum_increments_by_key(operations: Sequence[RedisPipelineIncrementOperation]) -> Mapping[str, float]:
by_key: Final = groupby(
sorted(operations, key=lambda operation: operation["key"]),
key=lambda operation: operation["key"],
)
return MappingProxyType({key: sum(operation["increment_value"] for operation in group) for key, group in by_key})
class RouterBudgetLimiting(CustomLogger):
@ -109,6 +112,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
@ -392,17 +398,97 @@ class RouterBudgetLimiting(CustomLogger):
- Increments the spend in memory cache (so spend instantly updated in memory)
- Queues the increment operation to Redis Pipeline (using batched pipeline to optimize performance. Using Redis for multi instance environment of LiteLLM)
"""
await self.dual_cache.in_memory_cache.async_increment(
key=spend_key,
value=response_cost,
ttl=ttl,
)
increment_op: Final = RedisPipelineIncrementOperation(
key=spend_key,
increment_value=response_cost,
ttl=ttl,
)
self.redis_increment_operation_queue.append(increment_op)
async with self._get_redis_increment_queue_lock():
await self.dual_cache.in_memory_cache.async_increment(
key=spend_key,
value=response_cost,
ttl=ttl,
)
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)
if not increment_operations_to_flush:
return increment_operations_to_flush
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
operations: Final = (*detached_increment_operations, *self.redis_increment_operation_queue)
grouped_operations: Final = (
(key, tuple(group))
for key, group in groupby(
sorted(operations, key=lambda operation: operation["key"]),
key=lambda operation: operation["key"],
)
)
self.redis_increment_operation_queue = [
RedisPipelineIncrementOperation(
key=key,
increment_value=sum(operation["increment_value"] for operation in group),
ttl=group[-1]["ttl"],
)
for key, group in grouped_operations
]
self._detached_increment_operations = None
async def _flush_queued_increment_operations(self, redis_cache: RedisCache) -> bool:
flush_task: Final = asyncio.create_task(self._write_queued_increment_operations(redis_cache))
return await self._await_flush_task(flush_task)
async def _await_flush_task(self, flush_task: asyncio.Task[bool]) -> bool:
try:
return await asyncio.shield(flush_task)
except asyncio.CancelledError:
while not flush_task.done():
try:
await asyncio.shield(flush_task)
except asyncio.CancelledError:
continue
flush_task.result()
raise
async def _write_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:
await self._clear_detached_increment_operations()
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
)
try:
await redis_cache.async_increment_pipeline(increment_list=increment_list)
except Exception as error:
log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", error)
await self._requeue_detached_increment_operations()
return False
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"""
@ -528,29 +614,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,
)
queued: Final = self.redis_increment_operation_queue
self.redis_increment_operation_queue = []
if queued:
asyncio.create_task(_push_increments_to_redis(self.dual_cache.redis_cache, queued))
flush_task: Final = asyncio.create_task(self._flush_queued_increments_with_lock(redis_cache))
return await self._await_flush_task(flush_task)
except Exception as e:
log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e)
async def _flush_queued_increments_with_lock(self, redis_cache: RedisCache) -> bool:
async with self._redis_increment_flush_lock:
return await self._flush_queued_increment_operations(redis_cache)
async def _sync_in_memory_spend_with_redis(self):
"""
@ -569,44 +651,51 @@ class RouterBudgetLimiting(CustomLogger):
# No need to sync if Redis cache is not initialized
if self.dual_cache.redis_cache is None:
return
# 1. Push all provider spend increments to Redis
await self._push_in_memory_increments_to_redis()
# 2. Fetch all current provider spend from Redis to update in-memory cache
cache_keys: Final = []
if self.provider_budget_config is not None:
for provider, config in self.provider_budget_config.items():
if config is None:
continue
cache_keys.append(f"provider_spend:{provider}:{config.budget_duration}")
if self.deployment_budget_config is not None:
for model_id, config in self.deployment_budget_config.items():
if config is None:
continue
cache_keys.append(f"deployment_spend:{model_id}:{config.budget_duration}")
if self.tag_budget_config is not None:
for tag, config in self.tag_budget_config.items():
if config is None:
continue
cache_keys.append(f"tag_spend:{tag}:{config.budget_duration}")
# Batch fetch current spend values from Redis
redis_values: Final = await self.dual_cache.redis_cache.async_batch_get_cache(key_list=cache_keys)
# Update in-memory cache with Redis values
if isinstance(redis_values, dict): # Check if redis_values is a dictionary
for key, value in redis_values.items():
if value is not None:
await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=float(value))
verbose_router_logger.debug("Updated in-memory cache for %s: %s", key, value)
async with self._redis_increment_flush_lock:
await self._flush_increments_then_copy_redis_spend()
except Exception as e:
log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e)
async def _flush_increments_then_copy_redis_spend(self) -> None:
redis_cache: Final = self.dual_cache.redis_cache
if redis_cache is None or not await self._flush_queued_increment_operations(redis_cache):
return
cache_keys: Final = []
if self.provider_budget_config is not None:
for provider, config in self.provider_budget_config.items():
if config is None:
continue
cache_keys.append(f"provider_spend:{provider}:{config.budget_duration}")
if self.deployment_budget_config is not None:
for model_id, config in self.deployment_budget_config.items():
if config is None:
continue
cache_keys.append(f"deployment_spend:{model_id}:{config.budget_duration}")
if self.tag_budget_config is not None:
for tag, config in self.tag_budget_config.items():
if config is None:
continue
cache_keys.append(f"tag_spend:{tag}:{config.budget_duration}")
redis_values: Final = await redis_cache.async_batch_get_cache(key_list=cache_keys)
if not isinstance(redis_values, dict):
return
async with self._get_redis_increment_queue_lock():
pending_spend_by_key: Final = _sum_increments_by_key(self.redis_increment_operation_queue)
updated_spend_by_key: Final = tuple(
(key, float(value) + pending_spend_by_key.get(key, 0.0))
for key, value in redis_values.items()
if value is not None
)
for key, updated_spend in updated_spend_by_key:
await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=updated_spend)
verbose_router_logger.debug("Updated in-memory cache for %s: %s", key, updated_spend)
def _get_budget_config_for_deployment(
self,
model_id: str,

450
tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py
View file

@ -1,6 +1,8 @@
import asyncio
import gc
import logging
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -9,6 +11,7 @@ import litellm
from litellm.caching.caching import DualCache
from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.router import LiteLLM_Params
from litellm.types.utils import BudgetConfig
@ -30,9 +33,7 @@ async def test_get_llm_provider_for_deployment_dict_does_not_require_litellm_par
):
class RaiseOnInit:
def __init__(self, *args, **kwargs):
raise AssertionError(
"LiteLLM_Params should not be instantiated in hot path"
)
raise AssertionError("LiteLLM_Params should not be instantiated in hot path")
monkeypatch.setattr(
"litellm.router_strategy.budget_limiter.LiteLLM_Params",
@ -99,9 +100,7 @@ async def test_get_llm_provider_for_deployment_dict_view_supports_mapping_and_at
@pytest.mark.asyncio
async def test_async_filter_deployments_resolves_provider_once_per_deployment(
disable_budget_sync, monkeypatch
):
async def test_async_filter_deployments_resolves_provider_once_per_deployment(disable_budget_sync, monkeypatch):
provider_budget = RouterBudgetLimiting(
dual_cache=DualCache(),
provider_budget_config={
@ -207,9 +206,7 @@ def _legacy_provider_resolution(deployment):
Reference implementation used before hot-path optimization.
"""
try:
_litellm_params = LiteLLM_Params(
**deployment.get("litellm_params", {"model": ""})
)
_litellm_params = LiteLLM_Params(**deployment.get("litellm_params", {"model": ""}))
_, custom_llm_provider, _, _ = litellm.get_llm_provider(
model=_litellm_params.model,
litellm_params=_litellm_params,
@ -228,9 +225,7 @@ def _legacy_provider_resolution(deployment):
],
)
@pytest.mark.asyncio
async def test_get_llm_provider_for_deployment_matches_legacy_behavior(
disable_budget_sync, deployment
):
async def test_get_llm_provider_for_deployment_matches_legacy_behavior(disable_budget_sync, deployment):
provider_budget = RouterBudgetLimiting(
dual_cache=DualCache(),
provider_budget_config={},
@ -242,9 +237,7 @@ async def test_get_llm_provider_for_deployment_matches_legacy_behavior(
assert current_provider == legacy_provider
def test_register_deployment_budget_for_runtime_added_deployment(
disable_budget_sync, monkeypatch
):
def test_register_deployment_budget_for_runtime_added_deployment(disable_budget_sync, monkeypatch):
import asyncio
monkeypatch.setattr(asyncio, "create_task", lambda coro: None)
@ -274,9 +267,7 @@ def test_register_deployment_budget_for_runtime_added_deployment(
assert budget_limiter._get_budget_config_for_deployment(model_id) is None
def test_router_add_deployment_registers_deployment_budget(
disable_budget_sync, monkeypatch
):
def test_router_add_deployment_registers_deployment_budget(disable_budget_sync, monkeypatch):
import asyncio
from litellm import Router
@ -304,9 +295,7 @@ def test_router_add_deployment_registers_deployment_budget(
budget_limiter = router._get_router_deployment_budget_limiter()
assert budget_limiter is not None
config = budget_limiter._get_budget_config_for_deployment(
"runtime-budget-deployment"
)
config = budget_limiter._get_budget_config_for_deployment("runtime-budget-deployment")
assert config is not None
assert config.max_budget == 0.000000000001
@ -338,7 +327,9 @@ async def test_sync_refused_by_the_open_circuit_breaker_is_quiet_and_leaks_no_ta
assert caplog.records == []
unretrieved.assert_not_called()
assert limiter.redis_increment_operation_queue == []
assert limiter.redis_increment_operation_queue == [
{"key": "provider_spend:openai:1d", "increment_value": 0.5, "ttl": 60}
]
assert redis_cache.async_increment_pipeline.await_count == 1
@ -353,32 +344,34 @@ async def _limiter_with_redis(redis_cache: MagicMock) -> RouterBudgetLimiting:
@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."""
async def test_push_waits_for_redis_before_completing(disable_budget_sync):
redis_started = asyncio.Event()
redis_answered = asyncio.Event()
async def wait_for_redis(**_: object) -> None:
redis_started.set()
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()
push_task = asyncio.create_task(limiter._push_in_memory_increments_to_redis())
await asyncio.wait_for(redis_started.wait(), timeout=1)
assert not push_task.done()
redis_answered.set()
assert await asyncio.wait_for(push_task, timeout=1) is True
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"))
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()
@ -396,3 +389,398 @@ async def test_push_task_failure_is_logged_once_and_not_leaked(disable_budget_sy
"Error syncing in-memory cache with Redis: Error 61 connecting to 127.0.0.1:6379"
]
unretrieved.assert_not_called()
_SPEND_KEY = "provider_spend:openai:1d"
def _increment(increment_value: float) -> RedisPipelineIncrementOperation:
return RedisPipelineIncrementOperation(key=_SPEND_KEY, increment_value=increment_value, ttl=86400)
class _ObservedLock(asyncio.Lock):
def __init__(self) -> None:
super().__init__()
self.waiter_started = asyncio.Event()
async def acquire(self) -> bool:
if self.locked():
self.waiter_started.set()
return await super().acquire()
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,
pipeline_completed: asyncio.Event | None = None,
read_started: asyncio.Event | None = None,
allow_read_to_complete: asyncio.Event | None = None,
) -> 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
self.pipeline_completed = pipeline_completed
self.read_started = read_started
self.allow_read_to_complete = allow_read_to_complete
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")
if self.pipeline_completed is not None:
self.pipeline_completed.set()
async def async_batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, float | None]:
self.events.append("batch_get")
snapshot = {key: self.values.get(key) for key in key_list}
if self.read_started is not None:
self.read_started.set()
if self.allow_read_to_complete is not None:
await self.allow_read_to_complete.wait()
return snapshot
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,
queue_lock: asyncio.Lock | None = None,
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 = queue_lock if queue_lock is not None else 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)]
@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(30.0)]
@pytest.mark.asyncio
async def test_failed_redis_flushes_coalesce_spend_by_key() -> None:
other_spend_key: Final = "provider_spend:other:1d"
redis_cache: Final = _MockRedisCache(
initial_values={_SPEND_KEY: 0.0, other_spend_key: 0.0}, should_fail_pipeline=True
)
in_memory_cache: Final = _MockInMemoryCache(initial_values={_SPEND_KEY: 0.0, other_spend_key: 0.0})
budget_limiter: Final = _new_router_budget_limiter(redis_cache=redis_cache, in_memory_cache=in_memory_cache)
for spend_key, response_cost, ttl in (
(_SPEND_KEY, 10.0, 90),
(other_spend_key, 4.0, 50),
(_SPEND_KEY, 20.0, 80),
(_SPEND_KEY, 30.0, 70),
):
await budget_limiter._increment_spend_in_current_window(spend_key, response_cost, ttl)
assert await budget_limiter._push_in_memory_increments_to_redis() is False
queued: Final = {operation["key"]: operation for operation in budget_limiter.redis_increment_operation_queue}
assert len(budget_limiter.redis_increment_operation_queue) == 2
assert queued[_SPEND_KEY] == RedisPipelineIncrementOperation(key=_SPEND_KEY, increment_value=60.0, ttl=70)
assert queued[other_spend_key] == RedisPipelineIncrementOperation(key=other_spend_key, increment_value=4.0, ttl=50)
redis_cache.should_fail_pipeline = False
assert await budget_limiter._push_in_memory_increments_to_redis() is True
assert redis_cache.values == {_SPEND_KEY: 60.0, other_spend_key: 4.0}
assert budget_limiter.redis_increment_operation_queue == []
@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 == []
@pytest.mark.asyncio
async def test_cancelled_push_waiting_for_flush_lock_still_writes_spend() -> None:
redis_cache = _MockRedisCache(initial_values={_SPEND_KEY: 0.0})
budget_limiter = _new_router_budget_limiter(
redis_cache=redis_cache,
redis_increment_operation_queue=[_increment(10.0)],
)
flush_lock = _ObservedLock()
budget_limiter._redis_increment_flush_lock = flush_lock
async with flush_lock:
push_task = asyncio.create_task(budget_limiter._push_in_memory_increments_to_redis())
await asyncio.wait_for(flush_lock.waiter_started.wait(), timeout=1)
push_task.cancel()
await asyncio.sleep(0)
assert not push_task.done()
with pytest.raises(asyncio.CancelledError):
await asyncio.wait_for(push_task, timeout=1)
assert redis_cache.values[_SPEND_KEY] == 10.0
assert budget_limiter.redis_increment_operation_queue == []
@pytest.mark.asyncio
async def test_empty_flush_does_not_block_later_increment_sync() -> None:
redis_cache = _MockRedisCache(initial_values={_SPEND_KEY: 100.0})
in_memory_cache = _MockInMemoryCache(initial_values={_SPEND_KEY: 100.0})
budget_limiter = _new_router_budget_limiter(
redis_cache=redis_cache,
in_memory_cache=in_memory_cache,
provider_budget_config={"openai": BudgetConfig(time_period="1d", budget_limit=500.0)},
)
empty_flush_succeeded = await budget_limiter._push_in_memory_increments_to_redis()
await budget_limiter._increment_spend_in_current_window(spend_key=_SPEND_KEY, response_cost=20.0, ttl=86400)
await budget_limiter._sync_in_memory_spend_with_redis()
assert empty_flush_succeeded is True
assert budget_limiter.redis_increment_operation_queue == []
assert redis_cache.values[_SPEND_KEY] == 120.0
assert in_memory_cache.values[_SPEND_KEY] == 120.0
assert redis_cache.events == [
"increment_pipeline:start",
"increment_pipeline:done",
"batch_get",
]
@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)]
@pytest.mark.asyncio
@pytest.mark.parametrize("pause_during", ["write", "read"])
async def test_sync_preserves_spend_recorded_during_redis_io(pause_during: str) -> None:
io_started = asyncio.Event()
allow_io_to_complete = asyncio.Event()
redis_cache = _MockRedisCache(
initial_values={_SPEND_KEY: 100.0},
pipeline_started=io_started if pause_during == "write" else None,
allow_pipeline_to_complete=allow_io_to_complete if pause_during == "write" else None,
read_started=io_started if pause_during == "read" else None,
allow_read_to_complete=allow_io_to_complete if pause_during == "read" else None,
)
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=175.0)},
)
sync_task = asyncio.create_task(budget_limiter._sync_in_memory_spend_with_redis())
await asyncio.wait_for(io_started.wait(), timeout=1)
await budget_limiter._increment_spend_in_current_window(_SPEND_KEY, 20.0, 86400)
allow_io_to_complete.set()
await sync_task
assert in_memory_cache.values[_SPEND_KEY] == 180.0
assert redis_cache.values[_SPEND_KEY] == 160.0
assert budget_limiter.redis_increment_operation_queue == [_increment(20.0)]
await budget_limiter._sync_in_memory_spend_with_redis()
assert in_memory_cache.values[_SPEND_KEY] == 180.0
assert redis_cache.values[_SPEND_KEY] == 180.0
assert budget_limiter.redis_increment_operation_queue == []
@pytest.mark.asyncio
@pytest.mark.parametrize("cancellations", [1, 2])
async def test_cancelled_flush_does_not_requeue_an_applied_batch(cancellations: int) -> None:
pipeline_started = asyncio.Event()
pipeline_completed = asyncio.Event()
allow_pipeline = asyncio.Event()
redis_cache = _MockRedisCache(
initial_values={_SPEND_KEY: 0.0},
pipeline_started=pipeline_started,
pipeline_completed=pipeline_completed,
allow_pipeline_to_complete=allow_pipeline,
)
queue_lock = _ObservedLock()
limiter = _new_router_budget_limiter(
redis_cache=redis_cache, queue_lock=queue_lock, redis_increment_operation_queue=[_increment(10.0)]
)
push_task = asyncio.create_task(limiter._push_in_memory_increments_to_redis())
await asyncio.wait_for(pipeline_started.wait(), timeout=1)
async with limiter._redis_increment_queue_lock:
allow_pipeline.set()
await asyncio.wait_for(pipeline_completed.wait(), timeout=1)
await asyncio.wait_for(queue_lock.waiter_started.wait(), timeout=1)
for _ in range(cancellations):
push_task.cancel()
await asyncio.sleep(0)
assert not push_task.done()
with pytest.raises(asyncio.CancelledError):
await push_task
await limiter._push_in_memory_increments_to_redis()
assert redis_cache.values[_SPEND_KEY] == 10.0
assert limiter.redis_increment_operation_queue == []

18
tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py
View file

@ -19,6 +19,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
@ -487,6 +488,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,