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

Merge 53478e8651 into b6143b3711

Browse source
This commit is contained in:
Emerson Gomes 2026-09-16 15:56:38 -04:00 • committed by GitHub
commit 2c0489ef41
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 517 additions and 73 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
@ -268,6 +269,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(

192
litellm/router_strategy/budget_limiter.py
View file

@ -93,13 +93,6 @@ 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)
class RouterBudgetLimiting(CustomLogger):
def __init__(
self,
@ -109,6 +102,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 +388,84 @@ 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
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 _flush_queued_increment_operations(self, redis_cache: RedisCache) -> bool:
flush_task: Final = asyncio.create_task(self._write_queued_increment_operations(redis_cache))
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 +591,21 @@ 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))
except Exception as e:
log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e)
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 +624,53 @@ 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():
for key, value in redis_values.items():
if value is None:
continue
pending_spend = sum( # rebind-ok: each cache key has independent queued spend
operation["increment_value"]
for operation in self.redis_increment_operation_queue
if operation["key"] == key
)
updated_spend = float(value) + pending_spend # rebind-ok: each cache key has an independent total
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,

18
tests/proxy_unit_tests/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,

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

@ -0,0 +1,354 @@
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 _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)]
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_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._detached_increment_operations is None
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)]
assert budget_limiter._detached_increment_operations is None
@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 == []
assert limiter._detached_increment_operations is None

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

@ -338,7 +338,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,25 +355,27 @@ 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()
assert limiter._detached_increment_operations is not None
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()))
assert limiter._detached_increment_operations is None
@pytest.mark.asyncio