diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 4e98823ba77..be5b307b5ae 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,7 +12,6 @@ import json import os import random import time -import traceback from collections.abc import AsyncGenerator, Callable, Coroutine, Iterable, Mapping, Sequence from contextlib import asynccontextmanager from contextvars import ContextVar @@ -90,6 +89,7 @@ from litellm.proxy.spend_tracking.savings import ( ) from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.repositories.prisma_protocols import BatchTable +from litellm.responses.streaming_iterator import spawn_logging_task from litellm.types.utils import CallTypes if TYPE_CHECKING: @@ -313,6 +313,38 @@ def _spend_tables_left_to_send( return cast(DBSpendUpdateTransactions, remaining) if any(remaining.values()) else None +def _report_spend_update_step_failure(step_name: str, *, payload: Mapping[str, object], exc: Exception) -> None: + spend_log_error( + "Spend tracking - %s failed, this request is missing from that spend rollup. request_id=%s api_key=%s", + step_name, + payload.get("request_id"), + payload.get("api_key"), + exc=exc, + ) + + +_EMPTY_SPEND_LOG_METADATA: Final = cast(SpendLogsMetadata, MappingProxyType({})) # cast-ok: empty read-only metadata + + +def _parse_spend_log_metadata(payload: Mapping[str, object]) -> SpendLogsMetadata: + """Malformed metadata costs the request its metadata-derived fields, never its daily rollup row.""" + raw: Final = payload.get("metadata") + if not isinstance(raw, str): + return _EMPTY_SPEND_LOG_METADATA + try: + parsed: Final[SpendLogsMetadata] = json.loads(raw) + except json.JSONDecodeError as e: + spend_log_error( + "Spend tracking - unreadable spend log metadata, rollups for this request are written without " + "metadata-derived fields. request_id=%s api_key=%s", + payload.get("request_id"), + payload.get("api_key"), + exc=e, + ) + return _EMPTY_SPEND_LOG_METADATA + return parsed + + def _is_billed_retry_attempt(metadata: SpendLogsMetadata) -> bool: """An earlier billed retry attempt's own spend row, not a request the caller sent.""" return BILLED_RETRY_ATTEMPT_METADATA_KEY in metadata and metadata[BILLED_RETRY_ATTEMPT_METADATA_KEY] is True @@ -647,8 +679,8 @@ class DBSpendUpdateWriter: "disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur." ) - # Single task replaces 11 create_task() calls - asyncio.create_task( + spawn_logging_task( + asyncio.get_running_loop(), self._batch_database_updates( response_cost=response_cost, user_id=user_id, @@ -661,7 +693,8 @@ class DBSpendUpdateWriter: litellm_proxy_budget_name=litellm_proxy_budget_name, payload=payload, request_model_access_groups=get_request_model_access_groups(kwargs), - ) + ), + task_name="spend rollups", ) self._enqueue_tool_registry_upsert( @@ -1086,91 +1119,52 @@ class DBSpendUpdateWriter: project_id: str | None = None, ): """ - Runs all 13 spend-update helpers sequentially inside a single asyncio task. - - Each helper is wrapped in try/except so one failure doesn't prevent the others. - - The deepcopy runs here, off the awaited request path, so the daily spend - helpers get a payload isolated from the spend-log queue entry and the caller. + Runs every spend-update helper for one request. Each helper goes through + _run_spend_update_step so one failure never skips the sibling writes. The deepcopy + isolates the daily spend helpers' payload from the spend-log queue entry and the caller. """ payload_copy: Final = copy.deepcopy(payload) request_tags: Final = payload_copy.get("request_tags") - try: - await self._update_user_db( + await self._run_spend_update_step( + "_update_user_db", + self._update_user_db( response_cost=response_cost, user_id=user_id, prisma_client=prisma_client, litellm_proxy_budget_name=litellm_proxy_budget_name, end_user_id=end_user_id, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_user_db failed: %s", - traceback.format_exc(), - ) - - try: - await self._update_key_db( - response_cost=response_cost, - hashed_token=hashed_token, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_key_db failed: %s", - traceback.format_exc(), - ) - - try: - await self._update_team_db( - response_cost=response_cost, - team_id=team_id, - user_id=user_id, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_team_db failed: %s", - traceback.format_exc(), - ) - - try: - await self._update_org_db( - response_cost=response_cost, - org_id=org_id, - user_id=user_id, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_org_db failed: %s", - traceback.format_exc(), - ) - - try: - await self._update_project_db( - response_cost=response_cost, - project_id=project_id, - prisma_client=prisma_client, - ) - except Exception: # noqa: BLE001 # a project enqueue failure must not skip the sibling spend writes - verbose_proxy_logger.debug( - "_batch_database_updates: _update_project_db failed: %s", - traceback.format_exc(), - ) - - try: - await self._update_tag_db( - response_cost=response_cost, - request_tags=request_tags, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_tag_db failed: %s", - traceback.format_exc(), - ) - + ), + payload=payload_copy, + ) + await self._run_spend_update_step( + "_update_key_db", + self._update_key_db(response_cost=response_cost, hashed_token=hashed_token, prisma_client=prisma_client), + payload=payload_copy, + ) + await self._run_spend_update_step( + "_update_team_db", + self._update_team_db( + response_cost=response_cost, team_id=team_id, user_id=user_id, prisma_client=prisma_client + ), + payload=payload_copy, + ) + await self._run_spend_update_step( + "_update_org_db", + self._update_org_db( + response_cost=response_cost, org_id=org_id, user_id=user_id, prisma_client=prisma_client + ), + payload=payload_copy, + ) + await self._run_spend_update_step( + "_update_project_db", + self._update_project_db(response_cost=response_cost, project_id=project_id, prisma_client=prisma_client), + payload=payload_copy, + ) + await self._run_spend_update_step( + "_update_tag_db", + self._update_tag_db(response_cost=response_cost, request_tags=request_tags, prisma_client=prisma_client), + payload=payload_copy, + ) await self._update_model_access_group_db( response_cost=response_cost, request_model_access_groups=request_model_access_groups, @@ -1178,96 +1172,64 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, router=get_llm_router(), ) - - _agent_id_for_spend: Final = payload_copy.get("agent_id") - try: - await self._update_agent_db( - response_cost=response_cost, - agent_id=_agent_id_for_spend, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_agent_db failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_user_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_user_transaction failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_end_user_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_end_user_transaction failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_agent_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_agent_transaction failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_team_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_team_transaction failed: %s", - traceback.format_exc(), - ) + await self._run_spend_update_step( + "_update_agent_db", + self._update_agent_db( + response_cost=response_cost, agent_id=payload_copy.get("agent_id"), prisma_client=prisma_client + ), + payload=payload_copy, + ) + await self._run_spend_update_step( + "add_spend_log_transaction_to_daily_user_transaction", + self.add_spend_log_transaction_to_daily_user_transaction(payload=payload_copy, prisma_client=prisma_client), + payload=payload_copy, + ) + await self._run_spend_update_step( + "add_spend_log_transaction_to_daily_end_user_transaction", + self.add_spend_log_transaction_to_daily_end_user_transaction( + payload=payload_copy, prisma_client=prisma_client + ), + payload=payload_copy, + ) + await self._run_spend_update_step( + "add_spend_log_transaction_to_daily_agent_transaction", + self.add_spend_log_transaction_to_daily_agent_transaction( + payload=payload_copy, prisma_client=prisma_client + ), + payload=payload_copy, + ) + await self._run_spend_update_step( + "add_spend_log_transaction_to_daily_team_transaction", + self.add_spend_log_transaction_to_daily_team_transaction(payload=payload_copy, prisma_client=prisma_client), + payload=payload_copy, + ) try: self._record_request_error(payload=payload_copy, prisma_client=prisma_client) - except Exception: # noqa: BLE001 # the rollup must never skip the sibling spend writes - verbose_proxy_logger.debug( - "_batch_database_updates: _record_request_error failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_org_transaction( - payload=payload_copy, - org_id=org_id, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_org_transaction failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_tag_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_tag_transaction failed: %s", - traceback.format_exc(), - ) - + except Exception as e: # noqa: BLE001 # the rollup must never skip the sibling spend writes + _report_spend_update_step_failure("_record_request_error", payload=payload_copy, exc=e) + await self._run_spend_update_step( + "add_spend_log_transaction_to_daily_org_transaction", + self.add_spend_log_transaction_to_daily_org_transaction( + payload=payload_copy, org_id=org_id, prisma_client=prisma_client + ), + payload=payload_copy, + ) + await self._run_spend_update_step( + "add_spend_log_transaction_to_daily_tag_transaction", + self.add_spend_log_transaction_to_daily_tag_transaction(payload=payload_copy, prisma_client=prisma_client), + payload=payload_copy, + ) await self._enqueue_model_usage_transaction(payload=payload_copy, prisma_client=prisma_client) + @staticmethod + async def _run_spend_update_step( + step_name: str, step: Coroutine[object, object, object], *, payload: SpendLogsPayload + ) -> None: + try: + await step + except Exception as e: # noqa: BLE001 # one helper failing must not skip the sibling spend writes + _report_spend_update_step_failure(step_name, payload=payload, exc=e) + async def _update_key_db( self, response_cost: float | None, @@ -2810,7 +2772,7 @@ class DBSpendUpdateWriter: date: Final = start_time.isoformat() if isinstance(start_time, datetime) else str(start_time or "") if not date: return - metadata: Final[SpendLogsMetadata] = json.loads(payload["metadata"]) + metadata: Final = _parse_spend_log_metadata(payload) internal_origin: Final = bool(metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY)) # pyright: ignore[reportUnknownMemberType] # TypedDict.get overloads carry Any defaults request_error_accumulator.record( payload=payload, @@ -2861,7 +2823,7 @@ class DBSpendUpdateWriter: # spend, token and per-entity columns written here stay either way. request_status: Final = prisma_client.get_request_status(payload) verbose_proxy_logger.debug("Logged request status: %s", request_status) - _metadata: Final[SpendLogsMetadata] = json.loads(payload["metadata"]) + _metadata: Final = _parse_spend_log_metadata(payload) usage_obj: Final = _metadata.get("usage_object", {}) or {} if isinstance(payload["startTime"], datetime): start_time: Final = payload["startTime"].isoformat() diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index 7437042e4bd..9b2a96bbc58 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -36,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.responses.streaming_iterator import drain_logging_tasks from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation from tests.unit.proxy.db.fake_prisma_engine import engine_call @@ -2158,46 +2159,122 @@ async def test_commit_spend_updates_to_db_increments_key_total_spend_alongside_s @pytest.mark.asyncio -async def test_update_database_creates_single_task(): +async def test_shutdown_drain_finishes_the_rollups_update_database_started(): """ - Test that update_database() fires exactly 1 asyncio.create_task() call - (the batched task) instead of the previous 11. + Regression for LIT-6027. The key, team and daily rollups ran in an unreferenced asyncio task, + so the shutdown sequence (flush, drain_logging_tasks, flush) never waited for them and spend + logs landed without matching daily rows. Once update_database returns, drain_logging_tasks + must leave every rollup for that request in its queue. """ - db_writer = DBSpendUpdateWriter() - - # Mock all helpers so nothing real runs + db_writer: Final = DBSpendUpdateWriter() db_writer._insert_spend_log_to_db = AsyncMock() - db_writer._batch_database_updates = AsyncMock() - + prisma: Final = MagicMock() + prisma.get_request_status = MagicMock(return_value="success") + fake_payload: Final = { + "request_id": "req-lit-6027", + "api_key": "hashed-key", + "user": "test-user", + "team_id": "test-team", + "agent_id": None, + "startTime": "2024-01-01T00:00:00", + "endTime": "2024-01-01T00:00:01", + "model": "gpt-4o-mini", + "custom_llm_provider": "openai", + "model_group": "gpt-4o-mini", + "prompt_tokens": 10, + "completion_tokens": 5, + "spend": 0.0, + "metadata": '{"usage_object": {}}', + } with ( patch("litellm.proxy.proxy_server.disable_spend_logs", False), - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.prisma_client", prisma), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), - patch("litellm.proxy.db.db_spend_update_writer.asyncio.create_task") as mock_create_task, + patch("litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", return_value=fake_payload), ): await db_writer.update_database( token="test-token", user_id="test-user", - end_user_id="test-end-user", + end_user_id=None, + team_id="test-team", + org_id=None, + kwargs={"model": "gpt-4o-mini", "custom_llm_provider": "openai"}, + completion_response=MagicMock(), start_time=datetime.now(), end_time=datetime.now(), - team_id="test-team", - org_id="test-org", - completion_response=MagicMock(), response_cost=0.1, - kwargs={"model": "gpt-4", "custom_llm_provider": "openai"}, ) + await drain_logging_tasks() - # Exactly 1 create_task call (the batch), not 11 - assert mock_create_task.call_count == 1 + spend: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert spend["key_list_transactions"] == {"test-token": 0.1} + assert spend["team_list_transactions"] == {"test-team": 0.1} + daily_user: Final = ( + await db_writer.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + daily_team: Final = ( + await db_writer.daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + assert daily_user is not None and daily_team is not None + (user_row,) = daily_user.values() + (team_row,) = daily_team.values() + assert (user_row["api_requests"], user_row["spend"]) == (1, 0.1) + assert (team_row["api_requests"], team_row["spend"]) == (1, 0.1) @pytest.mark.asyncio -async def test_batch_database_updates_isolation_on_failure(): +async def test_shutdown_drain_waits_for_an_in_flight_rollup_enqueue(): + """ + Regression for LIT-6027. A rollup enqueue that is still in flight when the shutdown drain + starts (here: blocked on a full spend queue) must be finished by the time drain_logging_tasks + returns, so the second shutdown flush commits it. + """ + enqueue_started: Final = asyncio.Event() + release_enqueue: Final = asyncio.Event() + + class _BlockedSpendQueue(SpendUpdateQueue): + async def add_update(self, update: SpendUpdateQueueItem): + enqueue_started.set() + await release_enqueue.wait() + await super().add_update(update) + + db_writer: Final = DBSpendUpdateWriter() + db_writer.spend_update_queue = _BlockedSpendQueue() + db_writer._insert_spend_log_to_db = AsyncMock() + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", True), + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + ): + await db_writer.update_database( + token="t1", + user_id=None, + end_user_id=None, + team_id="team-1", + org_id=None, + kwargs={"model": "gpt-4o-mini", "custom_llm_provider": "openai"}, + completion_response=MagicMock(), + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.1, + ) + await enqueue_started.wait() + drain: Final = asyncio.create_task(drain_logging_tasks()) + release_enqueue.set() + await drain + spend: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + + assert spend["key_list_transactions"] == {"t1": 0.1} + assert spend["team_list_transactions"] == {"team-1": 0.1} + + +@pytest.mark.asyncio +async def test_batch_database_updates_isolation_on_failure(caplog: pytest.LogCaptureFixture): """ Test that if one helper inside _batch_database_updates raises, - all other helpers still execute. + all other helpers still execute and the failure is reported at ERROR with the request id. """ db_writer = DBSpendUpdateWriter() @@ -2217,6 +2294,7 @@ async def test_batch_database_updates_isolation_on_failure(): db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock() + caplog.set_level(logging.ERROR, logger=verbose_proxy_logger.name) await db_writer._batch_database_updates( response_cost=0.1, user_id="u1", @@ -2226,7 +2304,7 @@ async def test_batch_database_updates_isolation_on_failure(): end_user_id="eu1", prisma_client=MagicMock(), litellm_proxy_budget_name="budget", - payload={"key": "value"}, + payload={"request_id": "req-iso", "key": "value"}, ) # _update_key_db raised, but all others should still have been called @@ -2242,6 +2320,48 @@ async def test_batch_database_updates_isolation_on_failure(): db_writer.add_spend_log_transaction_to_daily_team_transaction.assert_awaited_once() db_writer.add_spend_log_transaction_to_daily_org_transaction.assert_awaited_once() db_writer.add_spend_log_transaction_to_daily_tag_transaction.assert_awaited_once() + assert any( + "_update_key_db" in record.getMessage() and "req-iso" in record.getMessage() + for record in caplog.records + if record.levelno >= logging.ERROR + ) + + +@pytest.mark.asyncio +async def test_malformed_metadata_keeps_the_daily_rollup_row_and_is_reported(caplog: pytest.LogCaptureFixture): + """ + Regression for LIT-6027. A spend log whose metadata is not JSON used to vanish from every + daily rollup while the spend log itself persisted, with nothing above debug in the logs. + """ + writer: Final = DBSpendUpdateWriter() + prisma: Final = MagicMock() + prisma.get_request_status = MagicMock(return_value="success") + payload: Final = { + "request_id": "req-bad-metadata", + "user": "test-user", + "startTime": "2024-01-01T00:00:00", + "endTime": "2024-01-01T00:00:01", + "api_key": "test-key", + "model": "gpt-4o-mini", + "custom_llm_provider": "openai", + "model_group": "gpt-4o-mini", + "prompt_tokens": 100, + "completion_tokens": 50, + "spend": 0.05, + "metadata": "{not-json", + } + caplog.set_level(logging.ERROR, logger=verbose_proxy_logger.name) + + await writer.add_spend_log_transaction_to_daily_user_transaction( # pyright: ignore[reportUnknownMemberType] # prisma_client param is untyped on the writer + payload=payload, prisma_client=prisma + ) + + daily: Final = await writer.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + (row,) = daily.values() + assert (row["api_requests"], row["spend"], row["prompt_tokens"]) == (1, 0.05, 100) + assert any( + "req-bad-metadata" in record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR + ) @pytest.mark.asyncio @@ -2775,12 +2895,12 @@ async def test_update_database_does_not_deepcopy_on_request_path(): """ Regression for LIT-4088: copy.deepcopy must not run while the caller awaits update_database(). The deepcopy used to isolate the daily-spend helpers is - relocated into the _batch_database_updates background task, and the spend-log + relocated into the _batch_database_updates logging task, and the spend-log insert receives the payload directly (all consumers are read-only). Asserts: - zero copy.deepcopy calls happen on the awaited request path - - the batch background task still hands the daily helpers an isolated copy + - the tracked rollup task still hands the daily helpers an isolated copy (mutating the original after the task ran does not bleed into it) - the spend-log insert receives the payload on the request path with the correct content @@ -2856,16 +2976,13 @@ async def test_update_database_does_not_deepcopy_on_request_path(): response_cost=0.1, ) - # Request path is clean: nothing was deepcopied while the caller awaited. assert len(deepcopy_calls) == 0 # The spend-log insert ran inline on the request path with the real payload. assert captured_spend_log["ref"] is fake_payload assert captured_spend_log["model_at_call"] == "gpt-4" assert fake_payload["spend"] == 0.1 - - # Now let the batch background task run; the deepcopy happens here. - await asyncio.sleep(0) + await drain_logging_tasks() assert len(deepcopy_calls) >= 1 assert len(captured_batch_payloads) == 1