From bff83d9fc0aa3110b121e2c9f795cb416490e719 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Tue, 20 Oct 2026 16:35:08 -0700 Subject: [PATCH] fix(batches): retry confirmed pre-accounting failures (#47625) Backport of #47625 to stable/1.105.x. Cherry-picked from 6564a490f40fba668201fa60d378cb6c165bb03c (main). Adapted to the line: the router budget limiter keeps the line's clock and window TTL (budget_window_seconds from #45797 is not on the line) and its websocket skip, the charge-once helpers move to batch_utils from where the line already defines them, the write-started check drops apply_search_free_tier (#47494, not on the line) and its test, and tests use the line's _ProxyDBLogger name with only this PR's tests appended. --- .../proxy/common_utils/check_batch_cost.py | 36 ++-- .../proxy/hooks/model_max_budget_limiter.py | 15 +- .../proxy/hooks/proxy_track_cost_callback.py | 56 +++++- litellm/router_strategy/budget_limiter.py | 35 +++- litellm/router_utils/batch_utils.py | 13 ++ .../common_utils/test_check_batch_cost.py | 170 ++++++++++++++++-- .../hooks/test_proxy_track_cost_callback.py | 77 +++++++- .../router_strategy/test_budget_limiter.py | 91 ++++++++++ 8 files changed, 441 insertions(+), 52 deletions(-) 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 afc1dd47c92..adf3f74bb90 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -971,20 +971,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: @@ -1111,7 +1117,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/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index d2db5145fce..9bafb0add62 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -6,8 +6,6 @@ from dataclasses import dataclass from types import MappingProxyType from typing import Final -from openai.types import Batch - import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache @@ -16,7 +14,7 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth from litellm.router_strategy.budget_limiter import RouterBudgetLimiting -from litellm.router_utils.batch_utils import is_batch_retrieve_call_type +from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import BudgetConfig, StandardLoggingPayload @@ -121,17 +119,6 @@ def model_budget_start_time_cache_key( return f"{_BUDGET_START_TIME_KEY_PREFIXES[entity_type]}:{entity_id}:{budget_model}:{budget_duration}" -def batch_charged_once_marker_key(spend_key: str, batch_id: str) -> str: - return f"{spend_key}:batch:{batch_id}" - - -def batch_id_to_charge_once(call_type: object, response_obj: object, response_cost: float) -> str | None: - """A finished batch reports its whole cost on every poll, so its id is charged once per counter.""" - if response_cost <= 0 or not is_batch_retrieve_call_type(call_type): - return None - return response_obj.id if isinstance(response_obj, Batch) else None - - def resolve_model_budget(model: str, model_max_budget: Mapping[str, object]) -> ResolvedModelBudget | None: """Find the `model_max_budget` entry that governs `model`, or None.""" for candidate in _budget_model_candidates(model): diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index dc2723c4267..3494dddf34c 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, Protocol, cast +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import batch_cost_is_final @@ -63,6 +67,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, @@ -304,6 +339,11 @@ class _ProxyDBLogger(CustomLogger): update_cache_read_keys, ) + 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( @@ -324,6 +364,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 @@ -474,6 +516,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 = ( @@ -795,7 +839,15 @@ async def _update_database_and_spend_counters_in_batch( from litellm.proxy.proxy_server import arm_update_cache_read 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/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index a792654e2a9..0a68bd8c253 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -27,6 +27,8 @@ from itertools import groupby from types import MappingProxyType from typing import Any, Final +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache @@ -37,6 +39,7 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.router_strategy.tag_based_routing import _get_tags_from_request_kwargs +from litellm.router_utils.batch_utils import batch_charged_once_marker_key, batch_id_to_charge_once from litellm.router_utils.cooldown_callbacks import ( _get_prometheus_logger_from_callbacks, ) @@ -46,6 +49,8 @@ from litellm.types.utils import BudgetConfig, GenericBudgetConfigType, StandardL from litellm.types.utils import BudgetConfig as GenericBudgetInfo DEFAULT_REDIS_SYNC_INTERVAL: Final = 1 +_CALLBACK_FIELDS: Final = TypeAdapter(dict[str, object]) +_BATCH_CHARGE_VALUES: Final = TypeAdapter(tuple[object, float]) class _LiteLLMParamsDictView: @@ -491,17 +496,22 @@ class RouterBudgetLimiting(CustomLogger): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Original method now uses helper functions""" verbose_router_logger.debug("in RouterBudgetLimiting.async_log_success_event") + callback_fields: Final = _CALLBACK_FIELDS.validate_python(kwargs) # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. - if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): + if callback_fields.get("call_type") in ("_aresponses_websocket", "_arealtime"): return standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_payload is None: raise ValueError("standard_logging_payload is required") - response_cost: Final[float] = standard_logging_payload.get("response_cost", 0) + batch_response, response_cost = _BATCH_CHARGE_VALUES.validate_python( + (response_obj, _CALLBACK_FIELDS.validate_python(standard_logging_payload).get("response_cost", 0)) + ) model_id: Final[str] = str(standard_logging_payload.get("model_id", "")) custom_llm_provider: Final[str | None] = standard_logging_payload.get("custom_llm_provider") + batch_id: Final = batch_id_to_charge_once(callback_fields.get("call_type"), batch_response, response_cost) + budget_config: Final = ( self._get_budget_config_for_provider(custom_llm_provider) if custom_llm_provider is not None else None ) @@ -514,6 +524,7 @@ class RouterBudgetLimiting(CustomLogger): spend_key=spend_key, start_time_key=start_time_key, response_cost=response_cost, + batch_id=batch_id, ) deployment_budget_config: Final = self._get_budget_config_for_deployment(model_id) @@ -526,11 +537,12 @@ class RouterBudgetLimiting(CustomLogger): spend_key=deployment_spend_key, start_time_key=deployment_start_time_key, response_cost=response_cost, + batch_id=batch_id, ) request_tags: Final = _get_tags_from_request_kwargs( - kwargs, - metadata_variable_name=get_metadata_variable_name_from_kwargs(kwargs or {}), + callback_fields, + metadata_variable_name=get_metadata_variable_name_from_kwargs(callback_fields), ) if len(request_tags) > 0: for _tag in request_tags: @@ -543,6 +555,7 @@ class RouterBudgetLimiting(CustomLogger): spend_key=_tag_spend_key, start_time_key=_tag_start_time_key, response_cost=response_cost, + batch_id=batch_id, ) async def _increment_spend_for_key( @@ -551,12 +564,15 @@ class RouterBudgetLimiting(CustomLogger): spend_key: str, start_time_key: str, response_cost: float, + batch_id: str | None = None, ): if budget_config.budget_duration is None: return current_time: Final = datetime.now(timezone.utc).timestamp() ttl_seconds: Final = duration_in_seconds(budget_config.budget_duration) + if batch_id is not None and not await self._claim_batch_charge(spend_key, batch_id, ttl_seconds): + return budget_start = await self._get_or_set_budget_start_time( start_time_key=start_time_key, @@ -594,6 +610,17 @@ class RouterBudgetLimiting(CustomLogger): verbose_router_logger.debug("Incremented spend for %s by %s", spend_key, response_cost) + async def _claim_batch_charge(self, spend_key: str, batch_id: str, ttl_seconds: int) -> bool: + marker_key: Final = batch_charged_once_marker_key(spend_key=spend_key, batch_id=batch_id) + redis_cache: Final = self.dual_cache.redis_cache + if redis_cache is not None: + polls: Final = await redis_cache.async_increment(key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True) + return polls == 1 + local_polls: Final = await self.dual_cache.in_memory_cache.async_increment( + key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True + ) + return local_polls == 1 + async def periodic_sync_in_memory_spend_with_redis(self): """ Handler that triggers sync_in_memory_spend_with_redis every DEFAULT_REDIS_SYNC_INTERVAL seconds diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index 386a2135239..50085a055b2 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -3,6 +3,8 @@ import json from os import PathLike from typing import Final +from openai.types import Batch + from litellm._logging import verbose_logger from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose from litellm.types.utils import CallTypes @@ -189,3 +191,14 @@ def is_batch_retrieve_call_type(call_type: object) -> bool: routing strategies' own state, must not be fed from it. """ return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES + + +def batch_charged_once_marker_key(spend_key: str, batch_id: str) -> str: + return f"{spend_key}:batch:{batch_id}" + + +def batch_id_to_charge_once(call_type: object, response_obj: object, response_cost: float) -> str | None: + """A finished batch reports its whole cost on every poll, so its id is charged once per counter.""" + if response_cost <= 0 or not is_batch_retrieve_call_type(call_type): + return None + return response_obj.id if isinstance(response_obj, Batch) else None diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 1edf313c647..16e033f42fa 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3223,6 +3223,151 @@ 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.caching.caching import DualCache + from litellm.integrations.custom_logger import CustomLogger + from litellm.router_strategy.budget_limiter import RouterBudgetLimiting + 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") + + budget_cache: Final = DualCache() + RouterBudgetLimiting( + dual_cache=budget_cache, + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + model_list=[ + { + "model_name": "batch-alias", + "litellm_params": {"model": "openai/batch-test-model", "max_budget": 100, "budget_duration": "1d"}, + "model_info": {"id": "model-123"}, + } + ], + ) + 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 + assert await budget_cache.async_get_cache("provider_spend:openai:1d") == pytest.approx(100 * input_rate) + assert await budget_cache.async_get_cache("deployment_spend:model-123:1d") == pytest.approx(100 * input_rate) + 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"), @@ -3508,30 +3653,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/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 28376be64b6..595f784fec0 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -1,13 +1,16 @@ 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 httpx 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 @@ -2877,3 +2880,75 @@ def test_autonomous_agent_cost_tracking_needs_no_human_or_virtual_key(agent_id: assert _should_track_cost_callback( user_api_key=None, user_id=None, team_id=None, end_user_id=None, call_type="acompletion", agent_id=agent_id ) is expected + + +_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 diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 62de1586fdd..8882c18c73d 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -6,10 +6,13 @@ anthropic_messages, embedding and rerank surfaces leave it unset, which used to the callback raise before any spend was recorded, so those budgets never moved. """ +import asyncio from typing import Final import pytest +import litellm + from litellm.caching.caching import DualCache from litellm.router_strategy.budget_limiter import RouterBudgetLimiting @@ -135,3 +138,91 @@ async def test_deployment_budget_tracked_when_provider_is_unresolvable(disable_b ) assert await limiter.dual_cache.async_get_cache("deployment_spend:deployment-1:1d") == 0.25 + + +def _completed_budget_batch(batch_id: str): + from litellm.types.utils import LiteLLMBatch + + return LiteLLMBatch( + id=batch_id, + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-budget", + object="batch", + status="completed", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ("aretrieve_batch", "retrieve_batch")) +async def test_batch_polls_charge_each_router_budget_once(monkeypatch, call_type): + monkeypatch.setattr(litellm, "tag_budget_config", {"release": {"budget_limit": 100, "time_period": "1d"}}) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + cache: Final = DualCache() + limiter: Final = RouterBudgetLimiting( + dual_cache=cache, + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + model_list=[ + { + "model_name": "batch-budget", + "litellm_params": {"model": "openai/batch-budget", "max_budget": 100, "budget_duration": "1d"}, + "model_info": {"id": "deployment-1"}, + } + ], + ) + kwargs: Final = { + **_success_kwargs(provider_in_litellm_params=None, provider_in_payload="openai", call_type=call_type), + "metadata": {"tags": ["release"]}, + } + response: Final = _completed_budget_batch("batch-first") + await asyncio.gather(*(limiter.async_log_success_event(kwargs, response, None, None) for _ in range(10))) + keys: Final = ("provider_spend:openai:1d", "deployment_spend:deployment-1:1d", "tag_spend:release:1d") + assert tuple([await cache.async_get_cache(key) for key in keys]) == (0.25, 0.25, 0.25) + await limiter.async_log_success_event(kwargs, _completed_budget_batch("batch-second"), None, None) + await limiter.async_log_success_event({**kwargs, "call_type": "acompletion"}, response, None, None) + await limiter.async_log_success_event({**kwargs, "call_type": "acompletion"}, response, None, None) + assert tuple([await cache.async_get_cache(key) for key in keys]) == (1.0, 1.0, 1.0) + await limiter.async_log_success_event({**kwargs, "call_type": "_arealtime"}, response, None, None) + assert tuple([await cache.async_get_cache(key) for key in keys]) == (1.0, 1.0, 1.0) + + +@pytest.mark.asyncio +async def test_batch_charge_is_shared_across_workers_and_requires_redis_confirmation(): + from fakeredis import FakeServer, FakeRedis as FakeSyncRedis + from redis.exceptions import ConnectionError as RedisConnectionError + from fakeredis.aioredis import FakeRedis + from litellm.caching.redis_cache import RedisCache + + server: Final = FakeServer() + client: Final = FakeRedis(server=server) + sync_client: Final = FakeSyncRedis(server=server) + redis_cache: Final = RedisCache(host="router-batch-fixture", connection_pool=sync_client.connection_pool) + litellm.in_memory_llm_clients_cache.set_cache(key=redis_cache._get_async_client_cache_key(), value=client) + first: Final = RouterBudgetLimiting( + dual_cache=DualCache(redis_cache=redis_cache), + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + ) + second: Final = RouterBudgetLimiting( + dual_cache=DualCache(redis_cache=redis_cache), + provider_budget_config={"openai": {"budget_limit": 100, "time_period": "1d"}}, + ) + kwargs: Final = _success_kwargs( + provider_in_litellm_params=None, provider_in_payload="openai", call_type="aretrieve_batch" + ) + response: Final = _completed_budget_batch("batch-shared") + await asyncio.gather( + first.async_log_success_event(kwargs, response, None, None), + second.async_log_success_event(kwargs, response, None, None), + ) + await first._push_in_memory_increments_to_redis() + await second._push_in_memory_increments_to_redis() + assert float(await client.get("provider_spend:openai:1d")) == 0.25 + server.connected = False + with pytest.raises(RedisConnectionError): + await second.async_log_success_event(kwargs, _completed_budget_batch("batch-unconfirmed"), None, None) + server.connected = True + await second.async_log_success_event(kwargs, _completed_budget_batch("batch-unconfirmed"), None, None) + await second._push_in_memory_increments_to_redis() + assert float(await client.get("provider_spend:openai:1d")) == 0.5 + await client.aclose()