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

Backport of the router charge-once part of #47625 to stable/1.103.x, from main 6564a490f4 via the stable/1.104.x port 91cb4a55b5.
The batch accounting retry already landed in the earlier #47625 commit; this adds the charge-once helpers moved into router_utils.batch_utils and RouterBudgetLimiting's per-batch marker using the line's duration_in_seconds window. The shared-worker test waits for this line's background Redis push before reading spend.

litellm/proxy/hooks/model_max_budget_limiter.py

@@ -5,8 +5,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
@@ -15,7 +13,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
 
@@ -120,17 +118,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/router_strategy/budget_limiter.py

@@ -25,6 +25,8 @@ from collections.abc import Mapping
 from datetime import datetime, timedelta, timezone
 from typing import Any, Final
 
+from pydantic import TypeAdapter
+
 import litellm
 from litellm._logging import verbose_router_logger
 from litellm.caching.caching import DualCache
@@ -35,6 +37,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,
 )
@@ -44,6 +47,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:
@@ -414,10 +419,14 @@ class RouterBudgetLimiting(CustomLogger):
         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(kwargs.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
         )
@@ -430,6 +439,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)
@@ -442,6 +452,7 @@ 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(
@@ -459,6 +470,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(
@@ -467,12 +479,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,
@@ -510,6 +525,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

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
@@ -188,3 +190,14 @@ def is_batch_retrieve_call_type(call_type: object) -> bool:
     finished batch. Per-minute usage counters 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/router_unit_tests/test_router_batch_utils.py

@@ -3,7 +3,7 @@ import pytest
 
 import json
 from io import BytesIO
-from typing import Dict, List
+from typing import Dict, List, Final
 from litellm.router_utils.batch_utils import (
     replace_model_in_jsonl,
     _get_router_metadata_variable_name,
@@ -317,3 +317,46 @@ def test_replace_model_in_jsonl_with_embedded_newlines():
         == "This is a message\nwith multiple\nlines"
     )
     assert result_json["custom_id"] == "test123"
+
+
+def test_is_batch_retrieve_call_type_matches_only_batch_retrieves():
+    from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
+    from litellm.types.utils import CallTypes
+
+    assert is_batch_retrieve_call_type(CallTypes.aretrieve_batch.value) is True
+    assert is_batch_retrieve_call_type(CallTypes.retrieve_batch.value) is True
+
+    for call_type in CallTypes:
+        if call_type in (CallTypes.aretrieve_batch, CallTypes.retrieve_batch):
+            continue
+        assert is_batch_retrieve_call_type(call_type.value) is False
+
+    assert is_batch_retrieve_call_type(None) is False
+
+
+@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

tests/test_litellm/router_strategy/test_budget_limiter.py

@@ -6,10 +6,12 @@ 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 +137,109 @@ 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)
+
+
+async def _drain_background_pushes() -> None:
+    pending: Final = [
+        task for task in asyncio.all_tasks() if task.get_coro().__name__ == "_push_increments_to_redis"
+    ]
+    await asyncio.gather(*pending, return_exceptions=True)
+
+
+@pytest.mark.asyncio
+async def test_batch_charge_is_shared_across_workers_and_requires_redis_confirmation(monkeypatch):
+    import importlib.util
+    import sys
+
+    loaded: Final = sys.modules.get("valkey")
+    spec: Final = importlib.util.find_spec("valkey") if loaded is None else None
+    if (loaded is not None and not hasattr(loaded, "Valkey")) or (spec is not None and spec.origin is None):
+        monkeypatch.setitem(sys.modules, "valkey", None)
+        for name in [m for m in sys.modules if m.startswith("fakeredis")]:
+            monkeypatch.delitem(sys.modules, name)
+    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()
+    await _drain_background_pushes()
+    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()
+    await _drain_background_pushes()
+    assert float(await client.get("provider_spend:openai:1d")) == 0.5
+    await client.aclose()