scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 199

fix(batches): retry confirmed pre-accounting failures (#47625)

Browse source 
Backport of #47625 to stable/1.103.x.
Cherry-picked from 6564a490f4 (main).
Adapted to the line: the batch accounting retry (track_batch_accounting in proxy_track_cost_callback and its use in check_batch_cost) is ported, while this PR's router-budget charge-once changes in router_utils/batch_utils, budget_limiter and model_max_budget_limiter are left out, since they refactor charge-once code from main (8a33b37c39, 9acf09f60d) that is not on this line. The write-started mark keys on a budget reservation only, as the search free tier is not on this line. Tests use the line's _ProxyDBLogger, and the router-budget and search-allowance assertions are dropped.
This commit is contained in:
joshua-berri 2026-10-20 16:37:09 -07:00 • committed by Mrinal Chanshetty
parent 2f3d44b469
commit 32bfdeaf31
4 changed files with 289 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

36
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -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

56
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -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,

154
tests/proxy_unit_tests/test_check_batch_cost.py
View file

@ -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):

77
tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -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