mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
perf(proxy): pipeline Redis RPUSH/LPOP in spend update cycle
Replace 14 sequential Redis round-trips (7 RPUSH + 7 LPOP) per spend update cycle with 2 pipelined calls (1 RPUSH pipeline + 1 LPOP pipeline). This reduces connection pool contention at scale (50+ pods). - Add RedisPipelineRpushOperation and RedisPipelineLpopOperation TypedDicts - Add async_rpush_pipeline() and async_lpop_pipeline() to RedisCache - Refactor store_in_memory_spend_updates_in_redis() to use pipeline - Add get_all_transactions_from_redis_buffer_pipeline() for batched drain - Update _commit_spend_updates_to_db_with_redis() to use pipeline drain - Existing individual methods preserved for backward compatibility
This commit is contained in:
Ryan Crabbe
Sameer Kankute
parent
62d90fc7a2
commit
1de3b90db2
7 changed files with 829 additions and 296 deletions
|
|
@ -22,7 +22,11 @@ from litellm._logging import print_verbose, verbose_logger
|
|||
from litellm.constants import DEFAULT_REDIS_MAJOR_VERSION
|
||||
from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
|
||||
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
from litellm.types.caching import (
|
||||
RedisPipelineIncrementOperation,
|
||||
RedisPipelineLpopOperation,
|
||||
RedisPipelineRpushOperation,
|
||||
)
|
||||
from litellm.types.services import ServiceTypes
|
||||
|
||||
from .base_cache import BaseCache
|
||||
|
|
@ -1320,6 +1324,75 @@ class RedisCache(BaseCache):
|
|||
)
|
||||
raise e
|
||||
|
||||
async def _pipeline_rpush_helper(
|
||||
self,
|
||||
pipe: pipeline,
|
||||
rpush_list: List[RedisPipelineRpushOperation],
|
||||
) -> List[int]:
|
||||
"""Helper function for pipeline rpush operations"""
|
||||
for rpush_op in rpush_list:
|
||||
pipe.rpush(rpush_op["key"], *rpush_op["values"])
|
||||
results = await pipe.execute()
|
||||
# Preserve positional correspondence — raise on per-command errors
|
||||
for r in results:
|
||||
if isinstance(r, Exception):
|
||||
raise r
|
||||
return results
|
||||
|
||||
async def async_rpush_pipeline(
|
||||
self,
|
||||
rpush_list: List[RedisPipelineRpushOperation],
|
||||
) -> List[int]:
|
||||
"""
|
||||
Use Redis Pipelines for bulk RPUSH operations
|
||||
|
||||
Args:
|
||||
rpush_list: List of RedisPipelineRpushOperation dicts containing:
|
||||
- key: str
|
||||
- values: List[Any]
|
||||
|
||||
Returns:
|
||||
List[int]: List lengths after each push
|
||||
"""
|
||||
if len(rpush_list) == 0:
|
||||
return []
|
||||
|
||||
_redis_client: Any = self.init_async_client()
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
async with _redis_client.pipeline(transaction=False) as pipe:
|
||||
results = await self._pipeline_rpush_helper(pipe, rpush_list)
|
||||
|
||||
## LOGGING ##
|
||||
end_time = time.time()
|
||||
_duration = end_time - start_time
|
||||
asyncio.create_task(
|
||||
self.service_logger_obj.async_service_success_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=_duration,
|
||||
call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
return results
|
||||
except Exception as e:
|
||||
## LOGGING ##
|
||||
end_time = time.time()
|
||||
_duration = end_time - start_time
|
||||
asyncio.create_task(
|
||||
self.service_logger_obj.async_service_failure_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=_duration,
|
||||
error=e,
|
||||
call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS %s",
|
||||
str(e),
|
||||
)
|
||||
raise e
|
||||
|
||||
async def handle_lpop_count_for_older_redis_versions(
|
||||
self, pipe: pipeline, key: str, count: int
|
||||
) -> List[bytes]:
|
||||
|
|
@ -1400,3 +1473,115 @@ class RedisCache(BaseCache):
|
|||
f"LiteLLM Redis Cache LPOP: - Got exception from REDIS : {str(e)}"
|
||||
)
|
||||
raise e
|
||||
|
||||
async def _pipeline_lpop_helper(
|
||||
self,
|
||||
pipe: pipeline,
|
||||
lpop_list: List[RedisPipelineLpopOperation],
|
||||
) -> List[Optional[List[str]]]:
|
||||
"""Helper function for pipeline lpop operations.
|
||||
|
||||
For Redis >= 7, queues one LPOP(key, count) per operation.
|
||||
For Redis < 7, queues `count` individual LPOP(key) commands per operation.
|
||||
"""
|
||||
major_version = self._parse_redis_major_version()
|
||||
|
||||
if major_version >= 7:
|
||||
for lpop_op in lpop_list:
|
||||
pipe.lpop(lpop_op["key"], lpop_op["count"])
|
||||
raw_results = await pipe.execute()
|
||||
else:
|
||||
# For Redis < 7, LPOP doesn't support count param.
|
||||
# Issue `count` individual LPOP commands per key, all in one pipeline.
|
||||
counts: List[int] = []
|
||||
for lpop_op in lpop_list:
|
||||
count = lpop_op["count"] or 1
|
||||
counts.append(count)
|
||||
for _ in range(count):
|
||||
pipe.lpop(lpop_op["key"])
|
||||
flat_results = await pipe.execute()
|
||||
|
||||
# Re-group the flat results back into per-key lists
|
||||
raw_results = []
|
||||
offset = 0
|
||||
for count in counts:
|
||||
key_results = [
|
||||
r for r in flat_results[offset : offset + count] if r is not None
|
||||
]
|
||||
raw_results.append(key_results if key_results else None)
|
||||
offset += count
|
||||
|
||||
# Decode bytes -> str for each result set
|
||||
decoded_results: List[Optional[List[str]]] = []
|
||||
for r in raw_results:
|
||||
if r is None:
|
||||
decoded_results.append(None)
|
||||
elif isinstance(r, list):
|
||||
try:
|
||||
decoded_results.append(
|
||||
[
|
||||
item.decode("utf-8") if isinstance(item, bytes) else item
|
||||
for item in r
|
||||
if item is not None
|
||||
]
|
||||
or None
|
||||
)
|
||||
except Exception:
|
||||
decoded_results.append(r) # type: ignore
|
||||
else:
|
||||
decoded_results.append(None)
|
||||
return decoded_results
|
||||
|
||||
async def async_lpop_pipeline(
|
||||
self,
|
||||
lpop_list: List[RedisPipelineLpopOperation],
|
||||
) -> List[Optional[List[str]]]:
|
||||
"""
|
||||
Use Redis Pipelines for bulk LPOP operations
|
||||
|
||||
Args:
|
||||
lpop_list: List of RedisPipelineLpopOperation dicts containing:
|
||||
- key: str
|
||||
- count: Optional[int]
|
||||
|
||||
Returns:
|
||||
List[Optional[List[str]]]: Decoded results per key, None if key was empty
|
||||
"""
|
||||
if len(lpop_list) == 0:
|
||||
return []
|
||||
|
||||
_redis_client: Any = self.init_async_client()
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
async with _redis_client.pipeline(transaction=False) as pipe:
|
||||
results = await self._pipeline_lpop_helper(pipe, lpop_list)
|
||||
|
||||
## LOGGING ##
|
||||
end_time = time.time()
|
||||
_duration = end_time - start_time
|
||||
asyncio.create_task(
|
||||
self.service_logger_obj.async_service_success_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=_duration,
|
||||
call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
return results
|
||||
except Exception as e:
|
||||
## LOGGING ##
|
||||
end_time = time.time()
|
||||
_duration = end_time - start_time
|
||||
asyncio.create_task(
|
||||
self.service_logger_obj.async_service_failure_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=_duration,
|
||||
error=e,
|
||||
call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS %s",
|
||||
str(e),
|
||||
)
|
||||
raise e
|
||||
|
|
|
|||
|
|
@ -666,9 +666,16 @@ class DBSpendUpdateWriter:
|
|||
verbose_proxy_logger.debug("acquired lock for spend updates")
|
||||
|
||||
try:
|
||||
db_spend_update_transactions = (
|
||||
await self.redis_update_buffer.get_all_update_transactions_from_redis_buffer()
|
||||
)
|
||||
(
|
||||
db_spend_update_transactions,
|
||||
daily_spend_update_transactions,
|
||||
daily_team_spend_update_transactions,
|
||||
daily_org_spend_update_transactions,
|
||||
daily_end_user_spend_update_transactions,
|
||||
daily_agent_spend_update_transactions,
|
||||
daily_tag_spend_update_transactions,
|
||||
) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
|
||||
if db_spend_update_transactions is not None:
|
||||
verbose_proxy_logger.info(
|
||||
"Spend tracking - committing spend updates from Redis to DB: "
|
||||
|
|
@ -688,9 +695,6 @@ class DBSpendUpdateWriter:
|
|||
db_spend_update_transactions=db_spend_update_transactions,
|
||||
)
|
||||
|
||||
daily_spend_update_transactions = (
|
||||
await self.redis_update_buffer.get_all_daily_spend_update_transactions_from_redis_buffer()
|
||||
)
|
||||
if daily_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter.update_daily_user_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
|
|
@ -698,9 +702,6 @@ class DBSpendUpdateWriter:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=daily_spend_update_transactions,
|
||||
)
|
||||
daily_team_spend_update_transactions = (
|
||||
await self.redis_update_buffer.get_all_daily_team_spend_update_transactions_from_redis_buffer()
|
||||
)
|
||||
if daily_team_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter.update_daily_team_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
|
|
@ -709,9 +710,6 @@ class DBSpendUpdateWriter:
|
|||
daily_spend_transactions=daily_team_spend_update_transactions,
|
||||
)
|
||||
|
||||
daily_org_spend_update_transactions = (
|
||||
await self.redis_update_buffer.get_all_daily_org_spend_update_transactions_from_redis_buffer()
|
||||
)
|
||||
if daily_org_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter.update_daily_org_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
|
|
@ -720,9 +718,6 @@ class DBSpendUpdateWriter:
|
|||
daily_spend_transactions=daily_org_spend_update_transactions,
|
||||
)
|
||||
|
||||
daily_tag_spend_update_transactions = (
|
||||
await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer()
|
||||
)
|
||||
if daily_tag_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter.update_daily_tag_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
|
|
@ -730,9 +725,6 @@ class DBSpendUpdateWriter:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=daily_tag_spend_update_transactions,
|
||||
)
|
||||
daily_end_user_spend_update_transactions = (
|
||||
await self.redis_update_buffer.get_all_daily_end_user_spend_update_transactions_from_redis_buffer()
|
||||
)
|
||||
if daily_end_user_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter.update_daily_end_user_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
|
|
@ -740,9 +732,6 @@ class DBSpendUpdateWriter:
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=daily_end_user_spend_update_transactions,
|
||||
)
|
||||
daily_agent_spend_update_transactions = (
|
||||
await self.redis_update_buffer.get_all_daily_agent_spend_update_transactions_from_redis_buffer()
|
||||
)
|
||||
if daily_agent_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter.update_daily_agent_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ This is to prevent deadlocks and improve reliability
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching import RedisCache
|
||||
|
|
@ -36,6 +36,7 @@ from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
|
|||
)
|
||||
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation
|
||||
from litellm.types.services import ServiceTypes
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -209,47 +210,44 @@ class RedisUpdateBuffer:
|
|||
"ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions
|
||||
)
|
||||
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=db_spend_update_transactions,
|
||||
redis_key=REDIS_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_SPEND_UPDATE_QUEUE,
|
||||
# Build a list of rpush operations, skipping empty/None transaction sets
|
||||
_queue_configs: List[Tuple[Any, str, ServiceTypes]] = [
|
||||
(db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_SPEND_UPDATE_QUEUE),
|
||||
(daily_spend_update_transactions, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE),
|
||||
(daily_team_spend_update_transactions, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE),
|
||||
(daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE),
|
||||
(daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE),
|
||||
(daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE),
|
||||
(daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE),
|
||||
]
|
||||
|
||||
rpush_list: List[RedisPipelineRpushOperation] = []
|
||||
service_types: 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:
|
||||
return
|
||||
|
||||
result_lengths = await self.redis_cache.async_rpush_pipeline(
|
||||
rpush_list=rpush_list,
|
||||
)
|
||||
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=daily_spend_update_transactions,
|
||||
redis_key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE,
|
||||
)
|
||||
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=daily_team_spend_update_transactions,
|
||||
redis_key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE,
|
||||
)
|
||||
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=daily_org_spend_update_transactions,
|
||||
redis_key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE,
|
||||
)
|
||||
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=daily_end_user_spend_update_transactions,
|
||||
redis_key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE,
|
||||
)
|
||||
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=daily_agent_spend_update_transactions,
|
||||
redis_key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE,
|
||||
)
|
||||
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=daily_tag_spend_update_transactions,
|
||||
redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE,
|
||||
)
|
||||
# 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],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _number_of_transactions_to_store_in_redis(
|
||||
|
|
@ -338,6 +336,77 @@ class RedisUpdateBuffer:
|
|||
|
||||
return combined_transaction
|
||||
|
||||
async def get_all_transactions_from_redis_buffer_pipeline(
|
||||
self,
|
||||
) -> Tuple[
|
||||
Optional[DBSpendUpdateTransactions],
|
||||
Optional[Dict[str, DailyUserSpendTransaction]],
|
||||
Optional[Dict[str, DailyTeamSpendTransaction]],
|
||||
Optional[Dict[str, DailyOrganizationSpendTransaction]],
|
||||
Optional[Dict[str, DailyEndUserSpendTransaction]],
|
||||
Optional[Dict[str, DailyAgentSpendTransaction]],
|
||||
Optional[Dict[str, DailyTagSpendTransaction]],
|
||||
]:
|
||||
"""
|
||||
Drains all 7 Redis buffer queues in a single pipeline round-trip.
|
||||
|
||||
Returns a 7-tuple of parsed results in this order:
|
||||
0: DBSpendUpdateTransactions
|
||||
1: daily user spend
|
||||
2: daily team spend
|
||||
3: daily org spend
|
||||
4: daily end-user spend
|
||||
5: daily agent spend
|
||||
6: daily tag spend
|
||||
"""
|
||||
if self.redis_cache is None:
|
||||
return None, None, None, None, None, None, None
|
||||
|
||||
lpop_list: 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_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT),
|
||||
]
|
||||
|
||||
raw_results = 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: Optional[DBSpendUpdateTransactions] = None
|
||||
if raw_results[0] is not None:
|
||||
parsed = self._parse_list_of_transactions(raw_results[0])
|
||||
if len(parsed) > 0:
|
||||
db_spend = self._combine_list_of_transactions(parsed)
|
||||
|
||||
# Slots 1-6: daily spend categories
|
||||
daily_results: List[Optional[Dict[str, Any]]] = []
|
||||
for slot in range(1, 7):
|
||||
if raw_results[slot] is None:
|
||||
daily_results.append(None)
|
||||
else:
|
||||
list_of_daily = [json.loads(t) for t in raw_results[slot]] # type: ignore
|
||||
aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(
|
||||
list_of_daily
|
||||
)
|
||||
daily_results.append(aggregated)
|
||||
|
||||
return (
|
||||
db_spend,
|
||||
cast(Optional[Dict[str, DailyUserSpendTransaction]], daily_results[0]),
|
||||
cast(Optional[Dict[str, DailyTeamSpendTransaction]], daily_results[1]),
|
||||
cast(Optional[Dict[str, DailyOrganizationSpendTransaction]], daily_results[2]),
|
||||
cast(Optional[Dict[str, DailyEndUserSpendTransaction]], daily_results[3]),
|
||||
cast(Optional[Dict[str, DailyAgentSpendTransaction]], daily_results[4]),
|
||||
cast(Optional[Dict[str, DailyTagSpendTransaction]], daily_results[5]),
|
||||
)
|
||||
|
||||
async def get_all_daily_spend_update_transactions_from_redis_buffer(
|
||||
self,
|
||||
) -> Optional[Dict[str, DailyUserSpendTransaction]]:
|
||||
|
|
|
|||
|
|
@ -52,6 +52,24 @@ class RedisPipelineSetOperation(TypedDict):
|
|||
ttl: Optional[int]
|
||||
|
||||
|
||||
class RedisPipelineRpushOperation(TypedDict):
|
||||
"""
|
||||
TypedDict for 1 Redis Pipeline RPUSH Operation
|
||||
"""
|
||||
|
||||
key: str
|
||||
values: List[Any]
|
||||
|
||||
|
||||
class RedisPipelineLpopOperation(TypedDict):
|
||||
"""
|
||||
TypedDict for 1 Redis Pipeline LPOP Operation
|
||||
"""
|
||||
|
||||
key: str
|
||||
count: Optional[int]
|
||||
|
||||
|
||||
DynamicCacheControl = TypedDict(
|
||||
"DynamicCacheControl",
|
||||
{
|
||||
|
|
|
|||
|
|
@ -122,6 +122,219 @@ async def test_handle_lpop_count_for_older_redis_versions(monkeypatch):
|
|||
assert mock_pipeline.execute.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_rpush_pipeline_executes_all_operations(monkeypatch, redis_no_ping):
|
||||
"""Verify that multiple rpush ops are batched into a single pipeline execute"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline)
|
||||
mock_pipeline.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_pipeline.rpush = MagicMock()
|
||||
mock_pipeline.execute = AsyncMock(return_value=[3, 5, 1])
|
||||
mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline)
|
||||
|
||||
from litellm.types.caching import RedisPipelineRpushOperation
|
||||
|
||||
rpush_list = [
|
||||
RedisPipelineRpushOperation(key="key1", values=["a", "b"]),
|
||||
RedisPipelineRpushOperation(key="key2", values=["c"]),
|
||||
RedisPipelineRpushOperation(key="key3", values=["d", "e", "f"]),
|
||||
]
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
result = await redis_cache.async_rpush_pipeline(rpush_list=rpush_list)
|
||||
|
||||
assert result == [3, 5, 1]
|
||||
assert mock_pipeline.rpush.call_count == 3
|
||||
mock_pipeline.rpush.assert_any_call("key1", "a", "b")
|
||||
mock_pipeline.rpush.assert_any_call("key2", "c")
|
||||
mock_pipeline.rpush.assert_any_call("key3", "d", "e", "f")
|
||||
mock_pipeline.execute.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_rpush_pipeline_empty_list_returns_empty(monkeypatch, redis_no_ping):
|
||||
"""Empty rpush_list should return empty list without touching Redis"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
result = await redis_cache.async_rpush_pipeline(rpush_list=[])
|
||||
|
||||
assert result == []
|
||||
mock_redis_instance.pipeline.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_rpush_pipeline_raises_on_redis_error(monkeypatch, redis_no_ping):
|
||||
"""Pipeline errors should propagate"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline)
|
||||
mock_pipeline.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_pipeline.rpush = MagicMock()
|
||||
mock_pipeline.execute = AsyncMock(side_effect=ConnectionError("Redis down"))
|
||||
mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline)
|
||||
|
||||
from litellm.types.caching import RedisPipelineRpushOperation
|
||||
|
||||
rpush_list = [RedisPipelineRpushOperation(key="key1", values=["a"])]
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
with pytest.raises(ConnectionError, match="Redis down"):
|
||||
await redis_cache.async_rpush_pipeline(rpush_list=rpush_list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_lpop_pipeline_single_round_trip(monkeypatch, redis_no_ping):
|
||||
"""Verify that multiple lpop ops are batched into a single pipeline execute"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
redis_cache.redis_version = "7.0.0"
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline)
|
||||
mock_pipeline.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_pipeline.lpop = MagicMock()
|
||||
mock_pipeline.execute = AsyncMock(return_value=[
|
||||
[b"val1", b"val2"], # key1 results
|
||||
None, # key2 empty
|
||||
[b"val3"], # key3 results
|
||||
])
|
||||
mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline)
|
||||
|
||||
from litellm.types.caching import RedisPipelineLpopOperation
|
||||
|
||||
lpop_list = [
|
||||
RedisPipelineLpopOperation(key="key1", count=10),
|
||||
RedisPipelineLpopOperation(key="key2", count=10),
|
||||
RedisPipelineLpopOperation(key="key3", count=5),
|
||||
]
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
results = await redis_cache.async_lpop_pipeline(lpop_list=lpop_list)
|
||||
|
||||
assert len(results) == 3
|
||||
assert results[0] == ["val1", "val2"]
|
||||
assert results[1] is None
|
||||
assert results[2] == ["val3"]
|
||||
mock_pipeline.execute.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_lpop_pipeline_redis_lt7_regroups_flat_results(monkeypatch, redis_no_ping):
|
||||
"""Verify Redis < 7 fallback issues individual LPOPs and regroups correctly"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
redis_cache.redis_version = "6.2.0"
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline)
|
||||
mock_pipeline.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_pipeline.lpop = MagicMock()
|
||||
|
||||
# With count=3 for key1 and count=2 for key2, we get 5 individual LPOP commands
|
||||
# Simulate: key1 has 2 values then None, key2 has 1 value then None
|
||||
mock_pipeline.execute = AsyncMock(return_value=[
|
||||
b"val1", b"val2", None, # 3 LPOPs for key1
|
||||
b"val3", None, # 2 LPOPs for key2
|
||||
])
|
||||
mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline)
|
||||
|
||||
from litellm.types.caching import RedisPipelineLpopOperation
|
||||
|
||||
lpop_list = [
|
||||
RedisPipelineLpopOperation(key="key1", count=3),
|
||||
RedisPipelineLpopOperation(key="key2", count=2),
|
||||
]
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
results = await redis_cache.async_lpop_pipeline(lpop_list=lpop_list)
|
||||
|
||||
assert len(results) == 2
|
||||
assert results[0] == ["val1", "val2"] # 2 values, None filtered out
|
||||
assert results[1] == ["val3"] # 1 value, None filtered out
|
||||
# All 5 individual LPOPs should be queued, but only 1 execute() call
|
||||
assert mock_pipeline.lpop.call_count == 5
|
||||
mock_pipeline.execute.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_rpush_pipeline_raises_on_per_command_error(monkeypatch, redis_no_ping):
|
||||
"""Verify that per-command errors in pipeline results are raised, not silently dropped"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline)
|
||||
mock_pipeline.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_pipeline.rpush = MagicMock()
|
||||
# Simulate: first RPUSH succeeds, second returns a per-command error
|
||||
mock_pipeline.execute = AsyncMock(return_value=[3, Exception("WRONGTYPE")])
|
||||
mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline)
|
||||
|
||||
from litellm.types.caching import RedisPipelineRpushOperation
|
||||
|
||||
rpush_list = [
|
||||
RedisPipelineRpushOperation(key="key1", values=["a"]),
|
||||
RedisPipelineRpushOperation(key="key2", values=["b"]),
|
||||
]
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
with pytest.raises(Exception, match="WRONGTYPE"):
|
||||
await redis_cache.async_rpush_pipeline(rpush_list=rpush_list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_lpop_pipeline_empty_list(monkeypatch, redis_no_ping):
|
||||
"""Empty lpop_list should return empty list without touching Redis"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
result = await redis_cache.async_lpop_pipeline(lpop_list=[])
|
||||
|
||||
assert result == []
|
||||
mock_redis_instance.pipeline.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_lpop_pipeline_propagates_redis_exception(monkeypatch, redis_no_ping):
|
||||
"""Pipeline errors should propagate"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
redis_cache.redis_version = "7.0.0"
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline)
|
||||
mock_pipeline.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_pipeline.lpop = MagicMock()
|
||||
mock_pipeline.execute = AsyncMock(side_effect=ConnectionError("Redis down"))
|
||||
mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline)
|
||||
|
||||
from litellm.types.caching import RedisPipelineLpopOperation
|
||||
|
||||
lpop_list = [RedisPipelineLpopOperation(key="key1", count=10)]
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
with pytest.raises(ConnectionError, match="Redis down"):
|
||||
await redis_cache.async_lpop_pipeline(lpop_list=lpop_list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"redis_version",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,194 @@
|
|||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
|
||||
from litellm.types.caching import RedisPipelineRpushOperation
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redis_cache():
|
||||
"""Create a mock RedisCache instance"""
|
||||
mock = AsyncMock()
|
||||
return mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis_update_buffer(mock_redis_cache):
|
||||
"""Create a RedisUpdateBuffer with a mock RedisCache"""
|
||||
return RedisUpdateBuffer(redis_cache=mock_redis_cache)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache):
|
||||
"""
|
||||
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 7 have data
|
||||
spend_update_queue = AsyncMock()
|
||||
spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(
|
||||
return_value={"key_list_transactions": {"key1": 1.0}}
|
||||
)
|
||||
|
||||
daily_spend_queue = AsyncMock()
|
||||
daily_spend_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={"user_key1": {"spend": 1.0}}
|
||||
)
|
||||
|
||||
daily_team_queue = AsyncMock()
|
||||
daily_team_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={"team_key1": {"spend": 2.0}}
|
||||
)
|
||||
|
||||
# Empty queues
|
||||
daily_org_queue = AsyncMock()
|
||||
daily_org_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={}
|
||||
)
|
||||
|
||||
daily_end_user_queue = AsyncMock()
|
||||
daily_end_user_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
|
||||
daily_agent_queue = AsyncMock()
|
||||
daily_agent_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={}
|
||||
)
|
||||
|
||||
daily_tag_queue = AsyncMock()
|
||||
daily_tag_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={}
|
||||
)
|
||||
|
||||
await redis_update_buffer.store_in_memory_spend_updates_in_redis(
|
||||
spend_update_queue=spend_update_queue,
|
||||
daily_spend_update_queue=daily_spend_queue,
|
||||
daily_team_spend_update_queue=daily_team_queue,
|
||||
daily_org_spend_update_queue=daily_org_queue,
|
||||
daily_end_user_spend_update_queue=daily_end_user_queue,
|
||||
daily_agent_spend_update_queue=daily_agent_queue,
|
||||
daily_tag_spend_update_queue=daily_tag_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
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_in_memory_spend_updates_all_empty_returns_early(
|
||||
redis_update_buffer, mock_redis_cache
|
||||
):
|
||||
"""
|
||||
When all queues are empty, pipeline should never be called.
|
||||
"""
|
||||
mock_redis_cache.async_rpush_pipeline = AsyncMock()
|
||||
|
||||
# All queues return empty
|
||||
empty_queue = AsyncMock()
|
||||
empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(
|
||||
return_value={}
|
||||
)
|
||||
empty_daily_queue = AsyncMock()
|
||||
empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={}
|
||||
)
|
||||
|
||||
await redis_update_buffer.store_in_memory_spend_updates_in_redis(
|
||||
spend_update_queue=empty_queue,
|
||||
daily_spend_update_queue=empty_daily_queue,
|
||||
daily_team_spend_update_queue=empty_daily_queue,
|
||||
daily_org_spend_update_queue=empty_daily_queue,
|
||||
daily_end_user_spend_update_queue=empty_daily_queue,
|
||||
daily_agent_spend_update_queue=empty_daily_queue,
|
||||
daily_tag_spend_update_queue=empty_daily_queue,
|
||||
)
|
||||
|
||||
mock_redis_cache.async_rpush_pipeline.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
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.
|
||||
"""
|
||||
# Simulate pipeline results: slot 0 = spend updates, slots 1-6 = daily categories
|
||||
db_spend_json = json.dumps(
|
||||
{
|
||||
"key_list_transactions": {"key1": 1.0, "key2": 2.0},
|
||||
"user_list_transactions": {"user1": 0.5},
|
||||
"end_user_list_transactions": {},
|
||||
"team_list_transactions": {},
|
||||
"team_member_list_transactions": {},
|
||||
"org_list_transactions": {},
|
||||
"tag_list_transactions": {},
|
||||
}
|
||||
)
|
||||
daily_user_json = json.dumps({"user_key1": {"spend": 1.0, "api_requests": 1}})
|
||||
daily_team_json = json.dumps({"team_key1": {"spend": 2.0, "api_requests": 2}})
|
||||
|
||||
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)
|
||||
None, # slot 6: daily tag (empty)
|
||||
]
|
||||
)
|
||||
|
||||
result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
|
||||
assert len(result) == 7
|
||||
db_spend, daily_user, daily_team, daily_org, daily_end_user, daily_agent, daily_tag = result
|
||||
|
||||
# Verify db spend was parsed correctly
|
||||
assert db_spend is not None
|
||||
assert db_spend["key_list_transactions"]["key1"] == 1.0
|
||||
assert db_spend["key_list_transactions"]["key2"] == 2.0
|
||||
assert db_spend["user_list_transactions"]["user1"] == 0.5
|
||||
|
||||
# Verify daily user was parsed
|
||||
assert daily_user is not None
|
||||
assert daily_user["user_key1"]["spend"] == 1.0
|
||||
|
||||
# Verify daily team was parsed
|
||||
assert daily_team is not None
|
||||
assert daily_team["team_key1"]["spend"] == 2.0
|
||||
|
||||
# Verify empty slots
|
||||
assert daily_org is None
|
||||
assert daily_end_user is None
|
||||
assert daily_agent is None
|
||||
assert daily_tag is None
|
||||
|
||||
# Verify pipeline was called once with correct keys
|
||||
mock_redis_cache.async_lpop_pipeline.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis():
|
||||
"""When redis_cache is None, should return all Nones"""
|
||||
buffer = RedisUpdateBuffer(redis_cache=None)
|
||||
result = await buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
assert result == (None, None, None, None, None, None, None)
|
||||
|
|
@ -1,4 +1,3 @@
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
|
@ -52,9 +51,6 @@ async def test_daily_spend_tracking_with_disabled_spend_logs():
|
|||
# Call the method
|
||||
await db_writer.update_database(**test_data)
|
||||
|
||||
# Let the single batched task run
|
||||
await asyncio.sleep(0)
|
||||
|
||||
# Verify that _insert_spend_log_to_db was NOT called (since disable_spend_logs is True)
|
||||
db_writer._insert_spend_log_to_db.assert_not_called()
|
||||
|
||||
|
|
@ -119,9 +115,7 @@ async def test_update_daily_spend_with_null_entity_id():
|
|||
|
||||
# Verify the where clause contains null entity_id
|
||||
call_args = mock_table.upsert.call_args[1]
|
||||
where_clause = call_args["where"][
|
||||
"user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint"
|
||||
]
|
||||
where_clause = call_args["where"]["user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint"]
|
||||
assert where_clause["user_id"] is None
|
||||
assert where_clause["date"] == "2024-01-01"
|
||||
assert where_clause["api_key"] == "test-api-key"
|
||||
|
|
@ -167,7 +161,7 @@ async def test_update_daily_spend_sorting():
|
|||
upsert_calls = []
|
||||
for i in range(50):
|
||||
daily_spend_transactions[f"test_key_{i}"] = {
|
||||
"user_id": f"user{60-i}", # user60 ... user11, reverse order
|
||||
"user_id": f"user{60-i}", # user60 ... user11, reverse order
|
||||
"date": "2024-01-01",
|
||||
"api_key": "test-api-key",
|
||||
"model": "gpt-4",
|
||||
|
|
@ -179,48 +173,46 @@ async def test_update_daily_spend_sorting():
|
|||
"successful_requests": 1,
|
||||
"failed_requests": 0,
|
||||
}
|
||||
upsert_calls.append(
|
||||
call(
|
||||
where={
|
||||
"user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint": {
|
||||
"user_id": f"user{i+11}", # user11 ... user60, sorted order
|
||||
"date": "2024-01-01",
|
||||
"api_key": "test-api-key",
|
||||
"model": "gpt-4",
|
||||
"custom_llm_provider": "openai",
|
||||
"mcp_namespaced_tool_name": "",
|
||||
"endpoint": "",
|
||||
}
|
||||
upsert_calls.append(call(
|
||||
where={
|
||||
"user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint": {
|
||||
"user_id": f"user{i+11}", # user11 ... user60, sorted order
|
||||
"date": "2024-01-01",
|
||||
"api_key": "test-api-key",
|
||||
"model": "gpt-4",
|
||||
"custom_llm_provider": "openai",
|
||||
"mcp_namespaced_tool_name": "",
|
||||
"endpoint": "",
|
||||
}
|
||||
},
|
||||
data={
|
||||
"create": {
|
||||
"user_id": f"user{i+11}",
|
||||
"date": "2024-01-01",
|
||||
"api_key": "test-api-key",
|
||||
"model": "gpt-4",
|
||||
"model_group": None,
|
||||
"mcp_namespaced_tool_name": "",
|
||||
"custom_llm_provider": "openai",
|
||||
"endpoint": "",
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 20,
|
||||
"spend": 0.1,
|
||||
"api_requests": 1,
|
||||
"successful_requests": 1,
|
||||
"failed_requests": 0,
|
||||
},
|
||||
data={
|
||||
"create": {
|
||||
"user_id": f"user{i+11}",
|
||||
"date": "2024-01-01",
|
||||
"api_key": "test-api-key",
|
||||
"model": "gpt-4",
|
||||
"model_group": None,
|
||||
"mcp_namespaced_tool_name": "",
|
||||
"custom_llm_provider": "openai",
|
||||
"endpoint": "",
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 20,
|
||||
"spend": 0.1,
|
||||
"api_requests": 1,
|
||||
"successful_requests": 1,
|
||||
"failed_requests": 0,
|
||||
},
|
||||
"update": {
|
||||
"prompt_tokens": {"increment": 10},
|
||||
"completion_tokens": {"increment": 20},
|
||||
"spend": {"increment": 0.1},
|
||||
"api_requests": {"increment": 1},
|
||||
"successful_requests": {"increment": 1},
|
||||
"failed_requests": {"increment": 0},
|
||||
"endpoint": "",
|
||||
},
|
||||
"update": {
|
||||
"prompt_tokens": {"increment": 10},
|
||||
"completion_tokens": {"increment": 20},
|
||||
"spend": {"increment": 0.1},
|
||||
"api_requests": {"increment": 1},
|
||||
"successful_requests": {"increment": 1},
|
||||
"failed_requests": {"increment": 0},
|
||||
"endpoint": "",
|
||||
},
|
||||
)
|
||||
)
|
||||
},
|
||||
))
|
||||
|
||||
# Call the method
|
||||
await DBSpendUpdateWriter._update_daily_spend(
|
||||
|
|
@ -283,7 +275,7 @@ async def test_update_daily_spend_tag_with_request_id():
|
|||
|
||||
# Verify that table.upsert was called
|
||||
mock_table.upsert.assert_called_once()
|
||||
|
||||
|
||||
# Verify request_id is in update_data
|
||||
call_args = mock_table.upsert.call_args[1]
|
||||
update_data = call_args["data"]["update"]
|
||||
|
|
@ -291,13 +283,15 @@ async def test_update_daily_spend_tag_with_request_id():
|
|||
assert update_data["request_id"] == "test-request-id-123"
|
||||
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_daily_spend_with_none_values_in_sorting_fields():
|
||||
"""
|
||||
Test that _update_daily_spend handles None values in sorting fields correctly.
|
||||
|
||||
|
||||
This test ensures that when fields like date, api_key, model, or custom_llm_provider
|
||||
are None, the sorting doesn't crash with TypeError: '<' not supported between
|
||||
are None, the sorting doesn't crash with TypeError: '<' not supported between
|
||||
instances of 'NoneType' and 'str'.
|
||||
"""
|
||||
# Setup
|
||||
|
|
@ -515,7 +509,6 @@ async def test_update_tag_db_without_prisma_client():
|
|||
|
||||
assert writer.spend_update_queue.add_update.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id():
|
||||
"""
|
||||
|
|
@ -525,7 +518,7 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i
|
|||
writer = DBSpendUpdateWriter()
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.get_request_status = MagicMock(return_value="success")
|
||||
|
||||
|
||||
request_id = "test-request-id-123"
|
||||
payload = {
|
||||
"request_id": request_id,
|
||||
|
|
@ -553,15 +546,13 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i
|
|||
|
||||
# Should be called twice (once for each tag)
|
||||
assert writer.daily_tag_spend_update_queue.add_update.call_count == 2
|
||||
|
||||
|
||||
# Check that request_id is included in both transactions
|
||||
for call in writer.daily_tag_spend_update_queue.add_update.call_args_list:
|
||||
transaction_dict = call[1]["update"]
|
||||
# Each transaction should have one key with the format tag_date_api_key_model_provider
|
||||
for key, transaction in transaction_dict.items():
|
||||
assert (
|
||||
transaction["request_id"] == request_id
|
||||
), f"request_id should be {request_id} but got {transaction.get('request_id')}"
|
||||
assert transaction["request_id"] == request_id, f"request_id should be {request_id} but got {transaction.get('request_id')}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -875,11 +866,11 @@ async def test_endpoint_field_is_correctly_mapped_from_call_type():
|
|||
call_args = writer.daily_spend_update_queue.add_update.call_args[1]
|
||||
update_dict = call_args["update"]
|
||||
assert len(update_dict) == 1
|
||||
|
||||
|
||||
for key, transaction in update_dict.items():
|
||||
# Verify endpoint is included in the key
|
||||
assert key == f"test-user_2024-01-01_test-key_gpt-4_openai_/chat/completions"
|
||||
|
||||
|
||||
# Verify endpoint is set in the transaction
|
||||
assert transaction["endpoint"] == "/chat/completions"
|
||||
assert transaction["user_id"] == "test-user"
|
||||
|
|
@ -896,7 +887,7 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure():
|
|||
This ensures proper debugging information is available for issues like unique constraint violations.
|
||||
"""
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
|
||||
# Setup
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_batcher = MagicMock()
|
||||
|
|
@ -904,13 +895,13 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure():
|
|||
mock_batch_context = MagicMock()
|
||||
mock_batch_context.__aenter__ = AsyncMock(return_value=mock_batcher)
|
||||
mock_batcher.litellm_dailyuserspend = mock_table
|
||||
|
||||
|
||||
# Make the batch context manager's exit raise an exception
|
||||
# This simulates a batch commit failure (e.g., unique constraint violation)
|
||||
test_exception = Exception("Unique constraint violation")
|
||||
mock_batch_context.__aexit__ = AsyncMock(side_effect=test_exception)
|
||||
mock_prisma_client.db.batch_.return_value = mock_batch_context
|
||||
|
||||
|
||||
# Create a transaction
|
||||
daily_spend_transactions = {
|
||||
"test_key": {
|
||||
|
|
@ -927,13 +918,13 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure():
|
|||
"failed_requests": 0,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
# Create a mock proxy_logging_obj with failure_handler as AsyncMock
|
||||
mock_proxy_logging = MagicMock()
|
||||
mock_proxy_logging.failure_handler = AsyncMock()
|
||||
|
||||
|
||||
# Mock the logger to capture exception calls
|
||||
with patch.object(verbose_proxy_logger, "exception") as mock_exception_logger:
|
||||
with patch.object(verbose_proxy_logger, 'exception') as mock_exception_logger:
|
||||
# Call the method and expect it to raise the exception
|
||||
with pytest.raises(Exception, match="Unique constraint violation"):
|
||||
await DBSpendUpdateWriter._update_daily_spend(
|
||||
|
|
@ -946,16 +937,13 @@ async def test_update_daily_spend_logs_detailed_error_on_batch_upsert_failure():
|
|||
table_name="litellm_dailyuserspend",
|
||||
unique_constraint_name="user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint",
|
||||
)
|
||||
|
||||
|
||||
# Verify that exception was logged with detailed information
|
||||
assert mock_exception_logger.called
|
||||
call_args = mock_exception_logger.call_args[0][0]
|
||||
assert "Daily user spend batch upsert failed" in call_args
|
||||
assert "Table: litellm_dailyuserspend" in call_args
|
||||
assert (
|
||||
"Constraint: user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint"
|
||||
in call_args
|
||||
)
|
||||
assert "Constraint: user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name_endpoint" in call_args
|
||||
assert "Batch size: 1" in call_args
|
||||
assert "Unique constraint violation" in call_args
|
||||
|
||||
|
|
@ -973,7 +961,7 @@ async def test_update_daily_spend_re_raises_exception_after_logging():
|
|||
mock_batch_context = MagicMock()
|
||||
mock_batch_context.__aenter__ = AsyncMock(return_value=mock_batcher)
|
||||
mock_batcher.litellm_dailyuserspend = mock_table
|
||||
|
||||
|
||||
# Create a transaction
|
||||
daily_spend_transactions = {
|
||||
"test_key": {
|
||||
|
|
@ -990,16 +978,16 @@ async def test_update_daily_spend_re_raises_exception_after_logging():
|
|||
"failed_requests": 0,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
# Create a custom exception to verify it's re-raised
|
||||
custom_exception = ValueError("Database connection lost")
|
||||
mock_batch_context.__aexit__ = AsyncMock(side_effect=custom_exception)
|
||||
mock_prisma_client.db.batch_.return_value = mock_batch_context
|
||||
|
||||
|
||||
# Create a mock proxy_logging_obj with failure_handler as AsyncMock
|
||||
mock_proxy_logging = MagicMock()
|
||||
mock_proxy_logging.failure_handler = AsyncMock()
|
||||
|
||||
|
||||
# Verify the exception is re-raised
|
||||
with pytest.raises(ValueError, match="Database connection lost"):
|
||||
await DBSpendUpdateWriter._update_daily_spend(
|
||||
|
|
@ -1030,12 +1018,10 @@ async def test_commit_key_spend_updates_includes_last_active():
|
|||
mock_transaction = AsyncMock()
|
||||
mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction)
|
||||
mock_transaction.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_transaction.batch_ = MagicMock(
|
||||
return_value=AsyncMock(
|
||||
__aenter__=AsyncMock(return_value=mock_batcher),
|
||||
__aexit__=AsyncMock(return_value=False),
|
||||
)
|
||||
)
|
||||
mock_transaction.batch_ = MagicMock(return_value=AsyncMock(
|
||||
__aenter__=AsyncMock(return_value=mock_batcher),
|
||||
__aexit__=AsyncMock(return_value=False),
|
||||
))
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db = MagicMock()
|
||||
|
|
@ -1063,7 +1049,9 @@ async def test_commit_key_spend_updates_includes_last_active():
|
|||
|
||||
before_call = datetime.now(timezone.utc)
|
||||
|
||||
with patch("litellm.proxy.utils._raise_failed_update_spend_exception"):
|
||||
with patch(
|
||||
"litellm.proxy.utils._raise_failed_update_spend_exception"
|
||||
):
|
||||
await db_writer._commit_spend_updates_to_db(
|
||||
prisma_client=mock_prisma_client,
|
||||
n_retry_times=0,
|
||||
|
|
@ -1091,166 +1079,43 @@ async def test_commit_key_spend_updates_includes_last_active():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_database_creates_single_task():
|
||||
async def test_commit_spend_updates_uses_pipeline():
|
||||
"""
|
||||
Test that update_database() fires exactly 1 asyncio.create_task() call
|
||||
(the batched task) instead of the previous 11.
|
||||
Verify that _commit_spend_updates_to_db_with_redis uses
|
||||
get_all_transactions_from_redis_buffer_pipeline instead of 7 individual calls.
|
||||
"""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
|
||||
# Mock all helpers so nothing real runs
|
||||
db_writer._insert_spend_log_to_db = AsyncMock()
|
||||
db_writer._batch_database_updates = AsyncMock()
|
||||
mock_redis_update_buffer = AsyncMock()
|
||||
mock_redis_update_buffer.store_in_memory_spend_updates_in_redis = AsyncMock()
|
||||
# Return all-None tuple (no data to commit)
|
||||
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
|
||||
return_value=(None, None, None, None, None, None, None)
|
||||
)
|
||||
db_writer.redis_update_buffer = mock_redis_update_buffer
|
||||
|
||||
with patch("litellm.proxy.proxy_server.disable_spend_logs", False), patch(
|
||||
"litellm.proxy.proxy_server.prisma_client", MagicMock()
|
||||
), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), patch(
|
||||
"litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"
|
||||
), patch(
|
||||
"litellm.proxy.db.db_spend_update_writer.asyncio.create_task"
|
||||
) as mock_create_task:
|
||||
await db_writer.update_database(
|
||||
token="test-token",
|
||||
user_id="test-user",
|
||||
end_user_id="test-end-user",
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
team_id="test-team",
|
||||
org_id="test-org",
|
||||
completion_response=MagicMock(),
|
||||
response_cost=0.1,
|
||||
kwargs={"model": "gpt-4", "custom_llm_provider": "openai"},
|
||||
)
|
||||
mock_pod_lock_manager = AsyncMock()
|
||||
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
||||
mock_pod_lock_manager.release_lock = AsyncMock()
|
||||
db_writer.pod_lock_manager = mock_pod_lock_manager
|
||||
|
||||
# Exactly 1 create_task call (the batch), not 11
|
||||
assert mock_create_task.call_count == 1
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_proxy_logging = MagicMock()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_database_updates_isolation_on_failure():
|
||||
"""
|
||||
Test that if one helper inside _batch_database_updates raises,
|
||||
all other helpers still execute.
|
||||
"""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
|
||||
# Make _update_key_db raise
|
||||
db_writer._update_key_db = AsyncMock(side_effect=RuntimeError("key db boom"))
|
||||
|
||||
# All other helpers are normal mocks
|
||||
db_writer._update_user_db = AsyncMock()
|
||||
db_writer._update_team_db = AsyncMock()
|
||||
db_writer._update_org_db = AsyncMock()
|
||||
db_writer._update_tag_db = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock()
|
||||
|
||||
await db_writer._batch_database_updates(
|
||||
response_cost=0.1,
|
||||
user_id="u1",
|
||||
hashed_token="t1",
|
||||
team_id="team1",
|
||||
org_id="org1",
|
||||
end_user_id="eu1",
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
litellm_proxy_budget_name="budget",
|
||||
payload_copy={"key": "value"},
|
||||
request_tags=None,
|
||||
await db_writer._commit_spend_updates_to_db_with_redis(
|
||||
prisma_client=mock_prisma_client,
|
||||
n_retry_times=1,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
)
|
||||
|
||||
# _update_key_db raised, but all others should still have been called
|
||||
db_writer._update_user_db.assert_awaited_once()
|
||||
db_writer._update_key_db.assert_awaited_once()
|
||||
db_writer._update_team_db.assert_awaited_once()
|
||||
db_writer._update_org_db.assert_awaited_once()
|
||||
db_writer._update_tag_db.assert_awaited_once()
|
||||
db_writer.add_spend_log_transaction_to_daily_user_transaction.assert_awaited_once()
|
||||
db_writer.add_spend_log_transaction_to_daily_end_user_transaction.assert_awaited_once()
|
||||
db_writer.add_spend_log_transaction_to_daily_agent_transaction.assert_awaited_once()
|
||||
db_writer.add_spend_log_transaction_to_daily_team_transaction.assert_awaited_once()
|
||||
db_writer.add_spend_log_transaction_to_daily_org_transaction.assert_awaited_once()
|
||||
db_writer.add_spend_log_transaction_to_daily_tag_transaction.assert_awaited_once()
|
||||
# Pipeline method should be called once
|
||||
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_daily_agent_receives_deepcopied_payload():
|
||||
"""
|
||||
Test that the daily agent handler receives a deepcopied payload (not the original).
|
||||
|
||||
Previously, add_spend_log_transaction_to_daily_agent_transaction received the raw
|
||||
payload without a deepcopy, which was a mutation bug. This test goes through
|
||||
update_database() to verify the production deepcopy path.
|
||||
"""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
|
||||
# Capture the payload object that get_logging_payload returns (the "original")
|
||||
# and the payload the agent handler receives (should be a deepcopy)
|
||||
original_payload_ref = {}
|
||||
captured_agent_payloads = []
|
||||
|
||||
async def capture_agent_payload(**kwargs):
|
||||
captured_agent_payloads.append(kwargs.get("payload"))
|
||||
|
||||
# Mock all helpers
|
||||
db_writer._insert_spend_log_to_db = AsyncMock()
|
||||
db_writer._update_user_db = AsyncMock()
|
||||
db_writer._update_key_db = AsyncMock()
|
||||
db_writer._update_team_db = AsyncMock()
|
||||
db_writer._update_org_db = AsyncMock()
|
||||
db_writer._update_tag_db = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock(
|
||||
side_effect=capture_agent_payload
|
||||
)
|
||||
db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock()
|
||||
db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock()
|
||||
|
||||
# Mock get_logging_payload to return a known dict and capture its identity
|
||||
fake_payload = {
|
||||
"startTime": "2024-01-01T00:00:00",
|
||||
"endTime": "2024-01-01T00:01:00",
|
||||
"model": "gpt-4",
|
||||
"custom_llm_provider": "openai",
|
||||
"spend": 0.0,
|
||||
"nested": {"a": 1},
|
||||
}
|
||||
original_payload_ref["obj"] = fake_payload # store reference to the original
|
||||
|
||||
with patch("litellm.proxy.proxy_server.disable_spend_logs", True), patch(
|
||||
"litellm.proxy.proxy_server.prisma_client", MagicMock()
|
||||
), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), patch(
|
||||
"litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"
|
||||
), patch(
|
||||
"litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload",
|
||||
return_value=fake_payload,
|
||||
):
|
||||
await db_writer.update_database(
|
||||
token="test-token",
|
||||
user_id="test-user",
|
||||
end_user_id="test-end-user",
|
||||
team_id="test-team",
|
||||
org_id="test-org",
|
||||
kwargs={"model": "gpt-4", "custom_llm_provider": "openai"},
|
||||
completion_response=MagicMock(),
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
response_cost=0.1,
|
||||
)
|
||||
|
||||
# Let the single batched task run
|
||||
await asyncio.sleep(0)
|
||||
|
||||
# The agent handler should have been called
|
||||
assert len(captured_agent_payloads) == 1
|
||||
# The payload must NOT be the same object as the original (deepcopy occurred)
|
||||
assert captured_agent_payloads[0] is not original_payload_ref["obj"]
|
||||
# But it should have equivalent content
|
||||
assert captured_agent_payloads[0]["model"] == "gpt-4"
|
||||
assert captured_agent_payloads[0]["spend"] == 0.1
|
||||
# Individual methods should NOT be called
|
||||
mock_redis_update_buffer.get_all_update_transactions_from_redis_buffer.assert_not_called()
|
||||
mock_redis_update_buffer.get_all_daily_spend_update_transactions_from_redis_buffer.assert_not_called()
|
||||
mock_redis_update_buffer.get_all_daily_team_spend_update_transactions_from_redis_buffer.assert_not_called()
|
||||
mock_redis_update_buffer.get_all_daily_org_spend_update_transactions_from_redis_buffer.assert_not_called()
|
||||
mock_redis_update_buffer.get_all_daily_end_user_spend_update_transactions_from_redis_buffer.assert_not_called()
|
||||
mock_redis_update_buffer.get_all_daily_agent_spend_update_transactions_from_redis_buffer.assert_not_called()
|
||||
mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer.assert_not_called()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue