From 6ad23ff70ddc2b757cad9ae6677b08a797efbb0a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:09:08 -0700 Subject: [PATCH] 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. --- litellm/router.py | 11 ++- litellm/router_strategy/least_busy.py | 49 +++++++--- .../test_least_busy_concurrent_decisions.py | 90 +++++++++++++++++++ tests/unit/router_strategy/test_least_busy.py | 53 ++++++++++- .../test_router_routing_groups.py | 31 +++++++ .../test_router/test_router_helper_utils.py | 28 ++++++ 6 files changed, 243 insertions(+), 19 deletions(-) create mode 100644 tests/integration/routing/test_least_busy_concurrent_decisions.py diff --git a/litellm/router.py b/litellm/router.py index 43a5a6fddb7..62024bc9e74 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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: diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 5d9a451770c..29fb99ff63d 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -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 diff --git a/tests/integration/routing/test_least_busy_concurrent_decisions.py b/tests/integration/routing/test_least_busy_concurrent_decisions.py new file mode 100644 index 00000000000..bb923534db4 --- /dev/null +++ b/tests/integration/routing/test_least_busy_concurrent_decisions.py @@ -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} + ) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 6a95e825e6c..46964bdb9f8 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -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 diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 5c0df8e4ff4..c43662a84e2 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -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" diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 513136ae150..f9eff9b5ae9 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -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(