scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83

Merge commit 'a5f47a271a' into litellm_window_spend_reader

Browse source
This commit is contained in:
ryan-crabbe-berri 2026-08-29 13:55:53 -07:00
commit b5ec80d903
4 changed files with 151 additions and 100 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
litellm/proxy/db/db_spend_update_writer.py
View file

@ -936,6 +936,7 @@ class DBSpendUpdateWriter:
"daily_org_spend_update_transactions": daily_org_spend_update_transactions,
"daily_end_user_spend_update_transactions": daily_end_user_spend_update_transactions,
"daily_agent_spend_update_transactions": daily_agent_spend_update_transactions,
"window_spend_update_transactions": window_spend_update_transactions,
}
if db_spend_update_transactions is not None:
@ -1008,6 +1009,7 @@ class DBSpendUpdateWriter:
prisma_client=prisma_client,
window_spend_transactions=window_spend_update_transactions,
)
uncommitted.pop("window_spend_update_transactions", None)
except Exception as e:
spend_log_error(
"Spend tracking - failed to commit spend updates from Redis to DB. "
@ -1129,10 +1131,20 @@ class DBSpendUpdateWriter:
await self.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions()
)
await DBSpendUpdateWriter._commit_window_spend_updates(
prisma_client=prisma_client,
window_spend_transactions=window_spend_update_transactions,
)
try:
await DBSpendUpdateWriter._commit_window_spend_updates(
prisma_client=prisma_client,
window_spend_transactions=window_spend_update_transactions,
)
except Exception as e: # noqa: BLE001 # the increments go back on the queue; the rest of the flush must run
spend_log_error(
"Spend tracking - failed to commit budget window spend updates. "
"Re-queued %d window increments for retry on next tick. Error: %s",
len(window_spend_update_transactions),
str(e),
exc=e,
)
await self.window_spend_update_queue.update_queue.put(window_spend_update_transactions)
################## Tool Registry Upserts ##################
await self._flush_tool_discovery_queue(prisma_client=prisma_client)
@ -1206,27 +1218,19 @@ class DBSpendUpdateWriter:
"""
Commit per-budget-window spend increments to LiteLLM_BudgetWindowSpend.
Failures are logged rather than raised: these rows exist so budget
enforcement can stop aggregating LiteLLM_SpendLogs, and the read path
falls back to that aggregate, so a failed commit must not abort the
entity and daily spend commits that share this scheduler tick.
Raises on failure so the caller re-queues the increments: budget
enforcement trusts a current row without reconciling it against
LiteLLM_SpendLogs, so a dropped increment would let the entity spend
past its window limit after the next counter reseed.
"""
from litellm.proxy.db.budget_window_spend_writer import (
commit_window_spend_updates,
)
try:
await commit_window_spend_updates(
prisma_client=prisma_client,
transactions=window_spend_transactions,
)
except Exception as e: # noqa: BLE001 # any DB failure here must stay contained to the window rows
spend_log_error(
"Spend tracking - failed to commit budget window spend updates. %d window increments lost. Error: %s",
len(window_spend_transactions),
str(e),
exc=e,
)
await commit_window_spend_updates(
prisma_client=prisma_client,
transactions=window_spend_transactions,
)
async def _drain_and_commit_daily_tag_spend_from_redis(
self,

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

@ -456,6 +456,7 @@ class RedisUpdateBuffer:
daily_end_user_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None,
daily_agent_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None,
daily_tag_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None,
window_spend_update_transactions: Sequence[WindowSpendTransaction] | None = None,
) -> None:
"""
Re-push transactions that were popped from Redis but not committed to the DB.
@ -477,6 +478,7 @@ class RedisUpdateBuffer:
(daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY),
(daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY),
(daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY),
(window_spend_update_transactions, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY),
)
rpush_list: Final = tuple(

140
tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py
View file

@ -1,4 +1,5 @@
import json
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -22,9 +23,7 @@ def redis_update_buffer(mock_redis_cache):
@pytest.mark.asyncio
async def test_store_in_memory_spend_updates_uses_pipeline(
redis_update_buffer, mock_redis_cache
):
async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache):
"""
Verify store_in_memory_spend_updates_in_redis calls async_rpush_pipeline once
with the correct operations and skips empty queues.
@ -33,35 +32,29 @@ async def test_store_in_memory_spend_updates_uses_pipeline(
# Create mock queues - only 3 of 6 have data
spend_update_queue = AsyncMock()
spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = (
AsyncMock(return_value={"key_list_transactions": {"key1": 1.0}})
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_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}})
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_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_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_agent_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,
@ -82,9 +75,7 @@ async def test_store_in_memory_spend_updates_uses_pipeline(
@pytest.mark.asyncio
async def test_store_in_memory_spend_updates_restores_on_rpush_failure(
redis_update_buffer, mock_redis_cache
):
async def test_store_in_memory_spend_updates_restores_on_rpush_failure(redis_update_buffer, mock_redis_cache):
"""
If async_rpush_pipeline raises, the already-drained transactions must be
put back into the in-memory queues so the next scheduler tick retries.
@ -98,9 +89,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure(
SpendUpdateQueue,
)
mock_redis_cache.async_rpush_pipeline = AsyncMock(
side_effect=ConnectionError("redis went away")
)
mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away"))
spend_queue = SpendUpdateQueue()
daily_user_queue = DailySpendUpdateQueue()
@ -145,16 +134,12 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure(
# After restore, the main spend queue should hold one item per
# (entity_type, entity_id) pair with the aggregated cost
restored_spend = (
await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
)
restored_spend = await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
assert restored_spend["key_list_transactions"] == {"key-abc": 1.5}
assert restored_spend["team_list_transactions"] == {"team-xyz": 2.5}
# Daily user queue should hold the same aggregated dict
restored_daily = (
await daily_user_queue.flush_and_get_aggregated_daily_spend_update_transactions()
)
restored_daily = await daily_user_queue.flush_and_get_aggregated_daily_spend_update_transactions()
assert restored_daily == {
"user1_day_model": {
"spend": 1.0,
@ -165,9 +150,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure(
@pytest.mark.asyncio
async def test_store_in_memory_spend_updates_all_empty_returns_early(
redis_update_buffer, mock_redis_cache
):
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.
"""
@ -175,13 +158,9 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early(
# All queues return empty
empty_queue = AsyncMock()
empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(
return_value={}
)
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={})
)
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,
@ -196,9 +175,7 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early(
@pytest.mark.asyncio
async def test_get_all_transactions_from_redis_buffer_pipeline(
redis_update_buffer, mock_redis_cache
):
async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buffer, mock_redis_cache):
"""
Verify get_all_transactions_from_redis_buffer_pipeline correctly parses
and aggregates results from async_lpop_pipeline.
@ -287,9 +264,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline(
mock_redis_cache.async_lpop_pipeline.assert_called_once()
from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
popped_keys = [
op["key"] for op in mock_redis_cache.async_lpop_pipeline.call_args.kwargs["lpop_list"]
]
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
@ -302,9 +277,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis():
@pytest.mark.asyncio
async def test_restore_transactions_to_redis_pushes_only_provided(
redis_update_buffer, mock_redis_cache
):
async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_buffer, mock_redis_cache):
"""
restore_transactions_to_redis re-pushes only the transaction sets it was
given, to their matching buffer keys, so uncommitted spend can be retried.
@ -338,9 +311,42 @@ async def test_restore_transactions_to_redis_pushes_only_provided(
@pytest.mark.asyncio
async def test_restore_transactions_to_redis_noop_when_empty(
redis_update_buffer, mock_redis_cache
):
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
in the store path's encoding, so the next drain returns the same increments."""
from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
build_window_spend_transaction,
)
window_transactions = (
build_window_spend_transaction(
entity_type="key",
entity_id="hashed-token",
window_duration="30d",
window_start=datetime(2026, 8, 1, tzinfo=timezone.utc),
spend=3.0,
request_id="req-1",
started_at=datetime(2026, 8, 10, 12, 0, tzinfo=timezone.utc),
),
)
mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1])
await redis_update_buffer.restore_transactions_to_redis(window_spend_update_transactions=window_transactions)
rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"]
assert [op["key"] for op in rpush_list] == [REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY]
mock_redis_cache.async_lpop_pipeline = AsyncMock(
return_value=[None, None, None, None, None, None, list(rpush_list[0]["values"])]
)
drained = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
assert drained[6] == window_transactions
@pytest.mark.asyncio
async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer, mock_redis_cache):
"""Nothing to restore -> no Redis call."""
mock_redis_cache.async_rpush_pipeline = AsyncMock()
await redis_update_buffer.restore_transactions_to_redis()
@ -348,15 +354,11 @@ async def test_restore_transactions_to_redis_noop_when_empty(
@pytest.mark.asyncio
async def test_restore_transactions_to_redis_swallows_redis_error(
redis_update_buffer, mock_redis_cache
):
async def test_restore_transactions_to_redis_swallows_redis_error(redis_update_buffer, mock_redis_cache):
"""A Redis failure during restore must not propagate to the caller's finally block."""
from redis.exceptions import RedisError
mock_redis_cache.async_rpush_pipeline = AsyncMock(
side_effect=RedisError("redis down")
)
mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=RedisError("redis down"))
await redis_update_buffer.restore_transactions_to_redis(
db_spend_update_transactions={"key_list_transactions": {"key1": 1.0}},
@ -472,9 +474,7 @@ def test_get_transaction_buffer_redis_cache_parses_string_flag(monkeypatch):
@pytest.mark.asyncio
async def test_store_in_memory_spend_updates_pushes_budget_window_spend(
redis_update_buffer, mock_redis_cache
):
async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_update_buffer, mock_redis_cache):
"""The budget window queue has to ride the same rpush as the daily queues,
otherwise multi-pod deployments never persist per-window spend."""
from datetime import datetime, timezone
@ -519,15 +519,17 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend(
assert len(rpush_list) == 1
assert rpush_list[0]["key"] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY
pushed = json.loads(rpush_list[0]["values"][0])
assert pushed == [{
"entity_type": "key",
"entity_id": "hashed-token",
"window_duration": "30d",
"window_start": "2026-08-01T00:00:00.000000",
"spend": 1.25,
"request_ids": ["req-1"],
"started_at": "2026-08-10T12:00:00.000000",
}]
assert pushed == [
{
"entity_type": "key",
"entity_id": "hashed-token",
"window_duration": "30d",
"window_start": "2026-08-01T00:00:00.000000",
"spend": 1.25,
"request_ids": ["req-1"],
"started_at": "2026-08-10T12:00:00.000000",
}
]
@pytest.mark.asyncio

65
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -2569,19 +2569,21 @@ async def test_window_spend_transactions_are_not_committed_without_the_pod_lock(
@pytest.mark.asyncio
async def test_failed_window_spend_commit_does_not_abort_the_rest_of_the_flush():
"""Window rows are an optimization over aggregating LiteLLM_SpendLogs, so a
failure must not take the tool registry flush down with it."""
async def test_failed_window_spend_commit_requeues_the_increments_and_continues_the_flush():
"""Budget enforcement trusts a current window row without reconciling it
against LiteLLM_SpendLogs, so a dropped increment would let the key spend
past its limit after the next reseed. The increments must go back on the
queue, and the tool registry flush must still run."""
db_writer = DBSpendUpdateWriter()
await db_writer.window_spend_update_queue.add_update(
build_window_spend_transaction(
entity_type="key",
entity_id="hashed-token",
window_duration="30d",
window_start=datetime(2026, 8, 1, tzinfo=timezone.utc),
spend=0.5,
)
transaction = build_window_spend_transaction(
entity_type="key",
entity_id="hashed-token",
window_duration="30d",
window_start=datetime(2026, 8, 1, tzinfo=timezone.utc),
spend=0.5,
request_id="req-1",
)
await db_writer.window_spend_update_queue.add_update(transaction)
db = _WindowSpendFakeDB()
db.query_raw = AsyncMock(side_effect=Exception("connection reset"))
db_writer._flush_tool_discovery_queue = AsyncMock()
@ -2593,6 +2595,47 @@ async def test_failed_window_spend_commit_does_not_abort_the_rest_of_the_flush()
)
db_writer._flush_tool_discovery_queue.assert_called_once()
requeued = await db_writer.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions()
assert requeued == (transaction,)
@pytest.mark.asyncio
async def test_failed_window_spend_commit_from_redis_is_restored_to_redis():
"""The Redis drain is destructive, so a failed window commit has to push
the popped increments back exactly like the other spend categories."""
db_writer = DBSpendUpdateWriter()
window_transactions = (
build_window_spend_transaction(
entity_type="team",
entity_id="team-1",
window_duration="7d",
window_start=datetime(2026, 8, 1, tzinfo=timezone.utc),
spend=2.0,
request_id="req-1",
),
)
mock_redis_update_buffer = AsyncMock()
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
return_value=(None, None, None, None, None, None, window_transactions)
)
mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock()
db_writer.redis_update_buffer = mock_redis_update_buffer
db_writer.pod_lock_manager = AsyncMock()
db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
db = _WindowSpendFakeDB()
db.query_raw = AsyncMock(side_effect=Exception("connection reset"))
await db_writer._commit_spend_updates_to_db_with_redis(
prisma_client=_WindowSpendFakePrisma(db),
n_retry_times=0,
proxy_logging_obj=MagicMock(),
)
assert _window_spend_upserts(db) == []
mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once_with(
window_spend_update_transactions=window_transactions
)
db_writer.pod_lock_manager.release_lock.assert_awaited_once()
@pytest.mark.asyncio