scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

refactor(redis): move the spend update buffers behind a typed store

Browse source
This commit is contained in:
Devin AI 2026-10-30 22:41:38 +00:00
parent ffae582a50
commit f969987849
4 changed files with 356 additions and 265 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

253
litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
View file

@ -20,17 +20,10 @@ from litellm._logging import verbose_proxy_logger
from litellm.caching import RedisCache
from litellm.constants import (
MAX_REDIS_BUFFER_DEQUEUE_COUNT,
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
REDIS_SPEND_LOGS_BUFFER_KEY,
REDIS_SPEND_LOGS_BUFFER_MAX_ROWS,
REDIS_TOOL_DISCOVERY_BUFFER_KEY,
REDIS_UPDATE_BUFFER_KEY,
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import (
@ -61,6 +54,13 @@ 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_update_buffer import (
SpendBuffer,
SpendUpdateBufferStore,
TransportSpendUpdateBufferStore,
spend_buffer_key,
)
from litellm.redis.transport.contract import RedisFailure, raise_redis_failure
from litellm.secret_managers.main import str_to_bool
from litellm.types.caching import (
RedisPipelineLpopOperation,
@ -107,6 +107,19 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = (
_ValueT = TypeVar("_ValueT")
_DECODED_JSON: Final = TypeAdapter(object)
_DRAINED_BUFFERS: Final[tuple[SpendBuffer, ...]] = (
"spend",
"daily_user",
"daily_team",
"daily_org",
"daily_end_user",
"daily_agent",
"window",
)
def _spend_buffers(redis_cache: RedisCache) -> SpendUpdateBufferStore:
return TransportSpendUpdateBufferStore(redis_cache.transport())
def _zero_for_non_finite_constant(constant: str) -> float:
@ -191,7 +204,7 @@ class RedisUpdateBuffer:
async def _store_transactions_in_redis(
self,
transactions: dict[str, BaseDailySpendTransaction] | None, # mutable-ok: add_update requires dict
redis_key: str,
buffer: SpendBuffer,
service_type: ServiceTypes,
*,
source_queue: DailySpendUpdateQueue,
@ -201,39 +214,34 @@ class RedisUpdateBuffer:
Args:
transactions: The transactions to store
redis_key: The Redis key to store under
buffer: The Redis spend buffer to store under
service_type: The service type for event emission
source_queue: The in-memory queue to restore transactions to if Redis push fails
"""
if transactions is None or len(transactions) == 0:
return
list_of_transactions: Final = [safe_dumps(transactions)]
if self.redis_cache is None:
return
try:
current_redis_buffer_size: Final = await self.redis_cache.async_rpush(
key=redis_key,
values=list_of_transactions,
)
except Exception as e:
current_redis_buffer_size: Final = await _spend_buffers(self.redis_cache).push(buffer, transactions)
if isinstance(current_redis_buffer_size, RedisFailure):
verbose_proxy_logger.error(
"Spend tracking - failed to push spend updates to Redis (redis_key=%s). "
"Restoring them to the in-memory queue for retry on next tick. Error: %s",
redis_key,
str(e),
spend_buffer_key(buffer),
current_redis_buffer_size.message,
)
await source_queue.add_update(transactions)
else:
verbose_proxy_logger.debug(
"Spend tracking - pushed spend updates to Redis buffer. redis_key=%s, buffer_size=%s",
redis_key,
current_redis_buffer_size,
)
await self._emit_new_item_added_to_redis_buffer_event(
queue_size=current_redis_buffer_size,
service=service_type,
)
return
verbose_proxy_logger.debug(
"Spend tracking - pushed spend updates to Redis buffer. redis_key=%s, buffer_size=%s",
spend_buffer_key(buffer),
current_redis_buffer_size,
)
await self._emit_new_item_added_to_redis_buffer_event(
queue_size=current_redis_buffer_size,
service=service_type,
)
@with_service_target(SPEND_QUEUE_TARGET)
async def store_in_memory_spend_updates_in_redis(
@ -323,65 +331,54 @@ class RedisUpdateBuffer:
verbose_proxy_logger.debug("ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions)
# Build a list of rpush operations, skipping empty/None transaction sets
_queue_configs: Final[list[tuple[BufferedSpendTransactions | None, str, ServiceTypes]]] = [
_queue_configs: Final[tuple[tuple[object, SpendBuffer, ServiceTypes], ...]] = (
(
db_spend_update_transactions,
REDIS_UPDATE_BUFFER_KEY,
"spend",
ServiceTypes.REDIS_SPEND_UPDATE_QUEUE,
),
(
daily_spend_update_transactions,
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
"daily_user",
ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE,
),
(
daily_team_spend_update_transactions,
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
"daily_team",
ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE,
),
(
daily_org_spend_update_transactions,
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
"daily_org",
ServiceTypes.REDIS_DAILY_ORG_SPEND_UPDATE_QUEUE,
),
(
daily_end_user_spend_update_transactions,
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
"daily_end_user",
ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE,
),
(
daily_agent_spend_update_transactions,
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
"daily_agent",
ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE,
),
(
tuple(map(to_wire_payload, window_spend_update_transactions)),
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
"window",
ServiceTypes.REDIS_WINDOW_SPEND_UPDATE_QUEUE,
),
]
)
rpush_list: Final[list[RedisPipelineRpushOperation]] = []
service_types: Final[list[ServiceTypes]] = []
for transactions, redis_key, service_type in _queue_configs:
if transactions is None or len(transactions) == 0:
continue
rpush_list.append(
RedisPipelineRpushOperation(
key=redis_key,
values=[safe_dumps(transactions)],
)
)
service_types.append(service_type)
if len(rpush_list) == 0:
pending: Final[tuple[tuple[SpendBuffer, object, ServiceTypes], ...]] = tuple(
(config[1], config[0], config[2]) for config in _queue_configs if config[0]
)
if len(pending) == 0:
return
try:
result_lengths: Final = await self.redis_cache.async_rpush_pipeline(
rpush_list=rpush_list,
)
except Exception as e:
result_lengths: Final = await _spend_buffers(self.redis_cache).push_all(
tuple((buffer, transactions) for buffer, transactions, _ in pending)
)
if isinstance(result_lengths, RedisFailure):
# The in-memory queues were already drained above. If we let the
# exception propagate without restoring, the aggregated spend is
# permanently lost. Re-enqueue so the next scheduler tick retries.
@ -389,8 +386,8 @@ class RedisUpdateBuffer:
"Spend tracking - failed to push aggregated spend updates to Redis. "
"Restoring %d transaction sets to in-memory queues for retry on next tick. "
"Error: %s",
len(rpush_list),
str(e),
len(pending),
result_lengths.message,
)
await self._restore_spend_updates_to_in_memory_queues(
db_spend_update_transactions=db_spend_update_transactions,
@ -411,12 +408,8 @@ class RedisUpdateBuffer:
return
# Emit gauge events for each queue
for i, queue_size in enumerate(result_lengths):
if i < len(service_types):
await self._emit_new_item_added_to_redis_buffer_event(
queue_size=queue_size,
service=service_types[i],
)
for queue_size, (_, _, service_type) in zip(result_lengths, pending, strict=False):
await self._emit_new_item_added_to_redis_buffer_event(queue_size=queue_size, service=service_type)
@staticmethod
async def _restore_spend_updates_to_in_memory_queues(
@ -542,42 +535,42 @@ class RedisUpdateBuffer:
if self.redis_cache is None:
return
restore_configs: Final = (
(db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY),
(daily_spend_update_transactions, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY),
(daily_team_spend_update_transactions, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY),
(daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY),
(daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY),
(daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY),
(daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY),
restore_configs: Final[tuple[tuple[object, SpendBuffer], ...]] = (
(db_spend_update_transactions, "spend"),
(daily_spend_update_transactions, "daily_user"),
(daily_team_spend_update_transactions, "daily_team"),
(daily_org_spend_update_transactions, "daily_org"),
(daily_end_user_spend_update_transactions, "daily_end_user"),
(daily_agent_spend_update_transactions, "daily_agent"),
(daily_tag_spend_update_transactions, "daily_tag"),
(
None
if window_spend_update_transactions is None
else tuple(map(to_wire_payload, window_spend_update_transactions)),
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
"window",
),
)
rpush_list: Final = tuple(
RedisPipelineRpushOperation(key=redis_key, values=(safe_dumps(transactions),))
for transactions, redis_key in restore_configs
if transactions
batches: Final[tuple[tuple[SpendBuffer, object], ...]] = tuple(
(buffer, transactions) for transactions, buffer in restore_configs if transactions
)
if len(rpush_list) == 0:
if len(batches) == 0:
return
try:
await self.redis_cache.async_rpush_pipeline(rpush_list=rpush_list)
restored: Final = await _spend_buffers(self.redis_cache).push_all(batches)
if not isinstance(restored, RedisFailure):
verbose_proxy_logger.info(
"Spend tracking - restored %d uncommitted transaction set(s) to Redis for retry on next tick.",
len(rpush_list),
)
except RedisError as e:
verbose_proxy_logger.error(
"Spend tracking - failed to restore uncommitted transactions to Redis. "
"These spend updates are lost. Error: %s",
str(e),
len(batches),
)
return
if not isinstance(restored.exception, RedisError):
raise_redis_failure(restored)
verbose_proxy_logger.error(
"Spend tracking - failed to restore uncommitted transactions to Redis. "
"These spend updates are lost. Error: %s",
restored.message,
)
@with_service_target(SPEND_QUEUE_TARGET)
async def store_tool_discoveries_in_redis(self, tool_discovery_queue: ToolDiscoveryQueue) -> None:
@ -734,22 +727,23 @@ class RedisUpdateBuffer:
"""
if self.redis_cache is None:
return None
list_of_transactions: Final[str | list[str] | None] = await self.redis_cache.async_lpop(
key=REDIS_UPDATE_BUFFER_KEY,
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
list_of_transactions: Final = await _spend_buffers(self.redis_cache).pop(
"spend", count=MAX_REDIS_BUFFER_DEQUEUE_COUNT
)
if isinstance(list_of_transactions, RedisFailure):
raise_redis_failure(list_of_transactions)
if list_of_transactions is None:
return None
verbose_proxy_logger.info(
"Spend tracking - popped %d spend update batches from Redis buffer (key=%s). "
"These items are now removed from Redis and must be committed to DB.",
len(list_of_transactions) if isinstance(list_of_transactions, list) else 1,
len(list_of_transactions),
REDIS_UPDATE_BUFFER_KEY,
)
# Parse the list of transactions from JSON strings
parsed_transactions: Final = self._parse_list_of_transactions(list_of_transactions)
parsed_transactions: Final = self._parse_list_of_transactions(list(list_of_transactions))
# If there are no transactions, return None
if len(parsed_transactions) == 0:
@ -787,44 +781,16 @@ class RedisUpdateBuffer:
if self.redis_cache is None:
return None, None, None, None, None, None, None
lpop_list: Final[list[RedisPipelineLpopOperation]] = [
RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT),
RedisPipelineLpopOperation(
key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
),
RedisPipelineLpopOperation(
key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
),
RedisPipelineLpopOperation(
key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
),
RedisPipelineLpopOperation(
key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
),
RedisPipelineLpopOperation(
key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
),
RedisPipelineLpopOperation(
key=REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
),
]
raw_results: Final = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list)
# Pad with None if pipeline returned fewer results than expected
while len(raw_results) < 7:
raw_results.append(None)
raw_results: Final = await _spend_buffers(self.redis_cache).pop_all(
_DRAINED_BUFFERS, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT
)
if isinstance(raw_results, RedisFailure):
raise_redis_failure(raw_results)
# Slot 0: DBSpendUpdateTransactions
db_spend: DBSpendUpdateTransactions | None = None
if raw_results[0] is not None:
parsed: Final = self._parse_list_of_transactions(raw_results[0])
parsed: Final = self._parse_list_of_transactions(list(raw_results[0]))
if len(parsed) > 0:
db_spend = self._combine_list_of_transactions(parsed)
@ -869,7 +835,7 @@ class RedisUpdateBuffer:
)
await self._store_transactions_in_redis(
transactions=daily_tag_spend_update_transactions,
redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
buffer="daily_tag",
service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE,
source_queue=daily_tag_spend_update_queue,
)
@ -877,17 +843,18 @@ class RedisUpdateBuffer:
@with_service_target(SPEND_QUEUE_TARGET)
async def _lpop_daily_spend_transactions(
self,
redis_key: str,
buffer: SpendBuffer,
) -> list[dict[str, BaseDailySpendTransaction]] | None:
"""
Drains a daily spend buffer key and parses each popped item as JSON.
"""
if self.redis_cache is None:
return None
list_of_transactions: Final[list[str] | None] = await self.redis_cache.async_lpop(
key=redis_key,
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
list_of_transactions: Final = await _spend_buffers(self.redis_cache).pop(
buffer, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT
)
if isinstance(list_of_transactions, RedisFailure):
raise_redis_failure(list_of_transactions)
if list_of_transactions is None:
return None
return [_load_buffered_transaction(transaction) for transaction in list_of_transactions]
@ -900,9 +867,7 @@ class RedisUpdateBuffer:
"""
if self.redis_cache is None:
return None
list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions(
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY
)
list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions("daily_user")
if list_of_daily_spend_update_transactions is None:
return None
return cast(
@ -920,9 +885,7 @@ class RedisUpdateBuffer:
"""
if self.redis_cache is None:
return None
list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions(
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY
)
list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions("daily_team")
if list_of_daily_spend_update_transactions is None:
return None
return cast(
@ -940,9 +903,7 @@ class RedisUpdateBuffer:
"""
if self.redis_cache is None:
return None
list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions(
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY
)
list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions("daily_org")
if list_of_daily_spend_update_transactions is None:
return None
return cast(
@ -960,9 +921,7 @@ class RedisUpdateBuffer:
"""
if self.redis_cache is None:
return None
list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions(
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY
)
list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions("daily_end_user")
if list_of_daily_spend_update_transactions is None:
return None
return cast(
@ -980,9 +939,7 @@ class RedisUpdateBuffer:
"""
if self.redis_cache is None:
return None
list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions(
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY
)
list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions("daily_agent")
if list_of_daily_spend_update_transactions is None:
return None
return cast(
@ -1000,9 +957,7 @@ class RedisUpdateBuffer:
"""
if self.redis_cache is None:
return None
list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions(
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY
)
list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions("daily_tag")
if list_of_daily_spend_update_transactions is None:
return None
return cast(

89
litellm/redis/stores/spend_update_buffer.py Normal file
View file

@ -0,0 +1,89 @@
from collections.abc import Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final, Literal, Protocol, TypeAlias
from litellm.constants import (
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
REDIS_UPDATE_BUFFER_KEY,
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.redis.transport.contract import Popped, RedisFailure, RedisTransport
SpendBuffer: TypeAlias = Literal[
"spend", "daily_user", "daily_team", "daily_org", "daily_end_user", "daily_agent", "daily_tag", "window"
]
_KEYS: Final = MappingProxyType[SpendBuffer, str](
{
"spend": REDIS_UPDATE_BUFFER_KEY,
"daily_user": REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
"daily_team": REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
"daily_org": REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
"daily_end_user": REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
"daily_agent": REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
"daily_tag": REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
"window": REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
}
)
def spend_buffer_key(buffer: SpendBuffer) -> str:
return _KEYS[buffer]
class SpendUpdateBufferStore(Protocol):
async def push(self, buffer: SpendBuffer, transactions: object) -> int | RedisFailure: ...
async def push_all(self, batches: Sequence[tuple[SpendBuffer, object]]) -> tuple[int, ...] | RedisFailure: ...
async def pop(self, buffer: SpendBuffer, *, count: int) -> tuple[str, ...] | None | RedisFailure: ...
async def pop_all(
self, buffers: Sequence[SpendBuffer], *, count: int
) -> tuple[tuple[str, ...] | None, ...] | RedisFailure: ...
@dataclass(frozen=True, slots=True)
class TransportSpendUpdateBufferStore:
transport: RedisTransport
async def push(self, buffer: SpendBuffer, transactions: object) -> int | RedisFailure:
return await self.transport.rpush(spend_buffer_key(buffer), safe_dumps(transactions))
async def push_all(self, batches: Sequence[tuple[SpendBuffer, object]]) -> tuple[int, ...] | RedisFailure:
batch: Final = self.transport.batch(name="spend_update_buffer")
pushed: Final = tuple(
batch.rpush(spend_buffer_key(buffer), safe_dumps(transactions)) for buffer, transactions in batches
)
lengths: Final = tuple([await reply for reply in pushed])
failure: Final = next((length for length in lengths if isinstance(length, RedisFailure)), None)
return failure if failure is not None else tuple(length for length in lengths if isinstance(length, int))
async def pop(self, buffer: SpendBuffer, *, count: int) -> tuple[str, ...] | None | RedisFailure:
return _decode(await self.transport.lpop(spend_buffer_key(buffer), count))
async def pop_all(
self, buffers: Sequence[SpendBuffer], *, count: int
) -> tuple[tuple[str, ...] | None, ...] | RedisFailure:
batch: Final = self.transport.batch(name="spend_update_buffer")
popped: Final = tuple(batch.lpop(spend_buffer_key(buffer), count) for buffer in buffers)
decoded: Final = tuple([_decode(await reply) for reply in popped])
failure: Final = next((entries for entries in decoded if isinstance(entries, RedisFailure)), None)
return failure if failure is not None else tuple(e for e in decoded if not isinstance(e, RedisFailure))
def _decode(popped: Popped | RedisFailure) -> tuple[str, ...] | None | RedisFailure:
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 spend update buffer entry: {error.reason}")

216
tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py
View file

@ -4,15 +4,28 @@ from datetime import datetime, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import fakeredis
import pytest
import pytest_asyncio
from litellm._logging import verbose_proxy_logger
from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY
from litellm.caching.redis_cache import RedisCache
from litellm.constants import (
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
REDIS_UPDATE_BUFFER_KEY,
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import DailyTagSpendTransaction
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue
from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
from litellm.proxy.proxy_server import ProxyStartupEvent
from tests._support.fake_redis import fake_redis_cache
def _daily_tag_spend_update_transactions() -> dict[str, DailyTagSpendTransaction]:
@ -45,9 +58,9 @@ def _daily_tag_spend_update_transactions() -> dict[str, DailyTagSpendTransaction
@pytest.mark.asyncio
async def test_daily_tag_spend_restored_to_in_memory_queue_on_rpush_failure(
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, fake_server
):
mock_redis_cache.async_rpush = AsyncMock(side_effect=ConnectionError("redis went away"))
fake_server.connected = False
queue: Final = DailySpendUpdateQueue()
original_transactions: Final = _daily_tag_spend_update_transactions()
await queue.add_update(original_transactions)
@ -59,31 +72,36 @@ async def test_daily_tag_spend_restored_to_in_memory_queue_on_rpush_failure(
@pytest.mark.asyncio
async def test_daily_tag_spend_restored_after_rpush_failure_is_pushed_once_on_next_tick(
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, fake_server, raw_redis
):
mock_redis_cache.async_rpush = AsyncMock(side_effect=ConnectionError("redis went away"))
fake_server.connected = False
queue: Final = DailySpendUpdateQueue()
original_transactions: Final = _daily_tag_spend_update_transactions()
await queue.add_update(original_transactions)
await redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue)
mock_redis_cache.async_rpush = AsyncMock(return_value=1)
fake_server.connected = True
await redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue)
mock_redis_cache.async_rpush.assert_called_once()
redis_key: Final = mock_redis_cache.async_rpush.call_args.kwargs["key"]
values: Final = mock_redis_cache.async_rpush.call_args.kwargs["values"]
assert redis_key == REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY
assert json.loads(values[0]) == json.loads(safe_dumps(original_transactions))
values: Final = await raw_redis.lrange(REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, 0, -1)
assert values == [safe_dumps(original_transactions).encode()]
assert await queue.flush_and_get_aggregated_daily_spend_update_transactions() == {}
@pytest.fixture
def mock_redis_cache():
"""Create a mock RedisCache instance"""
mock = AsyncMock()
return mock
def fake_server() -> fakeredis.FakeServer:
return fakeredis.FakeServer()
@pytest.fixture
def raw_redis(fake_server: fakeredis.FakeServer) -> fakeredis.FakeAsyncRedis:
return fakeredis.FakeAsyncRedis(server=fake_server)
@pytest_asyncio.fixture(loop_scope="function")
async def mock_redis_cache(fake_server: fakeredis.FakeServer) -> RedisCache:
return fake_redis_cache(fake_server)
@pytest.fixture
@ -93,13 +111,11 @@ def redis_update_buffer(mock_redis_cache):
@pytest.mark.asyncio
async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache):
async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache, raw_redis):
"""
Verify store_in_memory_spend_updates_in_redis calls async_rpush_pipeline once
with the correct operations and skips empty queues.
"""
mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[3, 5, 2])
# Create mock queues - only 3 of 6 have data
spend_update_queue = AsyncMock()
spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(
@ -135,17 +151,23 @@ async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer,
daily_agent_spend_update_queue=daily_agent_queue,
)
# Should be called exactly once (pipeline)
mock_redis_cache.async_rpush_pipeline.assert_called_once()
# Verify only 3 operations were included (empty ones skipped)
call_args = mock_redis_cache.async_rpush_pipeline.call_args
rpush_list = call_args.kwargs["rpush_list"]
assert len(rpush_list) == 3
assert await raw_redis.lrange(REDIS_UPDATE_BUFFER_KEY, 0, -1) == [b'{"key_list_transactions": {"key1": 1.0}}']
assert await raw_redis.lrange(REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, 0, -1) == [b'{"user_key1": {"spend": 1.0}}']
assert await raw_redis.lrange(REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, 0, -1) == [b'{"team_key1": {"spend": 2.0}}']
assert (
await raw_redis.exists(
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
)
== 0
)
@pytest.mark.asyncio
async def test_store_in_memory_spend_updates_restores_on_rpush_failure(redis_update_buffer, mock_redis_cache):
async def test_store_in_memory_spend_updates_restores_on_rpush_failure(
redis_update_buffer, mock_redis_cache, fake_server
):
"""
If async_rpush_pipeline raises, the already-drained transactions must be
put back into the in-memory queues so the next scheduler tick retries.
@ -159,7 +181,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure(redis_upd
SpendUpdateQueue,
)
mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away"))
fake_server.connected = False
spend_queue = SpendUpdateQueue()
daily_user_queue = DailySpendUpdateQueue()
@ -245,7 +267,7 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_updat
@pytest.mark.asyncio
async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buffer, mock_redis_cache):
async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buffer, mock_redis_cache, raw_redis):
"""
Verify get_all_transactions_from_redis_buffer_pipeline correctly parses
and aggregates results from async_lpop_pipeline.
@ -278,17 +300,10 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buff
]
)
mock_redis_cache.async_lpop_pipeline = AsyncMock(
return_value=[
[db_spend_json], # slot 0: db spend updates
[daily_user_json], # slot 1: daily user
[daily_team_json], # slot 2: daily team
None, # slot 3: daily org (empty)
None, # slot 4: daily end-user (empty)
None, # slot 5: daily agent (empty)
[window_spend_json, window_spend_json], # slot 6: budget window spend
]
)
await raw_redis.rpush(REDIS_UPDATE_BUFFER_KEY, db_spend_json)
await raw_redis.rpush(REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, daily_user_json)
await raw_redis.rpush(REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, daily_team_json)
await raw_redis.rpush(REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, window_spend_json, window_spend_json)
result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
@ -328,17 +343,12 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buff
assert daily_end_user is None
assert daily_agent is None
# Verify pipeline was called once with correct keys
mock_redis_cache.async_lpop_pipeline.assert_called_once()
from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
popped_keys = [op["key"] for op in mock_redis_cache.async_lpop_pipeline.call_args.kwargs["lpop_list"]]
assert popped_keys[6] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
assert await raw_redis.exists(REDIS_UPDATE_BUFFER_KEY, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY) == 0
@pytest.mark.asyncio
async def test_org_member_spend_is_summed_across_pods_and_restored_on_rpush_failure(
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, fake_server, raw_redis
):
from litellm.proxy._types import Litellm_EntityType
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
@ -350,16 +360,14 @@ async def test_org_member_spend_is_summed_across_pods_and_restored_on_rpush_fail
member_key: Final = "organization_id::org-1::user_id::user-1"
pod_json: Final = json.dumps({"org_member_list_transactions": {member_key: 0.25}})
mock_redis_cache.async_lpop_pipeline = AsyncMock(
return_value=[[pod_json, pod_json], None, None, None, None, None, None]
)
await raw_redis.rpush(REDIS_UPDATE_BUFFER_KEY, pod_json, pod_json)
(db_spend, *_rest) = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
assert db_spend is not None
assert db_spend["org_member_list_transactions"] == {member_key: 0.5}
mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away"))
fake_server.connected = False
spend_queue: Final = SpendUpdateQueue()
await spend_queue.add_update(
{
@ -383,10 +391,11 @@ async def test_org_member_spend_is_summed_across_pods_and_restored_on_rpush_fail
@pytest.mark.asyncio
async def test_a_single_parked_batch_with_a_nan_spend_drains_as_zero(
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, raw_redis
):
mock_redis_cache.async_lpop = AsyncMock(
return_value=json.dumps({"key_list_transactions": {"nan-key": float("nan"), "ok-key": 0.5}}, allow_nan=True)
await raw_redis.rpush(
REDIS_UPDATE_BUFFER_KEY,
json.dumps({"key_list_transactions": {"nan-key": float("nan"), "ok-key": 0.5}}, allow_nan=True),
)
drained: Final = await redis_update_buffer.get_all_update_transactions_from_redis_buffer()
@ -416,7 +425,7 @@ def _parked_daily_batch(spend: float) -> str:
@pytest.mark.asyncio
async def test_a_nan_spend_parked_by_an_older_version_drains_from_the_pipeline_as_zero(
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, raw_redis
):
entity_batch: Final = json.dumps({"key_list_transactions": {"hashed-key": float("nan")}}, allow_nan=True)
window_batches: Final = tuple(
@ -435,17 +444,11 @@ async def test_a_nan_spend_parked_by_an_older_version_drains_from_the_pipeline_a
)
for spend in (float("nan"), 3.0)
)
mock_redis_cache.async_lpop_pipeline = AsyncMock(
return_value=[
[entity_batch],
[_parked_daily_batch(float("nan")), _parked_daily_batch(0.0042)],
None,
None,
None,
None,
list(window_batches),
]
await raw_redis.rpush(REDIS_UPDATE_BUFFER_KEY, entity_batch)
await raw_redis.rpush(
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, _parked_daily_batch(float("nan")), _parked_daily_batch(0.0042)
)
await raw_redis.rpush(REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, *window_batches)
(
db_spend,
@ -465,10 +468,10 @@ async def test_a_nan_spend_parked_by_an_older_version_drains_from_the_pipeline_a
@pytest.mark.asyncio
async def test_a_nan_spend_parked_by_an_older_version_drains_from_a_daily_key_as_zero(
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, caplog: pytest.LogCaptureFixture
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock, caplog: pytest.LogCaptureFixture, raw_redis
):
mock_redis_cache.async_lpop = AsyncMock(
return_value=[_parked_daily_batch(float("inf")), _parked_daily_batch(0.0042)]
await raw_redis.rpush(
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, _parked_daily_batch(float("inf")), _parked_daily_batch(0.0042)
)
with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name):
@ -488,18 +491,11 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis():
@pytest.mark.asyncio
async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_buffer, mock_redis_cache):
async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_buffer, mock_redis_cache, raw_redis):
"""
restore_transactions_to_redis re-pushes only the transaction sets it was
given, to their matching buffer keys, so uncommitted spend can be retried.
"""
from litellm.constants import (
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
REDIS_UPDATE_BUFFER_KEY,
)
mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1, 1])
db_spend = {"key_list_transactions": {"key1": 1.0}}
daily_user = {"user_key1": {"spend": 1.0}}
@ -508,24 +504,19 @@ async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_b
daily_spend_update_transactions=daily_user,
)
mock_redis_cache.async_rpush_pipeline.assert_called_once()
rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"]
pushed_keys = {op["key"] for op in rpush_list}
assert pushed_keys == {
REDIS_UPDATE_BUFFER_KEY,
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
}
# Payloads round-trip through the same JSON encoding used on the store path
payloads = {op["key"]: json.loads(op["values"][0]) for op in rpush_list}
assert payloads[REDIS_UPDATE_BUFFER_KEY] == db_spend
assert payloads[REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY] == daily_user
assert sorted(await raw_redis.keys("*")) == sorted(
(REDIS_UPDATE_BUFFER_KEY.encode(), REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY.encode())
)
assert await raw_redis.lrange(REDIS_UPDATE_BUFFER_KEY, 0, -1) == [safe_dumps(db_spend).encode()]
assert await raw_redis.lrange(REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, 0, -1) == [safe_dumps(daily_user).encode()]
@pytest.mark.asyncio
async def test_restored_window_spend_transactions_drain_back_unchanged(redis_update_buffer, mock_redis_cache):
async def test_restored_window_spend_transactions_drain_back_unchanged(
redis_update_buffer, mock_redis_cache, raw_redis
):
"""A window commit that fails after the destructive lpop must be re-pushed
in the store path's encoding, so the next drain returns the same increments."""
from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
build_window_spend_transaction,
)
@ -540,16 +531,9 @@ async def test_restored_window_spend_transactions_drain_back_unchanged(redis_upd
started_at=datetime(2026, 8, 10, 12, 0, tzinfo=timezone.utc),
),
)
mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1])
await redis_update_buffer.restore_transactions_to_redis(window_spend_update_transactions=window_transactions)
rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"]
assert [op["key"] for op in rpush_list] == [REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY]
mock_redis_cache.async_lpop_pipeline = AsyncMock(
return_value=[None, None, None, None, None, None, list(rpush_list[0]["values"])]
)
assert await raw_redis.keys("*") == [REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY.encode()]
drained = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
assert drained[6] == window_transactions
@ -564,17 +548,18 @@ async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer
@pytest.mark.asyncio
async def test_restore_transactions_to_redis_swallows_redis_error(redis_update_buffer, mock_redis_cache):
async def test_restore_transactions_to_redis_swallows_redis_error(
redis_update_buffer, mock_redis_cache, fake_server, raw_redis
):
"""A Redis failure during restore must not propagate to the caller's finally block."""
from redis.exceptions import RedisError
mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=RedisError("redis down"))
fake_server.connected = False
await redis_update_buffer.restore_transactions_to_redis(
db_spend_update_transactions={"key_list_transactions": {"key1": 1.0}},
)
mock_redis_cache.async_rpush_pipeline.assert_called_once()
fake_server.connected = True
assert await raw_redis.llen(REDIS_UPDATE_BUFFER_KEY) == 0
def test_validate_redis_transaction_buffer_raises_without_redis():
@ -684,19 +669,18 @@ def test_get_transaction_buffer_redis_cache_parses_string_flag(monkeypatch):
@pytest.mark.asyncio
async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_update_buffer, mock_redis_cache):
async def test_store_in_memory_spend_updates_pushes_budget_window_spend(
redis_update_buffer, mock_redis_cache, raw_redis
):
"""The budget window queue has to ride the same rpush as the daily queues,
otherwise multi-pod deployments never persist per-window spend."""
from datetime import datetime, timezone
from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
WindowSpendUpdateQueue,
build_window_spend_transaction,
)
mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1])
empty_queue = AsyncMock()
empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={})
empty_daily_queue = AsyncMock()
@ -724,10 +708,9 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_up
window_spend_update_queue=window_queue,
)
rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"]
assert len(rpush_list) == 1
assert rpush_list[0]["key"] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
pushed = json.loads(rpush_list[0]["values"][0])
assert await raw_redis.keys("*") == [REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY.encode()]
(pushed_bytes,) = await raw_redis.lrange(REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, 0, -1)
pushed = json.loads(pushed_bytes)
assert pushed == [
{
"entity_type": "key",
@ -742,7 +725,9 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_up
@pytest.mark.asyncio
async def test_budget_window_payloads_keep_request_ids_for_older_workers(redis_update_buffer, mock_redis_cache):
async def test_budget_window_payloads_keep_request_ids_for_older_workers(
redis_update_buffer, mock_redis_cache, raw_redis
):
"""A leader from before the field was dropped indexes request_ids while
merging what it popped, and the pop is destructive, so a payload without
the key would cost a rolling deploy those increments."""
@ -753,7 +738,6 @@ async def test_budget_window_payloads_keep_request_ids_for_older_workers(redis_u
build_window_spend_transaction,
)
mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1])
window_queue = WindowSpendUpdateQueue()
await window_queue.add_update(
build_window_spend_transaction(
@ -769,14 +753,14 @@ async def test_budget_window_payloads_keep_request_ids_for_older_workers(redis_u
window_spend_update_transactions=await window_queue.flush_and_get_aggregated_window_spend_transactions(),
)
rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"]
restored = json.loads(rpush_list[0]["values"][0])
(restored_bytes,) = await raw_redis.lrange(REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, 0, -1)
restored = json.loads(restored_bytes)
assert [payload["request_ids"] for payload in restored] == [[]]
@pytest.mark.asyncio
async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpush_failure(
redis_update_buffer, mock_redis_cache
redis_update_buffer, mock_redis_cache, fake_server
):
"""The window queue is drained before the rpush, so a Redis hiccup would
silently drop per-window spend without the restore."""
@ -787,7 +771,7 @@ async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpu
build_window_spend_transaction,
)
mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away"))
fake_server.connected = False
empty_queue = AsyncMock()
empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={})
@ -859,7 +843,7 @@ async def test_store_spend_logs_in_redis_drops_oldest_rows_past_the_cap():
@pytest.mark.asyncio
@pytest.mark.parametrize("spend", [float("nan"), float("inf")])
async def test_a_non_finite_daily_spend_buffered_in_redis_drains_as_zero(spend: float):
buffer: Final = RedisUpdateBuffer(redis_cache=_ListRedis())
buffer: Final = RedisUpdateBuffer(redis_cache=fake_redis_cache(fakeredis.FakeServer()))
queue: Final = DailySpendUpdateQueue()
await queue.add_update(json.loads(_parked_daily_batch(spend)))
await queue.add_update(json.loads(_parked_daily_batch(0.0042)))

63
tests/unit/redis/stores/test_spend_update_buffer.py Normal file
View file

@ -0,0 +1,63 @@
from typing import Final
import fakeredis
import pytest
from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY
from litellm.redis.stores.spend_update_buffer import TransportSpendUpdateBufferStore
from litellm.redis.transport.contract import RedisFailure
from tests._support.fake_redis import fake_redis_cache
def _store(server: fakeredis.FakeServer) -> TransportSpendUpdateBufferStore:
return TransportSpendUpdateBufferStore(fake_redis_cache(server, namespace="ns").transport())
@pytest.mark.asyncio
async def test_batches_land_on_the_legacy_lists_as_safe_dumps_json(server: fakeredis.FakeServer) -> None:
store: Final = _store(server)
raw: Final = fakeredis.FakeAsyncRedis(server=server)
assert await store.push("daily_tag", {"k": {"spend": 1.5}}) == 1
assert await store.push_all((("spend", {"key_list_transactions": {"a": 0.5}}), ("daily_tag", {"k": {}}))) == (1, 2)
assert await raw.lrange(f"ns:{REDIS_UPDATE_BUFFER_KEY}", 0, -1) == [b'{"key_list_transactions": {"a": 0.5}}']
assert await raw.lrange(f"ns:{REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY}", 0, -1) == [
b'{"k": {"spend": 1.5}}',
b'{"k": {}}',
]
assert await store.pop("daily_tag", count=1) == ('{"k": {"spend": 1.5}}',)
assert await store.pop_all(("spend", "daily_tag", "window"), count=10) == (
('{"key_list_transactions": {"a": 0.5}}',),
('{"k": {}}',),
None,
)
assert await raw.exists(f"ns:{REDIS_UPDATE_BUFFER_KEY}", f"ns:{REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY}") == 0
@pytest.mark.asyncio
async def test_a_rejected_push_fails_the_batch_but_the_other_lists_keep_their_push(
server: fakeredis.FakeServer,
) -> None:
raw: Final = fakeredis.FakeAsyncRedis(server=server)
await raw.set(f"ns:{REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY}", "not a list")
store: Final = _store(server)
pushed: Final = await store.push_all((("spend", {"a": 1}), ("daily_tag", {"b": 2})))
assert isinstance(pushed, RedisFailure) and pushed.kind == "server"
assert await raw.lrange(f"ns:{REDIS_UPDATE_BUFFER_KEY}", 0, -1) == [b'{"a": 1}']
@pytest.mark.asyncio
async def test_undecodable_entries_and_outages_are_failures(server: fakeredis.FakeServer) -> None:
raw: Final = fakeredis.FakeAsyncRedis(server=server)
await raw.rpush(f"ns:{REDIS_UPDATE_BUFFER_KEY}", b"\xff\xfe")
store: Final = _store(server)
undecodable: Final = await store.pop("spend", count=10)
server.connected = False
outage: Final = await store.pop_all(("spend",), count=10)
assert isinstance(undecodable, RedisFailure) and undecodable.kind == "decode"
assert isinstance(outage, RedisFailure) and outage.kind == "transport"