scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-27 06:13:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 72

fix(proxy): restore daily tag spend to the in-memory queue when the Redis RPUSH fails (#47307)

Browse source 
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yassin Kortam 2026-10-20 14:17:37 -07:00 • committed by GitHub
parent 2f139a70ae
commit 06bebcc5bb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 82 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -168,9 +168,11 @@ class RedisUpdateBuffer:
@with_service_target(SPEND_QUEUE_TARGET)
async def _store_transactions_in_redis(
self,
transactions: Mapping[str, BaseDailySpendTransaction] | None,
transactions: dict[str, BaseDailySpendTransaction] | None, # mutable-ok: add_update requires dict
redis_key: str,
service_type: ServiceTypes,
*,
source_queue: DailySpendUpdateQueue,
) -> None:
"""
Helper method to store transactions in Redis and emit an event
@ -179,6 +181,7 @@ class RedisUpdateBuffer:
transactions: The transactions to store
redis_key: The Redis key to store under
service_type: The service type for event emission
source_queue: The in-memory queue to restore transactions to if Redis push fails
"""
if transactions is None or len(transactions) == 0:
return
@ -191,6 +194,15 @@ class RedisUpdateBuffer:
key=redis_key,
values=list_of_transactions,
)
except Exception as e:
verbose_proxy_logger.error(
"Spend tracking - failed to push spend updates to Redis (redis_key=%s). "
"Restoring them to the in-memory queue for retry on next tick. Error: %s",
redis_key,
str(e),
)
await source_queue.add_update(transactions)
else:
verbose_proxy_logger.debug(
"Spend tracking - pushed spend updates to Redis buffer. redis_key=%s, buffer_size=%s",
redis_key,
@ -200,12 +212,6 @@ class RedisUpdateBuffer:
queue_size=current_redis_buffer_size,
service=service_type,
)
except Exception as e:
verbose_proxy_logger.error(
"Spend tracking - failed to push spend updates to Redis (redis_key=%s). Error: %s",
redis_key,
str(e),
)
@with_service_target(SPEND_QUEUE_TARGET)
async def store_in_memory_spend_updates_in_redis(
@ -793,6 +799,7 @@ class RedisUpdateBuffer:
transactions=daily_tag_spend_update_transactions,
redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE,
source_queue=daily_tag_spend_update_queue,
)
@with_service_target(SPEND_QUEUE_TARGET)

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

@ -6,10 +6,78 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.constants import REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import DailyTagSpendTransaction
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import DailySpendUpdateQueue
from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
from litellm.proxy.proxy_server import ProxyStartupEvent
def _daily_tag_spend_update_transactions() -> dict[str, DailyTagSpendTransaction]:
return {
"research_2025-01-15_api-key-123_gpt-4o_openai_/chat/completions": DailyTagSpendTransaction(
tag="research",
request_id="request-123",
date="2025-01-15",
api_key="api-key-123",
model="gpt-4o",
model_group=None,
mcp_namespaced_tool_name=None,
custom_llm_provider="openai",
endpoint="/chat/completions",
prompt_tokens=12,
completion_tokens=7,
cache_read_input_tokens=0,
cache_creation_input_tokens=0,
compression_saved_tokens=0,
compression_savings_spend=0.0,
prompt_caching_savings_spend=0.0,
gateway_injected_caching_savings_spend=0.0,
spend=0.0042,
api_requests=1,
successful_requests=1,
failed_requests=0,
)
}
@pytest.mark.asyncio
async def test_daily_tag_spend_restored_to_in_memory_queue_on_rpush_failure(
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock
):
mock_redis_cache.async_rpush = AsyncMock(side_effect=ConnectionError("redis went away"))
queue: Final = DailySpendUpdateQueue()
original_transactions: Final = _daily_tag_spend_update_transactions()
await queue.add_update(original_transactions)
await redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue)
assert await queue.flush_and_get_aggregated_daily_spend_update_transactions() == original_transactions
@pytest.mark.asyncio
async def test_daily_tag_spend_restored_after_rpush_failure_is_pushed_once_on_next_tick(
redis_update_buffer: RedisUpdateBuffer, mock_redis_cache: AsyncMock
):
mock_redis_cache.async_rpush = AsyncMock(side_effect=ConnectionError("redis went away"))
queue: Final = DailySpendUpdateQueue()
original_transactions: Final = _daily_tag_spend_update_transactions()
await queue.add_update(original_transactions)
await redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue)
mock_redis_cache.async_rpush = AsyncMock(return_value=1)
await redis_update_buffer.store_in_memory_daily_tag_spend_updates_in_redis(queue)
mock_redis_cache.async_rpush.assert_called_once()
redis_key: Final = mock_redis_cache.async_rpush.call_args.kwargs["key"]
values: Final = mock_redis_cache.async_rpush.call_args.kwargs["values"]
assert redis_key == REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY
assert json.loads(values[0]) == json.loads(safe_dumps(original_transactions))
assert await queue.flush_and_get_aggregated_daily_spend_update_transactions() == {}
@pytest.fixture
def mock_redis_cache():
"""Create a mock RedisCache instance"""