From b067e836f8df15bb3dbefe345bf503b7d3811b9a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 18:58:11 -0700 Subject: [PATCH] fix(batches): claim the batch cost spend row in the database before charging The cost callback used to look for an existing `_batch_cost` row before charging a completed batch, which left a window where concurrent retrieves on any instance all charged the key, and it would honor a row any request had written under that id. The spend update writer now inserts the batch cost row itself with `create_many(skip_duplicates=True)` and only the retrieve whose insert lands charges the key, team, and user. An existing row only takes the charge when it is a successful `aretrieve_batch` row, so a client-chosen `x-litellm-call-id` on another endpoint cannot suppress billing. Batch cost rows no longer get their own immediate flush path `batch_cost_is_final` now treats the proxy's normalized `complete` status like `completed`, which the enterprise batch cost poller relies on when it decides whether a completed batch is safe to retire. Tests build that status with `model_copy` since the OpenAI `Batch` model rejects it The `test-quality-ok` markers sit on the `patch(` lines the gate keys on, and the logging tests no longer wrap the priced retrieve in `contextlib.suppress` --- litellm/batches/batch_utils.py | 5 +- litellm/proxy/db/db_spend_update_writer.py | 73 ++++++++-- .../proxy/hooks/proxy_track_cost_callback.py | 68 +++------- litellm/proxy/utils.py | 10 +- .../test_litellm/batches/test_batch_utils.py | 14 +- .../test_litellm_logging.py | 12 +- .../proxy/db/test_db_spend_update_writer.py | 128 +++++++++++++++++- .../hooks/test_proxy_track_cost_callback.py | 116 ++++------------ .../prisma_and_spend/test_spend_functions.py | 15 -- 9 files changed, 249 insertions(+), 192 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index eaac3bf0e9f..959c7498479 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -25,7 +25,8 @@ class BatchCostUsageResult: failed_requests: int -_TERMINAL_BATCH_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "expired"}) +_COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"}) +_TERMINAL_BATCH_STATUSES: Final = _COMPLETED_BATCH_STATUSES | frozenset({"failed", "cancelled", "expired"}) def batch_cost_is_final(batch: Batch) -> bool: @@ -39,7 +40,7 @@ def batch_cost_is_final(batch: Batch) -> bool: """ if batch.status not in _TERMINAL_BATCH_STATUSES: return False - if batch.status != "completed" or batch.output_file_id is not None: + if batch.status not in _COMPLETED_BATCH_STATUSES or batch.output_file_id is not None: return True request_counts: Final = batch.request_counts return request_counts is not None and request_counts.total > 0 and request_counts.completed == 0 diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index ff48b00dc70..3fad351224b 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -82,7 +82,10 @@ else: RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value}) -IMMEDIATE_FLUSH_CALL_TYPES: Final = RESPONSES_SESSION_CALL_TYPES | frozenset({CallTypes.aretrieve_batch.value}) + + +def _is_batch_cost_row(payload: SpendLogsPayload) -> bool: + return payload.get("call_type") == CallTypes.aretrieve_batch.value and payload.get("status") == "success" class _SpendBatch(Protocol): @@ -216,7 +219,12 @@ class DBSpendUpdateWriter: start_time: datetime | None, end_time: datetime | None, response_cost: float | None, - ) -> None: + ) -> bool: + """Record the request's spend, answering whether its cost still needs charging. + + False only for a batch retrieve whose cost row another retrieve already wrote, + so the caller leaves the key, team, and user counters alone (LIT-7048). + """ from litellm.proxy.proxy_server import ( disable_spend_logs, litellm_proxy_budget_name, @@ -233,7 +241,7 @@ class DBSpendUpdateWriter: team_id, ) if ProxyUpdateSpend.disable_spend_updates() is True: - return + return True if token is not None and isinstance(token, str) and token.startswith("sk-"): hashed_token = hash_token(token=token) else: @@ -264,10 +272,8 @@ class DBSpendUpdateWriter: payload["team_id"] = team_id if disable_spend_logs is False: - await self._insert_spend_log_to_db( - payload=payload, - prisma_client=prisma_client, - ) + if not await self._record_spend_log(payload=payload, prisma_client=prisma_client): + return False await self._enqueue_tool_usage_transaction( payload=payload, completion_response=completion_response, @@ -307,6 +313,7 @@ class DBSpendUpdateWriter: ) verbose_proxy_logger.debug("Runs spend update on all tables") + return True except Exception: spend_log_error( "Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue " @@ -319,7 +326,55 @@ class DBSpendUpdateWriter: org_id, end_user_id, ) - return + return True + + async def _record_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient | None") -> bool: + if prisma_client is None or not _is_batch_cost_row(payload): + await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) + return True + return await self._claim_batch_cost_spend_log(payload=payload, prisma_client=prisma_client) + + async def _claim_batch_cost_spend_log(self, payload: SpendLogsPayload, prisma_client: "PrismaClient") -> bool: + """Write the batch's cost row now, or learn that another retrieve already did. + + Every retrieve of one batch shares this row, so the insert that lands first owns + the charge and every later one finds the row and charges nothing (LIT-7048). Only + a row a successful retrieve wrote counts: a failed retrieve, or any request whose + client picked the batch id as its call id, cannot take the charge away. + """ + from litellm.repositories.table_repositories import SpendLogsRepository + + request_id: Final = payload["request_id"] + spend_logs: Final = SpendLogsRepository(prisma_client).table + try: + claimed: Final = await spend_logs.create_many( + data=[prisma_client.jsonify_object(payload)], # mutable-ok: prisma create_many takes a list + skip_duplicates=True, + ) + if claimed == 1: + return True + existing: Final = await spend_logs.find_unique( + where={"request_id": request_id} # mutable-ok: prisma where clause + ) + except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreachable DB queues the row like any other spend log + verbose_proxy_logger.warning( + "Could not claim spend row %s for a batch's cost, queueing it: %s", request_id, e + ) + await self._insert_spend_log_to_db(payload=payload, prisma_client=prisma_client) + return True + if ( + existing is not None + and existing.call_type == CallTypes.aretrieve_batch.value + and existing.status == "success" + ): + verbose_proxy_logger.debug("Cost tracking skipped: spend row %s already charged this batch", request_id) + return False + verbose_proxy_logger.warning( + "Spend row %s belongs to a %s request, so this batch's cost is charged without a row of its own", + request_id, + getattr(existing, "call_type", None), + ) + return True async def _enqueue_tool_usage_transaction( self, @@ -940,7 +995,7 @@ class DBSpendUpdateWriter: from litellm.proxy.utils import enqueue_spend_logs, request_spend_log_flush await enqueue_spend_logs(prisma_client, (payload,)) - if payload.get("call_type") in IMMEDIATE_FLUSH_CALL_TYPES: + if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES: request_spend_log_flush() else: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 95e61fdd98c..f0c889a8cb4 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -34,7 +34,6 @@ from litellm.proxy.spend_tracking.spend_log_error_logger import ( from litellm.proxy.spend_tracking.spend_tracking_utils import ( _sanitize_error_information_for_spend_logs, get_request_model_access_groups, - get_spend_logs_id, ) from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( @@ -46,9 +45,7 @@ from litellm.types.utils import ( from litellm.utils import get_end_user_id_for_cost_tracking if TYPE_CHECKING: - from prisma.types import LiteLLM_SpendLogsWhereUniqueInput - - from litellm.proxy.utils import PrismaClient, ProxyLogging + from litellm.proxy.utils import ProxyLogging _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset( { @@ -229,7 +226,6 @@ class _ProxyDBLogger(CustomLogger): ): from litellm.proxy.proxy_server import ( increment_spend_counters, - prisma_client, proxy_logging_obj, update_cache, ) @@ -257,15 +253,15 @@ class _ProxyDBLogger(CustomLogger): if ( isinstance(completion_response, LiteLLMBatch) and kwargs.get("call_type") == CallTypes.aretrieve_batch.value + and not batch_cost_is_final(completion_response) ): - batch_spend_log_id: Final = get_spend_logs_id( - CallTypes.aretrieve_batch.value, completion_response.model_dump(), kwargs + verbose_proxy_logger.debug( + "Cost tracking deferred for batch %s still in status %s", + completion_response.id, + completion_response.status, ) - if not await _batch_cost_is_trackable_now( - batch=completion_response, spend_log_id=batch_spend_log_id, prisma_client=prisma_client - ): - await _release_budget_reservation(budget_reservation=budget_reservation) - return + await _release_budget_reservation(budget_reservation=budget_reservation) + return user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None)) team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None)) org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None)) @@ -307,7 +303,7 @@ class _ProxyDBLogger(CustomLogger): call_type=call_type, ): ## UPDATE DATABASE - await _update_database_and_spend_counters( + charged: Final = await _update_database_and_spend_counters( proxy_logging_obj=proxy_logging_obj, increment_spend_counters=increment_spend_counters, user_api_key=user_api_key, @@ -324,6 +320,8 @@ class _ProxyDBLogger(CustomLogger): request_tags=tags, model_access_groups=model_access_groups, ) + if not charged: + return # update cache (fire-and-forget for backward compat: # cached object fields, soft budget alerts, etc.) @@ -509,42 +507,6 @@ def _write_spend_metadata_to_kwargs(kwargs: dict, metadata: dict) -> None: bucket[key] = value -async def _batch_cost_is_trackable_now( - batch: LiteLLMBatch, spend_log_id: str | None, prisma_client: "PrismaClient | None" -) -> bool: - """A batch is billed exactly once, from the first retrieve that sees it final. - - Every retrieve of one batch shares a single spend row (its id plus the batch cost - suffix), so a poll that lands before the output exists would write that row at $0 - and pin it there, and every retrieve after the first would add the cost to the - key, team, and user counters again. - """ - if not batch_cost_is_final(batch): - verbose_proxy_logger.debug("Cost tracking deferred for batch %s still in status %s", batch.id, batch.status) - return False - if prisma_client is None or spend_log_id is None: - return True - if not await _spend_log_already_recorded(prisma_client=prisma_client, request_id=spend_log_id): - return True - verbose_proxy_logger.debug( - "Cost tracking skipped for batch %s: spend row %s already recorded", batch.id, spend_log_id - ) - return False - - -async def _spend_log_already_recorded(prisma_client: "PrismaClient", request_id: str) -> bool: - from litellm.proxy.utils import spend_log_is_queued - - if await spend_log_is_queued(prisma_client, request_id): - return True - spend_log_row: Final[LiteLLM_SpendLogsWhereUniqueInput] = {"request_id": request_id} - try: - return await prisma_client.db.litellm_spendlogs.find_unique(where=spend_log_row) is not None - except Exception as e: # noqa: BLE001 # prisma raises its own hierarchy; an unreadable DB must not drop the batch's only spend row - verbose_proxy_logger.warning("Could not check for an existing spend row %s, tracking anyway: %s", request_id, e) - return False - - def _is_unbilled_interaction_response(completion_response: object) -> bool: from litellm.interactions.background_cost_polling import missing_usage_is_expected from litellm.types.interactions import InteractionsAPIResponse @@ -636,9 +598,9 @@ async def _update_database_and_spend_counters( budget_reservation: dict | None, request_tags: list[str] | None = None, model_access_groups: Sequence[str] | None = None, -) -> None: +) -> bool: try: - await proxy_logging_obj.db_spend_update_writer.update_database( + charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key, response_cost=response_cost, user_id=user_id, @@ -663,6 +625,9 @@ async def _update_database_and_spend_counters( "Failed to invalidate budget reservation counters after release failed" ) raise + if not charged: + await _release_budget_reservation(budget_reservation=budget_reservation) + return False try: await increment_spend_counters( @@ -688,6 +653,7 @@ async def _update_database_and_spend_counters( finally: budget_reservation["finalized"] = True raise + return True async def _release_budget_reservation(budget_reservation: dict | None) -> None: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f64b51bc6c6..accf7b720fb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -6251,9 +6251,7 @@ def request_spend_log_flush() -> None: The Responses API hands the client an id it can chain from straight away, and that lookup reads the DB, so the row cannot sit in this worker's queue for a poll interval. - A batch's cost row is what every other worker checks before charging the same batch - again, so it cannot wait either. Repeated requests coalesce into the monitor's next - pass, so the batching holds. + Repeated requests coalesce into the monitor's next pass, so the batching holds. """ PrismaClient.spend_log_flush_requested.set() @@ -6268,12 +6266,6 @@ async def _wait_for_spend_log_flush_request(interval: float) -> bool: return True -async def spend_log_is_queued(prisma_client: PrismaClient, request_id: str) -> bool: - """Whether a spend log with ``request_id`` is still waiting for the next flush.""" - async with prisma_client._spend_log_transactions_lock: - return any(row.get("request_id") == request_id for row in prisma_client.spend_log_transactions) - - async def dequeue_spend_logs(prisma_client: PrismaClient, limit: int) -> list[dict[str, object]]: """Take up to ``limit`` of the oldest queued spend logs off the queue. diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 8d4f68164b4..976a96f2db1 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1735,10 +1735,10 @@ def _retrieved_batch( endpoint="/v1/chat/completions", input_file_id="file-in", object="batch", - status=status, + status="validating", output_file_id=output_file_id, request_counts=counts, - ) + ).model_copy(update={"status": status}) class TestBatchCostIsFinal: @@ -1750,8 +1750,9 @@ class TestBatchCostIsFinal: def test_in_flight_batch_is_not_final(self, status): assert bu.batch_cost_is_final(_retrieved_batch(status)) is False - def test_completed_with_output_is_final(self): - assert bu.batch_cost_is_final(_retrieved_batch("completed", output_file_id="file-out")) is True + @pytest.mark.parametrize("status", ["completed", "complete"]) + def test_completed_with_output_is_final(self, status): + assert bu.batch_cost_is_final(_retrieved_batch(status, output_file_id="file-out")) is True def test_completed_without_output_and_unknown_counts_is_not_final(self): assert bu.batch_cost_is_final(_retrieved_batch("completed")) is False @@ -1764,9 +1765,10 @@ class TestBatchCostIsFinal: counts = BatchRequestCounts(total=2, completed=2, failed=0) assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False - def test_completed_without_output_and_every_line_failed_is_final(self): + @pytest.mark.parametrize("status", ["completed", "complete"]) + def test_completed_without_output_and_every_line_failed_is_final(self, status): counts = BatchRequestCounts(total=2, completed=0, failed=2) - assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is True + assert bu.batch_cost_is_final(_retrieved_batch(status, counts=counts)) is True @pytest.mark.parametrize("status", ["failed", "expired", "cancelled"]) def test_other_terminal_statuses_are_final(self, status): diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 4583429bd11..174efaa4679 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -665,14 +665,14 @@ class TestRetrieveBatchPricesOnlyFinalBatches: endpoint="/v1/chat/completions", input_file_id="file-in", object="batch", - status=status, + status="validating", output_file_id=output_file_id, - ) + ).model_copy(update={"status": status}) @pytest.mark.asyncio @pytest.mark.parametrize( ("status", "output_file_id"), - [("validating", None), ("in_progress", None), ("finalizing", None), ("completed", None)], + [("validating", None), ("in_progress", None), ("finalizing", None), ("completed", None), ("complete", None)], ) async def test_non_final_batch_is_not_priced(self, monkeypatch, status, output_file_id) -> None: from litellm.litellm_core_utils import litellm_logging as logging_module @@ -681,8 +681,7 @@ class TestRetrieveBatchPricesOnlyFinalBatches: monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch) batch = self._batch(status, output_file_id) - with contextlib.suppress(Exception): - await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) + await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) handle_completed_batch.assert_not_awaited() assert "response_cost" not in batch._hidden_params @@ -705,8 +704,7 @@ class TestRetrieveBatchPricesOnlyFinalBatches: monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch) batch = self._batch("completed", "file-out") - with contextlib.suppress(Exception): - await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) + await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) handle_completed_batch.assert_awaited_once() assert batch._hidden_params["response_cost"] == 8e-06 diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index e1b2d151c6d..500a0e7bb06 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -7,6 +7,7 @@ import re from collections.abc import Callable from contextlib import asynccontextmanager from datetime import datetime, timezone +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -2934,16 +2935,14 @@ async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkey @pytest.mark.asyncio @pytest.mark.parametrize( "call_type, expects_flush", - [("aresponses", True), ("responses", True), ("aretrieve_batch", True), ("acompletion", False)], + [("aresponses", True), ("responses", True), ("acompletion", False)], ) async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_workers_read_back( call_type: str, expects_flush: bool ): """ A `previous_response_id` chained straight off the previous turn reads the DB, so a - Responses row cannot sit in this worker's queue until the monitor's next poll. A - batch's cost row is what another worker checks before charging the same batch again - (LIT-7048), so it cannot wait either. + Responses row cannot sit in this worker's queue until the monitor's next poll. """ from litellm.proxy.utils import PrismaClient @@ -2961,6 +2960,127 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_rows_other_worker PrismaClient.spend_log_flush_requested.clear() +def _batch_cost_payload() -> dict: + return { + **_minimal_spend_payload(), + "request_id": "batch_abc_batch_cost", + "call_type": "aretrieve_batch", + "status": "success", + } + + +def _spend_logs_prisma(inserted: int, existing: object) -> MagicMock: + prisma = _tool_usage_prisma() + prisma.jsonify_object = lambda data: dict(data) + prisma.db.litellm_spendlogs.create_many = AsyncMock(return_value=inserted) + prisma.db.litellm_spendlogs.find_unique = AsyncMock(return_value=existing) + return prisma + + +async def _update_database_with(db_writer: DBSpendUpdateWriter, prisma: MagicMock, payload: dict) -> bool: + with ( + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.disable_spend_logs", False + ), + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.prisma_client", prisma + ), + patch( # test-quality-ok: update_database reads this proxy_server global at call time, no seam + "litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget" + ), + patch( # test-quality-ok: update_database imports the payload builder inside its body, no seam + "litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", + return_value=payload, + ), + ): + charged = await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id=None, + org_id=None, + kwargs={"model": "gpt-5.6-luna", "call_type": "aretrieve_batch"}, + completion_response=None, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc), + response_cost=0.25, + ) + await asyncio.sleep(0) + return charged + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("inserted", "existing", "charged"), + [ + (1, None, True), + (0, SimpleNamespace(call_type="aretrieve_batch", status="success"), False), + (0, SimpleNamespace(call_type="aretrieve_batch", status="failure"), True), + (0, SimpleNamespace(call_type="aembedding", status="success"), True), + (0, None, True), + ], + ids=[ + "first_retrieve_owns_the_row", + "another_retrieve_already_charged", + "failed_retrieve_holds_the_row", + "client_chosen_call_id_holds_the_row", + "row_gone_between_insert_and_lookup", + ], +) +async def test_update_database_charges_a_batch_only_from_the_retrieve_that_wrote_its_row( + inserted: int, existing: object, charged: bool +): + """ + Every retrieve of one batch shares one spend row, so the insert that lands first is + the charge and every later retrieve must leave the counters alone (LIT-7048). A row + written by anything but a successful retrieve, say a request whose client picked the + batch id as its call id, must not be able to take the charge away. + """ + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(inserted, existing) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is charged + + claimed_rows = prisma.db.litellm_spendlogs.create_many.await_args.kwargs + assert claimed_rows["skip_duplicates"] is True + assert [(row["request_id"], row["spend"]) for row in claimed_rows["data"]] == [("batch_abc_batch_cost", 0.25)] + assert prisma.spend_log_transactions == [] + assert db_writer._batch_database_updates.await_count == (1 if charged else 0) + + +@pytest.mark.asyncio +async def test_update_database_queues_a_batch_cost_row_it_could_not_claim(): + """An unreachable DB must not drop the batch's only spend row, nor its charge.""" + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(0, None) + prisma.db.litellm_spendlogs.create_many = AsyncMock(side_effect=RuntimeError("db unreachable")) + + assert await _update_database_with(db_writer, prisma, _batch_cost_payload()) is True + + assert [row["request_id"] for row in prisma.spend_log_transactions] == ["batch_abc_batch_cost"] + assert db_writer._batch_database_updates.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload", + [{**_batch_cost_payload(), "call_type": "acompletion"}, {**_batch_cost_payload(), "status": "failure"}], + ids=["not_a_batch_retrieve", "failed_batch_retrieve"], +) +async def test_update_database_queues_every_other_spend_row_for_the_next_flush(payload: dict): + db_writer = DBSpendUpdateWriter() + db_writer._batch_database_updates = AsyncMock() + prisma = _spend_logs_prisma(1, None) + + assert await _update_database_with(db_writer, prisma, payload) is True + + prisma.db.litellm_spendlogs.create_many.assert_not_called() + assert prisma.spend_log_transactions == [payload] + assert db_writer._batch_database_updates.await_count == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize( "injected_deployment, attributed", diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index b7037e8d621..2965f8b4006 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,3 +1,4 @@ +import asyncio from datetime import datetime from unittest.mock import AsyncMock, MagicMock, patch @@ -6,7 +7,6 @@ import pytest from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.proxy_track_cost_callback import ( - _batch_cost_is_trackable_now, _get_budget_reservation_from_metadata, _ProxyDBLogger, _should_track_cost_callback, @@ -783,110 +783,46 @@ def _retrieved_batch(status: str, output_file_id: str | None): ) -def _prisma_client_with(queued_request_ids: tuple[str, ...], stored_row: object) -> MagicMock: - import asyncio - - prisma_client = MagicMock() - prisma_client._spend_log_transactions_lock = asyncio.Lock() - prisma_client.spend_log_transactions = [{"request_id": request_id} for request_id in queued_request_ids] - prisma_client.db.litellm_spendlogs.find_unique = AsyncMock(return_value=stored_row) - return prisma_client - - @pytest.mark.asyncio @pytest.mark.parametrize( - ("status", "output_file_id", "spend_log_id", "prisma_client", "trackable"), + ("call_type", "status", "output_file_id", "row_claimed", "spend_written", "charged"), [ - ("in_progress", None, "batch_abc_batch_cost", None, False), - ("in_progress", None, "batch_abc_batch_cost", _prisma_client_with((), None), False), - ("completed", None, "batch_abc_batch_cost", _prisma_client_with((), None), False), - ("completed", "file-out", "batch_abc_batch_cost", None, True), - ("completed", "file-out", None, _prisma_client_with((), None), True), - ("completed", "file-out", "batch_abc_batch_cost", _prisma_client_with(("batch_abc_batch_cost",), None), False), - ( - "completed", - "file-out", - "batch_abc_batch_cost", - _prisma_client_with((), {"request_id": "batch_abc_batch_cost"}), - False, - ), - ("completed", "file-out", "batch_abc_batch_cost", _prisma_client_with((), None), True), - ("failed", None, "batch_abc_batch_cost", _prisma_client_with((), None), True), + ("aretrieve_batch", "in_progress", None, True, False, False), + ("aretrieve_batch", "completed", None, True, False, False), + ("aretrieve_batch", "completed", "file-out", False, True, False), + ("aretrieve_batch", "completed", "file-out", True, True, True), + ("aretrieve_batch", "failed", None, True, True, True), + ("acreate_batch", "validating", None, True, True, True), ], ids=[ - "in_progress_without_db", - "in_progress_never_consults_db", - "completed_without_output_yet", - "final_without_db", - "final_without_spend_log_id", - "final_row_queued_for_flush", - "final_row_already_stored", - "final_first_sighting", - "failed_first_sighting", + "retrieve_before_final", + "retrieve_completed_without_output_yet", + "retrieve_after_another_retrieve_charged", + "retrieve_first_final", + "retrieve_failed_batch", + "create_before_final", ], ) -async def test_batch_cost_is_trackable_now(status, output_file_id, spend_log_id, prisma_client, trackable): - """ - A batch is billed from the first retrieve that sees it final and never again: - a poll before that wrote the shared spend row at $0 and pinned it there, and - every completed retrieve after the first charged the key again (LIT-7048). - """ - assert ( - await _batch_cost_is_trackable_now( - batch=_retrieved_batch(status, output_file_id), spend_log_id=spend_log_id, prisma_client=prisma_client - ) - is trackable - ) - - -@pytest.mark.asyncio -async def test_batch_cost_is_trackable_now_when_the_spend_row_lookup_fails(): - """An unreadable spend log table must not drop the batch's only spend row.""" - prisma_client = _prisma_client_with((), None) - prisma_client.db.litellm_spendlogs.find_unique = AsyncMock(side_effect=RuntimeError("db unreachable")) - - assert ( - await _batch_cost_is_trackable_now( - batch=_retrieved_batch("completed", "file-out"), - spend_log_id="batch_abc_batch_cost", - prisma_client=prisma_client, - ) - is True - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("call_type", "status", "output_file_id", "stored_row", "charged"), - [ - ("aretrieve_batch", "in_progress", None, None, False), - ("aretrieve_batch", "completed", "file-out", {"request_id": "batch_abc_batch_cost"}, False), - ("aretrieve_batch", "completed", "file-out", None, True), - ("acreate_batch", "validating", None, None, True), - ], - ids=["retrieve_before_final", "retrieve_already_recorded", "retrieve_first_final", "create_before_final"], -) -async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # test-quality-ok: whether the spend writer runs and whether the poll's reservation is handed back is the whole observable contract of the gate - call_type, status, output_file_id, stored_row, charged +async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # test-quality-ok: whether the spend writer runs, whether the counters move, and whether the poll's reservation is handed back is the whole observable contract of the gate + call_type, status, output_file_id, row_claimed, spend_written, charged ): """ - Only retrieves are gated, since creating a batch is its own billable request. - A retrieve that writes nothing hands its budget reservation back instead. + A poll before the batch is final used to pin its shared spend row at $0, and every + completed retrieve after the first charged the key again (LIT-7048). Only retrieves + are gated, since creating a batch is its own billable request, and a retrieve that + charges nothing hands its budget reservation back instead. """ logger = _ProxyDBLogger() budget_reservation = None if charged else {"reserved_cost": 0.5, "entries": []} kwargs = _batch_retrieve_kwargs(call_type, reservation=budget_reservation) with ( - patch( # test-quality-ok: prisma_client is a proxy_server global the callback reads lazily, no seam - "litellm.proxy.proxy_server.prisma_client", _prisma_client_with((), stored_row) - ), patch( # test-quality-ok: increment_spend_counters is a proxy_server global the callback reads lazily, no seam "litellm.proxy.proxy_server.increment_spend_counters", new_callable=AsyncMock - ), + ) as mock_increment_spend_counters, patch( # test-quality-ok: update_cache is a proxy_server global the callback reads lazily, no seam "litellm.proxy.proxy_server.update_cache", new_callable=AsyncMock - ), + ) as mock_update_cache, patch( # test-quality-ok: callback imports proxy_logging_obj off proxy_server in its body, no seam "litellm.proxy.proxy_server.proxy_logging_obj" ) as mock_proxy_logging, @@ -895,7 +831,7 @@ async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # ) as mock_release_budget_reservation, ): mock_proxy_logging.failed_tracking_alert = AsyncMock() - mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock() + mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock(return_value=row_claimed) mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock() await logger._PROXY_track_cost_callback( @@ -904,13 +840,15 @@ async def test_track_cost_callback_charges_a_batch_once_and_only_when_final( # start_time=datetime.now(), end_time=datetime.now(), ) + await asyncio.sleep(0) mock_proxy_logging.failed_tracking_alert.assert_not_called() + assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (1 if spend_written else 0) + assert mock_increment_spend_counters.await_count == (1 if charged else 0) + assert mock_update_cache.await_count == (1 if charged else 0) if charged: - mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once() mock_release_budget_reservation.assert_not_awaited() else: - mock_proxy_logging.db_spend_update_writer.update_database.assert_not_called() mock_release_budget_reservation.assert_awaited_once_with(budget_reservation=budget_reservation) diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index fc97d760226..a1eb88a7834 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -6,7 +6,6 @@ Symbols pinned here: - ``update_spend_logs_job`` - ``_monitor_spend_logs_queue`` - ``_raise_failed_update_spend_exception`` - - ``spend_log_is_queued`` """ from __future__ import annotations @@ -23,7 +22,6 @@ from litellm.proxy.utils import ( _monitor_spend_logs_queue, _raise_failed_update_spend_exception, drain_spend_logs_queue, - spend_log_is_queued, update_daily_tag_spend, update_spend, update_spend_logs_job, @@ -631,16 +629,3 @@ def test_raise_failed_update_spend_exception_raises_original_error() -> None: with pytest.raises(ValueError, match="specific"): asyncio.run(_runner()) - - -@pytest.mark.asyncio -async def test_spend_log_is_queued_matches_only_rows_awaiting_flush( - mock_prisma_client: Any, make_spend_log_row: Any -) -> None: - mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="batch_abc_batch_cost")] - - assert await spend_log_is_queued(mock_prisma_client, "batch_abc_batch_cost") is True - assert await spend_log_is_queued(mock_prisma_client, "batch_abc") is False - - mock_prisma_client.spend_log_transactions = [] - assert await spend_log_is_queued(mock_prisma_client, "batch_abc_batch_cost") is False