diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index c5d10b2749b..4811d1bfeb2 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 @@ -260,6 +261,9 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): def __init__(self, dual_cache: DualCache): self.dual_cache = dual_cache self.redis_increment_operation_queue = [] + self._redis_increment_queue_lock = asyncio.Lock() + self._redis_increment_flush_lock = asyncio.Lock() + self._detached_increment_operations = None self.deployment_budget_config = None async def is_key_within_model_budget( diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index d57d7da0410..f747e99cec1 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -26,7 +26,7 @@ from typing import Any, Final import litellm from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisPipelineIncrementOperation +from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, @@ -100,6 +100,9 @@ class RouterBudgetLimiting(CustomLogger): ): self.dual_cache = dual_cache self.redis_increment_operation_queue: list[RedisPipelineIncrementOperation] = [] + self._redis_increment_queue_lock = asyncio.Lock() + self._redis_increment_flush_lock = asyncio.Lock() + self._detached_increment_operations: tuple[RedisPipelineIncrementOperation, ...] | None = None asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) self.provider_budget_config: GenericBudgetConfigType | None = provider_budget_config self.deployment_budget_config: GenericBudgetConfigType | None = None @@ -393,7 +396,78 @@ class RouterBudgetLimiting(CustomLogger): increment_value=response_cost, ttl=ttl, ) - self.redis_increment_operation_queue.append(increment_op) + async with self._get_redis_increment_queue_lock(): + self.redis_increment_operation_queue.append(increment_op) + + def _get_redis_increment_queue_lock(self) -> asyncio.Lock: + return self._redis_increment_queue_lock + + async def _detach_queued_increment_operations(self) -> tuple[RedisPipelineIncrementOperation, ...]: + async with self._get_redis_increment_queue_lock(): + if self._detached_increment_operations is not None: + return self._detached_increment_operations + increment_operations_to_flush: Final = tuple(self.redis_increment_operation_queue) + self.redis_increment_operation_queue = [] # mutable-ok: emptied queue must stay appendable + self._detached_increment_operations = increment_operations_to_flush + return increment_operations_to_flush + + async def _clear_detached_increment_operations(self) -> None: + async with self._get_redis_increment_queue_lock(): + self._detached_increment_operations = None + + async def _requeue_detached_increment_operations(self) -> None: + async with self._get_redis_increment_queue_lock(): + detached_increment_operations: Final = self._detached_increment_operations + if detached_increment_operations is None: + return + self.redis_increment_operation_queue = ( + list( # mutable-ok: restored flush batch must stay appendable + detached_increment_operations + ) + + self.redis_increment_operation_queue + ) + self._detached_increment_operations = None + + async def _finish_increment_pipeline_after_cancellation( + self, + pipeline_task: asyncio.Task[object], + ) -> None: + try: + await pipeline_task + except Exception: + await self._requeue_detached_increment_operations() + verbose_router_logger.exception("Error pushing queued Redis increment operations to Redis") + return + await self._clear_detached_increment_operations() + + async def _flush_queued_increment_operations(self, redis_cache: RedisCache) -> bool: + increment_operations_to_flush: Final = await self._detach_queued_increment_operations() + if len(increment_operations_to_flush) == 0: + return True + + verbose_router_logger.debug( + "Pushing Redis Increment Pipeline for queue: %s", + increment_operations_to_flush, + ) + increment_list: Final = list( # mutable-ok: Redis pipeline contract requires a list + increment_operations_to_flush + ) + pipeline_task: Final = asyncio.create_task( + redis_cache.async_increment_pipeline( + increment_list=increment_list, + ) + ) + try: + await asyncio.shield(pipeline_task) + except Exception: + await asyncio.shield(self._requeue_detached_increment_operations()) + verbose_router_logger.exception("Error pushing queued Redis increment operations to Redis") + return False + except asyncio.CancelledError: + await asyncio.shield(self._finish_increment_pipeline_after_cancellation(pipeline_task)) + raise + await self._clear_detached_increment_operations() + return True async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Original method now uses helper functions""" @@ -519,33 +593,25 @@ class RouterBudgetLimiting(CustomLogger): DEFAULT_REDIS_SYNC_INTERVAL ) # Still wait DEFAULT_REDIS_SYNC_INTERVAL seconds on error before retrying - async def _push_in_memory_increments_to_redis(self): + async def _push_in_memory_increments_to_redis(self) -> bool: """ How this works: - async_log_success_event collects all provider spend increments in `redis_increment_operation_queue` - This function pushes all increments to Redis in a batched pipeline to optimize performance - Only runs if Redis is initialized + Only runs if Redis is initialized. Returns False when the detached batch could not be + written, so callers must not treat Redis as up to date. """ - try: - if not self.dual_cache.redis_cache: - return # Redis is not initialized + redis_cache: Final = self.dual_cache.redis_cache + if redis_cache is None: + return True - verbose_router_logger.debug( - "Pushing Redis Increment Pipeline for queue: %s", - self.redis_increment_operation_queue, - ) - if len(self.redis_increment_operation_queue) > 0: - asyncio.create_task( - self.dual_cache.redis_cache.async_increment_pipeline( - increment_list=self.redis_increment_operation_queue, - ) - ) - - self.redis_increment_operation_queue = [] - - except Exception as e: - verbose_router_logger.error("Error syncing in-memory cache with Redis: %s", e) + async with self._redis_increment_flush_lock: + try: + return await self._flush_queued_increment_operations(redis_cache) + except asyncio.CancelledError: + await asyncio.shield(self._requeue_detached_increment_operations()) + raise async def _sync_in_memory_spend_with_redis(self): """ @@ -566,7 +632,9 @@ class RouterBudgetLimiting(CustomLogger): return # 1. Push all provider spend increments to Redis - await self._push_in_memory_increments_to_redis() + flush_succeeded: Final = await self._push_in_memory_increments_to_redis() + if not flush_succeeded: + return # 2. Fetch all current provider spend from Redis to update in-memory cache cache_keys: Final = [] diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index 3785ccdcfba..25152fdfc0b 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -16,6 +16,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( resolve_model_budget, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import BudgetConfig as GenericBudgetInfo @@ -484,6 +485,23 @@ async def test_async_log_success_event_pushes_redis_increments_when_redis_config mock_push.assert_awaited_once() +@pytest.mark.asyncio +async def test_model_budget_limiter_initializes_redis_increment_queue_lock(): + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + spend_key = "virtual_key_spend:test-key:gpt-4:1d" + + await limiter._increment_spend_in_current_window( + spend_key=spend_key, response_cost=0.01, ttl=86400 + ) + + assert limiter.redis_increment_operation_queue == [ + RedisPipelineIncrementOperation( + key=spend_key, increment_value=0.01, ttl=86400 + ) + ] + + @pytest.mark.asyncio async def test_get_fallback_model_within_budget_returns_none_without_fallbacks( budget_limiter, diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py new file mode 100644 index 00000000000..9ef677cff42 --- /dev/null +++ b/tests/test_litellm/router_strategy/test_budget_limiter.py @@ -0,0 +1,232 @@ +import asyncio +from types import SimpleNamespace + +import pytest + +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.utils import BudgetConfig + +_SPEND_KEY = "provider_spend:openai:1d" + + +def _increment(increment_value: float) -> RedisPipelineIncrementOperation: + return RedisPipelineIncrementOperation(key=_SPEND_KEY, increment_value=increment_value, ttl=86400) + + +class _MockRedisCache: + def __init__( + self, + initial_values: dict[str, float], + pipeline_started: asyncio.Event | None = None, + allow_pipeline_to_complete: asyncio.Event | None = None, + should_fail_pipeline: bool = False, + ) -> None: + self.values = initial_values + self.events: list[str] = [] + self.pipeline_started = pipeline_started + self.allow_pipeline_to_complete = allow_pipeline_to_complete + self.should_fail_pipeline = should_fail_pipeline + + async def async_increment_pipeline( + self, increment_list: list[RedisPipelineIncrementOperation], **kwargs: object + ) -> None: + self.events.append("increment_pipeline:start") + if self.pipeline_started is not None: + self.pipeline_started.set() + if self.allow_pipeline_to_complete is not None: + await self.allow_pipeline_to_complete.wait() + if self.should_fail_pipeline: + raise RuntimeError("redis down") + for op in increment_list: + key = op["key"] + current = float(self.values.get(key, 0.0) or 0.0) + self.values[key] = current + float(op["increment_value"]) + self.events.append("increment_pipeline:done") + + async def async_batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, float | None]: + self.events.append("batch_get") + return {key: self.values.get(key) for key in key_list} + + +class _MockInMemoryCache: + def __init__(self, initial_values: dict[str, float]) -> None: + self.values = initial_values + + async def async_increment(self, key: str, value: float, ttl: int, **kwargs: object) -> float: + current = float(self.values.get(key, 0.0) or 0.0) + self.values[key] = current + float(value) + return self.values[key] + + async def async_set_cache(self, key: str, value: float, **kwargs: object) -> None: + self.values[key] = float(value) + + +def _new_router_budget_limiter( + *, + redis_cache: object, + in_memory_cache: object | None = None, + redis_increment_operation_queue: list[RedisPipelineIncrementOperation] | None = None, + provider_budget_config: dict[str, BudgetConfig] | None = None, +) -> RouterBudgetLimiting: + budget_limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting) + budget_limiter.dual_cache = SimpleNamespace( + redis_cache=redis_cache, + in_memory_cache=in_memory_cache if in_memory_cache is not None else SimpleNamespace(), + ) + budget_limiter.provider_budget_config = provider_budget_config + budget_limiter.deployment_budget_config = None + budget_limiter.tag_budget_config = None + budget_limiter.redis_increment_operation_queue = ( + list(redis_increment_operation_queue) if redis_increment_operation_queue is not None else [] + ) + budget_limiter._redis_increment_queue_lock = asyncio.Lock() + budget_limiter._redis_increment_flush_lock = asyncio.Lock() + budget_limiter._detached_increment_operations = None + return budget_limiter + + +@pytest.mark.asyncio +async def test_should_await_redis_pipeline_before_sync_reads() -> None: + pipeline_started = asyncio.Event() + allow_pipeline_to_complete = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={_SPEND_KEY: 100.0}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + ) + in_memory_cache = _MockInMemoryCache(initial_values={_SPEND_KEY: 160.0}) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + in_memory_cache=in_memory_cache, + redis_increment_operation_queue=[_increment(60.0)], + provider_budget_config={"openai": BudgetConfig(time_period="1d", budget_limit=500.0)}, + ) + + sync_task = asyncio.create_task(budget_limiter._sync_in_memory_spend_with_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + assert "batch_get" not in redis_cache.events + allow_pipeline_to_complete.set() + await sync_task + + assert redis_cache.values[_SPEND_KEY] == 160.0 + assert in_memory_cache.values[_SPEND_KEY] == 160.0 + assert budget_limiter.redis_increment_operation_queue == [] + assert redis_cache.events == [ + "increment_pipeline:start", + "increment_pipeline:done", + "batch_get", + ] + + +@pytest.mark.asyncio +async def test_should_requeue_increments_when_redis_pipeline_fails() -> None: + redis_cache = _MockRedisCache(initial_values={}, should_fail_pipeline=True) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + redis_increment_operation_queue=[_increment(10.0)], + ) + + flush_succeeded = await budget_limiter._push_in_memory_increments_to_redis() + + assert flush_succeeded is False + assert budget_limiter.redis_increment_operation_queue == [_increment(10.0)] + assert budget_limiter._detached_increment_operations is None + + +@pytest.mark.asyncio +async def test_should_keep_new_increments_when_pipeline_flush_fails() -> None: + pipeline_started = asyncio.Event() + allow_pipeline_to_complete = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + should_fail_pipeline=True, + ) + in_memory_cache = _MockInMemoryCache(initial_values={_SPEND_KEY: 0.0}) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + in_memory_cache=in_memory_cache, + redis_increment_operation_queue=[_increment(10.0)], + ) + + push_task = asyncio.create_task(budget_limiter._push_in_memory_increments_to_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + await budget_limiter._increment_spend_in_current_window(spend_key=_SPEND_KEY, response_cost=20.0, ttl=86400) + allow_pipeline_to_complete.set() + await push_task + + assert budget_limiter.redis_increment_operation_queue == [_increment(10.0), _increment(20.0)] + + +@pytest.mark.asyncio +async def test_should_keep_in_memory_spend_when_redis_pipeline_fails() -> None: + redis_cache = _MockRedisCache(initial_values={_SPEND_KEY: 100.0}, should_fail_pipeline=True) + in_memory_cache = _MockInMemoryCache(initial_values={_SPEND_KEY: 160.0}) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + in_memory_cache=in_memory_cache, + redis_increment_operation_queue=[_increment(60.0)], + provider_budget_config={"openai": BudgetConfig(time_period="1d", budget_limit=500.0)}, + ) + + await budget_limiter._sync_in_memory_spend_with_redis() + + assert in_memory_cache.values[_SPEND_KEY] == 160.0 + assert redis_cache.values[_SPEND_KEY] == 100.0 + assert budget_limiter.redis_increment_operation_queue == [_increment(60.0)] + assert "batch_get" not in redis_cache.events + + +@pytest.mark.asyncio +async def test_should_keep_increments_when_flush_is_cancelled_after_success() -> None: + pipeline_started = asyncio.Event() + allow_pipeline_to_complete = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={_SPEND_KEY: 0.0}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + ) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + redis_increment_operation_queue=[_increment(10.0)], + ) + + push_task = asyncio.create_task(budget_limiter._push_in_memory_increments_to_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + push_task.cancel() + allow_pipeline_to_complete.set() + with pytest.raises(asyncio.CancelledError): + await push_task + + assert redis_cache.values[_SPEND_KEY] == 10.0 + assert budget_limiter.redis_increment_operation_queue == [] + assert budget_limiter._detached_increment_operations is None + + +@pytest.mark.asyncio +async def test_should_requeue_increments_when_flush_is_cancelled_and_redis_fails() -> None: + pipeline_started = asyncio.Event() + allow_pipeline_to_complete = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={_SPEND_KEY: 0.0}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + should_fail_pipeline=True, + ) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + redis_increment_operation_queue=[_increment(10.0)], + ) + + push_task = asyncio.create_task(budget_limiter._push_in_memory_increments_to_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + push_task.cancel() + allow_pipeline_to_complete.set() + with pytest.raises(asyncio.CancelledError): + await push_task + + assert redis_cache.values[_SPEND_KEY] == 0.0 + assert budget_limiter.redis_increment_operation_queue == [_increment(10.0)] + assert budget_limiter._detached_increment_operations is None