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

* fix(batches): retry confirmed pre-accounting failures

* fix(batches): avoid promising retries for unknown accounting outcomes

* test(batches): exercise real spend writer during retry recovery

* fix(budgets): charge retried batches once per router budget

* test(budgets): cover shared batch charge helpers

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -852,20 +852,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:
@@ -1186,7 +1192,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

litellm/proxy/hooks/model_max_budget_limiter.py

@@ -7,8 +7,6 @@ from datetime import datetime, timezone
 from types import MappingProxyType
 from typing import Final, cast  # noqa: TID251  # untyped callback payload requires explicit narrowing
 
-from openai.types import Batch
-
 import litellm
 from litellm._internal_context import with_service_target
 from litellm._logging import verbose_proxy_logger
@@ -19,7 +17,7 @@ from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSet
 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, budget_window_seconds
-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
 
@@ -133,17 +131,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):

litellm/proxy/hooks/proxy_track_cost_callback.py

@@ -1,6 +1,8 @@
 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, timezone
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, Protocol, cast
@@ -68,6 +70,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,
@@ -445,6 +478,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(
@@ -465,6 +503,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 or litellm.apply_search_free_tier):
+                accounting.write_started = True
             if (
                 isinstance(completion_response, LiteLLMBatch)
                 and kwargs.get("call_type") == CallTypes.aretrieve_batch.value
@@ -617,6 +657,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 = (
@@ -936,7 +978,15 @@ async def _update_database_and_spend_counters_in_batch(
     project_id: str | None,
 ) -> bool:
     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,

litellm/router_strategy/budget_limiter.py

@@ -28,6 +28,8 @@ from itertools import groupby
 from types import MappingProxyType
 from typing import Any, Final
 
+from pydantic import TypeAdapter
+
 import litellm
 from litellm._internal_context import with_service_target
 from litellm._logging import verbose_router_logger
@@ -39,6 +41,7 @@ from litellm.litellm_core_utils.core_helpers import (
 )
 from litellm.litellm_core_utils.duration_parser import duration_in_seconds_from
 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,
 )
@@ -48,6 +51,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:
@@ -505,16 +510,21 @@ 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")
-        if kwargs.get("call_type") == "_arealtime":
+        callback_fields: Final = _CALLBACK_FIELDS.validate_python(kwargs)
+        if callback_fields.get("call_type") == "_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
         )
@@ -527,6 +537,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)
@@ -539,11 +550,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:
@@ -556,6 +568,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(
@@ -564,12 +577,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 = self._clock()
         new_window_seconds: Final = budget_window_seconds(current_time, budget_config.budget_duration)
+        if batch_id is not None and not await self._claim_batch_charge(spend_key, batch_id, new_window_seconds):
+            return
 
         budget_start: Final = await self._get_or_set_budget_start_time(
             start_time_key=start_time_key,
@@ -595,6 +611,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
+
     @with_service_target("router_budgets")
     async def periodic_sync_in_memory_spend_with_redis(self):
         """

litellm/router_utils/batch_utils.py

@@ -6,6 +6,7 @@ from contextvars import ContextVar
 from os import PathLike
 from typing import Final, TypeVar
 
+from openai.types import Batch
 from pydantic import TypeAdapter, ValidationError
 
 from litellm._logging import verbose_logger
@@ -286,3 +287,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

tests/unit/proxy/common_utils/test_check_batch_cost.py

@@ -3104,6 +3104,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
+        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"),
@@ -3389,30 +3534,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):

tests/unit/proxy/hooks/test_proxy_track_cost_callback.py

@@ -3270,3 +3270,100 @@ async def test_track_cost_callback_frees_grounding_under_the_allowance_only_when
     else:
         assert kwargs["response_cost"] == 1.0 and "discount_amount" not in breakdown
         assert counter.totals == {}
+
+
+@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
+
+
+@pytest.mark.asyncio
+async def test_batch_accounting_does_not_retry_after_a_search_allowance_update(monkeypatch: pytest.MonkeyPatch) -> None:
+    from litellm.proxy.hooks.proxy_track_cost_callback import track_batch_accounting
+
+    monkeypatch.setattr(litellm, "apply_search_free_tier", True)
+    counter: Final = _SharedCounter()
+    dependency: Final = SimpleNamespace(
+        failed_tracking_alert=AsyncMock(),
+        internal_usage_cache=SimpleNamespace(dual_cache=SimpleNamespace(redis_cache=counter)),
+    )
+    kwargs: Final = {**_grounded_gemini_kwargs(response_cost=1.0), "litellm_call_id": "grounded-batch"}
+    response: Final = ModelResponse(
+        usage=Usage(
+            prompt_tokens=10,
+            completion_tokens=5,
+            total_tokens=15,
+            prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=5),
+        )
+    )
+    with patch("litellm.proxy.proxy_server.proxy_logging_obj", dependency):
+        with track_batch_accounting("grounded-batch") as outcome:
+            await ProxyDBLogger().async_log_success_event(kwargs, response, _CALL_START, _CALL_START)
+
+    assert sum(counter.totals.values()) == 5
+    assert outcome.failed is True
+    assert outcome.retryable_failure is False

tests/unit/router_strategy/test_budget_limiter.py

@@ -449,3 +449,91 @@ def cleanup_redis():
             redis_client.delete(key)
     except Exception as e:
         print(f"Error cleaning up Redis: {str(e)}")
+
+
+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()

tests/unit/router_utils/test_batch_utils.py

@@ -2,7 +2,7 @@ import asyncio
 import importlib
 import json
 from io import BytesIO
-from typing import Dict, List
+from typing import Dict, Final, List
 
 import pytest
 
@@ -380,3 +380,31 @@ def setup_and_teardown():
     yield
     loop.close()
     asyncio.set_event_loop(None)
+
+
+@pytest.mark.parametrize(
+    ("call_type", "cost", "eligible"),
+    (("aretrieve_batch", 0.25, True), ("retrieve_batch", 0.25, True), ("acompletion", 0.25, False),
+     ("aretrieve_batch", 0, False), ("aretrieve_batch", -0.25, False)),
+)
+def test_batch_charge_identity_requires_a_cost_bearing_retrieval(call_type: str, cost: float, eligible: bool) -> None:
+    from openai.types import Batch
+    from litellm.router_utils.batch_utils import batch_id_to_charge_once
+
+    batch: Final = Batch(
+        id="batch-budget", completion_window="24h", created_at=1, endpoint="/v1/chat/completions",
+        input_file_id="file-budget", object="batch", status="completed",
+    )
+    assert batch_id_to_charge_once(call_type, batch, cost) == (batch.id if eligible else None)
+    assert batch_id_to_charge_once(call_type, {"id": batch.id}, cost) is None
+
+
+def test_batch_charge_markers_keep_counters_and_batches_independent() -> None:
+    from litellm.caching.in_memory_cache import InMemoryCache
+    from litellm.router_utils.batch_utils import batch_charged_once_marker_key
+
+    cache: Final = InMemoryCache()
+    cache.set_cache(batch_charged_once_marker_key("provider-spend", "first-batch"), 1)
+    assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "first-batch")) == 1
+    assert cache.get_cache(batch_charged_once_marker_key("deployment-spend", "first-batch")) is None
+    assert cache.get_cache(batch_charged_once_marker_key("provider-spend", "second-batch")) is None