scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

Merge pull request #42062 from BerriAI/litellm_pr38499_batch_retrieve_model_group

Browse source 
fix(router): stamp model_group when retrieving a batch, so batch tokens are attributable (internal copy of #38499)
This commit is contained in:
Mateo Wang 2026-09-19 18:23:06 -07:00 • committed by GitHub
commit 93e39d5042
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 723 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -6,7 +6,7 @@
######################################################################
import asyncio
import os
from collections.abc import Mapping
from collections.abc import Mapping, MutableMapping
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, cast
@ -141,6 +141,20 @@ async def _raise_when_input_file_must_be_managed(model: str, credentials: Mappin
)
def _litellm_metadata_of(data: MutableMapping[str, object]) -> MutableMapping[str, object]:
"""The request's litellm_metadata mapping, created on the request when it carries none.
The success handler reads this mapping, so a flag or a model group set here has to live
inside it rather than beside it.
"""
existing: Final = data.get("litellm_metadata")
if isinstance(existing, MutableMapping):
return existing
created: Final[dict[str, object]] = {} # mutable-ok: the logging layer copies and extends this mapping
data["litellm_metadata"] = created # rebind-ok: the success handler reads the request's own mapping
return created
def _raise_not_found_when_openai_fallback_unservable(
requested_provider: "str | None",
data: Mapping[str, object],
@ -668,11 +682,7 @@ async def retrieve_batch(
poller_owns_accounting: Final = bool(unified_batch_id) and batch_cost_poller_is_active()
if poller_owns_accounting:
litellm_metadata = data.get("litellm_metadata")
if not isinstance(litellm_metadata, dict):
litellm_metadata = {} # mutable-ok: the suppression flag must live inside litellm_metadata for the success handler to read it, and this request carried no mapping to extend
data["litellm_metadata"] = litellm_metadata
litellm_metadata["batch_ignore_default_logging"] = True
_litellm_metadata_of(data)["batch_ignore_default_logging"] = True
# Retrieve from provider (for non-terminal states or if DB lookup failed)
# SCENARIO 1: Batch ID is encoded with model info
@ -697,6 +707,7 @@ async def retrieve_batch(
# it the call falls into the legacy provider switch and 400s.
data["model"] = model_from_id
add_deployment_model_info(data=data, llm_router=llm_router, model_id=model_from_id)
_litellm_metadata_of(data).setdefault("model_group", model_from_id)
# Retrieve batch using model credentials
response = await litellm.aretrieve_batch(

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)
polls: Final = await self.dual_cache.async_increment_cache(
key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True
)
return polls == 1

12
litellm/router.py
View file

@ -152,6 +152,7 @@ from litellm.router_utils.auto_router_model_naming import (
)
from litellm.router_utils.batch_utils import (
_get_router_metadata_variable_name,
is_batch_retrieve_call_type,
replace_model_in_jsonl,
should_replace_model_in_jsonl,
)
@ -6200,6 +6201,8 @@ class Router:
"""
try:
parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs)
requested_model_group: Final = model
metadata_variable_name: Final = _get_router_metadata_variable_name(function_name="aretrieve_batch")
if model is not None:
filtered_model_list: (
list[DeploymentTypedDict] | list[dict] | dict | None
@ -6236,6 +6239,9 @@ class Router:
kwargs=new_kwargs,
function_name="aretrieve_batch",
)
model_group: Final = requested_model_group or model_name["model_name"]
if not new_kwargs[metadata_variable_name].get("model_group"):
new_kwargs[metadata_variable_name]["model_group"] = model_group
new_kwargs.pop("custom_llm_provider", None)
data.pop("custom_llm_provider", None)
return await litellm.aretrieve_batch(
@ -7944,6 +7950,8 @@ class Router:
# WS session wrappers fire with result=None; per-turn costs tracked by inner calls.
if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"):
return
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return
standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
if standard_logging_object is None:
raise ValueError("standard_logging_object is None")
@ -8090,6 +8098,8 @@ class Router:
- key: str - The key used to increment the cache
- None: if no key is found
"""
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return None
id = None
if kwargs["litellm_params"].get("metadata") is None:
pass
@ -8218,6 +8228,8 @@ class Router:
"""
Update RPM usage for a deployment
"""
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return
deployment_name: Final = kwargs["litellm_params"]["metadata"].get(
"deployment", None
) # handles wildcard routes - by giving the original name sent to `litellm.completion`

3
litellm/router_strategy/least_busy.py
View file

@ -9,6 +9,7 @@ from litellm._logging import verbose_router_logger
from litellm.caching.caching import DualCache
from litellm.caching.redis_cache import log_redis_failure
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60
@ -48,6 +49,8 @@ def _request_count_key(model_group: str, deployment_id: str) -> str:
def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None:
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return None
try:
call: Final = _CALL_KWARGS.validate_python(kwargs)
except ValidationError:

5
litellm/router_strategy/lowest_cost.py
View file

@ -8,6 +8,7 @@ from litellm import ModelResponse, token_counter, verbose_logger
from litellm._logging import verbose_router_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
class LowestCostLoggingHandler(CustomLogger):
@ -19,6 +20,8 @@ class LowestCostLoggingHandler(CustomLogger):
self.router_cache = router_cache
def log_success_event(self, kwargs, response_obj, start_time, end_time):
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return
try:
"""
Update usage on success
@ -92,6 +95,8 @@ class LowestCostLoggingHandler(CustomLogger):
)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return
try:
"""
Update cost usage on success

7
litellm/router_strategy/lowest_latency.py
View file

@ -13,6 +13,7 @@ from litellm import ModelResponse, token_counter, verbose_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs, safe_divide_seconds
from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
from litellm.types.utils import LiteLLMPydanticObjectBase
if TYPE_CHECKING:
@ -58,6 +59,8 @@ class LowestLatencyLoggingHandler(CustomLogger):
self.routing_args = RoutingArgs(**routing_args)
def log_success_event(self, kwargs, response_obj, start_time, end_time):
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return
try:
"""
Update latency usage on success
@ -182,6 +185,8 @@ class LowestLatencyLoggingHandler(CustomLogger):
"""
Check if Timeout Error, if timeout set deployment latency -> 100
"""
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return
try:
metadata_field: Final = self._select_metadata_field(kwargs)
_exception: Final = kwargs.get("exception", None)
@ -236,6 +241,8 @@ class LowestLatencyLoggingHandler(CustomLogger):
)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return
try:
"""
Update latency usage on success

5
litellm/router_strategy/lowest_tpm_rpm.py
View file

@ -8,6 +8,7 @@ from litellm import token_counter
from litellm._logging import verbose_router_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
from litellm.types.utils import LiteLLMPydanticObjectBase
from litellm.utils import print_verbose
@ -27,6 +28,8 @@ class LowestTPMLoggingHandler(CustomLogger):
self.routing_args = RoutingArgs(**routing_args)
def log_success_event(self, kwargs, response_obj, start_time, end_time):
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return
try:
"""
Update TPM/RPM usage on success
@ -79,6 +82,8 @@ class LowestTPMLoggingHandler(CustomLogger):
verbose_router_logger.debug(traceback.format_exc())
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return
try:
"""
Update TPM/RPM usage on success

5
litellm/router_strategy/lowest_tpm_rpm_v2.py
View file

@ -12,6 +12,7 @@ from litellm._logging import verbose_logger, verbose_router_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
from litellm.types.router import RouterErrors
from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload
from litellm.utils import get_utc_datetime, print_verbose
@ -210,6 +211,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger):
return deployment # don't fail calls if eg. redis fails to connect
def log_success_event(self, kwargs, response_obj, start_time, end_time):
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return
try:
"""
Update TPM/RPM usage on success
@ -250,6 +253,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger):
)
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
if is_batch_retrieve_call_type(kwargs.get("call_type")):
return
try:
"""
Update TPM usage on success

19
litellm/router_utils/batch_utils.py
View file

@ -5,6 +5,7 @@ from typing import Final
from litellm._logging import verbose_logger
from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose
from litellm.types.utils import CallTypes
class InMemoryFile(io.BytesIO):
@ -170,3 +171,21 @@ def _get_router_metadata_variable_name(function_name: str | None) -> str:
return "litellm_metadata"
else:
return "metadata"
BATCH_RETRIEVE_CALL_TYPES: Final = frozenset(
{
CallTypes.aretrieve_batch.value,
CallTypes.retrieve_batch.value,
}
)
def is_batch_retrieve_call_type(call_type: object) -> bool:
"""
A batch retrieve reports the whole job's token usage, which the provider spent
asynchronously over the life of the batch, and reports it again on every poll of the
finished batch. The counters that measure live traffic, per-minute rate limits and the
routing strategies' own state, must not be fed from it.
"""
return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES

15
tests/router_unit_tests/test_router_batch_utils.py
View file

@ -317,3 +317,18 @@ 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

40
tests/test_litellm/proxy/batches_endpoints/test_endpoints.py
View file

@ -1438,9 +1438,11 @@ async def call_retrieve(
user: Optional[UserAPIKeyAuth] = None,
headers: Optional[Dict[str, str]] = None,
query: Optional[Dict[str, str]] = None,
enriched_data: Optional[Dict[str, Any]] = None,
):
# Mirror the real flow: data starts as RetrieveBatchRequest(batch_id=...).
harness.data["data"] = {"batch_id": batch_id}
# Mirror the real flow: data starts as RetrieveBatchRequest(batch_id=...),
# then pre-call enrichment adds key/team metadata to it.
harness.data["data"] = {"batch_id": batch_id, **(enriched_data or {})}
return await endpoints.retrieve_batch(
request=FakeRequest(headers=headers, query=query),
fastapi_response=Response(),
@ -1476,6 +1478,7 @@ async def test_retrieve__model_encoded_id(retrieve_harness):
"api_key": "sk-azure",
"api_base": "https://azure.test",
"model": "azure/gpt-4o",
"litellm_metadata": {"model_group": "azure/gpt-4o"},
}
# 4. OUTPUT SHAPE - ids re-encoded with the model for the round-trip.
@ -1522,6 +1525,39 @@ async def test_retrieve__model_encoded_id__forwards_decoded_model_not_deployment
assert retrieve_harness.aretrieve_kwargs()["model"] == "azure/gpt-4o"
@pytest.mark.asyncio
async def test_retrieve__model_encoded_id__stamps_model_group(retrieve_harness):
"""This path never goes through the router, so nothing else labels the call.
Without the stamp the spend log lands under a blank model group and the batch
disappears from per-model usage."""
await call_retrieve(retrieve_harness, AZURE_BATCH_ID)
litellm_metadata = retrieve_harness.aretrieve_kwargs()["litellm_metadata"]
assert litellm_metadata["model_group"] == "azure/gpt-4o"
@pytest.mark.asyncio
async def test_retrieve__model_encoded_id__stamps_model_group_beside_existing_metadata(
retrieve_harness,
):
"""The stamp joins the metadata pre-call enrichment already built. Replacing
that dict instead of adding to it drops the key and team labels the spend log
is attributed with."""
await call_retrieve(
retrieve_harness,
AZURE_BATCH_ID,
enriched_data={"litellm_metadata": {"user_api_key_alias": "team-a-key"}},
)
litellm_metadata = retrieve_harness.aretrieve_kwargs()["litellm_metadata"]
assert litellm_metadata == {
"user_api_key_alias": "team-a-key",
"model_group": "azure/gpt-4o",
}
@pytest.mark.asyncio
async def test_retrieve__model_encoded_id__encodes_output_and_error_ids(
retrieve_harness,

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

@ -0,0 +1,198 @@
import asyncio
import time
from collections.abc import Callable, Mapping
from types import MappingProxyType
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.caching import RedisPipelineIncrementOperation
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[str, object]:
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
def _local_spend(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, spend_key: str) -> float:
return limiter.dual_cache.in_memory_cache.get_cache(key=spend_key) or 0.0
class _Clock:
def __init__(self) -> None:
self.seconds = 0.0
def now(self) -> float:
return self.seconds
def advance(self, seconds: float) -> None:
self.seconds = self.seconds + seconds
class _SharedRedisDouble:
def __init__(self, now: Callable[[], float] = time.time) -> None:
self.now = now
self.entries: Mapping[str, tuple[float, float | None]] = MappingProxyType({})
def _live(self, key: str) -> tuple[float, float | None] | None:
entry: Final = self.entries.get(key)
if entry is None:
return None
expires_at: Final = entry[1]
if expires_at is not None and expires_at <= self.now():
return None
return entry
def _store(self, key: str, value: float, expires_at: float | None) -> None:
self.entries = MappingProxyType({**self.entries, key: (value, expires_at)})
async def async_get_cache(self, key: str, **kwargs: object) -> float | None:
await asyncio.sleep(0)
entry: Final = self._live(key)
return None if entry is None else entry[0]
async def async_set_cache(self, key: str, value: float, ttl: int | None = None, **kwargs: object) -> None:
await asyncio.sleep(0)
self._store(key, value, None if ttl is None else self.now() + ttl)
async def async_increment(
self,
key: str,
value: float,
ttl: int | None = None,
parent_otel_span: object = None,
refresh_ttl: bool = False,
) -> float:
await asyncio.sleep(0)
live: Final = self._live(key)
total: Final = value if live is None else live[0] + value
kept_expiry: Final = None if live is None else live[1]
expires_at: Final = (
kept_expiry if ttl is None or (kept_expiry is not None and not refresh_ttl) else self.now() + ttl
)
self._store(key, total, expires_at)
return total
async def async_increment_pipeline(self, increment_list: list[RedisPipelineIncrementOperation]) -> list[float]:
return [await self.async_increment(op["key"], op["increment_value"], ttl=op["ttl"]) for op in increment_list]
def _worker(redis: _SharedRedisDouble) -> _PROXY_VirtualKeyModelMaxBudgetLimiter:
return _PROXY_VirtualKeyModelMaxBudgetLimiter(
dual_cache=DualCache(redis_cache=redis) # pyright: ignore[reportArgumentType] # duck-typed Redis double
)
async def _drain_redis_pushes() -> None:
await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task()))
@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)
@pytest.mark.asyncio
async def test_two_workers_polling_the_same_finished_batch_at_once_charge_it_once():
redis: Final = _SharedRedisDouble()
worker_a: Final = _worker(redis)
worker_b: Final = _worker(redis)
finished: Final = _batch("batch_first", "completed")
await asyncio.gather(_poll(worker_a, finished, BATCH_COST), _poll(worker_b, finished, BATCH_COST))
await _drain_redis_pushes()
assert _local_spend(worker_a, KEY_SPEND_KEY) + _local_spend(worker_b, KEY_SPEND_KEY) == pytest.approx(BATCH_COST)
assert await redis.async_get_cache(KEY_SPEND_KEY) == pytest.approx(BATCH_COST)
@pytest.mark.asyncio
async def test_a_batch_polled_within_every_budget_window_is_never_charged_again():
clock: Final = _Clock()
limiter: Final = _worker(_SharedRedisDouble(now=clock.now))
finished: Final = _batch("batch_first", "completed")
await _poll(limiter, finished, BATCH_COST)
clock.advance(12 * 3600)
await _poll(limiter, finished, BATCH_COST)
clock.advance(18 * 3600)
await _poll(limiter, finished, BATCH_COST)
assert _local_spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST)

