mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(router): keep batch retrieves out of the per-minute tpm/rpm counters
Stamping model_group let both router deployment callbacks past their `model_group is None` early return for batch retrieves. A batch reports the whole job's token total on retrieve and reports it again on every poll of the finished batch, so those tokens are not load in the current minute: three polls of one completed 1,200 token batch pushed a tpm:1000 deployment to 3,600. The fan-out also probed unrelated deployments, adding an rpm tick to each.
This commit is contained in:
parent
01bdfb34aa
commit
9acf09f60d
3 changed files with 99 additions and 0 deletions
|
|
@ -124,6 +124,7 @@ from litellm.router_utils.auto_router_model_naming import (
|
||||||
)
|
)
|
||||||
from litellm.router_utils.batch_utils import (
|
from litellm.router_utils.batch_utils import (
|
||||||
_get_router_metadata_variable_name,
|
_get_router_metadata_variable_name,
|
||||||
|
is_batch_retrieve_call_type,
|
||||||
replace_model_in_jsonl,
|
replace_model_in_jsonl,
|
||||||
should_replace_model_in_jsonl,
|
should_replace_model_in_jsonl,
|
||||||
)
|
)
|
||||||
|
|
@ -7878,6 +7879,8 @@ class Router:
|
||||||
# WS session wrappers fire with result=None; per-turn costs tracked by inner calls.
|
# WS session wrappers fire with result=None; per-turn costs tracked by inner calls.
|
||||||
if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"):
|
if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"):
|
||||||
return
|
return
|
||||||
|
if is_batch_retrieve_call_type(kwargs.get("call_type")):
|
||||||
|
return
|
||||||
standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
|
standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
|
||||||
if standard_logging_object is None:
|
if standard_logging_object is None:
|
||||||
raise ValueError("standard_logging_object is None")
|
raise ValueError("standard_logging_object is None")
|
||||||
|
|
@ -8117,6 +8120,8 @@ class Router:
|
||||||
"""
|
"""
|
||||||
Update RPM usage for a deployment
|
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_name: Final = kwargs["litellm_params"]["metadata"].get(
|
||||||
"deployment", None
|
"deployment", None
|
||||||
) # handles wildcard routes - by giving the original name sent to `litellm.completion`
|
) # handles wildcard routes - by giving the original name sent to `litellm.completion`
|
||||||
|
|
|
||||||
|
|
@ -5,6 +5,7 @@ from typing import Final
|
||||||
|
|
||||||
from litellm._logging import verbose_logger
|
from litellm._logging import verbose_logger
|
||||||
from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose
|
from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose
|
||||||
|
from litellm.types.utils import CallTypes
|
||||||
|
|
||||||
|
|
||||||
class InMemoryFile(io.BytesIO):
|
class InMemoryFile(io.BytesIO):
|
||||||
|
|
@ -170,3 +171,20 @@ def _get_router_metadata_variable_name(function_name: str | None) -> str:
|
||||||
return "litellm_metadata"
|
return "litellm_metadata"
|
||||||
else:
|
else:
|
||||||
return "metadata"
|
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. Per-minute usage counters must not be fed from it.
|
||||||
|
"""
|
||||||
|
return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES
|
||||||
|
|
|
||||||
|
|
@ -1090,6 +1090,82 @@ async def test_arouter_aretrieve_batch_with_model_stamps_requested_model_group(m
|
||||||
assert payload["model_group"] == _BATCH_GROUP
|
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.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_arouter_aretrieve_file_content():
|
async def test_arouter_aretrieve_file_content():
|
||||||
"""
|
"""
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue