refactor(redis): move parked spend logs behind a typed store

This commit is contained in:
Devin AI 2026-10-30 22:45:15 +00:00
parent f969987849
commit df3cdb255f
6 changed files with 170 additions and 123 deletions

View file

@ -7,7 +7,6 @@ This is to prevent deadlocks and improve reliability
import asyncio
import json
from collections.abc import Mapping, Sequence
from datetime import datetime
from functools import partial, reduce
from itertools import chain
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast
@ -20,7 +19,6 @@ from litellm._logging import verbose_proxy_logger
from litellm.caching import RedisCache
from litellm.constants import (
MAX_REDIS_BUFFER_DEQUEUE_COUNT,
REDIS_SPEND_LOGS_BUFFER_KEY,
REDIS_SPEND_LOGS_BUFFER_MAX_ROWS,
REDIS_TOOL_DISCOVERY_BUFFER_KEY,
REDIS_UPDATE_BUFFER_KEY,
@ -54,6 +52,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
to_wire_payload,
)
from litellm.proxy.db.spend_log_batching import SpendLogRow
from litellm.redis.stores.spend_logs_parking import SpendLogsParkingStore, TransportSpendLogsParkingStore
from litellm.redis.stores.spend_update_buffer import (
SpendBuffer,
SpendUpdateBufferStore,
@ -122,6 +121,10 @@ def _spend_buffers(redis_cache: RedisCache) -> SpendUpdateBufferStore:
return TransportSpendUpdateBufferStore(redis_cache.transport())
def _parked_spend_logs(redis_cache: RedisCache) -> SpendLogsParkingStore:
return TransportSpendLogsParkingStore(redis_cache.transport())
def _zero_for_non_finite_constant(constant: str) -> float:
verbose_proxy_logger.warning(
"Redis spend buffer - reading a parked %s spend as 0, the spend it held is lost", constant
@ -132,14 +135,6 @@ def _zero_for_non_finite_constant(constant: str) -> float:
_load_buffered_transaction: Final = partial(json.loads, parse_constant=_zero_for_non_finite_constant)
def _spend_log_json_default(value: object) -> str:
return value.isoformat() if isinstance(value, datetime) else str(value)
def _encode_spend_log_row(row: SpendLogRow) -> str:
return json.dumps(row, default=_spend_log_json_default)
def _decode_spend_log_row(encoded: str) -> dict[str, object] | None:
decoded: Final = _DECODED_JSON.validate_python(_load_buffered_transaction(encoded))
return decoded if isinstance(decoded, dict) else None
@ -631,24 +626,19 @@ class RedisUpdateBuffer:
"""Park spend-log rows in Redis so they outlive this pod, dropping the oldest past ``max_rows``."""
if self.redis_cache is None or len(rows) == 0 or not self.should_commit_spend_updates_to_redis():
return False
try:
buffer_size: Final = await self.redis_cache.async_rpush_and_trim(
key=REDIS_SPEND_LOGS_BUFFER_KEY,
values=tuple(_encode_spend_log_row(row) for row in rows),
max_len=max_rows,
)
overflow: Final = buffer_size - max_rows
if overflow > 0:
verbose_proxy_logger.error(
"Spend tracking - Redis spend log buffer is at its %d row cap; dropped the %d oldest spend logs",
max_rows,
overflow,
)
except Exception as e: # noqa: BLE001 # the caller falls back to the in-memory queue on any Redis fault
buffer_size: Final = await _parked_spend_logs(self.redis_cache).park(rows, max_rows=max_rows)
if isinstance(buffer_size, RedisFailure):
verbose_proxy_logger.error(
"Spend tracking - failed to park %d spend log rows in Redis. Error: %s", len(rows), str(e)
"Spend tracking - failed to park %d spend log rows in Redis. Error: %s", len(rows), buffer_size.message
)
return False
overflow: Final = buffer_size - max_rows
if overflow > 0:
verbose_proxy_logger.error(
"Spend tracking - Redis spend log buffer is at its %d row cap; dropped the %d oldest spend logs",
max_rows,
overflow,
)
verbose_proxy_logger.info("Spend tracking - parked %d spend log rows in Redis for a later flush", len(rows))
return True
@ -657,14 +647,12 @@ class RedisUpdateBuffer:
"""Atomically take up to ``limit`` parked spend-log rows out of Redis."""
if self.redis_cache is None or not self.should_commit_spend_updates_to_redis():
return ()
popped: Final[str | list[str] | None] = await self.redis_cache.async_lpop(
key=REDIS_SPEND_LOGS_BUFFER_KEY,
count=limit,
)
popped: Final = await _parked_spend_logs(self.redis_cache).take(limit)
if isinstance(popped, RedisFailure):
raise_redis_failure(popped)
if popped is None:
return ()
encoded_rows: Final = tuple(popped) if isinstance(popped, list) else (popped,)
decoded_rows: Final = (_decode_spend_log_row(encoded) for encoded in encoded_rows)
decoded_rows: Final = (_decode_spend_log_row(encoded) for encoded in popped)
return tuple(row for row in decoded_rows if row is not None)
@staticmethod

View file

@ -0,0 +1,45 @@
import json
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime
from typing import Final, Protocol
from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY
from litellm.redis.transport.contract import RedisFailure, RedisTransport
def _json_default(value: object) -> str:
return value.isoformat() if isinstance(value, datetime) else str(value)
def encode_spend_log_row(row: Mapping[str, object]) -> str:
return json.dumps(row, default=_json_default)
class SpendLogsParkingStore(Protocol):
async def park(self, rows: Sequence[Mapping[str, object]], *, max_rows: int) -> int | RedisFailure: ...
async def take(self, limit: int) -> tuple[str, ...] | None | RedisFailure: ...
@dataclass(frozen=True, slots=True)
class TransportSpendLogsParkingStore:
transport: RedisTransport
async def park(self, rows: Sequence[Mapping[str, object]], *, max_rows: int) -> int | RedisFailure:
batch: Final = self.transport.batch(transaction=True, name="spend_logs_parking")
pushed: Final = batch.rpush(REDIS_SPEND_LOGS_BUFFER_KEY, *(encode_spend_log_row(row) for row in rows))
trimmed: Final = batch.ltrim(REDIS_SPEND_LOGS_BUFFER_KEY, -max_rows, -1)
length: Final = await pushed
trim: Final = await trimmed
return trim if isinstance(trim, RedisFailure) and not isinstance(length, RedisFailure) else length
async def take(self, limit: int) -> tuple[str, ...] | None | RedisFailure:
popped: Final = await self.transport.lpop(REDIS_SPEND_LOGS_BUFFER_KEY, limit)
if popped is None or isinstance(popped, RedisFailure):
return popped
entries: Final = (popped,) if isinstance(popped, bytes) else popped
try:
return tuple(entry.decode("utf-8") for entry in entries)
except UnicodeDecodeError as error:
return RedisFailure(kind="decode", message=f"Invalid parked spend log row: {error.reason}")

View file

@ -804,32 +804,9 @@ async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpu
assert [payload["entity_id"] for payload in restored] == ["team-1"]
class _ListRedis:
def __init__(self) -> None:
self.rows: list[str] = []
async def async_rpush(self, key: str, values: list[str]) -> int:
self.rows.extend(values)
return len(self.rows)
async def async_rpush_and_trim(self, key: str, values: list[str], max_len: int) -> int:
self.rows.extend(values)
pushed_len = len(self.rows)
del self.rows[:-max_len]
return pushed_len
async def async_lpop(self, key: str, count: int | None = None, **kwargs: object) -> list[str] | None:
if not self.rows:
return None
popped = self.rows[:count]
del self.rows[:count]
return popped
@pytest.mark.asyncio
async def test_store_spend_logs_in_redis_drops_oldest_rows_past_the_cap():
redis = _ListRedis()
buffer = RedisUpdateBuffer(redis_cache=redis)
buffer = RedisUpdateBuffer(redis_cache=fake_redis_cache(fakeredis.FakeServer()))
buffer.should_commit_spend_updates_to_redis = MagicMock(return_value=True)
assert await buffer.store_spend_logs_in_redis([{"request_id": "old"}, {"request_id": "mid"}], max_rows=2) is True
@ -858,7 +835,7 @@ async def test_a_non_finite_daily_spend_buffered_in_redis_drains_as_zero(spend:
@pytest.mark.asyncio
async def test_a_spend_log_row_parked_with_a_nan_spend_is_read_back_with_zero_spend():
buffer: Final = RedisUpdateBuffer(redis_cache=_ListRedis())
buffer: Final = RedisUpdateBuffer(redis_cache=fake_redis_cache(fakeredis.FakeServer()))
buffer.should_commit_spend_updates_to_redis = MagicMock(return_value=True)
assert await buffer.store_spend_logs_in_redis([{"request_id": "nan-row", "spend": float("nan")}]) is True
@ -876,12 +853,12 @@ async def test_store_spend_logs_in_redis_reports_failure_without_redis():
@pytest.mark.asyncio
async def test_store_spend_logs_in_redis_is_off_unless_transaction_buffering_is_enabled():
redis = _ListRedis()
buffer = RedisUpdateBuffer(redis_cache=redis)
server: Final = fakeredis.FakeServer()
buffer = RedisUpdateBuffer(redis_cache=fake_redis_cache(server))
buffer.should_commit_spend_updates_to_redis = MagicMock(return_value=False)
assert await buffer.store_spend_logs_in_redis([{"request_id": "a"}]) is False
assert redis.rows == []
assert await fakeredis.FakeAsyncRedis(server=server).keys("*") == []
@pytest.mark.asyncio

View file

@ -13,15 +13,23 @@ needing a generated Prisma client or a real database.
from __future__ import annotations
import asyncio
import json
import sys
import threading
from collections.abc import Callable, Iterator
from dataclasses import dataclass, field
from email.message import EmailMessage
from pathlib import Path
from typing import Any, Callable, Dict, Iterator, List, Optional
from typing import Any
from unittest.mock import AsyncMock, MagicMock
import fakeredis
import pytest
import pytest_asyncio
from litellm.caching.redis_cache import RedisCache
from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY
from tests._support.fake_redis import fake_redis_cache
sys.path.insert(0, str(Path(__file__).resolve().parents[5]))
@ -51,16 +59,13 @@ VOLATILE_KEYS = frozenset(
def normalize(data: Any, volatile: frozenset = VOLATILE_KEYS) -> Any:
"""Recursively replace values for volatile keys with '<VOLATILE>'."""
if isinstance(data, dict):
return {
k: ("<VOLATILE>" if k in volatile else normalize(v, volatile))
for k, v in data.items()
}
return {k: ("<VOLATILE>" if k in volatile else normalize(v, volatile)) for k, v in data.items()}
if isinstance(data, list):
return [normalize(v, volatile) for v in data]
return data
_PRISMA_TABLES: List[str] = [
_PRISMA_TABLES: list[str] = [
"litellm_verificationtoken",
"litellm_teamtable",
"litellm_usertable",
@ -156,7 +161,7 @@ def mock_dual_cache() -> MagicMock:
``cache._store`` directly in a test for assertion convenience.
"""
cache = MagicMock(name="MockDualCache")
cache._store: Dict[str, Any] = {}
cache._store: dict[str, Any] = {}
def _sync_get(key: str, **_: Any) -> Any:
return cache._store.get(key)
@ -191,6 +196,7 @@ def patched_prisma_import(monkeypatch: pytest.MonkeyPatch) -> Iterator[MagicMock
directly and restore in teardown.
"""
import prisma as _prisma_pkg
import litellm.proxy.utils as _utils_mod
fake_prisma = MagicMock(name="FakePrisma")
@ -261,7 +267,7 @@ class FakeClock:
"""
now: float = 0.0
sleep_calls: List[float] = field(default_factory=list)
sleep_calls: list[float] = field(default_factory=list)
def advance(self, seconds: float) -> None:
self.now += seconds
@ -290,7 +296,7 @@ def fake_clock(monkeypatch: pytest.MonkeyPatch) -> FakeClock:
@pytest.fixture
def make_spend_log_row() -> Callable[..., Dict[str, Any]]:
def make_spend_log_row() -> Callable[..., dict[str, Any]]:
"""Factory for fake LiteLLM_SpendLogs rows."""
def _make(
@ -298,7 +304,7 @@ def make_spend_log_row() -> Callable[..., Dict[str, Any]]:
spend: float = 0.01,
model: str = "gpt-4o-mini",
**overrides: Any,
) -> Dict[str, Any]:
) -> dict[str, Any]:
row = {
"request_id": request_id,
"spend": spend,
@ -316,38 +322,21 @@ def make_spend_log_row() -> Callable[..., Dict[str, Any]]:
return _make
@dataclass(frozen=True)
class FakeRedisList:
def __init__(self) -> None:
self.items: dict[str, list[str]] = {}
self.down = False
server: fakeredis.FakeServer
cache: RedisCache
raw: fakeredis.FakeAsyncRedis
def _check_up(self) -> None:
if self.down:
raise ConnectionError("redis unreachable")
async def async_rpush_and_trim(self, key: str, values: list[str], max_len: int) -> int:
self._check_up()
stored = self.items.setdefault(key, [])
stored.extend(str(v) for v in values)
pushed_len = len(stored)
del stored[:-max_len]
return pushed_len
async def async_lpop(self, key: str, count: int | None = None, **kwargs: object) -> str | list[str] | None:
self._check_up()
stored = self.items.get(key, [])
if not stored:
return None
if count is None:
return stored.pop(0)
popped = stored[:count]
del stored[:count]
return popped
async def parked_request_ids(self) -> list[str]:
rows = await self.raw.lrange(REDIS_SPEND_LOGS_BUFFER_KEY, 0, -1)
return [json.loads(row)["request_id"] for row in rows]
@pytest.fixture
def fake_redis() -> FakeRedisList:
return FakeRedisList()
@pytest_asyncio.fixture(loop_scope="function")
async def fake_redis() -> FakeRedisList:
server = fakeredis.FakeServer()
return FakeRedisList(server=server, cache=fake_redis_cache(server), raw=fakeredis.FakeAsyncRedis(server=server))
@pytest.fixture
@ -358,7 +347,7 @@ def proxy_logging_with_redis(fake_redis: FakeRedisList) -> MagicMock:
proxy_logging.failure_handler = AsyncMock()
proxy_logging.db_spend_update_writer = MagicMock()
proxy_logging.db_spend_update_writer.db_update_spend_transaction_handler = AsyncMock()
buffer = RedisUpdateBuffer(redis_cache=fake_redis)
buffer = RedisUpdateBuffer(redis_cache=fake_redis.cache)
buffer.should_commit_spend_updates_to_redis = MagicMock(return_value=True)
proxy_logging.db_spend_update_writer.redis_update_buffer = buffer
return proxy_logging
@ -366,12 +355,12 @@ def proxy_logging_with_redis(fake_redis: FakeRedisList) -> MagicMock:
@dataclass
class _SentMessage:
from_addr: Optional[str]
from_addr: str | None
to_addrs: Any
subject: Optional[str]
body: Optional[str]
subject: str | None
body: str | None
starttls_called: bool
login_args: Optional[tuple]
login_args: tuple | None
thread_ident: int
@ -379,9 +368,9 @@ class _SentMessage:
class InMemorySMTP:
"""Captures outbound SMTP traffic for ``send_email`` tests."""
sent: List[_SentMessage] = field(default_factory=list)
raise_on_send: Optional[Exception] = None
connection_kwargs: List[Dict[str, Any]] = field(default_factory=list)
sent: list[_SentMessage] = field(default_factory=list)
raise_on_send: Exception | None = None
connection_kwargs: list[dict[str, Any]] = field(default_factory=list)
def server_factory(self) -> Callable[..., Any]:
outer = self
@ -389,9 +378,9 @@ class InMemorySMTP:
class _Conn:
def __init__(self) -> None:
self._starttls_called = False
self._login_args: Optional[tuple] = None
self._login_args: tuple | None = None
def __enter__(self) -> "_Conn":
def __enter__(self) -> _Conn:
return self
def __exit__(self, *exc: Any) -> None:
@ -406,7 +395,7 @@ class InMemorySMTP:
def send_message(
self,
msg: EmailMessage,
from_addr: Optional[str] = None,
from_addr: str | None = None,
to_addrs: Any = None,
) -> None:
if outer.raise_on_send is not None:

View file

@ -11,7 +11,6 @@ Symbols pinned here:
from __future__ import annotations
import asyncio
import json
from collections.abc import Callable
from contextlib import suppress
from datetime import datetime, timezone
@ -20,12 +19,11 @@ from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY
from litellm.proxy.utils import (
MAX_SPEND_LOG_DRAIN_ITERATIONS,
drain_spend_logs_queue,
monitor_spend_logs_queue,
raise_failed_update_spend_exception,
drain_spend_logs_queue,
recover_parked_spend_logs,
update_daily_tag_spend,
update_spend,
@ -719,10 +717,6 @@ def _table_gone_error() -> Exception:
)
def _parked_request_ids(fake_redis: Any) -> list[str]:
return [json.loads(row)["request_id"] for row in fake_redis.items.get(REDIS_SPEND_LOGS_BUFFER_KEY, [])]
@pytest.mark.asyncio
async def test_drain_spend_logs_queue_parks_unwritable_rows_in_redis_on_shutdown(
mock_prisma_client: Any, make_spend_log_row: Any, proxy_logging_with_redis: MagicMock, fake_redis: Any
@ -743,7 +737,7 @@ async def test_drain_spend_logs_queue_parks_unwritable_rows_in_redis_on_shutdown
)
assert mock_prisma_client.spend_log_transactions == []
assert sorted(_parked_request_ids(fake_redis)) == ["r1", "r2"]
assert sorted(await fake_redis.parked_request_ids()) == ["r1", "r2"]
@pytest.mark.asyncio
@ -779,7 +773,7 @@ async def test_drain_spend_logs_queue_waits_for_an_in_flight_write_before_parkin
proxy_logging_obj=proxy_logging_with_redis,
)
assert _parked_request_ids(fake_redis) == ["in-flight"]
assert await fake_redis.parked_request_ids() == ["in-flight"]
assert mock_prisma_client.spend_log_transactions == []
await release
with suppress(Exception):
@ -813,7 +807,7 @@ async def test_drain_spend_logs_queue_parks_rows_left_after_max_passes(
)
assert mock_prisma_client.spend_log_transactions == []
assert _parked_request_ids(fake_redis) == ["late"]
assert await fake_redis.parked_request_ids() == ["late"]
@pytest.mark.asyncio
@ -822,7 +816,7 @@ async def test_drain_spend_logs_queue_keeps_rows_in_memory_when_redis_is_down(
) -> None:
from prisma.errors import TableNotFoundError
fake_redis.down = True
fake_redis.server.connected = False
mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="r1")]
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_table_gone_error())
@ -834,7 +828,8 @@ async def test_drain_spend_logs_queue_keeps_rows_in_memory_when_redis_is_down(
)
assert [row["request_id"] for row in mock_prisma_client.spend_log_transactions] == ["r1"]
assert fake_redis.items == {}
fake_redis.server.connected = True
assert await fake_redis.raw.keys("*") == []
@pytest.mark.asyncio
@ -863,7 +858,7 @@ async def test_update_spend_writes_rows_parked_in_redis_by_a_previous_pod(
written = mock_prisma_client.db.litellm_spendlogs.create_many.await_args.kwargs["data"]
assert [row["request_id"] for row in written] == ["parked"]
assert _parked_request_ids(fake_redis) == []
assert await fake_redis.parked_request_ids() == []
assert mock_prisma_client.spend_log_transactions == []
@ -879,14 +874,14 @@ async def test_recover_parked_spend_logs_re_parks_rows_when_the_enqueue_is_cance
recover_parked_spend_logs(prisma_client=mock_prisma_client, proxy_logging_obj=proxy_logging_with_redis)
)
await asyncio.sleep(0.01)
assert _parked_request_ids(fake_redis) == []
assert await fake_redis.parked_request_ids() == []
recovery.cancel()
with pytest.raises(asyncio.CancelledError):
await recovery
mock_prisma_client._spend_log_transactions_lock.release()
assert _parked_request_ids(fake_redis) == ["parked"]
assert await fake_redis.parked_request_ids() == ["parked"]
assert mock_prisma_client.spend_log_transactions == []

View file

@ -0,0 +1,53 @@
from datetime import datetime, timezone
from typing import Final
import fakeredis
import pytest
from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY
from litellm.redis.stores.spend_logs_parking import TransportSpendLogsParkingStore
from litellm.redis.transport.contract import RedisFailure
from tests._support.fake_redis import fake_redis_cache
KEY: Final = f"ns:{REDIS_SPEND_LOGS_BUFFER_KEY}"
def _store(server: fakeredis.FakeServer) -> TransportSpendLogsParkingStore:
return TransportSpendLogsParkingStore(fake_redis_cache(server, namespace="ns").transport())
@pytest.mark.asyncio
async def test_rows_are_parked_as_json_with_iso_datetimes_and_capped_to_the_newest(
server: fakeredis.FakeServer,
) -> None:
store: Final = _store(server)
raw: Final = fakeredis.FakeAsyncRedis(server=server)
started: Final = datetime(2026, 9, 19, 20, 0, 5, tzinfo=timezone.utc)
assert await store.park(({"request_id": "old"}, {"request_id": "mid", "startTime": started}), max_rows=2) == 2
assert await store.park(({"request_id": "new"},), max_rows=2) == 3
assert await raw.lrange(KEY, 0, -1) == [
b'{"request_id": "mid", "startTime": "2026-09-19T20:00:05+00:00"}',
b'{"request_id": "new"}',
]
assert await store.take(1) == ('{"request_id": "mid", "startTime": "2026-09-19T20:00:05+00:00"}',)
assert await store.take(10) == ('{"request_id": "new"}',)
assert await store.take(10) is None
@pytest.mark.asyncio
async def test_rejections_undecodable_rows_and_outages_are_failures(server: fakeredis.FakeServer) -> None:
raw: Final = fakeredis.FakeAsyncRedis(server=server)
store: Final = _store(server)
await raw.set(KEY, "not a list")
rejected: Final = await store.park(({"request_id": "a"},), max_rows=10)
await raw.delete(KEY)
await raw.rpush(KEY, b"\xff")
undecodable: Final = await store.take(10)
server.connected = False
outage: Final = await store.take(10)
assert isinstance(rejected, RedisFailure) and rejected.kind == "server"
assert isinstance(undecodable, RedisFailure) and undecodable.kind == "decode"
assert isinstance(outage, RedisFailure) and outage.kind == "transport"