333
tests/test_litellm/test_router.py
View file

@ -49,6 +49,7 @@ from litellm.router import (
from litellm.router_strategy import simple_shuffle
from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit
from litellm.router_utils.cooldown_handlers import _async_get_cooldown_deployments
from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute
from litellm.types.llms.openai import ChatCompletionRequest
from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo, PreRoutingHookResponse, RetryPolicy
@ -1009,6 +1010,338 @@ async def test_arouter_aretrieve_batch():
assert mock_aretrieve_batch.call_args.kwargs["api_base"] == "my-custom-base"
_BATCH_GROUP = "gemini-batch-group"
_BATCH_DEPLOYMENT_MODEL = "openai/gpt-4o-mini"
_BATCH_API_BASE = "http://localhost:4001/v1"
_BATCH_ID = "batch-1"
_BATCH_ROWS = 2
_BATCH_TOKENS_PER_ROW = 600
_BATCH_COMPLETED = {
"id": _BATCH_ID,
"object": "batch",
"endpoint": "/v1/chat/completions",
"errors": None,
"input_file_id": "file-in-1",
"completion_window": "24h",
"status": "completed",
"output_file_id": "file-out-1",
"error_file_id": None,
"created_at": 0,
"completed_at": 1,
"request_counts": {"total": _BATCH_ROWS, "completed": _BATCH_ROWS, "failed": 0},
"metadata": None,
}
_BATCH_OUTPUT_JSONL = "\n".join(
json.dumps(
{
"id": f"req-{row}",
"custom_id": f"row-{row}",
"response": {
"status_code": 200,
"body": {
"id": f"chatcmpl-{row}",
"object": "chat.completion",
"model": "gpt-4o-mini",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}
],
"usage": {
"prompt_tokens": 500,
"completion_tokens": 100,
"total_tokens": _BATCH_TOKENS_PER_ROW,
},
},
},
}
)
for row in range(_BATCH_ROWS)
)
class _BatchPayloadCollector(CustomLogger):
def __init__(self):
super().__init__()
self.payloads = []
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
self.payloads.append(kwargs.get("standard_logging_object"))
async def retrieve_batch_payload(self):
for _ in range(100):
for payload in self.payloads:
if payload and payload.get("call_type") == "aretrieve_batch":
return payload
await asyncio.sleep(0.05)
raise AssertionError(f"no aretrieve_batch payload was emitted: {self.payloads}")
def _batch_model_group_router():
return litellm.Router(
model_list=[
{
"model_name": _BATCH_GROUP,
"litellm_params": {
"model": _BATCH_DEPLOYMENT_MODEL,
"api_base": _BATCH_API_BASE,
"api_key": "sk-fake",
},
}
]
)
def _mock_batch_provider(respx_mock):
respx_mock.get(f"{_BATCH_API_BASE}/batches/{_BATCH_ID}").mock(
return_value=httpx.Response(200, json=_BATCH_COMPLETED)
)
respx_mock.get(f"{_BATCH_API_BASE}/files/file-out-1/content").mock(
return_value=httpx.Response(200, text=_BATCH_OUTPUT_JSONL)
)
@pytest.mark.asyncio
async def test_arouter_aretrieve_batch_without_model_stamps_model_group(monkeypatch: pytest.MonkeyPatch):
"""
The proxy retrieves a managed batch by id only - no `model` in the request.
The router fans out over its deployments, so the model group is only known
from the deployment that answered.
"""
import respx
collector = _BatchPayloadCollector()
monkeypatch.setattr(litellm, "callbacks", [collector])
router = _batch_model_group_router()
with respx.mock(assert_all_called=True) as respx_mock:
_mock_batch_provider(respx_mock)
response = await router.aretrieve_batch(batch_id=_BATCH_ID)
payload = await collector.retrieve_batch_payload()
assert response.id == _BATCH_ID
assert payload["total_tokens"] == _BATCH_ROWS * _BATCH_TOKENS_PER_ROW
assert payload["model"] == _BATCH_DEPLOYMENT_MODEL
assert payload["model_group"] == _BATCH_GROUP
@pytest.mark.asyncio
async def test_arouter_aretrieve_batch_with_model_stamps_requested_model_group(monkeypatch: pytest.MonkeyPatch):
"""An explicitly requested model group is what gets logged."""
import respx
collector = _BatchPayloadCollector()
monkeypatch.setattr(litellm, "callbacks", [collector])
router = _batch_model_group_router()
with respx.mock(assert_all_called=True) as respx_mock:
_mock_batch_provider(respx_mock)
await router.aretrieve_batch(model=_BATCH_GROUP, batch_id=_BATCH_ID)
payload = await collector.retrieve_batch_payload()
assert payload["model_group"] == _BATCH_GROUP
_UNRELATED_BATCH_GROUP = "unrelated-batch-group"
_UNRELATED_BATCH_API_BASE = "http://localhost:4002/v1"
_BATCH_NOT_FOUND = {
"error": {
"message": f"No batch found with id '{_BATCH_ID}'.",
"type": "invalid_request_error",
"code": "batch_not_found",
}
}
async def _router_usage_keys(router, timeout: float = 2.0) -> list[str]:
loop = asyncio.get_event_loop()
deadline = loop.time() + timeout
while loop.time() < deadline:
keys = sorted(k for k in router.cache.in_memory_cache.cache_dict if k.startswith("global_router:"))
if keys:
return keys
await asyncio.sleep(0.05)
return []
@pytest.mark.asyncio
async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(monkeypatch: pytest.MonkeyPatch):
"""
A batch reports the whole job's tokens on retrieve, and reports them again on every
poll of the finished batch, so they are not a measure of load in the current minute.
The fan-out also probes deployments the caller never named. Neither may reach the
per-minute tpm/rpm counters that gate live traffic.
"""
import respx
collector = _BatchPayloadCollector()
monkeypatch.setattr(litellm, "callbacks", [collector])
router = litellm.Router(
model_list=[
{
"model_name": _BATCH_GROUP,
"litellm_params": {
"model": _BATCH_DEPLOYMENT_MODEL,
"api_base": _BATCH_API_BASE,
"api_key": "sk-fake",
},
"model_info": {"id": "batch-dep"},
"tpm": 1000,
"rpm": 10,
},
{
"model_name": _UNRELATED_BATCH_GROUP,
"litellm_params": {
"model": _BATCH_DEPLOYMENT_MODEL,
"api_base": _UNRELATED_BATCH_API_BASE,
"api_key": "sk-fake",
},
"model_info": {"id": "unrelated-dep"},
"tpm": 1000,
"rpm": 10,
},
]
)
with respx.mock(assert_all_called=True) as respx_mock:
_mock_batch_provider(respx_mock)
respx_mock.get(f"{_UNRELATED_BATCH_API_BASE}/batches/{_BATCH_ID}").mock(
return_value=httpx.Response(404, json=_BATCH_NOT_FOUND)
)
response = await router.aretrieve_batch(batch_id=_BATCH_ID)
payload = await collector.retrieve_batch_payload()
usage_keys = await _router_usage_keys(router)
assert response.id == _BATCH_ID
assert payload["model_group"] == _BATCH_GROUP
assert usage_keys == []
@pytest.mark.parametrize(
("call_type", "expected_key", "expected_successes"),
[
("aretrieve_batch", None, 0),
("retrieve_batch", None, 0),
("acompletion", "batch-dep:successes", 1),
],
)
def test_sync_deployment_callback_on_success_skips_batch_retrieves(
call_type: str, expected_key: str | None, expected_successes: int
):
router = litellm.Router(
model_list=[
{
"model_name": _BATCH_GROUP,
"litellm_params": {"model": _BATCH_DEPLOYMENT_MODEL, "api_base": _BATCH_API_BASE, "api_key": "sk-fake"},
"model_info": {"id": "batch-dep"},
}
]
)
key = router.sync_deployment_callback_on_success(
kwargs={
"call_type": call_type,
"litellm_params": {"metadata": {"model_group": _BATCH_GROUP}, "model_info": {"id": "batch-dep"}},
},
completion_response=None,
start_time=datetime.now(),
end_time=datetime.now(),
)
assert key == expected_key
assert (
get_deployment_successes_for_current_minute(litellm_router_instance=router, deployment_id="batch-dep")
== expected_successes
)
_ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:")
async def _moved_routing_counters(router, timeout: float = 2.0) -> list[str]:
loop = asyncio.get_event_loop()
deadline = loop.time() + timeout
while loop.time() < deadline:
cache_dict = router.cache.in_memory_cache.cache_dict
moved = sorted(
f"{key}={cache_dict[key]}"
for key in cache_dict
if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS)
and cache_dict[key]
)
if moved:
return moved
await asyncio.sleep(0.05)
return []
def _batch_fan_out_router(routing_strategy: str):
return litellm.Router(
routing_strategy=routing_strategy,
model_list=[
{
"model_name": _BATCH_GROUP,
"litellm_params": {
"model": _BATCH_DEPLOYMENT_MODEL,
"api_base": _BATCH_API_BASE,
"api_key": "sk-fake",
},
"model_info": {"id": "batch-dep"},
},
{
"model_name": _UNRELATED_BATCH_GROUP,
"litellm_params": {
"model": _BATCH_DEPLOYMENT_MODEL,
"api_base": _UNRELATED_BATCH_API_BASE,
"api_key": "sk-fake",
},
"model_info": {"id": "unrelated-dep"},
},
],
)
@pytest.mark.parametrize(
"routing_strategy",
[
"usage-based-routing",
"usage-based-routing-v2",
"latency-based-routing",
"cost-based-routing",
"least-busy",
],
)
@pytest.mark.asyncio
async def test_arouter_aretrieve_batch_does_not_feed_routing_strategies(
monkeypatch: pytest.MonkeyPatch, routing_strategy: str
):
"""
Every routing strategy picks a deployment from what recent live traffic did.
A batch retrieve reports the whole job on every poll and probes deployments the
caller never named, so polling a finished batch must not move the numbers that
decide where the next chat request goes.
"""
import respx
collector = _BatchPayloadCollector()
monkeypatch.setattr(litellm, "callbacks", [collector])
monkeypatch.setattr(litellm, "input_callback", [])
router = _batch_fan_out_router(routing_strategy)
with respx.mock(assert_all_called=True) as respx_mock:
_mock_batch_provider(respx_mock)
respx_mock.get(f"{_UNRELATED_BATCH_API_BASE}/batches/{_BATCH_ID}").mock(
return_value=httpx.Response(404, json=_BATCH_NOT_FOUND)
)
for _ in range(3):
response = await router.aretrieve_batch(batch_id=_BATCH_ID)
await collector.retrieve_batch_payload()
moved_counters = await _moved_routing_counters(router)
assert response.id == _BATCH_ID
assert moved_counters == []
@pytest.mark.asyncio
async def test_arouter_aretrieve_file_content():
"""