refactor(redis-store): enqueued-token store encodes and decodes the reservation record

litellm/proxy/hooks/batch_enqueued_tokens.py

@@ -13,7 +13,7 @@ import math
 import time
 import uuid
 from collections.abc import Callable, Mapping
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from datetime import timedelta
 from typing import TYPE_CHECKING, Annotated, Final, Literal, TypeAlias
 
@@ -25,6 +25,8 @@ from litellm.constants import BATCH_ENQUEUED_TOKEN_LIMIT_METADATA_KEY, BATCH_ENQ
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.redis.stores.enqueued_tokens import (
     AlreadyPopped,
+    BatchEnqueuedTokenReservation,
+    BatchEnqueuedTokenScope,
     EnqueuedTokensStore,
     OverLimit,
     Popped,
@@ -48,27 +50,6 @@ BATCH_ENQUEUED_REFUND_STATUSES: Final[frozenset[str]] = frozenset(
     {"completed", "complete", "failed", "expired", "cancelled", "cancelling"}
 )
 
-ScopeKey: TypeAlias = Literal["api_key", "team"]
-
-
-@dataclass(frozen=True, slots=True)
-class BatchEnqueuedTokenScope:
-    key: ScopeKey
-    value: str
-    limit: int
-
-
-ReservationBackend: TypeAlias = Literal["redis", "memory"]
-
-
-@dataclass(frozen=True, slots=True)
-class BatchEnqueuedTokenReservation:
-    tokens: int
-    scopes: tuple[BatchEnqueuedTokenScope, ...]
-    backend: ReservationBackend = "redis"
-    owner: str = ""
-    reserved_at_monotonic: float = field(default_factory=time.monotonic, compare=False)
-
 
 @dataclass(frozen=True, slots=True)
 class BatchEnqueuedTokenOverLimit:
@@ -80,7 +61,6 @@ BatchEnqueuedTokenOutcome: TypeAlias = BatchEnqueuedTokenReservation | BatchEnqu
 
 _LIMIT_ADAPTER: Final = TypeAdapter(Annotated[int, Field(gt=0)])
 _STORED_COUNTER_ADAPTER: Final = TypeAdapter(int | None)
-_RESERVATION_ADAPTER: Final = TypeAdapter(BatchEnqueuedTokenReservation)
 _LOCAL_RECORD_ADAPTER: Final = TypeAdapter(object)
 
 
@@ -309,16 +289,15 @@ class BatchEnqueuedTokenStore:
         reservation: BatchEnqueuedTokenReservation,
         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._store is not None:
-            saved: Final = await self._store.save_record(batch_id, serialized, ttl=timedelta(seconds=ttl))
+            saved: Final = await self._store.save_record(batch_id, reservation, ttl=timedelta(seconds=ttl))
             if saved is None:
                 return
         await self.internal_usage_cache.async_set_cache(
             key=self._record_key(batch_id),
-            value=serialized,
+            value=reservation,
             ttl=ttl,
             litellm_parent_otel_span=litellm_parent_otel_span,
             local_only=True,
@@ -337,28 +316,23 @@ class BatchEnqueuedTokenStore:
             # left behind by a save that raised after landing must not refund again.
             await self._pop_local_record(batch_id, litellm_parent_otel_span)
             return None
-        raw: Final = (
-            popped.document
-            if isinstance(popped, Popped)
-            else await self._pop_local_record(batch_id, litellm_parent_otel_span)
-        )
-        if raw is None:
-            return None
-        try:
-            if isinstance(raw, (str, bytes)):
-                return _RESERVATION_ADAPTER.validate_json(raw)
-            return _RESERVATION_ADAPTER.validate_python(raw)
-        except ValidationError:
+        if isinstance(popped, Popped):
+            return popped.reservation
+        if isinstance(popped, RedisFailure):
             verbose_proxy_logger.warning("Discarding malformed enqueued-token reservation record for %s", batch_id)
             return None
+        local: Final = await self._pop_local_record(batch_id, litellm_parent_otel_span)
+        return local if isinstance(local, BatchEnqueuedTokenReservation) else None
 
-    async def _pop_redis_record(self, batch_id: str) -> Popped | AlreadyPopped | None:
+    async def _pop_redis_record(self, batch_id: str) -> Popped | AlreadyPopped | RedisFailure | None:
         if self._store is None:
             return None
         popped: Final = await self._store.pop_record(
             batch_id, tombstone_ttl=timedelta(seconds=BATCH_ENQUEUED_TOKEN_TTL_SECONDS)
         )
-        return None if isinstance(popped, RedisFailure) else popped
+        if isinstance(popped, RedisFailure) and popped.kind != "decode":
+            return None
+        return popped
 
     async def _pop_local_record(self, batch_id: str, span: "Span | None") -> object:
         async with self._lock:

litellm/redis/stores/enqueued_tokens.py

@@ -1,9 +1,34 @@
-from dataclasses import dataclass
+import time
+from dataclasses import dataclass, field
 from datetime import timedelta
-from typing import Final, Protocol
+from typing import Final, Literal, Protocol, TypeAlias
+
+from pydantic import TypeAdapter, ValidationError
 
 from litellm.redis.transport.contract import EvalReply, RedisFailure, RedisTransport
 
+ScopeKey: TypeAlias = Literal["api_key", "team"]
+
+
+@dataclass(frozen=True, slots=True)
+class BatchEnqueuedTokenScope:
+    key: ScopeKey
+    value: str
+    limit: int
+
+
+ReservationBackend: TypeAlias = Literal["redis", "memory"]
+
+
+@dataclass(frozen=True, slots=True)
+class BatchEnqueuedTokenReservation:
+    tokens: int
+    scopes: tuple[BatchEnqueuedTokenScope, ...]
+    backend: ReservationBackend = "redis"
+    owner: str = ""
+    reserved_at_monotonic: float = field(default_factory=time.monotonic, compare=False)
+
+
 _RESERVE: Final = """
 local amount = tonumber(ARGV[1])
 local ttl = tonumber(ARGV[2])
@@ -25,6 +50,8 @@ end
 return 1
 """
 
+_RESERVATION: Final = TypeAdapter(BatchEnqueuedTokenReservation)
+
 _SAVE_RECORD: Final = """
 redis.call('SET', KEYS[1], ARGV[1], 'EX', tonumber(ARGV[2]))
 return 1
@@ -59,7 +86,7 @@ class OverLimit:
 
 @dataclass(frozen=True, slots=True)
 class Popped:
-    document: bytes
+    reservation: BatchEnqueuedTokenReservation
 
 
 @dataclass(frozen=True, slots=True)
@@ -74,7 +101,9 @@ class EnqueuedTokensStore(Protocol):
 
     async def refund(self, scope_key: str, scope_value: str, tokens: int) -> None | RedisFailure: ...
 
-    async def save_record(self, batch_id: str, document: str, *, ttl: timedelta) -> None | RedisFailure: ...
+    async def save_record(
+        self, batch_id: str, reservation: BatchEnqueuedTokenReservation, *, ttl: timedelta
+    ) -> None | RedisFailure: ...
 
     async def pop_record(
         self, batch_id: str, *, tombstone_ttl: timedelta
@@ -105,7 +134,10 @@ class TransportEnqueuedTokensStore:
         reply: Final = await self.transport.eval(_REFUND, (counter_key(scope_key, scope_value),), (tokens,))
         return reply if isinstance(reply, RedisFailure) else None
 
-    async def save_record(self, batch_id: str, document: str, *, ttl: timedelta) -> None | RedisFailure:
+    async def save_record(
+        self, batch_id: str, reservation: BatchEnqueuedTokenReservation, *, ttl: timedelta
+    ) -> None | RedisFailure:
+        document: Final = _RESERVATION.dump_json(reservation).decode("utf-8")
         reply: Final = await self.transport.eval(_SAVE_RECORD, (record_key(batch_id),), (document, _seconds(ttl)))
         return reply if isinstance(reply, RedisFailure) else None
 
@@ -119,7 +151,7 @@ class TransportEnqueuedTokensStore:
             case b"":
                 return AlreadyPopped()
             case bytes():
-                return Popped(reply)
+                return _decode_record(reply)
             case _:
                 return _decode_failure("enqueued-token reservation record", reply)
 
@@ -130,3 +162,10 @@ def _seconds(ttl: timedelta) -> int:
 
 def _decode_failure(what: str, reply: EvalReply) -> RedisFailure:
     return RedisFailure(kind="decode", message=f"unexpected {what} reply: {reply!r}")
+
+
+def _decode_record(document: bytes) -> Popped | RedisFailure:
+    try:
+        return Popped(_RESERVATION.validate_json(document))
+    except ValidationError as e:
+        return RedisFailure(kind="decode", message=f"enqueued-token reservation record does not decode: {e}")

tests/unit/proxy/hooks/test_batch_enqueued_tokens.py

@@ -17,6 +17,7 @@ from typing import Final
 
 import fakeredis
 import pytest
+from pydantic import TypeAdapter
 from redis.exceptions import ConnectionError as RedisConnectionError
 from redis.exceptions import NoScriptError
 from redis.exceptions import TimeoutError as RedisTimeoutError
@@ -26,7 +27,6 @@ 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 (
-    _RESERVATION_ADAPTER,
     BatchEnqueuedTokenOverLimit,
     BatchEnqueuedTokenReservation,
     BatchEnqueuedTokenScope,
@@ -38,6 +38,8 @@ from litellm.proxy.hooks.batch_enqueued_tokens import (
 from litellm.proxy.utils import InternalUsageCache
 from tests._support.fake_redis import fake_redis_cache
 
+_RESERVATION_ADAPTER: Final = TypeAdapter(BatchEnqueuedTokenReservation)
+
 
 def _in_memory_store() -> BatchEnqueuedTokenStore:
     return BatchEnqueuedTokenStore(internal_usage_cache=InternalUsageCache(DualCache(default_in_memory_ttl=60)))
@@ -235,7 +237,19 @@ async def test_a_reservation_record_round_trips_with_old_pods():
     assert await raw.get(_record("batch_old")) == b""
 
     await store.save_reservation("batch_new", reservation)
-    assert _RESERVATION_ADAPTER.validate_json(await raw.get(_record("batch_new"))) == reservation
+    assert await raw.get(_record("batch_new")) == _RESERVATION_ADAPTER.dump_json(reservation)
+
+
+@pytest.mark.asyncio
+async def test_a_reservation_record_that_does_not_decode_is_discarded(caplog):
+    store, raw = _redis_store(fakeredis.FakeServer())
+    await raw.set(_record("batch_bad"), "not a reservation", ex=60)
+
+    with caplog.at_level(logging.WARNING):
+        assert await store.pop_reservation("batch_bad") is None
+
+    assert any("Discarding malformed" in record.getMessage() for record in caplog.records)
+    assert await raw.get(_record("batch_bad")) == b""
 
 
 @pytest.mark.asyncio
@@ -361,11 +375,9 @@ async def test_failed_redis_refund_leaves_local_counters_untouched():
 
 @pytest.mark.asyncio
 async def test_pop_reservation_defaults_legacy_records_to_redis_backend():
-    store = _in_memory_store()
+    store, raw = _redis_store(fakeredis.FakeServer())
     legacy = '{"tokens": 5, "scopes": [{"key": "api_key", "value": "k", "limit": 10}]}'
-    store.internal_usage_cache.dual_cache.in_memory_cache.set_cache(
-        key="batch_enqueued_token_reservation:batch_legacy", value=legacy
-    )
+    await raw.set(_record("batch_legacy"), legacy, ex=60)
     popped = await store.pop_reservation("batch_legacy")
     assert popped == BatchEnqueuedTokenReservation(
         tokens=5, scopes=(BatchEnqueuedTokenScope(key="api_key", value="k", limit=10),), backend="redis"