diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index cfa54ae01a2..7ce50bf5ead 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -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( diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 3e094df7ac8..4e84bded9de 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -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, diff --git a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py index 4cc8fe78811..a2c38a898e9 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py @@ -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 == [] diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index efe41e1da9a..f194e43c74a 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -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,