diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 67577a68f5c..cfa54ae01a2 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -610,7 +610,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): async def _claim_batch_charge(self, spend_key: str, batch_id: str, ttl_seconds: int) -> bool: marker_key: Final = batch_charged_once_marker_key(spend_key=spend_key, batch_id=batch_id) - if await self.dual_cache.async_get_cache(key=marker_key) is not None: - return False - await self.dual_cache.async_set_cache(key=marker_key, value=1, ttl=ttl_seconds) - return True + polls: Final = await self.dual_cache.async_increment_cache( + key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True + ) + return polls == 1 diff --git a/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py index 47b14438212..ffb60fb4651 100644 --- a/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py @@ -1,3 +1,7 @@ +import asyncio +import time +from collections.abc import Callable, Mapping +from types import MappingProxyType from typing import Final import pytest @@ -6,6 +10,7 @@ from litellm.caching.caching import DualCache from litellm.proxy.hooks.model_max_budget_limiter import ( _PROXY_VirtualKeyModelMaxBudgetLimiter, ) +from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import LiteLLMBatch, Usage KEY_HASH: Final = "key-hash-batch" @@ -30,7 +35,7 @@ def _batch(batch_id: str, status: str) -> LiteLLMBatch: ) -def _event(call_type: str, response_cost: float) -> dict: +def _event(call_type: str, response_cost: float) -> dict[str, object]: return { "call_type": call_type, "standard_logging_object": { @@ -56,13 +61,88 @@ async def _poll(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, batch: LiteLLMB async def _chat(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter) -> None: - await limiter.async_log_success_event(_event("acompletion", CHAT_COST), response_obj=None, start_time=None, end_time=None) + await limiter.async_log_success_event( + _event("acompletion", CHAT_COST), response_obj=None, start_time=None, end_time=None + ) async def _spend(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, spend_key: str) -> float: return await limiter.dual_cache.async_get_cache(key=spend_key) or 0.0 +def _local_spend(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, spend_key: str) -> float: + return limiter.dual_cache.in_memory_cache.get_cache(key=spend_key) or 0.0 + + +class _Clock: + def __init__(self) -> None: + self.seconds = 0.0 + + def now(self) -> float: + return self.seconds + + def advance(self, seconds: float) -> None: + self.seconds = self.seconds + seconds + + +class _SharedRedisDouble: + def __init__(self, now: Callable[[], float] = time.time) -> None: + self.now = now + self.entries: Mapping[str, tuple[float, float | None]] = MappingProxyType({}) + + def _live(self, key: str) -> tuple[float, float | None] | None: + entry: Final = self.entries.get(key) + if entry is None: + return None + expires_at: Final = entry[1] + if expires_at is not None and expires_at <= self.now(): + return None + return entry + + def _store(self, key: str, value: float, expires_at: float | None) -> None: + self.entries = MappingProxyType({**self.entries, key: (value, expires_at)}) + + async def async_get_cache(self, key: str, **kwargs: object) -> float | None: + await asyncio.sleep(0) + entry: Final = self._live(key) + return None if entry is None else entry[0] + + async def async_set_cache(self, key: str, value: float, ttl: int | None = None, **kwargs: object) -> None: + await asyncio.sleep(0) + self._store(key, value, None if ttl is None else self.now() + ttl) + + async def async_increment( + self, + key: str, + value: float, + ttl: int | None = None, + parent_otel_span: object = None, + refresh_ttl: bool = False, + ) -> float: + await asyncio.sleep(0) + live: Final = self._live(key) + total: Final = value if live is None else live[0] + value + kept_expiry: Final = None if live is None else live[1] + expires_at: Final = ( + kept_expiry if ttl is None or (kept_expiry is not None and not refresh_ttl) else self.now() + ttl + ) + self._store(key, total, expires_at) + return total + + async def async_increment_pipeline(self, increment_list: list[RedisPipelineIncrementOperation]) -> list[float]: + return [await self.async_increment(op["key"], op["increment_value"], ttl=op["ttl"]) for op in increment_list] + + +def _worker(redis: _SharedRedisDouble) -> _PROXY_VirtualKeyModelMaxBudgetLimiter: + return _PROXY_VirtualKeyModelMaxBudgetLimiter( + dual_cache=DualCache(redis_cache=redis) # pyright: ignore[reportArgumentType] # duck-typed Redis double + ) + + +async def _drain_redis_pushes() -> None: + await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) + + @pytest.mark.asyncio async def test_polls_of_a_finished_batch_charge_each_per_model_budget_once(): limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) @@ -87,3 +167,32 @@ async def test_a_second_batch_and_chat_requests_still_charge_the_budget(): await _chat(limiter) assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(2 * BATCH_COST + 2 * CHAT_COST) + + +@pytest.mark.asyncio +async def test_two_workers_polling_the_same_finished_batch_at_once_charge_it_once(): + redis: Final = _SharedRedisDouble() + worker_a: Final = _worker(redis) + worker_b: Final = _worker(redis) + finished: Final = _batch("batch_first", "completed") + + await asyncio.gather(_poll(worker_a, finished, BATCH_COST), _poll(worker_b, finished, BATCH_COST)) + await _drain_redis_pushes() + + assert _local_spend(worker_a, KEY_SPEND_KEY) + _local_spend(worker_b, KEY_SPEND_KEY) == pytest.approx(BATCH_COST) + assert await redis.async_get_cache(KEY_SPEND_KEY) == pytest.approx(BATCH_COST) + + +@pytest.mark.asyncio +async def test_a_batch_polled_within_every_budget_window_is_never_charged_again(): + clock: Final = _Clock() + limiter: Final = _worker(_SharedRedisDouble(now=clock.now)) + finished: Final = _batch("batch_first", "completed") + + await _poll(limiter, finished, BATCH_COST) + clock.advance(12 * 3600) + await _poll(limiter, finished, BATCH_COST) + clock.advance(18 * 3600) + await _poll(limiter, finished, BATCH_COST) + + assert _local_spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST)