diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 95d2f6d3856..af36ec562a6 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -974,20 +974,26 @@ class CheckBatchCost: ) return None - try: - await logging_obj.async_success_handler( - result=response, - batch_cost=batch_result.cost, - batch_usage=batch_result.usage, - batch_models=batch_result.models, - batch_successful_requests=batch_result.successful_requests, - batch_failed_requests=batch_result.failed_requests, - batch_prompt_cost=batch_result.prompt_cost, - batch_completion_cost=batch_result.completion_cost, - ) - except Exception: - await self._release_job_claim(job, original_status, original_file_object) - raise + from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting + + with track_batch_accounting(logging_obj.litellm_call_id) as accounting: + try: + await logging_obj.async_success_handler( + result=response, + batch_cost=batch_result.cost, + batch_usage=batch_result.usage, + batch_models=batch_result.models, + batch_successful_requests=batch_result.successful_requests, + batch_failed_requests=batch_result.failed_requests, + batch_prompt_cost=batch_result.prompt_cost, + batch_completion_cost=batch_result.completion_cost, + ) + finally: + if accounting.retryable_failure: + await self._release_job_claim(job, original_status, original_file_object) + + if accounting.retryable_failure: + return None # Record batch duration (completed_at - created_at) if prom_logger and response.completed_at and response.created_at: @@ -1117,7 +1123,7 @@ class CheckBatchCost: continue verbose_proxy_logger.error( f"CheckBatchCost: failed to track cost for batch {batch_id} " - f"(job {job.id}); leaving it unprocessed so the next poll retries: {tracking_err}" + f"(job {job.id}): {tracking_err}" ) self._record_error(prom_logger, "cost_tracking_error") continue diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 425466396f4..898c58519dc 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1,9 +1,13 @@ import asyncio import traceback -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Generator, Mapping, Sequence +from contextlib import contextmanager +from contextvars import ContextVar from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import batch_cost_is_final @@ -59,6 +63,37 @@ from litellm.utils import get_end_user_id_for_cost_tracking if TYPE_CHECKING: from litellm.proxy.utils import ProxyLogging + +class BatchAccountingOutcome: + def __init__(self, call_id: str) -> None: + self.call_id: Final = call_id + self.write_started = False + self.failed = False + + @property + def retryable_failure(self) -> bool: + return self.failed and not self.write_started + + +_BATCH_ACCOUNTING: Final[ContextVar[BatchAccountingOutcome | None]] = ContextVar("batch_accounting", default=None) +_BATCH_CALLBACK_FIELDS: Final = TypeAdapter(dict[str, object]) + + +@contextmanager +def track_batch_accounting(call_id: str) -> Generator[BatchAccountingOutcome, None, None]: + outcome: Final = BatchAccountingOutcome(call_id) + token: Final = _BATCH_ACCOUNTING.set(outcome) + try: + yield outcome + finally: + _BATCH_ACCOUNTING.reset(token) + + +def _batch_accounting_for(call_id: object) -> BatchAccountingOutcome | None: + outcome: Final = _BATCH_ACCOUNTING.get() + return outcome if outcome is not None and call_id == outcome.call_id else None + + _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset( { CallTypes.pass_through.value, @@ -284,6 +319,11 @@ class _ProxyDBLogger(CustomLogger): update_cache, ) + accounting: Final = ( + _batch_accounting_for(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get("litellm_call_id")) + if _BATCH_ACCOUNTING.get() is not None + else None + ) verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback") try: verbose_proxy_logger.debug( @@ -304,6 +344,8 @@ class _ProxyDBLogger(CustomLogger): ) _write_spend_metadata_to_kwargs(kwargs=kwargs, metadata=metadata) budget_reservation: Final = _get_budget_reservation_from_metadata(metadata=metadata) + if accounting is not None and budget_reservation is not None: + accounting.write_started = True if ( isinstance(completion_response, LiteLLMBatch) and kwargs.get("call_type") == CallTypes.aretrieve_batch.value @@ -446,6 +488,8 @@ class _ProxyDBLogger(CustomLogger): f"Cost tracking failed for model={model}.\nDebug info - {cost_tracking_failure_debug_info}\nAdd custom pricing - https://docs.litellm.ai/docs/proxy/custom_pricing" ) except Exception as e: + if accounting is not None: + accounting.failed = True failing_model: Final = kwargs.get("model", "") failing_call_type: Final = kwargs.get("call_type", "") error_msg: Final = ( @@ -682,7 +726,15 @@ async def _update_database_and_spend_counters( budget_reservation=budget_reservation, response_cost=response_cost ) try: - charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database( + writer: Final = proxy_logging_obj.db_spend_update_writer + accounting: Final = ( + _batch_accounting_for(_BATCH_CALLBACK_FIELDS.validate_python(kwargs).get("litellm_call_id")) + if _BATCH_ACCOUNTING.get() is not None + else None + ) + if accounting is not None: + accounting.write_started = True + charged: Final = await writer.update_database( token=user_api_key, response_cost=response_cost, user_id=user_id, diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index a7886229cbd..6f1a2be7d21 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -3217,6 +3217,135 @@ class TestMultiPodBatchCostClaim: ): yield logging_obj + @pytest.mark.asyncio + @pytest.mark.parametrize("failure_phase", ("before_write", "after_write", "external")) + @pytest.mark.parametrize("collector_enabled", (False, True)) + async def test_builtin_accounting_retries_only_before_a_possible_charge( + self, failure_phase: str, collector_enabled: bool + ) -> None: + import base64 + from datetime import datetime + + import httpx + import litellm + import respx + + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger as ProxyDBLogger + from litellm.proxy.utils import PrismaClient + from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo + from litellm.types.utils import SpecialEnums + + class AccountingDependency: + def __init__(self) -> None: + self.failure_phase = failure_phase + self.writer: Final = DBSpendUpdateWriter() + self.failed_tracking_alert = AsyncMock() + self.slack_alerting_instance = MagicMock() + self.slack_alerting_instance.customer_spend_alert = AsyncMock( + side_effect=ConnectionError("Alert destination unavailable after charge") + if failure_phase == "after_write" + else None + ) + + @property + def db_spend_update_writer(self) -> DBSpendUpdateWriter: + if self.failure_phase == "before_write": + raise ConnectionError("Accounting dependency temporarily unavailable") + return self.writer + + class ExternalCallback(CustomLogger): + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + if failure_phase == "external": + raise ConnectionError("External destination unavailable") + + dependency: Final = AccountingDependency() + spend_prisma: Final = PrismaClient( + database_url="postgresql://unused:unused@localhost:5432/unused", proxy_logging_obj=dependency + ) + create_spend_rows: Final = AsyncMock(return_value=1) + input_rate: Final = 0.000002 + litellm.register_model( + { + "batch-test-model": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token_batches": input_rate, + "output_cost_per_token_batches": 0.000006, + } + } + ) + content: Final = json.dumps( + { + "response": { + "status_code": 200, + "body": { + "model": "batch-test-model", + "usage": {"prompt_tokens": 100, "completion_tokens": 0, "total_tokens": 100}, + }, + }, + "error": None, + } + ).encode() + row: Final = _FakeManagedObjectRow() + original_status: Final = row.status + original_file_object: Final = row.file_object + row.unified_object_id = ( + base64.urlsafe_b64encode( + SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format("model-123", "batch-456").encode() + ) + .decode() + .rstrip("=") + ) + prisma: Final = self._prisma(row, []) + router: Final = self._router() + router.get_deployment.return_value = Deployment( + model_name="batch-alias", + litellm_params=LiteLLM_Params(model="openai/batch-test-model"), + model_info=ModelInfo(id="model-123"), + ) + from litellm.proxy.spend_tracking.spend_event_producer import SpendEventProducer + + producer: Final = MagicMock(spec=SpendEventProducer) if collector_enabled else None + litellm.logging_callback_manager.add_litellm_async_success_callback( + ProxyDBLogger(spend_event_producer=producer) + ) + litellm.logging_callback_manager.add_litellm_async_success_callback(ExternalCallback()) + with ( + respx.mock, + patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency), + patch("litellm.proxy.proxy_server.prisma_client", spend_prisma), + patch.object( + spend_prisma.writer_db._original_prisma, + "litellm_spendlogs", + MagicMock(create_many=create_spend_rows), + ), + ): + respx.get(f"https://api.openai.com/v1/files/{_CLAIM_OUTPUT_FILE_ID}/content").mock( + return_value=httpx.Response(200, content=content) + ) + worker: Final = self._instance(prisma, router) + await worker.check_batch_cost() + assert row.batch_processed is (failure_phase != "before_write") + if failure_phase == "before_write": + assert row.status == original_status + assert json.loads(row.file_object) == original_file_object + assert create_spend_rows.await_count == (0 if failure_phase == "before_write" else 1) + dependency.failure_phase = "recovered" + await worker.check_batch_cost() + await worker.check_batch_cost() + + assert row.batch_processed is True + create_spend_rows.assert_awaited_once() + assert create_spend_rows.call_args.kwargs["data"][0]["spend"] == pytest.approx(100 * input_rate) + if producer is not None: + assert all( + json.loads(call.args[0])["call_type"] != "aretrieve_batch" for call in producer.publish.call_args_list + ) + @pytest.mark.asyncio @pytest.mark.parametrize( ("deployment_model", "successes", "failures", "expected_model"), @@ -3515,30 +3644,23 @@ class TestMultiPodBatchCostClaim: @pytest.mark.asyncio @pytest.mark.parametrize("stored_as_json_string", (True, False)) - async def test_a_failed_spend_log_write_releases_the_claim(self, stored_as_json_string: bool) -> None: - """A transient failure while billing a claimed batch must hand the row back, or its - spend is silently lost instead of being retried on the next cycle.""" - row = _FakeManagedObjectRow() - original_file_object: Final = row.file_object if stored_as_json_string else json.loads(row.file_object) - row.file_object = original_file_object - journal = [] - prisma = self._prisma(row, journal) + async def test_an_unclassified_logging_failure_keeps_the_claim(self, stored_as_json_string: bool) -> None: + row: Final = _FakeManagedObjectRow() + row.file_object = row.file_object if stored_as_json_string else json.loads(row.file_object) + journal: Final = [] + prisma: Final = self._prisma(row, journal) - with self._billing_patches(journal, bill_error=Exception("spend log write failed")): + with self._billing_patches(journal, bill_error=Exception("Accounting outcome unknown")): await self._instance(prisma, self._router()).check_batch_cost() - assert journal == ["fetch", "claim", "bill", "release"] - assert row.batch_processed is False - assert row.status == "in_progress" - assert json.loads(row.file_object) == original_file_object + assert journal == ["fetch", "claim", "bill"] + assert row.batch_processed is True with self._billing_patches(journal) as logging_obj: await self._instance(prisma, self._router()).check_batch_cost() await self._instance(prisma, self._router()).check_batch_cost() - assert row.batch_processed is True - assert row.status == "complete" - logging_obj.async_success_handler.assert_awaited_once() + logging_obj.async_success_handler.assert_not_awaited() @pytest.mark.asyncio async def test_a_worker_interrupted_mid_costing_leaves_the_batch_billable(self): 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 0e9c336a9eb..0fb98590be2 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,12 +1,15 @@ import asyncio import json import logging -from datetime import datetime +from datetime import datetime, timezone +from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +import litellm + from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth @@ -2618,3 +2621,75 @@ async def test_track_cost_callback_failure_alert_never_carries_request_metadata_ assert "headers" in failure_debug_lines[0] else: assert failure_debug_lines == [] + + +_CALL_START: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + + +@pytest.mark.asyncio +async def test_batch_accounting_failures_stay_with_their_originating_call() -> None: + from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting + + logger: Final = _ProxyDBLogger() + dependency: Final = SimpleNamespace(failed_tracking_alert=AsyncMock()) + + async def fail_before_accounting(call_id: str) -> None: + await logger.async_log_success_event( + kwargs={ + "litellm_call_id": call_id, + "call_type": "aretrieve_batch", + "model": "batch-model", + "stream": False, + "response_cost": None, + }, + response_obj=None, + start_time=_CALL_START, + end_time=_CALL_START, + ) + + async def independent_batch() -> bool: + with track_batch_accounting("independent") as outcome: + await fail_before_accounting("independent") + return outcome.retryable_failure + + with patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency): + with track_batch_accounting("outer") as outer: + await fail_before_accounting("unrelated") + assert outer.retryable_failure is False + with track_batch_accounting("inner") as inner: + await fail_before_accounting("inner") + assert inner.retryable_failure is True + assert outer.retryable_failure is False + independent, _ = await asyncio.gather(independent_batch(), fail_before_accounting("outer")) + assert independent is True + assert outer.retryable_failure is True + with pytest.raises(asyncio.CancelledError): + with track_batch_accounting("cancelled") as cancelled: + raise asyncio.CancelledError + await fail_before_accounting("cancelled") + assert cancelled.retryable_failure is False + + +@pytest.mark.asyncio +async def test_batch_accounting_does_not_retry_after_settling_a_budget_reservation() -> None: + from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting + + reservation: Final[dict[str, object]] = {"reserved_cost": 0.5, "entries": ()} + dependency: Final = SimpleNamespace(failed_tracking_alert=AsyncMock()) + with patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency): + with track_batch_accounting("reserved") as outcome: + await _ProxyDBLogger().async_log_success_event( + kwargs={ + "litellm_call_id": "reserved", + "call_type": "aretrieve_batch", + "model": "batch-model", + "stream": False, + "response_cost": None, + "litellm_params": {"metadata": {"user_api_key_budget_reservation": reservation}}, + }, + response_obj=None, + start_time=_CALL_START, + end_time=_CALL_START, + ) + assert reservation["finalized"] is True + assert outcome.retryable_failure is False