From 56b0f9c3c226cfe349ffeb64dc8077e8226f2a65 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Tue, 20 Oct 2026 17:10:50 -0700 Subject: [PATCH] fix(batches): retry confirmed pre-accounting failures (#47625) Backport of the router charge-once part of #47625 to stable/1.103.x, from main 6564a490f40fba668201fa60d378cb6c165bb03c 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. --- .../proxy/hooks/model_max_budget_limiter.py | 15 +-- litellm/router_strategy/budget_limiter.py | 28 ++++- litellm/router_utils/batch_utils.py | 13 +++ .../test_router_batch_utils.py | 45 +++++++- .../router_strategy/test_budget_limiter.py | 108 ++++++++++++++++++ 5 files changed, 193 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 67577a68f5c..745302abd03 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/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): diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index b9f7e96d7e7..5fd59109f54 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/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 diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index 6e110b586fb..c0f6f73f523 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 @@ -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 diff --git a/tests/router_unit_tests/test_router_batch_utils.py b/tests/router_unit_tests/test_router_batch_utils.py index c9f19731372..e77f97bed9c 100644 --- a/tests/router_unit_tests/test_router_batch_utils.py +++ b/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 diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py index 62de1586fdd..d2c4a3767d9 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter.py +++ b/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()