From 867a4df347a2632f8d93af52a9154f1679545ca7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 17:41:19 -0700 Subject: [PATCH] fix(proxy): claim a finished batch's per-model budget charge atomically A finished batch reports its whole cost on every poll. The charge-once marker is now taken with one atomic increment on the shared cache, so two workers polling the same batch at once cannot both charge it, and the marker's TTL is refreshed on every poll so a batch polled within every budget window is never charged again after the marker's first expiry. --- .../proxy/hooks/model_max_budget_limiter.py | 8 +- .../hooks/test_model_max_budget_limiter.py | 113 +++++++++++++++++- 2 files changed, 115 insertions(+), 6 deletions(-) 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)