scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(proxy): reserve tiered budget all-or-nothing across all deployments

Browse source 
Alibaba Model Studio (Dashscope) tiered pricing is all-or-nothing: the tier is
selected by a request's total input tokens and every token, input and output, is
billed at that one tier's rate. The reservation path used graduated slicing and,
worse, picked the output tier from the output-token count, so a long-context
request with a large output allowance reserved far less than the provider charges
and could slip past a depleted budget. Select the tier from input tokens and apply
its rates to all input and output tokens.

Reservation also read tiered pricing from only the first deployment in a model
group. A caller could hit an alias whose cheaper deployment was listed first and
exceed the budget once routed to a costlier sibling. Estimate against every
eligible deployment's pricing and reserve the maximum.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Shivam Rawat 2026-07-11 14:22:15 -07:00
parent 8548816ae4
commit c132a03ca7
3 changed files with 231 additions and 59 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

40
litellm/litellm_core_utils/llm_cost_calc/tiered_pricing.py
View file

@ -97,3 +97,43 @@ def calculate_tiered_cost(
total_cost += remaining_tokens * _coerce_cost_per_token(cost_per_token)
return total_cost
def select_tier_for_input(
tiered_pricing: List[dict],
input_tokens: int,
) -> Optional[dict]:
"""
Select the pricing tier for a request based on its total input token count.
Alibaba Model Studio (Dashscope) tiered pricing is all-or-nothing: the tier is
chosen by the total input tokens of a single request and every token in the
request (input and output) is billed at that one tier's rate, rather than
graduated income-tax-style slicing. A tier matches when
``range_start < input_tokens <= range_end`` (so a request of exactly
``range_end`` tokens stays in the lower tier, matching the official
``0 < Token <= 32K`` phrasing). Requests above the highest declared range fall
back to the last (most expensive) tier.
"""
if not tiered_pricing or input_tokens <= 0:
return None
sorted_tiers = sorted(tiered_pricing, key=lambda t: t.get("range", [0, 0])[0])
valid_tiers = [tier for tier in sorted_tiers if len(tier.get("range", [])) == 2]
if not valid_tiers:
return None
matching = [tier for tier in valid_tiers if tier["range"][0] < input_tokens <= tier["range"][1]]
if matching:
return matching[0]
return valid_tiers[-1]
def tier_rate(
tier: dict,
cost_key: str,
fallback_cost_key: Optional[str] = None,
) -> float:
"""Read a per-token rate from a tier, coercing YAML string costs to float."""
raw = tier.get(cost_key) or tier.get(fallback_cost_key, 0)
return _coerce_cost_per_token(raw)

166
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -10,7 +10,7 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import calculate_tiered_cost
from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate
from litellm.proxy._types import (
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
@ -925,9 +925,25 @@ def _estimate_request_input_cost_for_model(
model: str,
llm_router: Router | None,
) -> float | None:
model_info = _get_model_cost_info(model=model, llm_router=llm_router)
if model_info is None:
return None
estimates = [
_input_cost_for_cost_info(
request_body=request_body,
route=route,
model=model,
model_info=model_info,
)
for model_info in _get_model_cost_infos(model=model, llm_router=llm_router)
]
valid_estimates = [estimate for estimate in estimates if estimate is not None]
return max(valid_estimates) if valid_estimates else None
def _input_cost_for_cost_info(
request_body: dict,
route: str,
model: str,
model_info: Dict[str, Any],
) -> Optional[float]:
input_tokens = _estimate_input_tokens(
request_body=request_body,
route=route,
@ -938,11 +954,9 @@ def _estimate_request_input_cost_for_model(
return None
tiered_pricing = model_info.get("tiered_pricing")
if isinstance(tiered_pricing, list) and tiered_pricing:
return calculate_tiered_cost(
tokens=input_tokens,
tiered_pricing=tiered_pricing,
cost_key="input_cost_per_token",
)
tier = select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=input_tokens)
if tier is not None:
return input_tokens * tier_rate(tier, "input_cost_per_token")
input_cost_per_token = _to_float(model_info.get("input_cost_per_token"))
if input_cost_per_token is None:
return None
@ -955,10 +969,25 @@ def _estimate_request_max_cost_for_model(
model: str,
llm_router: Optional[Router],
) -> Optional[float]:
model_info = _get_model_cost_info(model=model, llm_router=llm_router)
if model_info is None:
return None
estimates = [
_max_cost_for_cost_info(
request_body=request_body,
route=route,
model=model,
model_info=model_info,
)
for model_info in _get_model_cost_infos(model=model, llm_router=llm_router)
]
valid_estimates = [estimate for estimate in estimates if estimate is not None]
return max(valid_estimates) if valid_estimates else None
def _max_cost_for_cost_info(
request_body: dict,
route: str,
model: str,
model_info: Dict[str, Any],
) -> Optional[float]:
image_cost = _estimate_image_generation_cost(
request_body=request_body,
model_info=model_info,
@ -966,8 +995,6 @@ def _estimate_request_max_cost_for_model(
if image_cost is not None:
return image_cost
input_cost_per_token = _to_float(model_info.get("input_cost_per_token"))
output_cost_per_token = _to_float(model_info.get("output_cost_per_token"))
input_tokens = _estimate_input_tokens(
request_body=request_body,
route=route,
@ -985,16 +1012,14 @@ def _estimate_request_max_cost_for_model(
output_multiplier = _get_output_multiplier(request_body=request_body)
tiered_pricing = model_info.get("tiered_pricing")
if isinstance(tiered_pricing, list) and tiered_pricing:
return calculate_tiered_cost(
tokens=input_tokens,
tiered_pricing=tiered_pricing,
cost_key="input_cost_per_token",
) + calculate_tiered_cost(
tokens=output_tokens * output_multiplier,
tiered_pricing=tiered_pricing,
cost_key="output_cost_per_token",
)
tier = select_tier_for_input(tiered_pricing=tiered_pricing, input_tokens=input_tokens)
if tier is not None:
return (input_tokens * tier_rate(tier, "input_cost_per_token")) + (
output_tokens * output_multiplier * tier_rate(tier, "output_cost_per_token")
)
input_cost_per_token = _to_float(model_info.get("input_cost_per_token"))
output_cost_per_token = _to_float(model_info.get("output_cost_per_token"))
cost = 0.0
if input_cost_per_token is not None:
cost += input_tokens * input_cost_per_token
@ -1052,46 +1077,69 @@ def _get_model_cost_info(
llm_router: Optional[Router],
) -> Optional[Dict[str, Any]]:
if llm_router is not None:
try:
model_group_info = llm_router.get_model_group_info(model_group=model)
if model_group_info is not None:
model_group_cost_info = model_group_info.model_dump()
# Reservation runs before routing, so the concrete deployment is
# unknown here. We assume deployments in a group share pricing and
# use the first tiered table we find as the estimate. This is a
# best-effort admission gate; post-request billing is exact against
# the deployment that actually served the request. Mixing tiered and
# flat deployments in one group can therefore over- or under-estimate
# at reservation time only.
deployments = llm_router.get_model_list(model_name=model) or []
for deployment in deployments:
model_id = deployment.get("model_info", {}).get("id")
backend_model = deployment.get("litellm_params", {}).get("model")
if not isinstance(model_id, str) or not isinstance(backend_model, str):
continue
deployment_model_info = llm_router.get_deployment_model_info(
model_id=model_id,
model_name=backend_model,
)
if deployment_model_info is None:
continue
tiered_pricing = deployment_model_info.get("tiered_pricing")
if isinstance(tiered_pricing, list) and tiered_pricing:
return {
**model_group_cost_info,
"tiered_pricing": tiered_pricing,
}
return model_group_cost_info
except Exception:
verbose_proxy_logger.debug(
"Unable to load router model group info for budget reservation",
exc_info=True,
)
model_group_info = llm_router.get_model_group_info(model_group=model)
if model_group_info is not None:
return model_group_info.model_dump()
return dict(litellm.get_model_info(model=model))
def _get_model_cost_infos(
model: str,
llm_router: Optional[Router],
) -> List[Dict[str, Any]]:
"""Cost-info candidates to estimate a request against for one model group.
Reservation runs before routing, so the deployment that will serve the request
is unknown. Rather than guess, we estimate the cost against every eligible
pricing shape in the group (the group's flat rates plus each deployment's
tiered table) and let the caller reserve the maximum, so a cheaper sibling
deployment can never leave the request under-reserved.
"""
try:
return dict(litellm.get_model_info(model=model))
base = _get_model_cost_info(model=model, llm_router=llm_router)
if base is None:
return []
tiered_tables = _get_deployment_tiered_pricing_tables(model=model, llm_router=llm_router)
except Exception:
verbose_proxy_logger.debug(
"Unable to load model cost info for budget reservation",
exc_info=True,
)
return []
if not tiered_tables:
return [base]
return [base, *({**base, "tiered_pricing": table} for table in tiered_tables)]
def _deployment_tiered_pricing_table(
deployment: Dict[str, Any],
llm_router: Router,
) -> Optional[List[dict]]:
model_id = deployment.get("model_info", {}).get("id")
backend_model = deployment.get("litellm_params", {}).get("model")
if not isinstance(model_id, str) or not isinstance(backend_model, str):
return None
deployment_model_info = llm_router.get_deployment_model_info(model_id=model_id, model_name=backend_model)
if deployment_model_info is None:
return None
tiered_pricing = deployment_model_info.get("tiered_pricing")
if isinstance(tiered_pricing, list) and tiered_pricing:
return tiered_pricing
return None
def _get_deployment_tiered_pricing_tables(
model: str,
llm_router: Optional[Router],
) -> List[List[dict]]:
if llm_router is None:
return []
deployments = llm_router.get_model_list(model_name=model) or []
return [
table
for deployment in deployments
if (table := _deployment_tiered_pricing_table(deployment, llm_router)) is not None
]
def _estimate_input_tokens(

84
tests/test_litellm/proxy/test_budget_reservation.py
View file

@ -824,6 +824,90 @@ async def test_should_reserve_tiered_pricing_cost(spend_counter_state):
await release_budget_reservation(reservation)
def test_tiered_reservation_is_all_or_nothing_with_output_tier_from_input_length():
"""Dashscope tiered pricing is all-or-nothing: the tier is chosen by the total
input tokens and every token (input and output) is billed at that tier's rate.
A long-context request with a large output allowance must reserve the output at
the input-selected tier, not at the cheapest tier picked from the output volume.
The earlier graduated calculation under-reserved such requests, letting a caller
slip past a depleted budget."""
tiered_pricing = [
{"range": [0, 32000], "input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06},
{"range": [32000, 128000], "input_cost_per_token": 4e-06, "output_cost_per_token": 8e-06},
]
input_tokens = 100000 # falls entirely in the second tier
output_tokens = 1000
with (
patch(
"litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info",
return_value={"tiered_pricing": tiered_pricing, "max_output_tokens": 200000},
),
patch(
"litellm.proxy.spend_tracking.budget_reservation._estimate_input_tokens",
return_value=input_tokens,
),
patch(
"litellm.proxy.spend_tracking.budget_reservation._estimate_output_tokens",
return_value=output_tokens,
),
):
estimated = estimate_request_max_cost(
request_body=_request_body(),
route="/chat/completions",
llm_router=None,
)
expected = (input_tokens * 4e-06) + (output_tokens * 8e-06)
assert estimated == pytest.approx(expected)
# What the old graduated math (with the output tier taken from output volume)
# would have reserved. The all-or-nothing estimate must be strictly larger.
graduated_under_reserve = (32000 * 1e-06) + (68000 * 4e-06) + (output_tokens * 2e-06)
assert estimated > graduated_under_reserve
def test_reservation_uses_most_expensive_deployment_in_group():
"""When a model group mixes deployments with different tiered rates, reservation
must estimate against the most expensive one. Reserving the cheaper sibling would
let a caller repeatedly hit the alias and exceed the budget once routed to the
costlier deployment."""
cheap = [{"range": [0, 32000], "input_cost_per_token": 1e-06, "output_cost_per_token": 2e-06}]
expensive = [{"range": [0, 32000], "input_cost_per_token": 5e-06, "output_cost_per_token": 1e-05}]
input_tokens = 1000
output_tokens = 10
with (
patch(
"litellm.proxy.spend_tracking.budget_reservation._get_model_cost_info",
return_value={"max_output_tokens": 200000},
),
patch(
"litellm.proxy.spend_tracking.budget_reservation._get_deployment_tiered_pricing_tables",
return_value=[cheap, expensive],
),
patch(
"litellm.proxy.spend_tracking.budget_reservation._estimate_input_tokens",
return_value=input_tokens,
),
patch(
"litellm.proxy.spend_tracking.budget_reservation._estimate_output_tokens",
return_value=output_tokens,
),
):
estimated = estimate_request_max_cost(
request_body=_request_body(),
route="/chat/completions",
llm_router=None,
)
expected_expensive = (input_tokens * 5e-06) + (output_tokens * 1e-05)
expected_cheap = (input_tokens * 1e-06) + (output_tokens * 2e-06)
assert expected_expensive > expected_cheap
assert estimated == pytest.approx(expected_expensive)
@pytest.mark.asyncio
async def test_should_clamp_reservation_to_model_ceiling_when_caller_overrequests(
spend_counter_state,