From 2a771caf022c5ad6195190da2f67dbf7b833d481 Mon Sep 17 00:00:00 2001 From: hiraku-miyoshi Date: Wed, 19 Aug 2026 19:12:22 -0700 Subject: [PATCH] fix(proxy): clamp reservation record TTL so stale records never outlive their counters --- litellm/proxy/hooks/batch_enqueued_tokens.py | 34 ++++++++++++++----- .../proxy/hooks/test_batch_enqueued_tokens.py | 20 +++++++++++ 2 files changed, 46 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py index 82d4e7c66fc..32bccca2ab0 100644 --- a/litellm/proxy/hooks/batch_enqueued_tokens.py +++ b/litellm/proxy/hooks/batch_enqueued_tokens.py @@ -9,9 +9,11 @@ the reservation is refunded when the batch reaches a terminal state """ 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 @@ class BatchEnqueuedTokenStore: 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 @@ class BatchEnqueuedTokenStore: 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,6 +308,7 @@ class BatchEnqueuedTokenStore: 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): @@ -302,7 +316,9 @@ class BatchEnqueuedTokenStore: 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 @@ class BatchEnqueuedTokenStore: 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 @@ class BatchEnqueuedTokenStore: 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, ) diff --git a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py index 440b3050a39..edc921a40a3 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py +++ b/tests/test_litellm/proxy/hooks/test_batch_enqueued_tokens.py @@ -16,6 +16,7 @@ from typing import Final 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 @@ class _SingleKeyRedisFake: 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 @@ class _SingleKeyRedisFake: 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()