mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
fix(spend_tracking): re-key spend log rows a reused provider response id would drop
LiteLLM_SpendLogs.request_id is the provider's response id and the flush inserts with skip_duplicates, so a self-hosted OpenAI-compatible server that answers every request with the same completion id had every request after the first served and charged but never logged. The flush now reads back, from the writer, the identity of the rows create_many skipped and re-inserts the ones that belong to a different request keyed on their own litellm_call_id, with the provider id kept in metadata.response_id. Rows the flush already stored (a replay after a transport retry, under either key) and rows keyed on a stored object's id (batch polls, file uploads, the zero-priced object reads) stay skipped, and a read-back the database rejects leaves the rows skipped instead of requeueing the flush. request_id keeps meaning the id the caller was shown, so GET /spend/logs?request_id= by response id, the Logs page search, and Responses API previous_response_id session lookups behave as before.
This commit is contained in:
Yucheng He
parent
af6dc1db08
commit
e7442c7232
7 changed files with 500 additions and 8 deletions
|
|
@ -3876,6 +3876,7 @@ class SpendLogsMetadata(TypedDict):
|
|||
autorouter_savings: ReadOnly[float | None] # stamped by the logging payload; None = not auto-routed
|
||||
litellm_gateway_injected_cache: ReadOnly[str | None]
|
||||
router_metadata: ReadOnly[SpendLogsRouterMetadata | None] # None = deployment not flagged internal_router_model
|
||||
response_id: ReadOnly[str | None] # the provider's response id; request_id unless the row was re-keyed
|
||||
|
||||
|
||||
class SpendLogsPayload(TypedDict):
|
||||
|
|
|
|||
|
|
@ -144,6 +144,7 @@ def _get_spend_logs_metadata(
|
|||
litellm_call_id: str | None = None,
|
||||
autorouter_savings: float | None = None,
|
||||
router_metadata: SpendLogsRouterMetadata | None = None,
|
||||
response_id: str | None = None,
|
||||
) -> SpendLogsMetadata:
|
||||
if metadata is None:
|
||||
return SpendLogsMetadata(
|
||||
|
|
@ -184,6 +185,7 @@ def _get_spend_logs_metadata(
|
|||
litellm_gateway_injected_cache=None,
|
||||
litellm_call_id=litellm_call_id,
|
||||
router_metadata=router_metadata,
|
||||
response_id=response_id,
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
"getting payload for SpendLogs, available keys in metadata: " + str(list(metadata.keys()))
|
||||
|
|
@ -191,8 +193,13 @@ def _get_spend_logs_metadata(
|
|||
|
||||
# Filter the metadata dictionary to include only the specified keys
|
||||
clean_metadata: Final = SpendLogsMetadata(
|
||||
**{key: metadata.get(key) for key in SpendLogsMetadata.__annotations__ if key != "router_metadata"},
|
||||
**{
|
||||
key: metadata.get(key)
|
||||
for key in SpendLogsMetadata.__annotations__
|
||||
if key not in ("router_metadata", "response_id")
|
||||
},
|
||||
router_metadata=router_metadata,
|
||||
response_id=response_id,
|
||||
)
|
||||
_raw_key: Final = clean_metadata.get("user_api_key")
|
||||
_trusted_hash: Final = metadata.get("user_api_key_hash")
|
||||
|
|
@ -387,6 +394,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
|
|||
usage = _combined_usage.model_dump()
|
||||
|
||||
id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs)
|
||||
raw_response_id: Final = response_obj_dict.get("id")
|
||||
standard_logging_payload: Final = cast(StandardLoggingPayload | None, kwargs.get("standard_logging_object", None))
|
||||
|
||||
end_user_id = get_end_user_id_for_cost_tracking(litellm_params)
|
||||
|
|
@ -522,6 +530,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
|
|||
standard_logging_payload.get("autorouter_savings", None) if standard_logging_payload is not None else None
|
||||
),
|
||||
litellm_call_id=litellm_call_id,
|
||||
response_id=raw_response_id if isinstance(raw_response_id, str) else None,
|
||||
router_metadata=_get_router_metadata_for_spend_log(
|
||||
metadata=metadata,
|
||||
requested_model=_model_group,
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ from litellm.constants import (
|
|||
DEFAULT_MODEL_CREATED_AT_TIME,
|
||||
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
|
||||
MAX_TEAM_LIST_LIMIT,
|
||||
NON_INFERENCE_CALL_TYPES,
|
||||
SPEND_LOG_QUEUE_MAX_BYTES,
|
||||
SPEND_LOG_WRITE_BATCH_MAX_BYTES,
|
||||
SPEND_LOG_WRITE_BATCH_MAX_ROWS,
|
||||
|
|
@ -142,7 +143,7 @@ from litellm.proxy.db.prisma_client import (
|
|||
PrismaWrapper,
|
||||
parse_iam_endpoint_from_url,
|
||||
)
|
||||
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
|
||||
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient
|
||||
from litellm.proxy.db.spend_log_batching import (
|
||||
spend_log_queue_within_budget,
|
||||
spend_log_row_bytes,
|
||||
|
|
@ -7234,11 +7235,102 @@ async def _monitor_spend_logs_queue(
|
|||
|
||||
MAX_SPEND_LOG_ISOLATION_FAILURES_PER_BATCH: Final = 256
|
||||
|
||||
_OBJECT_ID_KEYED_SPEND_LOG_CALL_TYPES: Final = NON_INFERENCE_CALL_TYPES | frozenset(
|
||||
{
|
||||
CallTypes.create_batch.value,
|
||||
CallTypes.acreate_batch.value,
|
||||
CallTypes.retrieve_batch.value,
|
||||
CallTypes.aretrieve_batch.value,
|
||||
CallTypes.create_file.value,
|
||||
CallTypes.acreate_file.value,
|
||||
}
|
||||
)
|
||||
"""Rows keyed on the id of the object the call addressed rather than on a fresh response id:
|
||||
the zero-priced reads and management calls of a stored object, batch polls (every poll of one
|
||||
batch collapses into its single cost row, the claim in ``_claim_batch_cost_spend_log``) and file
|
||||
uploads. A second row for one of these collapses on purpose."""
|
||||
|
||||
|
||||
def _is_transient_spend_log_write_error(e: Exception) -> bool:
|
||||
return PrismaDBExceptionHandler.is_database_transport_error(e) or PrismaDBExceptionHandler.is_deadlock_error(e)
|
||||
|
||||
|
||||
def _spend_log_row_can_rekey_on_call_id(row: Mapping[str, object]) -> bool:
|
||||
call_id: Final = row.get("litellm_call_id")
|
||||
return isinstance(call_id, str) and call_id != "" and call_id != row.get("request_id")
|
||||
|
||||
|
||||
def _spend_log_row_is_stored(row: Mapping[str, object], stored_identities: frozenset[tuple[object, object]]) -> bool:
|
||||
"""A stored row with this row's identity is this row or a replay of it, under either key."""
|
||||
call_id: Final = row.get("litellm_call_id")
|
||||
return (row.get("request_id"), call_id) in stored_identities or (call_id, call_id) in stored_identities
|
||||
|
||||
|
||||
async def _spend_logs_rekeyed_after_duplicate_skip(
|
||||
repo: SpendLogsRepository, rows: Sequence[Mapping[str, object]]
|
||||
) -> tuple[Mapping[str, object], ...]:
|
||||
"""The rows ``create_many(skip_duplicates=True)`` skipped, re-keyed on their own ``litellm_call_id``.
|
||||
|
||||
``request_id`` is the provider's response id, so a provider that reuses one completion id
|
||||
(a self-hosted OpenAI-compatible server, commonly) had every request after the first
|
||||
served and charged but never logged (LIT-6666). The insert reports only how many rows
|
||||
landed, so the rows are read back by identity from the writer (a lagging read replica
|
||||
would report the rows this very insert landed as missing): a stored row carrying this
|
||||
row's ``request_id`` AND ``litellm_call_id``, or already keyed on its call id, is this row
|
||||
or a replay of it after a transport retry, and stays skipped. Rows keyed on an object id
|
||||
collapse on purpose, and a row already keyed on its call id has nothing left to fall back
|
||||
to; those are skipped and only logged.
|
||||
"""
|
||||
unverified: Final = tuple(row for row in rows if row.get("call_type") not in _OBJECT_ID_KEYED_SPEND_LOG_CALL_TYPES)
|
||||
if not unverified:
|
||||
return ()
|
||||
stored_identities: Final = await SpendLogsRepository(WriterPinnedClient(repo.prisma_client.db)).stored_identities(
|
||||
str(key) for row in unverified for key in (row.get("request_id"), row.get("litellm_call_id")) if key
|
||||
)
|
||||
dropped: Final = tuple(row for row in unverified if not _spend_log_row_is_stored(row, stored_identities))
|
||||
rekeyable: Final = tuple(row for row in dropped if _spend_log_row_can_rekey_on_call_id(row))
|
||||
for row in dropped:
|
||||
if not _spend_log_row_can_rekey_on_call_id(row):
|
||||
verbose_proxy_logger.error(
|
||||
"Spend tracking - dropping spend log row whose request_id %s is already taken and whose "
|
||||
"litellm_call_id %s offers no other key",
|
||||
row.get("request_id"),
|
||||
row.get("litellm_call_id"),
|
||||
)
|
||||
if rekeyable:
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - %d spend log row(s) shared a request_id with a stored row and were re-keyed on "
|
||||
"their litellm_call_id: %s",
|
||||
len(rekeyable),
|
||||
sorted(frozenset(str(row.get("request_id")) for row in rekeyable)),
|
||||
)
|
||||
return tuple(
|
||||
{**row, "request_id": row["litellm_call_id"]} # mutable-ok: prisma create_many takes dict rows
|
||||
for row in rekeyable
|
||||
)
|
||||
|
||||
|
||||
async def _spend_logs_rekeyed_or_left_skipped(
|
||||
repo: SpendLogsRepository, rows: Sequence[Mapping[str, object]]
|
||||
) -> tuple[Mapping[str, object], ...]:
|
||||
"""``_spend_logs_rekeyed_after_duplicate_skip``, leaving the rows skipped (the pre-existing
|
||||
outcome) when the read-back fails on its data rather than on transport, so one bad read
|
||||
cannot requeue a whole flush behind it forever."""
|
||||
try:
|
||||
return await _spend_logs_rekeyed_after_duplicate_skip(repo, rows)
|
||||
except Exception as e:
|
||||
if _is_transient_spend_log_write_error(e):
|
||||
raise
|
||||
spend_log_error(
|
||||
"Spend tracking - could not read back which of %d skipped spend log rows are stored; leaving them "
|
||||
"skipped. error=%s",
|
||||
len(rows),
|
||||
str(e),
|
||||
exc=e,
|
||||
)
|
||||
return ()
|
||||
|
||||
|
||||
async def _create_spend_logs_with_poison_isolation(
|
||||
repo: SpendLogsRepository,
|
||||
rows: Sequence[Mapping[str, object]],
|
||||
|
|
@ -7267,8 +7359,7 @@ async def _create_spend_logs_with_poison_isolation(
|
|||
persist. Returns the budget left after this subtree.
|
||||
"""
|
||||
try:
|
||||
await repo.table.create_many(data=rows, skip_duplicates=True)
|
||||
return failure_budget
|
||||
inserted: Final = await repo.table.create_many(data=rows, skip_duplicates=True)
|
||||
except Exception as e:
|
||||
if not PrismaDBExceptionHandler.is_prisma_data_error(e):
|
||||
raise
|
||||
|
|
@ -7303,6 +7394,11 @@ async def _create_spend_logs_with_poison_isolation(
|
|||
)
|
||||
return 0
|
||||
return await _create_spend_logs_with_poison_isolation(repo, rows[mid:], remaining)
|
||||
if isinstance(inserted, int) and inserted < len(rows):
|
||||
rekeyed: Final = await _spend_logs_rekeyed_or_left_skipped(repo, rows)
|
||||
if rekeyed:
|
||||
return await _create_spend_logs_with_poison_isolation(repo, rekeyed, failure_budget)
|
||||
return failure_budget
|
||||
|
||||
|
||||
def _raise_failed_update_spend_exception(e: Exception, start_time: float, proxy_logging_obj: ProxyLogging):
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ These are thin wrappers for tables that do not (yet) need domain-specific query
|
|||
methods; richer repositories live in their own modules.
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable
|
||||
from typing import TYPE_CHECKING, Any, Final, Generic
|
||||
|
||||
from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync
|
||||
|
|
@ -64,9 +65,19 @@ class OrganizationMembershipRepository(PrismaTableRepository["prisma_models.Lite
|
|||
table_name = "litellm_organizationmembership"
|
||||
|
||||
|
||||
_SPEND_LOG_IDENTITIES_SQL: Final = (
|
||||
'SELECT request_id, litellm_call_id FROM "LiteLLM_SpendLogs" WHERE request_id = ANY($1::text[])'
|
||||
)
|
||||
|
||||
|
||||
class SpendLogsRepository(PrismaTableRepository["prisma_models.LiteLLM_SpendLogs"]):
|
||||
table_name = "litellm_spendlogs"
|
||||
|
||||
async def stored_identities(self, request_ids: Iterable[str]) -> frozenset[tuple[object, object]]:
|
||||
"""``(request_id, litellm_call_id)`` of every stored row keyed on one of ``request_ids``."""
|
||||
rows: Final = await self.prisma_client.db.query_raw(_SPEND_LOG_IDENTITIES_SQL, sorted(frozenset(request_ids)))
|
||||
return frozenset((row.get("request_id"), row.get("litellm_call_id")) for row in rows)
|
||||
|
||||
|
||||
class BudgetWindowSpendRepository(PrismaTableRepository["prisma_models.LiteLLM_BudgetWindowSpend"]):
|
||||
table_name = "litellm_budgetwindowspend"
|
||||
|
|
|
|||
|
|
@ -643,6 +643,7 @@ ignored_keys = [
|
|||
"request_id",
|
||||
"litellm_call_id",
|
||||
"metadata.litellm_call_id",
|
||||
"metadata.response_id",
|
||||
"session_id",
|
||||
"startTime",
|
||||
"endTime",
|
||||
|
|
|
|||
|
|
@ -1356,6 +1356,43 @@ def test_get_logging_payload_populates_litellm_call_id_alongside_provider_reques
|
|||
assert payload["litellm_call_id"] == call_id
|
||||
|
||||
|
||||
def test_get_logging_payload_keeps_the_provider_response_id_in_metadata():
|
||||
"""LIT-6666: a row the flush re-keys on its call id (its provider id was already taken
|
||||
by another request) still has to say which response it logged."""
|
||||
payload = get_logging_payload(
|
||||
kwargs={
|
||||
"model": "gpt-4o-mini",
|
||||
"litellm_call_id": "call-1",
|
||||
"litellm_params": {"metadata": {"user_api_key": "test-key", "response_id": "caller-supplied"}},
|
||||
},
|
||||
response_obj=litellm.ModelResponse(
|
||||
id="chatcmpl-static-1",
|
||||
choices=[],
|
||||
usage=litellm.Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2),
|
||||
),
|
||||
start_time=datetime.datetime.now(timezone.utc),
|
||||
end_time=datetime.datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
assert json.loads(payload["metadata"])["response_id"] == "chatcmpl-static-1"
|
||||
|
||||
|
||||
def test_get_logging_payload_leaves_metadata_response_id_empty_without_a_response_id():
|
||||
"""A failure row has no provider response; a caller cannot fill the slot through request metadata."""
|
||||
payload = get_logging_payload(
|
||||
kwargs={
|
||||
"model": "gpt-4o-mini",
|
||||
"litellm_call_id": "call-1",
|
||||
"litellm_params": {"metadata": {"user_api_key": "test-key", "response_id": "caller-supplied"}},
|
||||
},
|
||||
response_obj=None,
|
||||
start_time=datetime.datetime.now(timezone.utc),
|
||||
end_time=datetime.datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
assert json.loads(payload["metadata"])["response_id"] is None
|
||||
|
||||
|
||||
@patch("litellm.proxy.proxy_server.master_key", None)
|
||||
@patch("litellm.proxy.proxy_server.general_settings", {})
|
||||
def test_get_logging_payload_includes_overhead_in_spend_logs_metadata():
|
||||
|
|
|
|||
|
|
@ -10,6 +10,9 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
|
||||
import httpx
|
||||
from collections.abc import Iterator
|
||||
from typing import Any, Dict, List
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
|
@ -17,6 +20,7 @@ from unittest.mock import AsyncMock, MagicMock
|
|||
import pytest
|
||||
|
||||
import litellm.proxy.utils as utils_mod
|
||||
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
|
||||
from litellm.proxy.db.spend_log_batching import spend_log_row_bytes
|
||||
from litellm.proxy.utils import PrismaClient, ProxyUpdateSpend, enqueue_spend_logs
|
||||
|
||||
|
|
@ -475,7 +479,7 @@ async def test_update_spend_logs_retries_and_requeues_batch_on_db_outage(
|
|||
|
||||
def _deadlock_error() -> Exception:
|
||||
return _data_error(
|
||||
'Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: None, '
|
||||
"Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: None, "
|
||||
'kind: QueryError(PostgresError { code: "40P01", message: "deadlock detected", severity: "ERROR" }) })'
|
||||
)
|
||||
|
||||
|
|
@ -508,9 +512,7 @@ async def test_update_spend_logs_retries_deadlock_and_keeps_every_row(
|
|||
logs_to_process=[make_spend_log_row(request_id="a"), make_spend_log_row(request_id="b")],
|
||||
)
|
||||
|
||||
attempts = tuple(
|
||||
tuple(row["request_id"] for row in call.kwargs["data"]) for call in create_many.await_args_list
|
||||
)
|
||||
attempts = tuple(tuple(row["request_id"] for row in call.kwargs["data"]) for call in create_many.await_args_list)
|
||||
assert attempts == (("a", "b"), ("a", "b"), ("a", "b"))
|
||||
assert mock_prisma_client.spend_log_transactions == []
|
||||
|
||||
|
|
@ -883,3 +885,338 @@ def test_disable_spend_updates_error_when_general_settings_unavailable(
|
|||
monkeypatch.delattr(proxy_server_mod, "general_settings", raising=False)
|
||||
with pytest.raises(ImportError):
|
||||
ProxyUpdateSpend.disable_spend_updates()
|
||||
|
||||
|
||||
class _SpendLogsTable:
|
||||
"""``LiteLLM_SpendLogs`` as the flush sees it: ``create_many`` is ``INSERT ... ON CONFLICT
|
||||
DO NOTHING`` on ``request_id`` and reports how many rows landed, and ``query_raw`` answers the
|
||||
identity read-back with the stored ``(request_id, litellm_call_id)`` pairs."""
|
||||
|
||||
def __init__(self, seeded: List[Dict[str, Any]] | None = None) -> None:
|
||||
self.rows: Dict[str, Dict[str, Any]] = {row["request_id"]: dict(row) for row in seeded or []}
|
||||
self.inserts: List[List[str]] = []
|
||||
self.query_raw_calls = 0
|
||||
|
||||
async def create_many(self, *, data: Any, skip_duplicates: bool) -> int:
|
||||
assert skip_duplicates is True
|
||||
self.inserts.append([row["request_id"] for row in data])
|
||||
inserted = 0
|
||||
for row in data:
|
||||
if row["request_id"] in self.rows:
|
||||
continue
|
||||
self.rows[row["request_id"]] = dict(row)
|
||||
inserted += 1
|
||||
return inserted
|
||||
|
||||
async def query_raw(self, sql: str, request_ids: Any) -> List[Dict[str, Any]]:
|
||||
self.query_raw_calls += 1
|
||||
assert "WHERE request_id = ANY($1::text[])" in sql
|
||||
return [
|
||||
{"request_id": rid, "litellm_call_id": row.get("litellm_call_id")}
|
||||
for rid, row in self.rows.items()
|
||||
if rid in request_ids
|
||||
]
|
||||
|
||||
|
||||
def _wire_spend_logs_table(mock_prisma_client: Any, seeded: List[Dict[str, Any]] | None = None) -> _SpendLogsTable:
|
||||
table = _SpendLogsTable(seeded)
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=table.create_many)
|
||||
mock_prisma_client.db.query_raw = AsyncMock(side_effect=table.query_raw)
|
||||
return table
|
||||
|
||||
|
||||
def _wire_routed_db(mock_prisma_client: Any, table: _SpendLogsTable) -> MagicMock:
|
||||
"""``prisma_client.db`` as a read-replica router: top-level reads go to the reader, which never
|
||||
sees this flush's writes, and ``writer`` is the engine the rows landed on."""
|
||||
routed = MagicMock(spec=RoutingPrismaWrapper)
|
||||
routed.writer_unavailable = False
|
||||
routed.litellm_spendlogs = MagicMock()
|
||||
routed.litellm_spendlogs.create_many = AsyncMock(side_effect=table.create_many)
|
||||
routed.query_raw = AsyncMock(return_value=[])
|
||||
routed.writer = MagicMock()
|
||||
routed.writer.query_raw = AsyncMock(side_effect=table.query_raw)
|
||||
mock_prisma_client.db = routed
|
||||
return routed
|
||||
|
||||
|
||||
async def _flush(mock_prisma_client: Any, logs: List[Dict[str, Any]]) -> None:
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=0,
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
logs_to_process=logs,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_rekeys_the_rows_a_reused_provider_response_id_would_drop(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any
|
||||
) -> None:
|
||||
"""Three requests answered with one provider id are three charged requests, so three rows
|
||||
must land: the first keeps the provider id, the other two are re-keyed on their own call id.
|
||||
Observed on a live proxy against a self-hosted server returning a fixed completion id: the
|
||||
duplicate-tolerant flush kept one row while key and daily spend counted all three."""
|
||||
table = _wire_spend_logs_table(mock_prisma_client)
|
||||
logs = [
|
||||
make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}", call_type="acompletion")
|
||||
for i in range(3)
|
||||
]
|
||||
|
||||
await _flush(mock_prisma_client, logs)
|
||||
|
||||
assert {rid: row["litellm_call_id"] for rid, row in table.rows.items()} == {
|
||||
"chatcmpl-static-1": "call-0",
|
||||
"call-1": "call-1",
|
||||
"call-2": "call-2",
|
||||
}
|
||||
assert table.inserts == [["chatcmpl-static-1"] * 3, ["call-1", "call-2"]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_reads_stored_identities_back_from_the_writer(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any
|
||||
) -> None:
|
||||
"""The read-back must see the rows this flush just landed, so it goes to the writer: with a
|
||||
read replica configured, ``db.query_raw`` is routed to the reader, and a lagging reader would
|
||||
report every just-inserted row as missing and re-insert it under its call id."""
|
||||
table = _wire_spend_logs_table(mock_prisma_client)
|
||||
routed = _wire_routed_db(mock_prisma_client, table)
|
||||
logs = [
|
||||
make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}", call_type="acompletion")
|
||||
for i in range(2)
|
||||
]
|
||||
|
||||
await _flush(mock_prisma_client, logs)
|
||||
|
||||
routed.writer.query_raw.assert_awaited_once()
|
||||
routed.query_raw.assert_not_awaited()
|
||||
assert sorted(table.rows) == ["call-1", "chatcmpl-static-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_rekeys_only_the_colliding_rows_of_a_mixed_batch(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any
|
||||
) -> None:
|
||||
"""Two providers reusing ids and one issuing unique ids share a flush: every row of the
|
||||
unique provider keeps its key and only the later rows of each reused id are re-keyed."""
|
||||
table = _wire_spend_logs_table(mock_prisma_client)
|
||||
logs = [
|
||||
make_spend_log_row(request_id="chatcmpl-unique-a", litellm_call_id="call-a", call_type="acompletion"),
|
||||
make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id="call-b", call_type="acompletion"),
|
||||
make_spend_log_row(request_id="chatcmpl-static-2", litellm_call_id="call-c", call_type="acompletion"),
|
||||
make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id="call-d", call_type="acompletion"),
|
||||
make_spend_log_row(request_id="chatcmpl-unique-e", litellm_call_id="call-e", call_type="acompletion"),
|
||||
make_spend_log_row(request_id="chatcmpl-static-2", litellm_call_id="call-f", call_type="acompletion"),
|
||||
]
|
||||
|
||||
await _flush(mock_prisma_client, logs)
|
||||
|
||||
assert {rid: row["litellm_call_id"] for rid, row in table.rows.items()} == {
|
||||
"chatcmpl-unique-a": "call-a",
|
||||
"chatcmpl-static-1": "call-b",
|
||||
"chatcmpl-static-2": "call-c",
|
||||
"call-d": "call-d",
|
||||
"chatcmpl-unique-e": "call-e",
|
||||
"call-f": "call-f",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_rekeys_a_row_whose_provider_id_an_earlier_flush_stored(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any
|
||||
) -> None:
|
||||
"""The colliding row usually landed in an earlier flush, or from another worker."""
|
||||
table = _wire_spend_logs_table(
|
||||
mock_prisma_client,
|
||||
seeded=[
|
||||
make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id="call-old", call_type="acompletion")
|
||||
],
|
||||
)
|
||||
|
||||
await _flush(
|
||||
mock_prisma_client,
|
||||
[make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id="call-new", call_type="acompletion")],
|
||||
)
|
||||
|
||||
assert {rid: row["litellm_call_id"] for rid, row in table.rows.items()} == {
|
||||
"chatcmpl-static-1": "call-old",
|
||||
"call-new": "call-new",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_does_not_rekey_a_replay_of_a_stored_row(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any
|
||||
) -> None:
|
||||
"""A transport retry replays the whole batch; a row already stored under its own call id is
|
||||
the same request, not a collision, and must not become a second row."""
|
||||
table = _wire_spend_logs_table(
|
||||
mock_prisma_client,
|
||||
seeded=[make_spend_log_row(request_id="chatcmpl-1", litellm_call_id="call-1", call_type="acompletion")],
|
||||
)
|
||||
|
||||
await _flush(
|
||||
mock_prisma_client,
|
||||
[make_spend_log_row(request_id="chatcmpl-1", litellm_call_id="call-1", call_type="acompletion")],
|
||||
)
|
||||
|
||||
assert list(table.rows) == ["chatcmpl-1"]
|
||||
assert len(table.inserts) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"call_type", ["acreate_batch", "aretrieve_batch", "acreate_file", "avector_store_retrieve", "avector_store_delete"]
|
||||
)
|
||||
async def test_update_spend_logs_keeps_object_keyed_rows_collapsed_on_their_object_id(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any, call_type: str
|
||||
) -> None:
|
||||
"""Every poll of one batch shares its cost row by design, and every read of a stored object
|
||||
is keyed on that object's id, so a duplicate here is not a lost row and the identity read-back
|
||||
is not even issued."""
|
||||
table = _wire_spend_logs_table(mock_prisma_client)
|
||||
logs = [
|
||||
make_spend_log_row(request_id="batch_abc_batch_cost", litellm_call_id=f"call-{i}", call_type=call_type)
|
||||
for i in range(2)
|
||||
]
|
||||
|
||||
await _flush(mock_prisma_client, logs)
|
||||
|
||||
assert list(table.rows) == ["batch_abc_batch_cost"]
|
||||
assert table.query_raw_calls == 0
|
||||
assert len(table.inserts) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_drops_and_logs_a_row_already_keyed_on_its_call_id(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""A row whose key already is its call id (a client pinning x-litellm-call-id on a failure
|
||||
row, say) has nothing to fall back to: it stays dropped, but loudly."""
|
||||
table = _wire_spend_logs_table(
|
||||
mock_prisma_client,
|
||||
seeded=[make_spend_log_row(request_id="call-pinned", litellm_call_id="call-other", call_type="acompletion")],
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.ERROR, logger=utils_mod.verbose_proxy_logger.name):
|
||||
await _flush(
|
||||
mock_prisma_client,
|
||||
[make_spend_log_row(request_id="call-pinned", litellm_call_id="call-pinned", call_type="acompletion")],
|
||||
)
|
||||
|
||||
assert list(table.rows) == ["call-pinned"]
|
||||
assert len(table.inserts) == 1
|
||||
dropped = [record.getMessage() for record in caplog.records if "dropping spend log row" in record.getMessage()]
|
||||
assert dropped == [
|
||||
"Spend tracking - dropping spend log row whose request_id call-pinned is already taken and whose "
|
||||
"litellm_call_id call-pinned offers no other key"
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_rekey_that_collides_again_stops(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any
|
||||
) -> None:
|
||||
"""A client pinning one x-litellm-call-id against a provider reusing one response id
|
||||
collides on both keys. The stored row already keyed on that call id reads as this row's
|
||||
replay, so the row stays skipped in the same pass, exactly as a pinned call id skips today."""
|
||||
table = _wire_spend_logs_table(
|
||||
mock_prisma_client,
|
||||
seeded=[
|
||||
make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id="call-old", call_type="acompletion"),
|
||||
make_spend_log_row(request_id="call-pinned", litellm_call_id="call-pinned", call_type="acompletion"),
|
||||
],
|
||||
)
|
||||
|
||||
await _flush(
|
||||
mock_prisma_client,
|
||||
[make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id="call-pinned", call_type="acompletion")],
|
||||
)
|
||||
|
||||
assert sorted(table.rows) == ["call-pinned", "chatcmpl-static-1"]
|
||||
assert table.inserts == [["chatcmpl-static-1"]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_treats_a_replay_of_a_rekeyed_row_as_stored(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""A transport retry replays a whole batch, including rows an earlier attempt already re-keyed.
|
||||
Those are stored under their call id, so the read-back finds them there and nothing is re-keyed
|
||||
or warned about a second time."""
|
||||
table = _wire_spend_logs_table(mock_prisma_client)
|
||||
logs = [
|
||||
make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}", call_type="acompletion")
|
||||
for i in range(2)
|
||||
]
|
||||
await _flush(mock_prisma_client, logs)
|
||||
inserts_after_first_flush = len(table.inserts)
|
||||
caplog.clear()
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger=utils_mod.verbose_proxy_logger.name):
|
||||
await _flush(mock_prisma_client, logs)
|
||||
|
||||
assert sorted(table.rows) == ["call-1", "chatcmpl-static-1"]
|
||||
assert len(table.inserts) == inserts_after_first_flush + 1
|
||||
assert [record for record in caplog.records if "re-keyed" in record.getMessage()] == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_leaves_rows_skipped_when_the_read_back_fails_on_its_data(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
"""A read-back the database rejects (not a transport fault) must not requeue the whole flush
|
||||
behind it: the skipped rows stay skipped, which is what happened before, and the failure is
|
||||
logged."""
|
||||
table = _wire_spend_logs_table(mock_prisma_client)
|
||||
mock_prisma_client.db.query_raw = AsyncMock(side_effect=_data_error("invalid input syntax for type text[]"))
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
logs = [
|
||||
make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}", call_type="acompletion")
|
||||
for i in range(2)
|
||||
]
|
||||
|
||||
with caplog.at_level(logging.ERROR, logger=utils_mod.verbose_proxy_logger.name):
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=0,
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
logs_to_process=logs,
|
||||
)
|
||||
|
||||
assert list(table.rows) == ["chatcmpl-static-1"]
|
||||
assert len(table.inserts) == 1
|
||||
assert mock_prisma_client.spend_log_transactions == []
|
||||
assert any("could not read back" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_spend_logs_retries_the_flush_when_the_read_back_hits_a_transport_fault(
|
||||
mock_prisma_client: Any, make_spend_log_row: Any
|
||||
) -> None:
|
||||
"""A transport fault during the read-back is the same outage as one during the insert, so
|
||||
the flush retries (and finally requeues) instead of leaving the rows behind."""
|
||||
table = _wire_spend_logs_table(mock_prisma_client)
|
||||
mock_prisma_client.db.query_raw = AsyncMock(side_effect=httpx.ReadError("network blip"))
|
||||
proxy_logging = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
logs = [
|
||||
make_spend_log_row(request_id="chatcmpl-static-1", litellm_call_id=f"call-{i}", call_type="acompletion")
|
||||
for i in range(2)
|
||||
]
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=0,
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
logs_to_process=logs,
|
||||
)
|
||||
|
||||
assert mock_prisma_client.spend_log_transactions == logs
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue