scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-27 01:22:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

fix(proxy): charge a finished batch once against per-model budgets

Browse source 
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.
This commit is contained in:
mateo-berri 2026-09-19 08:20:14 -07:00
parent 952cafea7d
commit 8a33b37c39
2 changed files with 155 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

80
litellm/proxy/hooks/model_max_budget_limiter.py
View file

@ -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

89
tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py Normal file
View file

@ -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)