fix(proxy): clamp reservation record TTL so stale records never outlive their counters

Author: hiraku-miyoshi

litellm/proxy/hooks/batch_enqueued_tokens.py

@@ -9,9 +9,11 @@
 """
 
 import asyncio
+import math
+import time
 import uuid
-from collections.abc import Awaitable, Mapping, Sequence
-from dataclasses import dataclass
+from collections.abc import Awaitable, Callable, Mapping, Sequence
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias
 
 from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
@@ -85,6 +87,7 @@ class BatchEnqueuedTokenReservation:
     scopes: tuple[BatchEnqueuedTokenScope, ...]
     backend: ReservationBackend = "redis"
     owner: str = ""
+    reserved_at_monotonic: float = field(default_factory=time.monotonic, compare=False)
 
 
 @dataclass(frozen=True, slots=True)
@@ -180,11 +183,18 @@ class BatchEnqueuedTokenStore:
     granted them, and in-memory grants also remember the granting worker, so a
     refund never debits counters the grant did not charge. Everything expires after
     ``BATCH_ENQUEUED_TOKEN_TTL_SECONDS`` so a crash between submission and the
-    terminal-state refund can never leak tokens forever.
+    terminal-state refund can never leak tokens forever, and reservation records
+    expire no later than the counters they would refund, so a stale record can
+    never debit an allowance re-granted after its counters expired.
     """
 
-    def __init__(self, internal_usage_cache: "InternalUsageCache") -> None:
+    def __init__(
+        self,
+        internal_usage_cache: "InternalUsageCache",
+        monotonic: Callable[[], float] = time.monotonic,
+    ) -> None:
         self.internal_usage_cache = internal_usage_cache
+        self._monotonic: Final = monotonic
         self._lock = asyncio.Lock()
         self._owner_token = uuid.uuid4().hex
         redis_cache = internal_usage_cache.dual_cache.redis_cache
@@ -235,6 +245,7 @@ async def _reserve_via_redis(
         tokens: int,
         scopes: tuple[BatchEnqueuedTokenScope, ...],
     ) -> BatchEnqueuedTokenOutcome:
+        started: Final = self._monotonic()
         for index, scope in enumerate(scopes):
             result = await self._run_reserve_script(
                 reserve_script,
@@ -246,7 +257,9 @@ async def _reserve_via_redis(
             if result[0] != 1:
                 await self._rollback_partial_reserve(refund_script, tokens=tokens, scopes=scopes[:index])
                 return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=result[1])
-        return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes, backend="redis")
+        return BatchEnqueuedTokenReservation(
+            tokens=tokens, scopes=scopes, backend="redis", reserved_at_monotonic=started
+        )
 
     async def _run_reserve_script(
         self,
@@ -295,14 +308,17 @@ async def _reserve_in_memory(
         scopes: tuple[BatchEnqueuedTokenScope, ...],
         span: "Span | None",
     ) -> BatchEnqueuedTokenOutcome:
+        started: Final = self._monotonic()
         async with self._lock:
             currents: Final = tuple([await self._get_local_counter(scope, span) for scope in scopes])
             for scope, current in zip(scopes, currents):
                 if current + tokens > scope.limit:
                     return BatchEnqueuedTokenOverLimit(scope=scope, enqueued=current)
             for scope, current in zip(scopes, currents):
                 await self._set_local_counter(scope, current + tokens, span)
-        return BatchEnqueuedTokenReservation(tokens=tokens, scopes=scopes, backend="memory", owner=self._owner_token)
+        return BatchEnqueuedTokenReservation(
+            tokens=tokens, scopes=scopes, backend="memory", owner=self._owner_token, reserved_at_monotonic=started
+        )
 
     async def refund(
         self,
@@ -349,11 +365,13 @@ async def save_reservation(
         litellm_parent_otel_span: "Span | None" = None,
     ) -> None:
         serialized: Final = _RESERVATION_ADAPTER.dump_json(reservation).decode("utf-8")
+        elapsed: Final = self._monotonic() - reservation.reserved_at_monotonic
+        ttl: Final = max(1, BATCH_ENQUEUED_TOKEN_TTL_SECONDS - math.ceil(elapsed))
         if self._save_script is not None:
             try:
                 await self._save_script(
                     (self._record_key(batch_id),),
-                    (serialized, BATCH_ENQUEUED_TOKEN_TTL_SECONDS),
+                    (serialized, ttl),
                 )
             except Exception as e:  # noqa: BLE001  # any Redis failure must fall back to the in-memory record
                 verbose_proxy_logger.warning(
@@ -364,7 +382,7 @@ async def save_reservation(
         await self.internal_usage_cache.async_set_cache(
             key=self._record_key(batch_id),
             value=serialized,
-            ttl=BATCH_ENQUEUED_TOKEN_TTL_SECONDS,
+            ttl=ttl,
             litellm_parent_otel_span=litellm_parent_otel_span,
             local_only=True,
         )

tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py

@@ -16,6 +16,7 @@
 import pytest
 
 from litellm.caching.caching import DualCache
+from litellm.constants import BATCH_ENQUEUED_TOKEN_TTL_SECONDS
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.hooks.batch_enqueued_tokens import (
     BatchEnqueuedTokenOverLimit,
@@ -136,6 +137,7 @@ def __init__(
         raise_after_landing_save_keys: frozenset[str] = frozenset(),
     ) -> None:
         self.script_calls: tuple[tuple[str, tuple[str, ...]], ...] = ()
+        self.save_ttls: tuple[int, ...] = ()
         self.counters: Mapping[str, int] = MappingProxyType({})
         self.records: Mapping[str, str] = MappingProxyType({})
         self.fail_reserve_keys = fail_reserve_keys
@@ -180,6 +182,7 @@ def _run(self, kind: str, keys: tuple[str, ...], args: tuple[str | bytes | int |
             if keys[0] in self.fail_save_keys:
                 raise ConnectionError(f"simulated redis failure for {keys[0]}")
             self.records = MappingProxyType({**self.records, keys[0]: str(args[0])})
+            self.save_ttls = (*self.save_ttls, int(args[1]))
             if keys[0] in self.raise_after_landing_save_keys:
                 raise TimeoutError(f"simulated redis timeout after landing for {keys[0]}")
             return 1
@@ -303,6 +306,23 @@ async def test_local_ghost_left_by_landed_save_never_refunds_twice():
     assert fake.counters[f"batch_enqueued_tokens:{scope.key}:{scope.value}"] == 100
 
 
+@pytest.mark.asyncio
+async def test_record_ttl_shrinks_by_elapsed_time_so_stale_records_never_outlive_their_counters():
+    scope = _scope(limit=100)
+    fake = _SingleKeyRedisFake()
+    ticks = iter((1_000.0, 1_030.5))
+    store = BatchEnqueuedTokenStore(
+        internal_usage_cache=InternalUsageCache(DualCache(redis_cache=fake, default_in_memory_ttl=60)),
+        monotonic=lambda: next(ticks),
+    )
+    reservation = await store.reserve(tokens=60, scopes=(scope,))
+    assert isinstance(reservation, BatchEnqueuedTokenReservation)
+    assert reservation.reserved_at_monotonic == 1_000.0
+    await store.save_reservation("batch_ttl_clamp", reservation)
+    assert fake.save_ttls == (BATCH_ENQUEUED_TOKEN_TTL_SECONDS - 31,)
+    assert await store.pop_reservation("batch_ttl_clamp") == reservation
+
+
 @pytest.mark.asyncio
 async def test_memory_refund_skips_reservations_granted_by_another_worker():
     store = _in_memory_store()