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.

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

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)