mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
refactor(redis): move parked spend logs behind a typed store
This commit is contained in:
parent
f969987849
commit
df3cdb255f
6 changed files with 170 additions and 123 deletions
|
|
@ -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
|
||||
|
|
|
|||
45
litellm/redis/stores/spend_logs_parking.py
Normal file
45
litellm/redis/stores/spend_logs_parking.py
Normal 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}")
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 == []
|
||||
|
||||
|
||||
|
|
|
|||
53
tests/unit/redis/stores/test_spend_logs_parking.py
Normal file
53
tests/unit/redis/stores/test_spend_logs_parking.py
Normal 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"
|
||||
Loading…
Add table
Reference in a new issue