From 3c1b5749d2911530eab4e265e8ca9bd53d8cd215 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 30 Oct 2026 11:04:21 -0700 Subject: [PATCH] refactor(redis-store): enqueued-token store encodes and decodes the reservation record --- litellm/proxy/hooks/batch_enqueued_tokens.py | 54 +++++-------------- litellm/redis/stores/enqueued_tokens.py | 51 +++++++++++++++--- .../proxy/hooks/test_batch_enqueued_tokens.py | 24 ++++++--- 3 files changed, 77 insertions(+), 52 deletions(-) diff --git a/litellm/proxy/hooks/batch_enqueued_tokens.py b/litellm/proxy/hooks/batch_enqueued_tokens.py index ea6b5b8da99..19e19de7e6f 100644 --- a/litellm/proxy/hooks/batch_enqueued_tokens.py +++ b/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: diff --git a/litellm/redis/stores/enqueued_tokens.py b/litellm/redis/stores/enqueued_tokens.py index 447808a953a..b3ad41e15c8 100644 --- a/litellm/redis/stores/enqueued_tokens.py +++ b/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}") diff --git a/tests/unit/proxy/hooks/test_batch_enqueued_tokens.py b/tests/unit/proxy/hooks/test_batch_enqueued_tokens.py index 074d1679373..59158063d10 100644 --- a/tests/unit/proxy/hooks/test_batch_enqueued_tokens.py +++ b/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"