From 8a33b37c39503419c50adad19d806a8f51fbe117 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 08:20:14 -0700 Subject: [PATCH] fix(proxy): charge a finished batch once against per-model budgets A completed batch reports its whole cost on every retrieve, and the per-model budget limiter added that cost to the key, user, team, and end-user counters on each poll. Stamping model_group on plain-id retrieves widened this from model-encoded batch ids to every poll, so a key ran out of a budget it never spent. A marker per counter and batch id now lets the first poll charge and later polls skip. --- .../proxy/hooks/model_max_budget_limiter.py | 80 ++++++++++++++--- .../hooks/test_model_max_budget_limiter.py | 89 +++++++++++++++++++ 2 files changed, 155 insertions(+), 14 deletions(-) create mode 100644 tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index bbfc7325f40..67577a68f5c 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -5,6 +5,8 @@ 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 @@ -13,6 +15,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.types.llms.openai import AllMessageValues from litellm.types.utils import BudgetConfig, StandardLoggingPayload @@ -117,6 +120,17 @@ 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): @@ -537,22 +551,18 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ) return + batch_id: Final = batch_id_to_charge_once( + call_type=kwargs.get("call_type"), + response_obj=response_obj, + response_cost=response_cost, + ) for entity_type, entity_id, resolved in resolved_budgets: - await self._increment_spend_for_key( - budget_config=resolved.budget_config, - spend_key=model_budget_spend_cache_key( - entity_type=entity_type, - entity_id=entity_id, - budget_model=resolved.budget_model, - budget_duration=resolved.budget_config.budget_duration, - ), - start_time_key=model_budget_start_time_cache_key( - entity_type=entity_type, - entity_id=entity_id, - budget_model=resolved.budget_model, - budget_duration=resolved.budget_config.budget_duration, - ), + await self._charge_entity( + entity_type=entity_type, + entity_id=entity_id, + resolved=resolved, response_cost=response_cost, + batch_id=batch_id, ) if self.dual_cache.redis_cache is not None: @@ -562,3 +572,45 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): "current state of in memory cache %s", json.dumps(self.dual_cache.in_memory_cache.cache_dict, indent=4, default=str), ) + + async def _charge_entity( + self, + entity_type: Litellm_EntityType, + entity_id: str | None, + resolved: ResolvedModelBudget, + response_cost: float, + batch_id: str | None, + ) -> None: + budget_duration: Final = resolved.budget_config.budget_duration + if budget_duration is None: + return + spend_key: Final = model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=resolved.budget_model, + budget_duration=budget_duration, + ) + if batch_id is not None and not await self._claim_batch_charge( + spend_key=spend_key, + batch_id=batch_id, + ttl_seconds=duration_in_seconds(budget_duration), + ): + return + await self._increment_spend_for_key( + budget_config=resolved.budget_config, + spend_key=spend_key, + start_time_key=model_budget_start_time_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=resolved.budget_model, + budget_duration=budget_duration, + ), + response_cost=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) + if await self.dual_cache.async_get_cache(key=marker_key) is not None: + return False + await self.dual_cache.async_set_cache(key=marker_key, value=1, ttl=ttl_seconds) + return True diff --git a/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py new file mode 100644 index 00000000000..47b14438212 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py @@ -0,0 +1,89 @@ +from typing import Final + +import pytest + +from litellm.caching.caching import DualCache +from litellm.proxy.hooks.model_max_budget_limiter import ( + _PROXY_VirtualKeyModelMaxBudgetLimiter, +) +from litellm.types.utils import LiteLLMBatch, Usage + +KEY_HASH: Final = "key-hash-batch" +USER_ID: Final = "user-batch" +MODEL_GROUP: Final = "batch-qa-primary" +BATCH_COST: Final = 2.925e-05 +CHAT_COST: Final = 0.001 +KEY_SPEND_KEY: Final = f"virtual_key_spend:{KEY_HASH}:{MODEL_GROUP}:1d" +USER_SPEND_KEY: Final = f"user_model_spend:{USER_ID}:{MODEL_GROUP}:1d" + + +def _batch(batch_id: str, status: str) -> LiteLLMBatch: + return LiteLLMBatch( + id=batch_id, + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-batch", + object="batch", + status=status, + usage=Usage(prompt_tokens=20, completion_tokens=18, total_tokens=38), + ) + + +def _event(call_type: str, response_cost: float) -> dict: + return { + "call_type": call_type, + "standard_logging_object": { + "call_type": call_type, + "response_cost": response_cost, + "model": "openai/gpt-5.4-mini", + "model_group": MODEL_GROUP, + "metadata": {"user_api_key_hash": KEY_HASH, "user_api_key_user_id": USER_ID}, + }, + "litellm_params": { + "metadata": { + "user_api_key_model_max_budget": {MODEL_GROUP: {"budget_limit": 0.0001, "time_period": "1d"}}, + "user_api_key_user_model_max_budget": {MODEL_GROUP: {"budget_limit": 0.0001, "time_period": "1d"}}, + } + }, + } + + +async def _poll(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, batch: LiteLLMBatch, response_cost: float) -> None: + await limiter.async_log_success_event( + _event("aretrieve_batch", response_cost), response_obj=batch, start_time=None, end_time=None + ) + + +async def _chat(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter) -> None: + await limiter.async_log_success_event(_event("acompletion", CHAT_COST), response_obj=None, start_time=None, end_time=None) + + +async def _spend(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, spend_key: str) -> float: + return await limiter.dual_cache.async_get_cache(key=spend_key) or 0.0 + + +@pytest.mark.asyncio +async def test_polls_of_a_finished_batch_charge_each_per_model_budget_once(): + limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + first: Final = _batch("batch_first", "completed") + + await _poll(limiter, _batch("batch_first", "in_progress"), response_cost=0) + for _ in range(3): + await _poll(limiter, first, response_cost=BATCH_COST) + + assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST) + assert await _spend(limiter, USER_SPEND_KEY) == pytest.approx(BATCH_COST) + + +@pytest.mark.asyncio +async def test_a_second_batch_and_chat_requests_still_charge_the_budget(): + limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + + await _poll(limiter, _batch("batch_first", "completed"), response_cost=BATCH_COST) + await _poll(limiter, _batch("batch_first", "completed"), response_cost=BATCH_COST) + await _poll(limiter, _batch("batch_second", "completed"), response_cost=BATCH_COST) + await _chat(limiter) + await _chat(limiter) + + assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(2 * BATCH_COST + 2 * CHAT_COST)