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

perf(spend): build the daily rollup row once per spend log instead of once per entity

Browse source 
_batch_database_updates queued the daily user, end user, agent, team, org and tag rows by calling _common_add_spend_log_transaction_to_daily_transaction once per entity, and each call re-parsed the metadata and re-ran compute_savings_spend for the same request. The shared fields do not depend on the entity, so they are now built once and handed to every daily rollup. The adders still build their own row when called without one
This commit is contained in:
narmstrong-berri 2026-10-19 20:39:39 +00:00
parent 2e427c6247
commit bd158c8108
No known key found for this signature in database
2 changed files with 172 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1131,10 +1131,12 @@ class DBSpendUpdateWriter:
traceback.format_exc(),
)
daily_spend_base: Final = await self._shared_daily_spend_base_transaction(payload_copy, prisma_client)
try:
await self.add_spend_log_transaction_to_daily_user_transaction(
payload=payload_copy,
prisma_client=prisma_client,
daily_spend_base=daily_spend_base,
)
except Exception:
verbose_proxy_logger.debug(
@ -1146,6 +1148,7 @@ class DBSpendUpdateWriter:
await self.add_spend_log_transaction_to_daily_end_user_transaction(
payload=payload_copy,
prisma_client=prisma_client,
daily_spend_base=daily_spend_base,
)
except Exception:
verbose_proxy_logger.debug(
@ -1157,6 +1160,7 @@ class DBSpendUpdateWriter:
await self.add_spend_log_transaction_to_daily_agent_transaction(
payload=payload_copy,
prisma_client=prisma_client,
daily_spend_base=daily_spend_base,
)
except Exception:
verbose_proxy_logger.debug(
@ -1168,6 +1172,7 @@ class DBSpendUpdateWriter:
await self.add_spend_log_transaction_to_daily_team_transaction(
payload=payload_copy,
prisma_client=prisma_client,
daily_spend_base=daily_spend_base,
)
except Exception:
verbose_proxy_logger.debug(
@ -1187,6 +1192,7 @@ class DBSpendUpdateWriter:
payload=payload_copy,
org_id=org_id,
prisma_client=prisma_client,
daily_spend_base=daily_spend_base,
)
except Exception:
verbose_proxy_logger.debug(
@ -1198,6 +1204,7 @@ class DBSpendUpdateWriter:
await self.add_spend_log_transaction_to_daily_tag_transaction(
payload=payload_copy,
prisma_client=prisma_client,
daily_spend_base=daily_spend_base,
)
except Exception:
verbose_proxy_logger.debug(
@ -2761,9 +2768,10 @@ class DBSpendUpdateWriter:
async def _common_add_spend_log_transaction_to_daily_transaction(
self,
payload: dict | SpendLogsPayload,
payload: SpendLogsPayload,
prisma_client: PrismaClient,
type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user",
daily_spend_base: BaseDailySpendTransaction | None = None,
) -> BaseDailySpendTransaction | None:
entity: Final = "tag" if type == "request_tags" else type
identity_payload: Final = (
@ -2771,6 +2779,31 @@ class DBSpendUpdateWriter:
)
if not daily_spend_entity_ids(identity_payload, entity):
return None
if daily_spend_base is not None:
return daily_spend_base
return await self._daily_spend_base_transaction(payload, prisma_client)
async def _shared_daily_spend_base_transaction(
self,
payload: SpendLogsPayload,
prisma_client: PrismaClient | None,
) -> BaseDailySpendTransaction | None:
if prisma_client is None:
return None
try:
return await self._daily_spend_base_transaction(payload, prisma_client)
except Exception: # noqa: BLE001 # each daily rollup falls back to building its own row, as before
verbose_proxy_logger.debug(
"_batch_database_updates: _daily_spend_base_transaction failed: %s",
traceback.format_exc(),
)
return None
async def _daily_spend_base_transaction(
self,
payload: SpendLogsPayload,
prisma_client: PrismaClient,
) -> BaseDailySpendTransaction | None:
expected_keys: Final = ("startTime", "api_key")
if not all(key in payload for key in expected_keys):
verbose_proxy_logger.debug(
@ -2806,7 +2839,7 @@ class DBSpendUpdateWriter:
if isinstance(payload["startTime"], datetime):
start_time: Final = payload["startTime"].isoformat()
date = start_time.split("T")[0]
elif isinstance(payload["startTime"], str):
elif isinstance(payload["startTime"], str): # pyright: ignore[reportUnnecessaryIsInstance] # spend log payloads are not validated against the TypedDict at runtime
date = payload["startTime"].split("T")[0]
else:
verbose_proxy_logger.debug(
@ -2874,8 +2907,9 @@ class DBSpendUpdateWriter:
async def add_spend_log_transaction_to_daily_user_transaction(
self,
payload: dict | SpendLogsPayload,
payload: SpendLogsPayload,
prisma_client: PrismaClient | None = None,
daily_spend_base: BaseDailySpendTransaction | None = None,
):
"""
Add a spend log transaction to the `daily_spend_update_queue`
@ -2889,7 +2923,7 @@ class DBSpendUpdateWriter:
return
base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction(
payload, prisma_client, "user"
payload, prisma_client, "user", daily_spend_base
)
if base_daily_transaction is None:
return
@ -2903,13 +2937,14 @@ class DBSpendUpdateWriter:
self,
payload: SpendLogsPayload,
prisma_client: PrismaClient | None = None,
daily_spend_base: BaseDailySpendTransaction | None = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")
return
base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction(
payload, prisma_client, "team"
payload, prisma_client, "team", daily_spend_base
)
if base_daily_transaction is None:
return
@ -2927,6 +2962,7 @@ class DBSpendUpdateWriter:
payload: SpendLogsPayload,
prisma_client: PrismaClient | None = None,
org_id: str | None = None,
daily_spend_base: BaseDailySpendTransaction | None = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")
@ -2945,7 +2981,7 @@ class DBSpendUpdateWriter:
)
base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction(
payload_with_org, prisma_client, "org"
payload_with_org, prisma_client, "org", daily_spend_base
)
if base_daily_transaction is None:
return
@ -2959,6 +2995,7 @@ class DBSpendUpdateWriter:
self,
payload: SpendLogsPayload,
prisma_client: PrismaClient | None = None,
daily_spend_base: BaseDailySpendTransaction | None = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")
@ -2978,7 +3015,7 @@ class DBSpendUpdateWriter:
)
base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction(
payload_with_end_user_id, prisma_client, "end_user"
payload_with_end_user_id, prisma_client, "end_user", daily_spend_base
)
if base_daily_transaction is None:
return
@ -2992,6 +3029,7 @@ class DBSpendUpdateWriter:
self,
payload: SpendLogsPayload,
prisma_client: PrismaClient | None = None,
daily_spend_base: BaseDailySpendTransaction | None = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")
@ -3006,7 +3044,7 @@ class DBSpendUpdateWriter:
},
)
base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction(
payload_with_agent_id, prisma_client, "agent"
payload_with_agent_id, prisma_client, "agent", daily_spend_base
)
if base_daily_transaction is None:
return
@ -3019,13 +3057,14 @@ class DBSpendUpdateWriter:
self,
payload: SpendLogsPayload,
prisma_client: PrismaClient | None = None,
daily_spend_base: BaseDailySpendTransaction | None = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")
return
base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction(
payload, prisma_client, "request_tags"
payload, prisma_client, "request_tags", daily_spend_base
)
if base_daily_transaction is None:
return

124
tests/unit/proxy/db/test_db_spend_update_writer.py
View file

@ -1,5 +1,6 @@
import asyncio
import copy
import itertools
import json
import logging
import re
@ -35,6 +36,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
build_window_spend_transaction,
)
from litellm.proxy.db.rollup_lock_timeout import ROLLUP_LOCK_TIMEOUT_SQL
from litellm.proxy.spend_tracking.savings import compute_savings_spend
from tests.unit.proxy.db.fake_prisma_engine import engine_call
@ -1440,6 +1442,128 @@ async def test_failed_project_enqueue_is_reported_and_does_not_drop_the_rest_of_
assert transactions["team_list_transactions"] == {"team-1": 0.25}
@pytest.mark.asyncio
async def test_batch_database_updates_prices_savings_once_for_every_daily_rollup():
db_writer: Final = DBSpendUpdateWriter()
prisma_client: Final = MagicMock()
prisma_client.get_request_status = MagicMock(return_value="success")
payload: Final = {
"request_id": "req-daily",
"user": "user-1",
"team_id": "team-1",
"end_user": "end-user-1",
"agent_id": "agent-1",
"request_tags": ["tag-a", "tag-b"],
"startTime": "2024-01-01T12:00:00",
"api_key": "key-1",
"model": "claude-sonnet-5",
"custom_llm_provider": "anthropic",
"model_group": "claude-sonnet-5",
"prompt_tokens": 1000,
"completion_tokens": 10,
"spend": 0.01,
"metadata": json.dumps({"usage_object": {"prompt_tokens": 1000, "cache_read_input_tokens": 800}}),
}
with patch("litellm.proxy.db.db_spend_update_writer.compute_savings_spend", wraps=compute_savings_spend) as savings:
await db_writer._batch_database_updates(
response_cost=0.01,
user_id="user-1",
hashed_token="key-1",
team_id="team-1",
org_id="org-1",
end_user_id="end-user-1",
prisma_client=prisma_client,
litellm_proxy_budget_name=None,
payload=payload,
)
assert savings.call_count == 1
queues: Final = {
"user_id": db_writer.daily_spend_update_queue,
"team_id": db_writer.daily_team_spend_update_queue,
"organization_id": db_writer.daily_org_spend_update_queue,
"end_user_id": db_writer.daily_end_user_spend_update_queue,
"agent_id": db_writer.daily_agent_spend_update_queue,
"tag": db_writer.daily_tag_spend_update_queue,
}
rows: Final = {
entity: tuple(
sorted((await queue.flush_and_get_aggregated_daily_spend_update_transactions()).values(), key=str)
)
for entity, queue in queues.items()
}
assert {entity: tuple(row[entity] for row in entity_rows) for entity, entity_rows in rows.items()} == {
"user_id": ("user-1",),
"team_id": ("team-1",),
"organization_id": ("org-1",),
"end_user_id": ("end-user-1",),
"agent_id": ("agent-1",),
"tag": ("tag-a", "tag-b"),
}
all_rows: Final = tuple(itertools.chain.from_iterable(rows.values()))
savings_columns: Final = {
(row["prompt_caching_savings_spend"], row["cache_read_input_tokens"], row["spend"]) for row in all_rows
}
assert len(savings_columns) == 1
assert next(iter(savings_columns))[0] > 0
@pytest.mark.asyncio
@pytest.mark.parametrize(
"overrides, dropped, has_prisma",
[
({"metadata": "{not json"}, (), True),
({}, (), False),
({}, ("api_key",), True),
({}, ("model",), True),
({}, ("model_group",), True),
],
)
async def test_batch_database_updates_skips_daily_rows_it_cannot_build_but_keeps_key_spend(
overrides: dict[str, str], dropped: tuple[str, ...], has_prisma: bool
):
db_writer: Final = DBSpendUpdateWriter()
prisma_client: Final = MagicMock() if has_prisma else None
if prisma_client is not None:
prisma_client.get_request_status = MagicMock(return_value="success")
complete_payload: Final = {
"request_id": "req-bad",
"user": "user-1",
"team_id": "team-1",
"startTime": "2024-01-01T12:00:00",
"api_key": "key-1",
"model": "gpt-4o-mini",
"custom_llm_provider": "openai",
"model_group": "gpt-4o-mini",
"prompt_tokens": 10,
"completion_tokens": 5,
"spend": 0.01,
"metadata": '{"usage_object": {}}',
**overrides,
}
payload: Final = {key: value for key, value in complete_payload.items() if key not in dropped}
await db_writer._batch_database_updates(
response_cost=0.01,
user_id="user-1",
hashed_token="key-1",
team_id="team-1",
org_id=None,
end_user_id=None,
prisma_client=prisma_client,
litellm_proxy_budget_name=None,
payload=payload,
)
assert await db_writer.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() == {}
assert (
await db_writer.daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() == {}
)
transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
assert transactions["key_list_transactions"] == ({"key-1": 0.01} if has_prisma else {})
@pytest.mark.asyncio
async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id():
"""