mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(router): count least-busy decisions calls and release reservations on cache hits
Least-busy now reads the model group from litellm_metadata before metadata, so /v1/decisions calls that carry caller metadata are counted in flight and release their reservation at pre-call. A response cache hit releases its reservation instead of decrementing another request's in-flight count, and Responses WebSocket turn admission selects without reserving. The reservation store is reached through a typed protocol.
This commit is contained in:
mateo-berri
parent
d499a032fc
commit
6ad23ff70d
6 changed files with 243 additions and 19 deletions
|
|
@ -2431,6 +2431,7 @@ class Router:
|
|||
messages: list[dict[str, str]] | None,
|
||||
input: str | list | None,
|
||||
request_kwargs: dict | None,
|
||||
reserve: bool,
|
||||
) -> Any | None:
|
||||
"""
|
||||
Asks the strategy selector for a deployment. Caller handles
|
||||
|
|
@ -2444,7 +2445,7 @@ class Router:
|
|||
return await selector.async_get_available_deployments(
|
||||
model_group=model,
|
||||
healthy_deployments=healthy_deployments,
|
||||
reserve=request_kwargs not in (None, {}),
|
||||
reserve=reserve,
|
||||
)
|
||||
case "usage-based-routing":
|
||||
# `LowestTPMLoggingHandler` (v1) only exposes the sync
|
||||
|
|
@ -2485,6 +2486,7 @@ class Router:
|
|||
messages: list[dict[str, str]] | None,
|
||||
input: str | list | None,
|
||||
request_kwargs: dict | None,
|
||||
reserve: bool,
|
||||
) -> Any | None:
|
||||
"""
|
||||
Sync sibling of `_select_deployment_async`. Caller handles
|
||||
|
|
@ -2501,7 +2503,7 @@ class Router:
|
|||
return selector.get_available_deployments(
|
||||
model_group=model,
|
||||
healthy_deployments=healthy_deployments,
|
||||
reserve=request_kwargs not in (None, {}),
|
||||
reserve=reserve,
|
||||
)
|
||||
case "usage-based-routing" | "usage-based-routing-v2":
|
||||
return selector.get_available_deployments(
|
||||
|
|
@ -10257,6 +10259,7 @@ class Router:
|
|||
messages=None,
|
||||
input=None,
|
||||
request_kwargs=routing_kwargs,
|
||||
reserve=False,
|
||||
)
|
||||
if selected is None:
|
||||
no_deployment: Final = await async_raise_no_deployment_exception(
|
||||
|
|
@ -15550,6 +15553,7 @@ class Router:
|
|||
messages=messages,
|
||||
input=input,
|
||||
request_kwargs=request_kwargs,
|
||||
reserve=request_kwargs not in (None, {}),
|
||||
)
|
||||
if deployment is None:
|
||||
exception: Final = await async_raise_no_deployment_exception(
|
||||
|
|
@ -15719,6 +15723,7 @@ class Router:
|
|||
messages=messages,
|
||||
input=input,
|
||||
request_kwargs=request_kwargs,
|
||||
reserve=request_kwargs not in (None, {}),
|
||||
)
|
||||
|
||||
if deployment is None:
|
||||
|
|
@ -16581,6 +16586,7 @@ class Router:
|
|||
messages=messages,
|
||||
input=input,
|
||||
request_kwargs=request_kwargs,
|
||||
reserve=request_kwargs not in (None, {}),
|
||||
)
|
||||
|
||||
if deployment is None:
|
||||
|
|
@ -16763,6 +16769,7 @@ class Router:
|
|||
messages=messages,
|
||||
input=input,
|
||||
request_kwargs=request_kwargs,
|
||||
reserve=request_kwargs not in (None, {}),
|
||||
)
|
||||
|
||||
if deployment is None:
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import logging
|
||||
import time
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from typing import Final
|
||||
from typing import Final, Protocol
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
|
@ -27,6 +27,7 @@ class _Metadata(TypedDict, total=False):
|
|||
|
||||
class _LitellmParams(TypedDict, total=False):
|
||||
metadata: ReadOnly[_Metadata | None]
|
||||
litellm_metadata: ReadOnly[_Metadata | None]
|
||||
model_info: ReadOnly[_ModelInfo | None]
|
||||
|
||||
|
||||
|
|
@ -44,14 +45,18 @@ class _Deployment(TypedDict):
|
|||
|
||||
_CALL_KWARGS: Final = TypeAdapter(_CallKwargs)
|
||||
_DEPLOYMENTS: Final = TypeAdapter(list[_Deployment])
|
||||
_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None)
|
||||
_RESERVATION_TIMES: Final = TypeAdapter(tuple[float, ...] | None)
|
||||
_MEMORY_COUNTS: Final[TypeAdapter[tuple[float | None, ...] | None]] = TypeAdapter(tuple[float | None, ...] | None)
|
||||
_RESERVATION_TIMES: Final[TypeAdapter[tuple[float, ...] | None]] = TypeAdapter(tuple[float, ...] | None)
|
||||
|
||||
|
||||
def _request_count_key(model_group: str, deployment_id: str) -> str:
|
||||
return f"{model_group}_request_count:{deployment_id}"
|
||||
|
||||
|
||||
def _model_group(metadata: _Metadata | None) -> str | None:
|
||||
return metadata.get("model_group") if metadata else None
|
||||
|
||||
|
||||
def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None:
|
||||
if is_batch_retrieve_call_type(kwargs.get("call_type")):
|
||||
return None
|
||||
|
|
@ -59,10 +64,11 @@ def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None:
|
|||
call: Final = _CALL_KWARGS.validate_python(kwargs)
|
||||
except ValidationError:
|
||||
return None
|
||||
litellm_params: Final = call.get("litellm_params")
|
||||
metadata: Final = litellm_params.get("metadata") if litellm_params else None
|
||||
model_info: Final = litellm_params.get("model_info") if litellm_params else None
|
||||
model_group: Final = metadata.get("model_group") if metadata else None
|
||||
litellm_params: Final = call.get("litellm_params") or _LitellmParams()
|
||||
model_info: Final = litellm_params.get("model_info")
|
||||
model_group: Final = _model_group(litellm_params.get("litellm_metadata")) or _model_group(
|
||||
litellm_params.get("metadata")
|
||||
)
|
||||
deployment_id: Final = model_info.get("id") if model_info else None
|
||||
if model_group is None or deployment_id is None:
|
||||
return None
|
||||
|
|
@ -118,6 +124,14 @@ def _warn_unwritable(key: str, error: Exception) -> None:
|
|||
)
|
||||
|
||||
|
||||
class _ReservationStore(Protocol):
|
||||
def get_cache(self, key: str) -> object: ...
|
||||
|
||||
def set_cache(self, key: str, value: tuple[float, ...], *, ttl: float) -> None: ...
|
||||
|
||||
def delete_cache(self, key: str) -> None: ...
|
||||
|
||||
|
||||
class LeastBusyLoggingHandler(CustomLogger):
|
||||
test_flag: bool = False
|
||||
logged_success: int = 0
|
||||
|
|
@ -128,16 +142,19 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
self.router_cache_id = str(id(router_cache))
|
||||
self._clock = clock
|
||||
|
||||
def _reservation_store(self) -> _ReservationStore:
|
||||
return self.router_cache.in_memory_cache
|
||||
|
||||
def _live_reservations(self, key: str, now: float) -> tuple[float, ...]:
|
||||
raw: Final = self.router_cache.get_cache(_reservation_key(key), local_only=True)
|
||||
raw: Final = self._reservation_store().get_cache(_reservation_key(key))
|
||||
return _unexpired(_RESERVATION_TIMES.validate_python(raw) or (), now)
|
||||
|
||||
def _store_reservations(self, key: str, times: tuple[float, ...]) -> None:
|
||||
# InMemoryCache keeps a live key's first TTL, so drop the key to restart it
|
||||
local: Final = self.router_cache.in_memory_cache
|
||||
local.delete_cache(_reservation_key(key))
|
||||
store: Final = self._reservation_store()
|
||||
store.delete_cache(_reservation_key(key))
|
||||
if times:
|
||||
local.set_cache(_reservation_key(key), times, ttl=SELECTION_RESERVATION_TTL_SECONDS)
|
||||
store.set_cache(_reservation_key(key), times, ttl=SELECTION_RESERVATION_TTL_SECONDS)
|
||||
|
||||
def _select(
|
||||
self,
|
||||
|
|
@ -172,7 +189,10 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
def log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
self._increment(kwargs, -1)
|
||||
if kwargs.get("cache_hit") is True:
|
||||
self._release_reservation(kwargs)
|
||||
else:
|
||||
self._increment(kwargs, -1)
|
||||
if self.test_flag:
|
||||
self.logged_success += 1
|
||||
|
||||
|
|
@ -188,7 +208,10 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
async def async_log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
await self._async_increment(kwargs, -1)
|
||||
if kwargs.get("cache_hit") is True:
|
||||
self._release_reservation(kwargs)
|
||||
else:
|
||||
await self._async_increment(kwargs, -1)
|
||||
if self.test_flag:
|
||||
self.logged_success += 1
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,90 @@
|
|||
import json
|
||||
import threading
|
||||
from collections import Counter
|
||||
from collections.abc import Callable
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, gateway_from_environment
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
pytestmark = pytest.mark.timeout(240)
|
||||
|
||||
_GROUP: Final = "least-busy-decider"
|
||||
_DEPLOYMENTS: Final = ("least-busy-decider-a", "least-busy-decider-b")
|
||||
_BURST: Final = 4
|
||||
_REPLY: Final = json.dumps(
|
||||
{
|
||||
"model": "Qwen/Qwen3-0.6B",
|
||||
"answers": {"billing": {"type": "noul", "noul": 0.91}},
|
||||
"usage": {"input_tokens": 367, "output_tokens": 3},
|
||||
}
|
||||
).encode()
|
||||
_DECISION: Final = {
|
||||
"model": _GROUP,
|
||||
"input": "I was charged twice for order 4417.",
|
||||
"questions": [{"type": "predicate", "name": "billing", "instructions": "Is this a billing issue?"}],
|
||||
}
|
||||
|
||||
|
||||
def _held_until_the_burst_arrives(arrived: threading.Barrier) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
arrived.wait()
|
||||
return Reply(body=_REPLY)
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _write_config(path: Path, provider_url: str) -> Path:
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"model_list": [
|
||||
{
|
||||
"model_name": _GROUP,
|
||||
"litellm_params": {"model": "hosted_vllm/Qwen/Qwen3-0.6B", "api_base": provider_url},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
for deployment_id in _DEPLOYMENTS
|
||||
],
|
||||
"general_settings": {
|
||||
"master_key": "os.environ/LITELLM_MASTER_KEY",
|
||||
"database_url": "os.environ/DATABASE_URL",
|
||||
"disable_model_info_refresh": True,
|
||||
},
|
||||
"router_settings": {"routing_strategy": "least-busy", "num_retries": 0, "disable_cooldowns": True},
|
||||
}
|
||||
)
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
def _decide(proxy: Gateway) -> httpx.Response:
|
||||
return proxy.client.post(
|
||||
"/v1/decisions", json=_DECISION, headers={"Authorization": f"Bearer {proxy.key}"}, timeout=120
|
||||
)
|
||||
|
||||
|
||||
def test_least_busy_spreads_a_burst_of_decisions_calls_across_the_group(tmp_path: Path) -> None:
|
||||
arrived: Final = threading.Barrier(_BURST, timeout=60)
|
||||
with (
|
||||
wire_server(_held_until_the_burst_arrives(arrived)) as upstream,
|
||||
gateway_from_environment() as rig,
|
||||
owned_proxy(
|
||||
rig,
|
||||
tmp_path,
|
||||
{"STORE_MODEL_IN_DB": "False"},
|
||||
config=_write_config(tmp_path / "least-busy.yaml", upstream.url),
|
||||
) as proxy,
|
||||
ThreadPoolExecutor(max_workers=_BURST) as pool,
|
||||
):
|
||||
responses: Final = tuple(pool.map(lambda _: _decide(proxy), range(_BURST)))
|
||||
assert [response.status_code for response in responses] == [200] * _BURST, [r.text for r in responses]
|
||||
assert Counter(response.headers["x-litellm-model-id"] for response in responses) == Counter(
|
||||
{deployment_id: _BURST // len(_DEPLOYMENTS) for deployment_id in _DEPLOYMENTS}
|
||||
)
|
||||
|
|
@ -102,6 +102,31 @@ def test_sync_pick_reads_the_shared_counts() -> None:
|
|||
assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
|
||||
|
||||
|
||||
def _routed_through_litellm_metadata(deployment_id: str) -> dict[str, object]:
|
||||
return {
|
||||
"litellm_params": {
|
||||
"metadata": {"tags": ["caller-tag"]},
|
||||
"litellm_metadata": {"model_group": GROUP},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def test_a_request_routed_through_litellm_metadata_counts_while_in_flight() -> None:
|
||||
shared: Final = SharedRedisCounters()
|
||||
calling_worker: Final = _worker(shared)
|
||||
picking_worker: Final = _worker(shared)
|
||||
|
||||
assert calling_worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
|
||||
calling_worker.log_pre_api_call(model="m", messages=[], kwargs=_routed_through_litellm_metadata("dep-a"))
|
||||
|
||||
assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
|
||||
|
||||
calling_worker.log_success_event(_routed_through_litellm_metadata("dep-a"), None, None, None)
|
||||
|
||||
assert calling_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
|
||||
|
||||
|
||||
def test_the_handler_never_pushes_a_counters_ttl_forward() -> None:
|
||||
"""A worker that dies mid-request leaves a +1 nobody will ever decrement. Redis expires that
|
||||
stuck count an hour after the key was created, which only works while nothing writes the TTL
|
||||
|
|
@ -211,7 +236,9 @@ class OpenBreakerRedis(SharedRedisCounters):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request(caplog: pytest.LogCaptureFixture) -> None:
|
||||
async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request(
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
worker: Final = _worker(OpenBreakerRedis())
|
||||
|
||||
with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"):
|
||||
|
|
@ -383,9 +410,7 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts
|
|||
if use_async
|
||||
else router.get_available_deployment(model=ROUTER_GROUP, messages=None)
|
||||
)
|
||||
response: Final = await router.acompletion(
|
||||
model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]
|
||||
)
|
||||
response: Final = await router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}])
|
||||
for _ in range(10):
|
||||
await asyncio.sleep(0)
|
||||
await GLOBAL_LOGGING_WORKER.flush()
|
||||
|
|
@ -488,3 +513,23 @@ def test_a_request_started_through_another_selector_on_the_same_cache_releases_t
|
|||
|
||||
assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
|
||||
assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_cache_hit_releases_the_pick_it_reserved() -> None:
|
||||
worker: Final = _worker(None)
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
|
||||
|
||||
assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
|
||||
await worker.async_log_success_event({**_call_kwargs("dep-b"), "cache_hit": True}, None, None, None)
|
||||
|
||||
assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
|
||||
|
||||
|
||||
def test_a_cache_hit_leaves_requests_in_flight_on_its_deployment_counted() -> None:
|
||||
worker: Final = _worker(None)
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
|
||||
|
||||
worker.log_success_event({**_call_kwargs("dep-a"), "cache_hit": True}, None, None, None)
|
||||
|
||||
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
|
||||
|
|
|
|||
|
|
@ -481,6 +481,33 @@ def test_two_least_busy_groups_count_a_request_once(monkeypatch):
|
|||
assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0
|
||||
|
||||
|
||||
|
||||
def _filtered_model_pick(router: Router) -> str:
|
||||
deployment: Final = router.get_available_deployment(
|
||||
"filtered-model", request_kwargs=_proxy_shaped_request(model="filtered-model")
|
||||
)
|
||||
return deployment["model_info"]["id"]
|
||||
|
||||
|
||||
def test_a_least_busy_group_pick_stops_counting_once_its_request_finishes(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
monkeypatch.setattr(litellm, "input_callback", [])
|
||||
router: Final = _build_router(
|
||||
routing_strategy="least-busy",
|
||||
routing_groups=[{"group_name": "fast", "models": ["filtered-model"], "routing_strategy": "least-busy"}],
|
||||
)
|
||||
|
||||
first: Final = _filtered_model_pick(router)
|
||||
kwargs: Final = {"litellm_params": {"metadata": {"model_group": "filtered-model"}, "model_info": {"id": first}}}
|
||||
for callback in litellm.input_callback:
|
||||
if isinstance(callback, CustomLogger):
|
||||
callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs)
|
||||
for callback in litellm.callbacks:
|
||||
if isinstance(callback, CustomLogger):
|
||||
callback.log_success_event(kwargs, None, None, None)
|
||||
|
||||
assert _filtered_model_pick(router) == first
|
||||
|
||||
def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch):
|
||||
"""
|
||||
Least-busy hangs its counting off litellm's global callback lists, and those lists keep one
|
||||
|
|
@ -613,6 +640,7 @@ async def test_select_deployment_async_returns_none_without_selector():
|
|||
messages=None,
|
||||
input=None,
|
||||
request_kwargs=None,
|
||||
reserve=False,
|
||||
)
|
||||
assert result is None
|
||||
|
||||
|
|
@ -627,6 +655,7 @@ def test_select_deployment_sync_returns_none_without_selector():
|
|||
messages=None,
|
||||
input=None,
|
||||
request_kwargs=None,
|
||||
reserve=False,
|
||||
)
|
||||
assert result is None
|
||||
|
||||
|
|
@ -661,6 +690,7 @@ def test_select_deployment_sync_does_not_dispatch_cost_based_routing():
|
|||
messages=None,
|
||||
input=None,
|
||||
request_kwargs=None,
|
||||
reserve=False,
|
||||
)
|
||||
assert result is None
|
||||
|
||||
|
|
@ -707,6 +737,7 @@ async def test_async_dispatch_falls_back_to_sync_for_usage_based_routing_v1():
|
|||
messages=[{"role": "user", "content": "hi"}],
|
||||
input=None,
|
||||
request_kwargs={},
|
||||
reserve=False,
|
||||
)
|
||||
|
||||
assert v1_spy.called, "async dispatch must route v1 strategy through sync method"
|
||||
|
|
|
|||
|
|
@ -597,6 +597,34 @@ async def test_responses_websocket_turn_is_refused_once_its_deployment_crosses_b
|
|||
await admit_turn()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_websocket_turns_do_not_make_their_deployment_look_busy_to_least_busy():
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "ws-least-busy",
|
||||
"litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake"},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
for deployment_id in ("ws-least-busy-a", "ws-least-busy-b")
|
||||
],
|
||||
routing_strategy="least-busy",
|
||||
)
|
||||
admit_turn: Final = (await _kwargs_routed_to(router, "ws-least-busy", "_aresponses_websocket"))["turn_admission"]
|
||||
assert callable(admit_turn)
|
||||
for _ in range(3):
|
||||
await admit_turn()
|
||||
|
||||
picks: Final = [
|
||||
(await router.async_get_available_deployment(model="ws-least-busy", request_kwargs={"input": "hi"}))[
|
||||
"model_info"
|
||||
]["id"]
|
||||
for _ in range(2)
|
||||
]
|
||||
|
||||
assert sorted(picks) == ["ws-least-busy-a", "ws-least-busy-b"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_websocket_turn_after_the_first_counts_toward_deployment_rpm():
|
||||
router = Router(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue