mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
fix(proxy): keep Redis spend buffer items recoverable after a drain reply timeout
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin
parent
eb114e638d
commit
6b55fa8901
11 changed files with 864 additions and 106 deletions
|
|
@ -20,7 +20,7 @@ import threading
|
|||
import time
|
||||
import weakref
|
||||
from collections import Counter
|
||||
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
|
||||
from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence
|
||||
from contextvars import ContextVar
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
|
|
@ -72,6 +72,43 @@ else:
|
|||
Span = Any
|
||||
|
||||
|
||||
class _AsyncRedisListPipeline(Protocol):
|
||||
async def __aenter__(self) -> "_AsyncRedisListPipeline": ...
|
||||
|
||||
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> None: ...
|
||||
|
||||
def lmove(self, first_list: str, second_list: str, src: str, dest: str) -> "_AsyncRedisListPipeline": ...
|
||||
|
||||
def lpop(self, name: str) -> "_AsyncRedisListPipeline": ...
|
||||
|
||||
def expire(self, name: str, time: int) -> "_AsyncRedisListPipeline": ...
|
||||
|
||||
async def execute(self) -> list[object]: ...
|
||||
|
||||
|
||||
_LPOP_TO_PROCESSING_SCRIPT: Final = """
|
||||
local moved = {}
|
||||
for i = 1, tonumber(ARGV[1]) do
|
||||
local v = redis.call('LPOP', KEYS[1])
|
||||
if not v then break end
|
||||
redis.call('RPUSH', KEYS[2], v)
|
||||
moved[#moved + 1] = v
|
||||
end
|
||||
redis.call('EXPIRE', KEYS[2], tonumber(ARGV[2]))
|
||||
return moved
|
||||
"""
|
||||
|
||||
|
||||
class _RedisServiceLogger(Protocol):
|
||||
def async_service_success_hook(
|
||||
self, service: ServiceTypes, call_type: str, duration: float, caller: str
|
||||
) -> Coroutine[object, object, None]: ...
|
||||
|
||||
def async_service_failure_hook(
|
||||
self, service: ServiceTypes, duration: float, error: Exception, call_type: str, caller: str
|
||||
) -> Coroutine[object, object, None]: ...
|
||||
|
||||
|
||||
class _AsyncRedisCommands(Protocol):
|
||||
"""Async redis commands this cache issues.
|
||||
|
||||
|
|
@ -91,6 +128,8 @@ class _AsyncRedisCommands(Protocol):
|
|||
|
||||
def lpop(self, name: str, count: int | None = None) -> Awaitable[object]: ...
|
||||
|
||||
def lrange(self, name: str, start: int, end: int) -> Awaitable[list[bytes | str]]: ...
|
||||
|
||||
def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ...
|
||||
|
||||
def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ...
|
||||
|
|
@ -135,6 +174,7 @@ _INCREMENT_WITH_FLOOR_LUA: Final = (
|
|||
)
|
||||
|
||||
_LUA_COUNT: Final = TypeAdapter(int)
|
||||
_LUA_LIST_RESULT: Final = TypeAdapter(list[bytes | str | None])
|
||||
_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...])
|
||||
|
||||
|
||||
|
|
@ -144,6 +184,14 @@ def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, .
|
|||
)
|
||||
|
||||
|
||||
def _decode_redis_list(values: Sequence[bytes | str | None]) -> list[str]:
|
||||
return [
|
||||
value.decode("utf-8", errors="replace") if isinstance(value, bytes) else value
|
||||
for value in values
|
||||
if value is not None
|
||||
]
|
||||
|
||||
|
||||
def _is_generic_caller_frame(frame: FrameType) -> bool:
|
||||
module: Final = frame.f_globals.get("__name__")
|
||||
return module in _GENERIC_CALLER_MODULES or frame.f_code.co_name in _GENERIC_CALLER_FRAME_NAMES
|
||||
|
|
@ -680,7 +728,7 @@ class RedisCache(BaseCache):
|
|||
self.redis_flush_size: int = 100
|
||||
else:
|
||||
self.redis_flush_size = redis_flush_size
|
||||
self.redis_version: str = "Unknown"
|
||||
self.redis_version: str | int | float = "Unknown"
|
||||
try:
|
||||
if not coroutine_checker.is_async_callable(self.redis_client):
|
||||
self.redis_version = self.redis_client.info()["redis_version"]
|
||||
|
|
@ -2406,6 +2454,162 @@ class RedisCache(BaseCache):
|
|||
decoded_results.append(None)
|
||||
return decoded_results
|
||||
|
||||
@functools.cached_property
|
||||
def _lpop_to_processing_script(self) -> Callable[..., Awaitable[object]]:
|
||||
return self.async_register_script(_LPOP_TO_PROCESSING_SCRIPT)
|
||||
|
||||
async def _async_move_list_to_processing(
|
||||
self,
|
||||
redis_client: _AsyncRedisCommands,
|
||||
namespaced_key: str,
|
||||
processing_key: str,
|
||||
count: int,
|
||||
ttl_seconds: int,
|
||||
) -> list[str]:
|
||||
major_version: Final = self._parse_redis_major_version()
|
||||
version_parts: Final = str(self.redis_version).strip().split(".")
|
||||
minor_part: Final = (
|
||||
version_parts[1].split("-", maxsplit=1)[0] if major_version == 6 and len(version_parts) > 1 else ""
|
||||
)
|
||||
minor_version: Final = int(minor_part) if minor_part.isdigit() else 0
|
||||
use_lmove: Final = major_version > 6 or (major_version == 6 and minor_version >= 2)
|
||||
if not use_lmove:
|
||||
script_results: Final = await self._lpop_to_processing_script(
|
||||
keys=(namespaced_key, processing_key),
|
||||
args=(count, ttl_seconds),
|
||||
)
|
||||
return _decode_redis_list(_LUA_LIST_RESULT.validate_python(script_results))
|
||||
|
||||
list_pipeline: Final = cast( # cast-ok: RedisCluster stubs omit list commands
|
||||
_AsyncRedisListPipeline,
|
||||
redis_client.pipeline(transaction=False),
|
||||
)
|
||||
async with list_pipeline as pipe:
|
||||
for _ in range(count):
|
||||
pipe.lmove(namespaced_key, processing_key, src="LEFT", dest="RIGHT")
|
||||
pipe.expire(processing_key, ttl_seconds)
|
||||
raw_results: Final = await pipe.execute()
|
||||
|
||||
filtered_results: Final = tuple(
|
||||
result for result in raw_results[:-1] if result is None or isinstance(result, (bytes, str))
|
||||
)
|
||||
return _decode_redis_list(filtered_results)
|
||||
|
||||
def _processing_list_key(self, namespaced_key: str, worker_id: str) -> str:
|
||||
return f"{namespaced_key}:processing:{worker_id}"
|
||||
|
||||
@_redis_circuit_breaker_guard
|
||||
async def async_drain_list_to_processing(self, key: str, worker_id: str, count: int, ttl_seconds: int) -> list[str]:
|
||||
_redis_client: Final = self._async_commands()
|
||||
namespaced_key: Final = self.check_and_fix_namespace(key)
|
||||
processing_key: Final = self._processing_list_key(
|
||||
namespaced_key=namespaced_key,
|
||||
worker_id=worker_id,
|
||||
)
|
||||
start_time: Final = time.time()
|
||||
service_logger: Final = cast( # cast-ok: logger has the required hooks
|
||||
_RedisServiceLogger,
|
||||
self.service_logger_obj,
|
||||
)
|
||||
|
||||
try:
|
||||
pending_items: Final = _decode_redis_list(await _redis_client.lrange(processing_key, 0, -1))
|
||||
drained_items: Final = (
|
||||
pending_items
|
||||
if pending_items
|
||||
else await self._async_move_list_to_processing(
|
||||
redis_client=_redis_client,
|
||||
namespaced_key=namespaced_key,
|
||||
processing_key=processing_key,
|
||||
count=count,
|
||||
ttl_seconds=ttl_seconds,
|
||||
)
|
||||
)
|
||||
end_time: Final = time.time()
|
||||
_duration: Final = end_time - start_time
|
||||
asyncio.create_task(
|
||||
service_logger.async_service_success_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=_duration,
|
||||
call_type="async_drain_list_to_processing",
|
||||
caller=_get_call_stack_info(),
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
failure_end_time: Final = time.time()
|
||||
failure_duration: Final = failure_end_time - start_time
|
||||
asyncio.create_task(
|
||||
service_logger.async_service_failure_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=failure_duration,
|
||||
error=e,
|
||||
call_type="async_drain_list_to_processing",
|
||||
caller=_get_call_stack_info(),
|
||||
)
|
||||
)
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async_drain_list_to_processing() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
raise
|
||||
else:
|
||||
return drained_items
|
||||
|
||||
@_redis_circuit_breaker_guard
|
||||
async def async_ack_processing_list(self, key: str, worker_id: str) -> None:
|
||||
_redis_client: Final = self._async_commands()
|
||||
namespaced_key: Final = self.check_and_fix_namespace(key)
|
||||
processing_key: Final = self._processing_list_key(
|
||||
namespaced_key=namespaced_key,
|
||||
worker_id=worker_id,
|
||||
)
|
||||
start_time: Final = time.time()
|
||||
service_logger: Final = cast( # cast-ok: logger has the required hooks
|
||||
_RedisServiceLogger,
|
||||
self.service_logger_obj,
|
||||
)
|
||||
|
||||
try:
|
||||
for attempt in range(3):
|
||||
try:
|
||||
await _redis_client.delete(processing_key)
|
||||
break
|
||||
except Exception:
|
||||
if attempt == 2:
|
||||
raise
|
||||
|
||||
end_time: Final = time.time()
|
||||
_duration: Final = end_time - start_time
|
||||
asyncio.create_task(
|
||||
service_logger.async_service_success_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=_duration,
|
||||
call_type="async_ack_processing_list",
|
||||
caller=_get_call_stack_info(),
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
failure_end_time: Final = time.time()
|
||||
failure_duration: Final = failure_end_time - start_time
|
||||
asyncio.create_task(
|
||||
service_logger.async_service_failure_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=failure_duration,
|
||||
error=e,
|
||||
call_type="async_ack_processing_list",
|
||||
caller=_get_call_stack_info(),
|
||||
)
|
||||
)
|
||||
log_redis_failure(
|
||||
verbose_logger,
|
||||
logging.ERROR,
|
||||
"LiteLLM Redis Caching: async_ack_processing_list() - Got exception from REDIS",
|
||||
e,
|
||||
)
|
||||
raise
|
||||
|
||||
@_redis_circuit_breaker_guard
|
||||
async def async_lpop_pipeline(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -126,6 +126,10 @@ class RedisClusterCache(RedisCache):
|
|||
return call()
|
||||
return self._connection_gate.run(call, self._get_async_client_cache_key())
|
||||
|
||||
def _processing_list_key(self, namespaced_key: str, worker_id: str) -> str:
|
||||
namespace_prefix: Final = f"{self.namespace}:" if self.namespace else ""
|
||||
return f"{namespace_prefix}{{{namespaced_key}}}:processing:{worker_id}"
|
||||
|
||||
def _run_redis_mget_operation(self, keys: list[str]) -> Sequence[bytes | str | None]:
|
||||
"""
|
||||
Overrides `_run_redis_mget_operation` in redis_cache.py
|
||||
|
|
|
|||
|
|
@ -453,6 +453,7 @@ REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_up
|
|||
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer"
|
||||
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer"
|
||||
MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = get_env_int("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)
|
||||
REDIS_BUFFER_PROCESSING_LIST_TTL_SECONDS: Final = 86400
|
||||
REDIS_SPEND_LOGS_BUFFER_KEY: Final = "litellm_spend_logs_buffer"
|
||||
REDIS_SPEND_LOGS_BUFFER_MAX_ROWS: Final = 100000
|
||||
REDIS_SPEND_LOGS_BUFFER_DEQUEUE_COUNT: Final = 1000
|
||||
|
|
|
|||
|
|
@ -218,6 +218,8 @@ _SERVICE_VERB_BY_CALL_TYPE: Final[dict[str, str]] = {
|
|||
"async_delete_cache": "delete",
|
||||
"async_rpush": "rpush",
|
||||
"async_lpop": "lpop",
|
||||
"async_drain_list_to_processing": "lmove",
|
||||
"async_ack_processing_list": "delete",
|
||||
"async_scan_iter": "scan",
|
||||
"async_lpop_pipeline": "lpop",
|
||||
"async_rpush_pipeline": "rpush",
|
||||
|
|
|
|||
|
|
@ -6,9 +6,12 @@ This is to prevent deadlocks and improve reliability
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from contextlib import AsyncExitStack
|
||||
from datetime import datetime
|
||||
from functools import reduce
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast
|
||||
|
||||
from redis.exceptions import RedisError
|
||||
|
|
@ -18,6 +21,7 @@ from litellm._logging import verbose_proxy_logger
|
|||
from litellm.caching import RedisCache
|
||||
from litellm.constants import (
|
||||
MAX_REDIS_BUFFER_DEQUEUE_COUNT,
|
||||
REDIS_BUFFER_PROCESSING_LIST_TTL_SECONDS,
|
||||
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
|
||||
|
|
@ -55,7 +59,6 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
|||
from litellm.proxy.db.spend_log_batching import SpendLogRow
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.caching import (
|
||||
RedisPipelineLpopOperation,
|
||||
RedisPipelineRpushOperation,
|
||||
)
|
||||
from litellm.types.services import ServiceTypes
|
||||
|
|
@ -145,6 +148,26 @@ class RedisUpdateBuffer:
|
|||
redis_cache: RedisCache | None = None,
|
||||
):
|
||||
self.redis_cache = redis_cache
|
||||
self._drain_worker_id: Final = str(uuid.uuid4())
|
||||
self._pipeline_drain_keys: Final = (
|
||||
REDIS_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
|
||||
)
|
||||
self._drain_locks: Final = MappingProxyType(
|
||||
{
|
||||
key: asyncio.Lock()
|
||||
for key in (
|
||||
*self._pipeline_drain_keys,
|
||||
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_SPEND_LOGS_BUFFER_KEY,
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def should_commit_spend_updates_to_redis() -> bool:
|
||||
|
|
@ -503,11 +526,11 @@ class RedisUpdateBuffer:
|
|||
window_spend_update_transactions: Sequence[WindowSpendTransaction] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Re-push transactions that were popped from Redis but not committed to the DB.
|
||||
Re-push transactions drained from Redis but not committed to the DB.
|
||||
|
||||
The leader drains the buffers with a destructive ``lpop`` before committing to
|
||||
the database. When a commit fails after its retries are exhausted, the popped
|
||||
transactions must be pushed back so a later scheduler tick can retry them;
|
||||
The leader removes drained transactions from Redis before committing to the
|
||||
database. When a commit fails after its retries are exhausted, the transactions
|
||||
must be pushed back so a later scheduler tick can retry them;
|
||||
otherwise the aggregated spend is lost permanently. The re-pushed payloads use
|
||||
the same JSON encoding as the store path, so the next drain parses them normally.
|
||||
"""
|
||||
|
|
@ -581,19 +604,40 @@ class RedisUpdateBuffer:
|
|||
verbose_proxy_logger.info("Spend tracking - parked %d spend log rows in Redis for a later flush", len(rows))
|
||||
return True
|
||||
|
||||
async def _drain_to_processing(
|
||||
self,
|
||||
key: str,
|
||||
count: int,
|
||||
) -> list[str] | None: # mutable-ok: preserve the Redis list drain contract
|
||||
if self.redis_cache is None:
|
||||
return None
|
||||
drained_items: Final = await self.redis_cache.async_drain_list_to_processing(
|
||||
key=key,
|
||||
worker_id=self._drain_worker_id,
|
||||
count=count,
|
||||
ttl_seconds=REDIS_BUFFER_PROCESSING_LIST_TTL_SECONDS,
|
||||
)
|
||||
return drained_items if drained_items else None
|
||||
|
||||
async def _drain(self, key: str, count: int) -> list[str] | None:
|
||||
if self.redis_cache is None:
|
||||
return None
|
||||
async with self._drain_locks[key]:
|
||||
drained_items: Final = await self._drain_to_processing(key=key, count=count)
|
||||
if drained_items is None:
|
||||
return None
|
||||
await self.redis_cache.async_ack_processing_list(key=key, worker_id=self._drain_worker_id)
|
||||
return drained_items
|
||||
|
||||
@with_service_target(SPEND_QUEUE_TARGET)
|
||||
async def get_spend_logs_from_redis_buffer(self, limit: int) -> tuple[dict[str, object], ...]:
|
||||
"""Atomically take up to ``limit`` parked spend-log rows out of Redis."""
|
||||
"""Drain up to ``limit`` parked spend-log rows from 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,
|
||||
)
|
||||
if popped is None:
|
||||
drained_items: Final = await self._drain(key=REDIS_SPEND_LOGS_BUFFER_KEY, count=limit)
|
||||
if drained_items 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 drained_items)
|
||||
return tuple(row for row in decoded_rows if row is not None)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -656,7 +700,7 @@ class RedisUpdateBuffer:
|
|||
"""
|
||||
if self.redis_cache is None:
|
||||
return None
|
||||
list_of_transactions: Final[str | list[str] | None] = await self.redis_cache.async_lpop(
|
||||
list_of_transactions: Final = await self._drain(
|
||||
key=REDIS_UPDATE_BUFFER_KEY,
|
||||
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
|
||||
)
|
||||
|
|
@ -666,7 +710,7 @@ class RedisUpdateBuffer:
|
|||
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,
|
||||
)
|
||||
|
||||
|
|
@ -695,7 +739,7 @@ class RedisUpdateBuffer:
|
|||
tuple[WindowSpendTransaction, ...] | None,
|
||||
]:
|
||||
"""
|
||||
Drains the main 7 Redis buffer queues in a single pipeline round-trip.
|
||||
Concurrently drains the main 7 Redis buffer queues.
|
||||
|
||||
Returns a 7-tuple of parsed results in this order:
|
||||
0: DBSpendUpdateTransactions
|
||||
|
|
@ -709,63 +753,77 @@ 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)
|
||||
|
||||
# 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])
|
||||
if len(parsed) > 0:
|
||||
db_spend = self._combine_list_of_transactions(parsed)
|
||||
|
||||
# Slots 1-5: daily spend categories
|
||||
daily_results: Final[list[dict[str, BaseDailySpendTransaction] | None]] = []
|
||||
for slot in range(1, 6):
|
||||
slot_result = raw_results[slot]
|
||||
if slot_result is None:
|
||||
daily_results.append(None)
|
||||
else:
|
||||
list_of_daily = [json.loads(t) for t in slot_result]
|
||||
aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(list_of_daily)
|
||||
daily_results.append(aggregated)
|
||||
async with AsyncExitStack() as drain_lock_stack:
|
||||
for key in self._pipeline_drain_keys:
|
||||
await drain_lock_stack.enter_async_context(self._drain_locks[key])
|
||||
drain_results_with_errors: Final = tuple(
|
||||
await asyncio.gather(
|
||||
*(
|
||||
self._drain_to_processing(key=key, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT)
|
||||
for key in self._pipeline_drain_keys
|
||||
),
|
||||
return_exceptions=True,
|
||||
)
|
||||
)
|
||||
drain_errors: Final = tuple(
|
||||
result for result in drain_results_with_errors if isinstance(result, BaseException)
|
||||
)
|
||||
if drain_errors:
|
||||
raise drain_errors[0]
|
||||
raw_results: Final = tuple(
|
||||
result for result in drain_results_with_errors if not isinstance(result, BaseException)
|
||||
)
|
||||
acknowledgement_indexes: Final = tuple(
|
||||
index for index, items in enumerate(raw_results) if items is not None
|
||||
)
|
||||
acknowledgement_results: Final = tuple(
|
||||
await asyncio.gather(
|
||||
*(
|
||||
self.redis_cache.async_ack_processing_list(
|
||||
key=self._pipeline_drain_keys[index],
|
||||
worker_id=self._drain_worker_id,
|
||||
)
|
||||
for index in acknowledgement_indexes
|
||||
),
|
||||
return_exceptions=True,
|
||||
)
|
||||
)
|
||||
failed_acknowledgements: Final = tuple(
|
||||
(index, result)
|
||||
for index, result in zip(acknowledgement_indexes, acknowledgement_results)
|
||||
if isinstance(result, BaseException)
|
||||
)
|
||||
for index, error in failed_acknowledgements:
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - failed to acknowledge Redis buffer drain (key=%s). Items will be retried: %s",
|
||||
self._pipeline_drain_keys[index],
|
||||
error,
|
||||
)
|
||||
failed_acknowledgement_indexes: Final = frozenset(index for index, _ in failed_acknowledgements)
|
||||
results_ready_for_parse: Final = tuple(
|
||||
None if index in failed_acknowledgement_indexes else items for index, items in enumerate(raw_results)
|
||||
)
|
||||
|
||||
parsed_transactions: Final = (
|
||||
self._parse_list_of_transactions(results_ready_for_parse[0])
|
||||
if results_ready_for_parse[0] is not None
|
||||
else ()
|
||||
)
|
||||
db_spend: Final = self._combine_list_of_transactions(parsed_transactions) if parsed_transactions else None
|
||||
daily_results: Final = tuple(
|
||||
None
|
||||
if transactions is None
|
||||
else DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(
|
||||
[json.loads(transaction) for transaction in transactions]
|
||||
)
|
||||
for transactions in results_ready_for_parse[1:6]
|
||||
)
|
||||
window_transactions: Final = results_ready_for_parse[6]
|
||||
window_spend: Final = (
|
||||
WindowSpendUpdateQueue.get_aggregated_window_spend_transactions(
|
||||
tuple(json.loads(transaction) for transaction in raw_results[6])
|
||||
tuple(json.loads(transaction) for transaction in window_transactions)
|
||||
)
|
||||
if raw_results[6] is not None
|
||||
if window_transactions is not None
|
||||
else None
|
||||
)
|
||||
|
||||
|
|
@ -801,11 +859,11 @@ class RedisUpdateBuffer:
|
|||
redis_key: str,
|
||||
) -> list[dict[str, BaseDailySpendTransaction]] | None:
|
||||
"""
|
||||
Drains a daily spend buffer key and parses each popped item as JSON.
|
||||
Drains a daily spend buffer key and parses each item as JSON.
|
||||
"""
|
||||
if self.redis_cache is None:
|
||||
return None
|
||||
list_of_transactions: Final[list[str] | None] = await self.redis_cache.async_lpop(
|
||||
list_of_transactions: Final = await self._drain(
|
||||
key=redis_key,
|
||||
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -8,6 +8,8 @@ from typing import Final, cast
|
|||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fakeredis import FakeServer
|
||||
from fakeredis.aioredis import FakeRedis
|
||||
|
||||
import litellm
|
||||
from litellm._service_logger import ServiceLogging
|
||||
|
|
@ -23,6 +25,22 @@ def redis_no_ping():
|
|||
yield
|
||||
|
||||
|
||||
def _cache_with_fakeredis(
|
||||
monkeypatch: pytest.MonkeyPatch, namespace: str, redis_version: str | float = "7.0.0"
|
||||
) -> tuple[RedisCache, FakeRedis]:
|
||||
redis_client: Final = FakeRedis(server=FakeServer())
|
||||
sync_redis_client: Final = MagicMock()
|
||||
sync_redis_client.info.return_value = {"redis_version": redis_version}
|
||||
with (
|
||||
patch("asyncio.get_running_loop", side_effect=RuntimeError("No running event loop")),
|
||||
patch("litellm._redis.get_redis_client", return_value=sync_redis_client),
|
||||
patch("litellm._redis.get_redis_connection_pool", return_value=MagicMock()),
|
||||
):
|
||||
redis_cache: Final = RedisCache(namespace=namespace)
|
||||
monkeypatch.setattr(redis_cache, "init_async_client", lambda: redis_client)
|
||||
return redis_cache, redis_client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sync_batch_redis_cache(redis_no_ping):
|
||||
with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point
|
||||
|
|
@ -59,6 +77,101 @@ def test_check_and_fix_namespace_prefixes_keys_sharing_the_namespace_prefix(
|
|||
assert redis_cache.check_and_fix_namespace(key=key) == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_drain_list_to_processing_respects_count_pending_order_and_namespace(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
redis_cache, redis_client = _cache_with_fakeredis(monkeypatch, namespace="tenant")
|
||||
await redis_client.rpush("tenant:buffer", "first", "second", "third")
|
||||
|
||||
drained: Final = await redis_cache.async_drain_list_to_processing(
|
||||
key="buffer",
|
||||
worker_id="worker-1",
|
||||
count=2,
|
||||
ttl_seconds=60,
|
||||
)
|
||||
|
||||
processing_key: Final = "tenant:buffer:processing:worker-1"
|
||||
assert drained == ["first", "second"]
|
||||
assert await redis_client.lrange(processing_key, 0, -1) == [b"first", b"second"]
|
||||
await redis_client.rpush("tenant:buffer", "newer")
|
||||
|
||||
pending: Final = await redis_cache.async_drain_list_to_processing(
|
||||
key="buffer",
|
||||
worker_id="worker-1",
|
||||
count=2,
|
||||
ttl_seconds=60,
|
||||
)
|
||||
|
||||
assert pending == ["first", "second"]
|
||||
assert await redis_client.lrange("tenant:buffer", 0, -1) == [b"third", b"newer"]
|
||||
redis_keys: Final = await redis_client.keys("*")
|
||||
assert set(redis_keys) == {b"tenant:buffer", processing_key.encode()}
|
||||
assert all(redis_key.startswith(b"tenant:") for redis_key in redis_keys)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_drain_list_to_processing_replaces_invalid_utf8_and_acknowledges(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
redis_cache, redis_client = _cache_with_fakeredis(monkeypatch, namespace="tenant")
|
||||
await redis_client.rpush("tenant:buffer", b"\xff", b"valid")
|
||||
|
||||
drained: Final = await redis_cache.async_drain_list_to_processing(
|
||||
key="buffer",
|
||||
worker_id="worker-utf8",
|
||||
count=2,
|
||||
ttl_seconds=60,
|
||||
)
|
||||
assert drained == ["\ufffd", "valid"]
|
||||
|
||||
await redis_cache.async_ack_processing_list(key="buffer", worker_id="worker-utf8")
|
||||
assert (
|
||||
await redis_cache.async_drain_list_to_processing(
|
||||
key="buffer",
|
||||
worker_id="worker-utf8",
|
||||
count=2,
|
||||
ttl_seconds=60,
|
||||
)
|
||||
== []
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("redis_version", "expected_items", "expected_processing", "expected_buffer"),
|
||||
[
|
||||
("6.2.0", ("first", "second"), (b"first", b"second"), (b"third",)),
|
||||
(7.2, ("first", "second"), (b"first", b"second"), (b"third",)),
|
||||
(6.2, ("first", "second"), (b"first", b"second"), (b"third",)),
|
||||
],
|
||||
)
|
||||
async def test_async_drain_list_to_processing_preserves_pending_items_for_redis_version(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
redis_version: str | float,
|
||||
expected_items: tuple[str, ...],
|
||||
expected_processing: tuple[bytes, ...],
|
||||
expected_buffer: tuple[bytes, ...],
|
||||
) -> None:
|
||||
redis_cache, redis_client = _cache_with_fakeredis(
|
||||
monkeypatch,
|
||||
namespace="tenant",
|
||||
redis_version=redis_version,
|
||||
)
|
||||
await redis_client.rpush("tenant:buffer", "first", "second", "third")
|
||||
|
||||
drained: Final = await redis_cache.async_drain_list_to_processing(
|
||||
key="buffer",
|
||||
worker_id="worker-1",
|
||||
count=2,
|
||||
ttl_seconds=60,
|
||||
)
|
||||
|
||||
assert drained == list(expected_items)
|
||||
assert tuple(await redis_client.lrange("tenant:buffer:processing:worker-1", 0, -1)) == expected_processing
|
||||
assert tuple(await redis_client.lrange("tenant:buffer", 0, -1)) == expected_buffer
|
||||
|
||||
|
||||
@pytest.mark.parametrize("namespace", [None, "litellm"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_delete_cache_applies_namespace(namespace, monkeypatch, redis_no_ping):
|
||||
|
|
@ -1033,8 +1146,9 @@ def test_prometheus_services_logger_starts_from_the_breakers_already_built(monke
|
|||
monkeypatch.setattr(litellm, "service_callback", ["prometheus_system"])
|
||||
gc.collect()
|
||||
ServiceLogging()
|
||||
assert REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"}) == (
|
||||
redis_breaker_state_counts()["closed"]
|
||||
assert (
|
||||
REGISTRY.get_sample_value("litellm_redis_circuit_breaker_state", {"state": "closed"})
|
||||
== (redis_breaker_state_counts()["closed"])
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -5,13 +5,14 @@ import ssl
|
|||
import weakref
|
||||
from collections.abc import Awaitable, Callable
|
||||
from importlib import import_module
|
||||
from typing import cast
|
||||
from typing import Final, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from redis.asyncio import Redis, RedisCluster
|
||||
from redis.asyncio.cluster import ClusterNode
|
||||
from redis.asyncio.connection import SSLConnection
|
||||
from redis.crc import key_slot
|
||||
|
||||
from litellm.caching.evicted_client_closer import EvictedClientCloser
|
||||
from litellm.caching.llm_caching_handler import LLMClientCache
|
||||
|
|
@ -44,6 +45,31 @@ def test_redis_cluster_batch_get(mock_init_redis_cluster):
|
|||
assert not mock_redis.mget.called
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("namespace", "expected_source_key", "expected_processing_prefix"),
|
||||
[
|
||||
(None, "buffer", "{buffer}:"),
|
||||
("tenant", "tenant:buffer", "tenant:{tenant:buffer}:"),
|
||||
],
|
||||
)
|
||||
def test_redis_cluster_processing_list_key_preserves_namespace_and_hash_slot(
|
||||
namespace: str | None,
|
||||
expected_source_key: str,
|
||||
expected_processing_prefix: str,
|
||||
) -> None:
|
||||
with patch("litellm._redis.init_redis_cluster", return_value=MagicMock()):
|
||||
cache: Final = RedisClusterCache(
|
||||
startup_nodes=[{"host": "localhost", "port": 6379}],
|
||||
namespace=namespace,
|
||||
)
|
||||
|
||||
source_namespaced_key: Final = cache.check_and_fix_namespace("buffer")
|
||||
processing_key: Final = cache._processing_list_key(source_namespaced_key, "worker-1")
|
||||
assert source_namespaced_key == expected_source_key
|
||||
assert processing_key.startswith(expected_processing_prefix)
|
||||
assert key_slot(processing_key.encode()) == key_slot(source_namespaced_key.encode())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch("litellm._redis.init_redis_cluster")
|
||||
async def test_redis_cluster_async_batch_get(mock_init_redis_cluster):
|
||||
|
|
|
|||
|
|
@ -1,11 +1,17 @@
|
|||
import asyncio
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from itertools import count
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fakeredis import FakeServer
|
||||
from fakeredis.aioredis import FakeRedis
|
||||
from redis.asyncio.client import Pipeline
|
||||
from redis.exceptions import TimeoutError as RedisTimeoutError
|
||||
|
||||
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
|
||||
from litellm.proxy.proxy_server import ProxyStartupEvent
|
||||
|
||||
|
|
@ -23,6 +29,65 @@ def redis_update_buffer(mock_redis_cache):
|
|||
return RedisUpdateBuffer(redis_cache=mock_redis_cache)
|
||||
|
||||
|
||||
def _cache_with_fakeredis(monkeypatch: pytest.MonkeyPatch) -> tuple[RedisCache, FakeRedis]:
|
||||
redis_client: Final = FakeRedis(server=FakeServer())
|
||||
sync_redis_client: Final = MagicMock()
|
||||
sync_redis_client.info.return_value = {"redis_version": "7.0.0"}
|
||||
with (
|
||||
patch("asyncio.get_running_loop", side_effect=RuntimeError("no running event loop")),
|
||||
patch("litellm._redis.get_redis_client", return_value=sync_redis_client),
|
||||
patch("litellm._redis.get_redis_connection_pool", return_value=MagicMock()),
|
||||
):
|
||||
redis_cache: Final = RedisCache(namespace="litellm")
|
||||
monkeypatch.setattr(redis_cache, "init_async_client", lambda: redis_client)
|
||||
return redis_cache, redis_client
|
||||
|
||||
|
||||
def _timeout_first_pipeline_execute(monkeypatch: pytest.MonkeyPatch, redis_client: FakeRedis) -> None:
|
||||
original_pipeline: Final = redis_client.pipeline
|
||||
pipeline_calls: Final = count()
|
||||
|
||||
def pipeline_with_timeout(transaction: bool = True) -> Pipeline:
|
||||
pipeline: Final = original_pipeline(transaction=transaction)
|
||||
if next(pipeline_calls) == 0:
|
||||
original_execute: Final = pipeline.execute
|
||||
|
||||
async def execute_then_timeout() -> list[object]:
|
||||
await original_execute()
|
||||
raise RedisTimeoutError("simulated response timeout")
|
||||
|
||||
monkeypatch.setattr(pipeline, "execute", execute_then_timeout)
|
||||
return pipeline
|
||||
|
||||
monkeypatch.setattr(redis_client, "pipeline", pipeline_with_timeout)
|
||||
|
||||
|
||||
def _hold_first_pipeline_until_release(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
redis_client: FakeRedis,
|
||||
move_executed: asyncio.Event,
|
||||
release_move: asyncio.Event,
|
||||
) -> None:
|
||||
original_pipeline: Final = redis_client.pipeline
|
||||
pipeline_calls: Final = count()
|
||||
|
||||
def pipeline_with_timeout(transaction: bool = True) -> Pipeline:
|
||||
pipeline: Final = original_pipeline(transaction=transaction)
|
||||
if next(pipeline_calls) == 0:
|
||||
original_execute: Final = pipeline.execute
|
||||
|
||||
async def execute_after_release() -> list[object]:
|
||||
results: Final = await original_execute()
|
||||
move_executed.set()
|
||||
await release_move.wait()
|
||||
return results
|
||||
|
||||
monkeypatch.setattr(pipeline, "execute", execute_after_release)
|
||||
return pipeline
|
||||
|
||||
monkeypatch.setattr(redis_client, "pipeline", pipeline_with_timeout)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache):
|
||||
"""
|
||||
|
|
@ -179,7 +244,7 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_updat
|
|||
async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buffer, mock_redis_cache):
|
||||
"""
|
||||
Verify get_all_transactions_from_redis_buffer_pipeline correctly parses
|
||||
and aggregates results from async_lpop_pipeline.
|
||||
and aggregates results from each buffer drain.
|
||||
"""
|
||||
# Simulate pipeline results: slot 0 = spend updates, slots 1-5 = daily categories,
|
||||
# slot 6 = budget window spend
|
||||
|
|
@ -209,17 +274,18 @@ 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
|
||||
]
|
||||
mock_redis_cache.async_drain_list_to_processing = AsyncMock(
|
||||
side_effect=(
|
||||
[db_spend_json],
|
||||
[daily_user_json],
|
||||
[daily_team_json],
|
||||
[],
|
||||
[],
|
||||
[],
|
||||
[window_spend_json, window_spend_json],
|
||||
)
|
||||
)
|
||||
mock_redis_cache.async_ack_processing_list = AsyncMock()
|
||||
|
||||
result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
|
||||
|
|
@ -259,12 +325,29 @@ 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
|
||||
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_TEAM_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_UPDATE_BUFFER_KEY,
|
||||
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
|
||||
drained_keys: Final = tuple(
|
||||
call.kwargs["key"] for call in mock_redis_cache.async_drain_list_to_processing.await_args_list
|
||||
)
|
||||
assert drained_keys == (
|
||||
REDIS_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY,
|
||||
)
|
||||
assert mock_redis_cache.async_ack_processing_list.await_count == 4
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -281,8 +364,8 @@ 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]
|
||||
mock_redis_cache.async_drain_list_to_processing = AsyncMock(
|
||||
side_effect=([pod_json, pod_json], [], [], [], [], [], [])
|
||||
)
|
||||
|
||||
(db_spend, *_rest) = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
|
|
@ -356,7 +439,7 @@ async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_b
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_restored_window_spend_transactions_drain_back_unchanged(redis_update_buffer, mock_redis_cache):
|
||||
"""A window commit that fails after the destructive lpop must be re-pushed
|
||||
"""A window commit that fails after a drain 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 (
|
||||
|
|
@ -380,8 +463,8 @@ async def test_restored_window_spend_transactions_drain_back_unchanged(redis_upd
|
|||
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"])]
|
||||
mock_redis_cache.async_drain_list_to_processing = AsyncMock(
|
||||
side_effect=([], [], [], [], [], [], list(rpush_list[0]["values"]))
|
||||
)
|
||||
drained = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
|
||||
|
|
@ -410,6 +493,190 @@ async def test_restore_transactions_to_redis_swallows_redis_error(redis_update_b
|
|||
mock_redis_cache.async_rpush_pipeline.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pipeline_drain_recovers_spend_after_move_reply_timeout(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from litellm.proxy._types import Litellm_EntityType
|
||||
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue
|
||||
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
|
||||
|
||||
redis_cache, redis_client = _cache_with_fakeredis(monkeypatch)
|
||||
buffer: Final = RedisUpdateBuffer(redis_cache=redis_cache)
|
||||
spend_queue: Final = SpendUpdateQueue()
|
||||
await spend_queue.add_update({"entity_type": Litellm_EntityType.KEY, "entity_id": "key-abc", "response_cost": 0.25})
|
||||
await spend_queue.add_update({"entity_type": Litellm_EntityType.KEY, "entity_id": "key-abc", "response_cost": 0.5})
|
||||
await spend_queue.add_update(
|
||||
{"entity_type": Litellm_EntityType.TEAM, "entity_id": "team-xyz", "response_cost": 1.25}
|
||||
)
|
||||
await buffer.store_in_memory_spend_updates_in_redis(
|
||||
spend_update_queue=spend_queue,
|
||||
daily_spend_update_queue=DailySpendUpdateQueue(),
|
||||
daily_team_spend_update_queue=DailySpendUpdateQueue(),
|
||||
daily_org_spend_update_queue=DailySpendUpdateQueue(),
|
||||
daily_end_user_spend_update_queue=DailySpendUpdateQueue(),
|
||||
daily_agent_spend_update_queue=DailySpendUpdateQueue(),
|
||||
)
|
||||
|
||||
_timeout_first_pipeline_execute(monkeypatch, redis_client)
|
||||
with pytest.raises(RedisTimeoutError, match="simulated response timeout"):
|
||||
await buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
|
||||
recovered: Final = await buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
assert recovered[0] is not None
|
||||
assert recovered[0]["key_list_transactions"] == {"key-abc": 0.75}
|
||||
assert recovered[0]["team_list_transactions"] == {"team-xyz": 1.25}
|
||||
assert recovered[1:] == (None, None, None, None, None, None)
|
||||
assert await buffer.get_all_transactions_from_redis_buffer_pipeline() == (None, None, None, None, None, None, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pipeline_drain_keeps_all_buffers_pending_after_partial_move_failure(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.constants import (
|
||||
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_UPDATE_BUFFER_KEY,
|
||||
)
|
||||
|
||||
redis_cache, redis_client = _cache_with_fakeredis(monkeypatch)
|
||||
buffer: Final = RedisUpdateBuffer(redis_cache=redis_cache)
|
||||
db_spend: Final = json.dumps({"key_list_transactions": {"key-partial": 1.25}})
|
||||
daily_spend: Final = json.dumps({"user-partial": {"spend": 2.5, "api_requests": 1}})
|
||||
await redis_cache.async_rpush(key=REDIS_UPDATE_BUFFER_KEY, values=(db_spend,))
|
||||
await redis_cache.async_rpush(key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, values=(daily_spend,))
|
||||
|
||||
_timeout_first_pipeline_execute(monkeypatch, redis_client)
|
||||
with pytest.raises(RedisTimeoutError, match="simulated response timeout"):
|
||||
await buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
|
||||
recovered: Final = await buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
assert recovered[0] is not None
|
||||
assert recovered[0]["key_list_transactions"] == {"key-partial": 1.25}
|
||||
assert recovered[1] == {"user-partial": {"spend": 2.5, "api_requests": 1}}
|
||||
assert recovered[2:] == (None, None, None, None, None)
|
||||
assert await buffer.get_all_transactions_from_redis_buffer_pipeline() == (None, None, None, None, None, None, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pipeline_drain_returns_only_keys_with_successful_acknowledgements(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.constants import (
|
||||
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
|
||||
REDIS_UPDATE_BUFFER_KEY,
|
||||
)
|
||||
|
||||
redis_cache, redis_client = _cache_with_fakeredis(monkeypatch)
|
||||
buffer: Final = RedisUpdateBuffer(redis_cache=redis_cache)
|
||||
db_spend: Final = json.dumps({"key_list_transactions": {"key-ack": 1.5}})
|
||||
daily_spend: Final = json.dumps({"user-ack": {"spend": 2.75, "api_requests": 1}})
|
||||
await redis_cache.async_rpush(key=REDIS_UPDATE_BUFFER_KEY, values=(db_spend,))
|
||||
await redis_cache.async_rpush(key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, values=(daily_spend,))
|
||||
|
||||
original_delete: Final = redis_client.delete
|
||||
|
||||
async def fail_main_buffer_delete(key: str) -> int:
|
||||
if REDIS_UPDATE_BUFFER_KEY in key:
|
||||
raise RedisTimeoutError("DEL failed")
|
||||
return await original_delete(key)
|
||||
|
||||
monkeypatch.setattr(redis_client, "delete", fail_main_buffer_delete)
|
||||
first_result: Final = await buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
assert first_result[0] is None
|
||||
assert first_result[1] == {"user-ack": {"spend": 2.75, "api_requests": 1}}
|
||||
assert first_result[2:] == (None, None, None, None, None)
|
||||
|
||||
monkeypatch.setattr(redis_client, "delete", original_delete)
|
||||
second_result: Final = await buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
assert second_result[0] is not None
|
||||
assert second_result[0]["key_list_transactions"] == {"key-ack": 1.5}
|
||||
assert second_result[1:] == (None, None, None, None, None, None)
|
||||
assert await buffer.get_all_transactions_from_redis_buffer_pipeline() == (None, None, None, None, None, None, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spend_logs_are_recovered_once_after_move_reply_timeout(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
redis_cache, redis_client = _cache_with_fakeredis(monkeypatch)
|
||||
buffer: Final = RedisUpdateBuffer(redis_cache=redis_cache)
|
||||
buffer.should_commit_spend_updates_to_redis = MagicMock(return_value=True)
|
||||
rows: Final = ({"request_id": "req-a"}, {"request_id": "req-b"})
|
||||
assert await buffer.store_spend_logs_in_redis(rows) is True
|
||||
|
||||
_timeout_first_pipeline_execute(monkeypatch, redis_client)
|
||||
with pytest.raises(RedisTimeoutError, match="simulated response timeout"):
|
||||
await buffer.get_spend_logs_from_redis_buffer(limit=10)
|
||||
|
||||
assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == rows
|
||||
assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_overlapping_spend_log_drains_serialize_and_recover_rows_once(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
redis_cache, redis_client = _cache_with_fakeredis(monkeypatch)
|
||||
buffer: Final = RedisUpdateBuffer(redis_cache=redis_cache)
|
||||
buffer.should_commit_spend_updates_to_redis = MagicMock(return_value=True)
|
||||
rows: Final = ({"request_id": "req-a"}, {"request_id": "req-b"})
|
||||
assert await buffer.store_spend_logs_in_redis(rows) is True
|
||||
|
||||
move_executed: Final = asyncio.Event()
|
||||
release_move: Final = asyncio.Event()
|
||||
_hold_first_pipeline_until_release(
|
||||
monkeypatch,
|
||||
redis_client,
|
||||
move_executed,
|
||||
release_move,
|
||||
)
|
||||
|
||||
first_call: Final = asyncio.create_task(buffer.get_spend_logs_from_redis_buffer(limit=10))
|
||||
await move_executed.wait()
|
||||
|
||||
second_started: Final = asyncio.Event()
|
||||
|
||||
async def start_second_drain() -> tuple[dict[str, object], ...]:
|
||||
second_started.set()
|
||||
return await buffer.get_spend_logs_from_redis_buffer(limit=10)
|
||||
|
||||
second_call: Final = asyncio.create_task(start_second_drain())
|
||||
await second_started.wait()
|
||||
|
||||
release_move.set()
|
||||
first_result: Final = await first_call
|
||||
second_result: Final = await second_call
|
||||
request_ids: Final = tuple(str(row["request_id"]) for row in first_result + second_result)
|
||||
assert sorted(request_ids) == ["req-a", "req-b"]
|
||||
assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ack_failure_leaves_items_for_the_next_drain(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
redis_cache, redis_client = _cache_with_fakeredis(monkeypatch)
|
||||
buffer: Final = RedisUpdateBuffer(redis_cache=redis_cache)
|
||||
payload: Final = json.dumps(
|
||||
{
|
||||
"key_list_transactions": {"key-ack": 1.5},
|
||||
"user_list_transactions": {},
|
||||
"end_user_list_transactions": {},
|
||||
"team_list_transactions": {},
|
||||
"team_member_list_transactions": {},
|
||||
"org_list_transactions": {},
|
||||
"tag_list_transactions": {},
|
||||
}
|
||||
)
|
||||
await redis_cache.async_rpush(key="litellm_spend_update_buffer", values=(payload,))
|
||||
|
||||
original_delete: Final = redis_client.delete
|
||||
failed_delete: Final = AsyncMock(side_effect=RedisTimeoutError("DEL failed"))
|
||||
monkeypatch.setattr(redis_client, "delete", failed_delete)
|
||||
with pytest.raises(RedisTimeoutError, match="DEL failed"):
|
||||
await buffer.get_all_update_transactions_from_redis_buffer()
|
||||
|
||||
monkeypatch.setattr(redis_client, "delete", original_delete)
|
||||
recovered: Final = await buffer.get_all_update_transactions_from_redis_buffer()
|
||||
assert recovered is not None
|
||||
assert recovered["key_list_transactions"] == {"key-ack": 1.5}
|
||||
|
||||
|
||||
def test_validate_redis_transaction_buffer_raises_without_redis():
|
||||
"""
|
||||
When use_redis_transaction_buffer=true but no Redis cache is configured,
|
||||
|
|
@ -577,7 +844,7 @@ 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):
|
||||
"""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
|
||||
merging what it drained, and the drain removes the payload, so one without
|
||||
the key would cost a rolling deploy those increments."""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
|
|
@ -656,6 +923,7 @@ async def test_store_in_memory_spend_updates_restores_budget_window_spend_on_rpu
|
|||
class _ListRedis:
|
||||
def __init__(self) -> None:
|
||||
self.rows: list[str] = []
|
||||
self.processing: list[str] = []
|
||||
|
||||
async def async_rpush_and_trim(self, key: str, values: list[str], max_len: int) -> int:
|
||||
self.rows.extend(values)
|
||||
|
|
@ -663,12 +931,15 @@ class _ListRedis:
|
|||
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
|
||||
async def async_drain_list_to_processing(self, key: str, worker_id: str, count: int, ttl_seconds: int) -> list[str]:
|
||||
if self.processing:
|
||||
return self.processing
|
||||
self.processing = self.rows[:count]
|
||||
self.rows = self.rows[count:]
|
||||
return self.processing
|
||||
|
||||
async def async_ack_processing_list(self, key: str, worker_id: str) -> None:
|
||||
self.processing = []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ import threading
|
|||
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, Callable, Dict, Final, Iterator, List, Optional
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
|
@ -344,6 +344,29 @@ class FakeRedisList:
|
|||
del stored[:count]
|
||||
return popped
|
||||
|
||||
async def async_drain_list_to_processing(
|
||||
self,
|
||||
key: str,
|
||||
worker_id: str,
|
||||
count: int,
|
||||
ttl_seconds: int,
|
||||
) -> list[str]:
|
||||
self._check_up()
|
||||
processing_key: Final = f"{key}:processing:{worker_id}"
|
||||
pending: Final = self.items.get(processing_key, [])
|
||||
if pending:
|
||||
return list(pending)
|
||||
source: Final = self.items.get(key, [])
|
||||
moved: Final = source[:count]
|
||||
if moved:
|
||||
self.items[key] = source[len(moved) :]
|
||||
self.items[processing_key] = moved
|
||||
return list(moved)
|
||||
|
||||
async def async_ack_processing_list(self, key: str, worker_id: str) -> None:
|
||||
self._check_up()
|
||||
self.items.pop(f"{key}:processing:{worker_id}", None)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_redis() -> FakeRedisList:
|
||||
|
|
|
|||
|
|
@ -887,6 +887,9 @@ async def test_recover_parked_spend_logs_re_parks_rows_when_the_enqueue_is_cance
|
|||
mock_prisma_client._spend_log_transactions_lock.release()
|
||||
|
||||
assert _parked_request_ids(fake_redis) == ["parked"]
|
||||
recovered_rows: Final = await buffer.get_spend_logs_from_redis_buffer(limit=10)
|
||||
assert [row["request_id"] for row in recovered_rows] == ["parked"]
|
||||
assert await buffer.get_spend_logs_from_redis_buffer(limit=10) == ()
|
||||
assert mock_prisma_client.spend_log_transactions == []
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
import inspect
|
||||
import json
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
|
@ -36,6 +38,7 @@ from litellm.caching.redis_cache import RedisCache
|
|||
from litellm.caching.redis_cluster_cache import RedisClusterCache
|
||||
from litellm.constants import REDIS_CLUSTER_HEALTH_CHECK_INTERVAL
|
||||
from litellm.proxy._types import CoordinationRedisParams
|
||||
from tests.integration._support.redis_process import owned_redis
|
||||
|
||||
|
||||
class _StubCredentialProvider(CredentialProvider):
|
||||
|
|
@ -1960,3 +1963,52 @@ def test_async_sentinel_keeps_the_credential_provider_off_the_monitors(markers,
|
|||
master_kwargs = mock_sentinel_cls.return_value.master_for.call_args[1]
|
||||
assert isinstance(master_kwargs["credential_provider"], provider_cls)
|
||||
assert "password" not in master_kwargs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.skipif(shutil.which("redis-server") is None, reason="Lua drain needs a real redis-server")
|
||||
@pytest.mark.parametrize("redis_version", ["6.0", 6.1])
|
||||
async def test_async_drain_list_to_processing_keeps_legacy_redis_rows_pending_in_order(
|
||||
redis_version: str | float,
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
with owned_redis(tmp_path) as redis_server:
|
||||
redis_cache: Final = RedisCache(host=redis_server.host, port=redis_server.port, namespace="drain-test")
|
||||
redis_cache.redis_version = redis_version
|
||||
redis_client: Final = async_redis.Redis(host=redis_server.host, port=redis_server.port)
|
||||
try:
|
||||
rows: Final = ("row-1", "row-2", "row-3", "row-4", "row-5")
|
||||
await redis_client.rpush("drain-test:buffer", *rows)
|
||||
|
||||
first_batch: Final = await redis_cache.async_drain_list_to_processing(
|
||||
key="buffer",
|
||||
worker_id="worker-1",
|
||||
count=2,
|
||||
ttl_seconds=60,
|
||||
)
|
||||
processing_key: Final = "drain-test:buffer:processing:worker-1"
|
||||
assert first_batch == ["row-1", "row-2"]
|
||||
assert await redis_client.llen(processing_key) == 2
|
||||
assert await redis_client.lrange("drain-test:buffer", 0, -1) == [b"row-3", b"row-4", b"row-5"]
|
||||
|
||||
pending_batch: Final = await redis_cache.async_drain_list_to_processing(
|
||||
key="buffer",
|
||||
worker_id="worker-1",
|
||||
count=2,
|
||||
ttl_seconds=60,
|
||||
)
|
||||
assert pending_batch == ["row-1", "row-2"]
|
||||
assert await redis_client.lrange(processing_key, 0, -1) == [b"row-1", b"row-2"]
|
||||
|
||||
await redis_cache.async_ack_processing_list(key="buffer", worker_id="worker-1")
|
||||
next_batch: Final = await redis_cache.async_drain_list_to_processing(
|
||||
key="buffer",
|
||||
worker_id="worker-1",
|
||||
count=2,
|
||||
ttl_seconds=60,
|
||||
)
|
||||
assert next_batch == ["row-3", "row-4"]
|
||||
assert await redis_client.lrange("drain-test:buffer", 0, -1) == [b"row-5"]
|
||||
finally:
|
||||
await redis_client.aclose()
|
||||
await redis_cache.disconnect()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue