mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
fix(least-busy): free tier-group picks, stop lookups reserving, and count only started attempts
Reservations are keyed by deployment id, so a complexity, auto, adaptive or quality router frees its tier pick when the call starts. Each started attempt is recorded under its call id, and only a recorded start lowers the count when the attempt ends, so a fallback hop rejected before the provider frees its pick. Whether a pick reserves now comes from a context variable set by the routed call sites, so get_available_deployment and the deprecated async_get_available_deployment no longer hold a deployment, and custom routing strategies and Router subclasses keep their method signatures.
This commit is contained in:
mateo-berri
parent
15f8086684
commit
2c4bae982b
7 changed files with 238 additions and 120 deletions
|
|
@ -219,7 +219,12 @@ from litellm.router_strategy.complexity_router.context_compaction import (
|
|||
reject_recursive_compactor,
|
||||
surface_for_call,
|
||||
)
|
||||
from litellm.router_strategy.least_busy import SELECTION_RESERVATION_MAX_KEYS, LeastBusyLoggingHandler
|
||||
from litellm.router_strategy.least_busy import (
|
||||
SELECTION_RESERVATION_MAX_KEYS,
|
||||
LeastBusyLoggingHandler,
|
||||
pick_reserves,
|
||||
picks_reserve,
|
||||
)
|
||||
from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler
|
||||
from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler
|
||||
from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler
|
||||
|
|
@ -1381,14 +1386,14 @@ class Router(DeprecatedRouterMembers):
|
|||
(``pass_through``).
|
||||
"""
|
||||
if kind == "pick":
|
||||
return await self._core_async_get_available_deployment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts
|
||||
model=model,
|
||||
request_kwargs=request_kwargs if request_kwargs is not None else {},
|
||||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
reserve=False,
|
||||
)
|
||||
with picks_reserve(False):
|
||||
return await self._core_async_get_available_deployment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts
|
||||
model=model,
|
||||
request_kwargs=request_kwargs if request_kwargs is not None else {},
|
||||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
)
|
||||
if kind == "healthy":
|
||||
return await self._core_async_get_healthy_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts
|
||||
model=model,
|
||||
|
|
@ -3404,12 +3409,13 @@ class Router(DeprecatedRouterMembers):
|
|||
)
|
||||
|
||||
# pick the one that is available (lowest TPM/RPM)
|
||||
deployment = self.get_available_deployment(
|
||||
model=model,
|
||||
messages=messages,
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
request_kwargs=kwargs,
|
||||
)
|
||||
with picks_reserve(True):
|
||||
deployment = self.get_available_deployment(
|
||||
model=model,
|
||||
messages=messages,
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
request_kwargs=kwargs,
|
||||
)
|
||||
self._drop_unsupported_classifier_reasoning_effort(
|
||||
deployment=cast(DeploymentTypedDict, deployment), # cast-ok: selection returns a router deployment
|
||||
model=model,
|
||||
|
|
@ -5884,12 +5890,13 @@ class Router(DeprecatedRouterMembers):
|
|||
model_name: Final = ""
|
||||
try:
|
||||
verbose_router_logger.debug("Inside _image_generation()- model: %s", model)
|
||||
deployment: Final = self.get_available_deployment(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "prompt"}],
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
request_kwargs=kwargs,
|
||||
)
|
||||
with picks_reserve(True):
|
||||
deployment: Final = self.get_available_deployment(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "prompt"}],
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
request_kwargs=kwargs,
|
||||
)
|
||||
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
|
||||
data: Final = deployment["litellm_params"].copy()
|
||||
|
||||
|
|
@ -6237,7 +6244,7 @@ class Router(DeprecatedRouterMembers):
|
|||
kwargs.setdefault("metadata", {}).update({"model_group": model})
|
||||
|
||||
# pick the one that is available (lowest TPM/RPM)
|
||||
with endpoint_selection("text_completion", model=model):
|
||||
with endpoint_selection("text_completion", model=model), picks_reserve(True):
|
||||
deployment: Final = self.get_available_deployment(
|
||||
model=model,
|
||||
messages=messages,
|
||||
|
|
@ -7378,7 +7385,11 @@ class Router(DeprecatedRouterMembers):
|
|||
kwargs=kwargs,
|
||||
metadata_variable_name=metadata_variable_name,
|
||||
)
|
||||
with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name):
|
||||
with (
|
||||
self._video_call_pin(handler_name, kwargs),
|
||||
self._route_capability_pin(handler_name),
|
||||
picks_reserve(True),
|
||||
):
|
||||
deployment: Final = self.get_available_deployment(
|
||||
model=model,
|
||||
messages=kwargs.get("messages", None),
|
||||
|
|
@ -7448,7 +7459,7 @@ class Router(DeprecatedRouterMembers):
|
|||
model_name = None
|
||||
try:
|
||||
verbose_router_logger.debug("Inside embedding()- model: %s", model)
|
||||
with embedding_deployment_selection():
|
||||
with embedding_deployment_selection(), picks_reserve(True):
|
||||
deployment: Final = self.get_available_deployment(
|
||||
model=model,
|
||||
input=input,
|
||||
|
|
@ -8059,13 +8070,13 @@ class Router(DeprecatedRouterMembers):
|
|||
try:
|
||||
verbose_router_logger.debug("Inside _acancel_batch()- model: %s", model)
|
||||
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
|
||||
deployment: Final = await self._core_async_get_available_deployment(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "batch-api-fake-text"}],
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
request_kwargs=kwargs,
|
||||
reserve=False,
|
||||
)
|
||||
with picks_reserve(False):
|
||||
deployment: Final = await self._core_async_get_available_deployment(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "batch-api-fake-text"}],
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
request_kwargs=kwargs,
|
||||
)
|
||||
|
||||
selected_deployment_id: Final = (deployment.get("model_info") or {}).get("id")
|
||||
data: Final = deployment["litellm_params"].copy()
|
||||
|
|
@ -15697,7 +15708,6 @@ class Router(DeprecatedRouterMembers):
|
|||
messages: list[dict[str, str]] | None = None,
|
||||
input: str | list | None = None,
|
||||
specific_deployment: bool | None = False,
|
||||
reserve: bool = True,
|
||||
):
|
||||
"""
|
||||
Async implementation of 'get_available_deployments'.
|
||||
|
|
@ -15711,7 +15721,7 @@ class Router(DeprecatedRouterMembers):
|
|||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
reserve=reserve,
|
||||
reserve=pick_reserves(default=True),
|
||||
)
|
||||
|
||||
async def _async_get_available_deployment(
|
||||
|
|
@ -16864,7 +16874,7 @@ class Router(DeprecatedRouterMembers):
|
|||
messages=messages,
|
||||
input=input,
|
||||
request_kwargs=request_kwargs,
|
||||
reserve=request_kwargs not in (None, {}),
|
||||
reserve=pick_reserves(default=False) and request_kwargs not in (None, {}),
|
||||
)
|
||||
|
||||
if deployment is None:
|
||||
|
|
|
|||
|
|
@ -27,11 +27,12 @@ import inspect
|
|||
import sys
|
||||
import warnings
|
||||
from collections import Counter
|
||||
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
|
||||
from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Generic, TypeVar, overload
|
||||
from typing import TYPE_CHECKING, Final, Generic, ParamSpec, TypeVar, overload
|
||||
|
||||
from litellm.router_api import SETTING_FIELDS, ChangeSet
|
||||
from litellm.router_strategy.least_busy import picks_reserve
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from typing import Any # noqa: TID251 # only to repeat two old signatures that had Any, unchanged
|
||||
|
|
@ -58,6 +59,7 @@ if TYPE_CHECKING:
|
|||
from litellm.types.utils import ModelInfo as ModelMapInfo
|
||||
|
||||
_T: Final = TypeVar("_T")
|
||||
_P: Final = ParamSpec("_P")
|
||||
|
||||
DEPRECATED_ROUTER_MEMBER_USES: Final[Counter[str]] = Counter() # mutable-ok: the live tally _note increments
|
||||
|
||||
|
|
@ -163,6 +165,18 @@ def _attr(replacement: str, read: Callable[["Router"], _T]) -> _DeprecatedAttrib
|
|||
return _DeprecatedAttribute(replacement, read)
|
||||
|
||||
|
||||
def _without_reserving(lookup: Callable[_P, Awaitable[_T]]) -> Callable[_P, Coroutine[object, object, _T]]:
|
||||
@functools.wraps(lookup)
|
||||
async def looked_up(
|
||||
*args: _P.args,
|
||||
**kwargs: _P.kwargs, # kwargs-ok: forwards the wrapped lookup's ParamSpec unchanged
|
||||
) -> _T:
|
||||
with picks_reserve(False):
|
||||
return await lookup(*args, **kwargs)
|
||||
|
||||
return looked_up
|
||||
|
||||
|
||||
# (id(router), member name) pairs whose subclass override is running in this context.
|
||||
_OVERRIDES_RUNNING: Final[contextvars.ContextVar[frozenset[tuple[int, str]]]] = contextvars.ContextVar(
|
||||
"router_overrides_running", default=frozenset()
|
||||
|
|
@ -325,7 +339,7 @@ class DeprecatedRouterMembers:
|
|||
deployment_latency_map = _attr("health().latency", lambda r: r._core_deployment_latency_map)
|
||||
async_get_available_deployment = _attr(
|
||||
"select(model, request_kwargs)",
|
||||
lambda r: r._core_async_get_available_deployment,
|
||||
lambda r: _without_reserving(r._core_async_get_available_deployment),
|
||||
)
|
||||
|
||||
# -- observers ------------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -1,7 +1,9 @@
|
|||
import contextlib
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from collections.abc import Callable, Generator, Mapping, Sequence
|
||||
from contextvars import ContextVar
|
||||
from typing import Final, Protocol
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
|
@ -12,7 +14,6 @@ from litellm._logging import verbose_router_logger
|
|||
from litellm.caching.caching import DualCache
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.caching.redis_cache import log_redis_failure
|
||||
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
|
||||
|
||||
|
|
@ -20,6 +21,20 @@ IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60
|
|||
SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0
|
||||
SELECTION_RESERVATION_MAX_KEYS: Final = 10_000
|
||||
_RESERVATIONS_LOCK: Final = threading.Lock()
|
||||
_PICK_RESERVES: Final[ContextVar[bool]] = ContextVar("least_busy_pick_reserves")
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def picks_reserve(reserve: bool) -> Generator[None, None, None]:
|
||||
token: Final = _PICK_RESERVES.set(reserve)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_PICK_RESERVES.reset(token)
|
||||
|
||||
|
||||
def pick_reserves(default: bool) -> bool:
|
||||
return _PICK_RESERVES.get(default)
|
||||
|
||||
|
||||
class _ModelInfo(TypedDict, total=False):
|
||||
|
|
@ -97,23 +112,24 @@ def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None:
|
|||
return model_group, str(deployment_id)
|
||||
|
||||
|
||||
def _counted_in_flight(kwargs: Mapping[str, object]) -> bool:
|
||||
return kwargs.get("cache_hit") is not True and (
|
||||
kwargs.get("api_call_start_time") is not None or LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL in kwargs
|
||||
)
|
||||
def _reservation_key(deployment_id: str) -> str:
|
||||
return f"least_busy_reservations:{deployment_id}"
|
||||
|
||||
|
||||
def _reservation_key(request_count_key: str) -> str:
|
||||
return f"{request_count_key}:reservations"
|
||||
def _started_attempts_key(call_id: str, ref: tuple[str, str]) -> str:
|
||||
return f"least_busy_started:{call_id}:{_request_count_key(*ref)}"
|
||||
|
||||
|
||||
def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
|
||||
def _deployment_ids(healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
_request_count_key(model_group, str(deployment["model_info"]["id"]))
|
||||
for deployment in _DEPLOYMENTS.validate_python(healthy_deployments)
|
||||
str(deployment["model_info"]["id"]) for deployment in _DEPLOYMENTS.validate_python(healthy_deployments)
|
||||
)
|
||||
|
||||
|
||||
def _request_count_keys(model_group: str, deployment_ids: tuple[str, ...]) -> tuple[str, ...]:
|
||||
return tuple(_request_count_key(model_group, deployment_id) for deployment_id in deployment_ids)
|
||||
|
||||
|
||||
def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]:
|
||||
return tuple(0 if value is None else int(value) for value in values)
|
||||
|
||||
|
|
@ -180,20 +196,28 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
else InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS, clock=clock)
|
||||
)
|
||||
|
||||
def _live_reservations(self, key: str, now: float) -> tuple[float, ...]:
|
||||
raw: Final = self._reservations.get_cache(_reservation_key(key))
|
||||
return _unexpired(_RESERVATION_TIMES.validate_python(raw) or (), now)
|
||||
def _times(self, key: str) -> tuple[float, ...]:
|
||||
return _RESERVATION_TIMES.validate_python(self._reservations.get_cache(key)) or ()
|
||||
|
||||
def _store_reservations(self, key: str, times: tuple[float, ...]) -> None:
|
||||
def _store_times(self, key: str, times: tuple[float, ...], ttl: float) -> None:
|
||||
# InMemoryCache keeps a live key's first TTL, so drop the key to restart it
|
||||
self._reservations.delete_cache(_reservation_key(key))
|
||||
self._reservations.delete_cache(key)
|
||||
if times:
|
||||
self._reservations.set_cache(_reservation_key(key), times, ttl=SELECTION_RESERVATION_TTL_SECONDS)
|
||||
self._reservations.set_cache(key, times, ttl=ttl)
|
||||
|
||||
def _live_reservations(self, deployment_id: str, now: float) -> tuple[float, ...]:
|
||||
return _unexpired(self._times(_reservation_key(deployment_id)), now)
|
||||
|
||||
def _store_reservations(self, deployment_id: str, times: tuple[float, ...]) -> None:
|
||||
self._store_times(_reservation_key(deployment_id), times, SELECTION_RESERVATION_TTL_SECONDS)
|
||||
|
||||
def _drop_oldest_reservation(self, deployment_id: str, now: float) -> None:
|
||||
self._store_reservations(deployment_id, self._live_reservations(deployment_id, now)[1:])
|
||||
|
||||
def _select(
|
||||
self,
|
||||
healthy_deployments: Sequence[Mapping[str, object]],
|
||||
keys: tuple[str, ...],
|
||||
deployment_ids: tuple[str, ...],
|
||||
counts: tuple[int, ...],
|
||||
reserve: bool,
|
||||
) -> Mapping[str, object] | None:
|
||||
|
|
@ -201,25 +225,39 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
return None
|
||||
with _RESERVATIONS_LOCK:
|
||||
now: Final = self._clock()
|
||||
live: Final = tuple(self._live_reservations(key, now) for key in keys)
|
||||
live: Final = tuple(self._live_reservations(deployment_id, now) for deployment_id in deployment_ids)
|
||||
index: Final = _least_busy_index(tuple(count + len(r) for count, r in zip(counts, live, strict=True)))
|
||||
if reserve:
|
||||
self._store_reservations(keys[index], (*live[index], now))
|
||||
self._store_reservations(deployment_ids[index], (*live[index], now))
|
||||
return healthy_deployments[index]
|
||||
|
||||
def _release_reservation(self, kwargs: Mapping[str, object]) -> None:
|
||||
ref: Final = _deployment_ref(kwargs)
|
||||
if ref is None:
|
||||
return
|
||||
key: Final = _request_count_key(*ref)
|
||||
def _start_attempt(self, ref: tuple[str, str], call_id: object) -> None:
|
||||
with _RESERVATIONS_LOCK:
|
||||
now: Final = self._clock()
|
||||
self._store_reservations(key, self._live_reservations(key, now)[1:])
|
||||
self._drop_oldest_reservation(ref[1], now)
|
||||
if isinstance(call_id, str):
|
||||
key: Final = _started_attempts_key(call_id, ref)
|
||||
self._store_times(key, (*self._times(key), now), IN_FLIGHT_COUNT_TTL_SECONDS)
|
||||
|
||||
def _end_attempt(self, ref: tuple[str, str], call_id: object) -> bool:
|
||||
if not isinstance(call_id, str):
|
||||
return True
|
||||
key: Final = _started_attempts_key(call_id, ref)
|
||||
with _RESERVATIONS_LOCK:
|
||||
started: Final = self._times(key)
|
||||
if started:
|
||||
self._store_times(key, started[1:], IN_FLIGHT_COUNT_TTL_SECONDS)
|
||||
else:
|
||||
self._drop_oldest_reservation(ref[1], self._clock())
|
||||
return bool(started)
|
||||
|
||||
@with_service_target("router_usage")
|
||||
def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
|
||||
self._release_reservation(kwargs)
|
||||
self._increment(kwargs, 1)
|
||||
ref: Final = _deployment_ref(kwargs)
|
||||
if ref is None:
|
||||
return
|
||||
self._start_attempt(ref, kwargs.get("litellm_call_id"))
|
||||
self._increment(_request_count_key(*ref), 1)
|
||||
|
||||
@with_service_target("router_usage")
|
||||
def log_success_event(
|
||||
|
|
@ -257,7 +295,8 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
def get_available_deployments(
|
||||
self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]], reserve: bool = False
|
||||
) -> Mapping[str, object] | None:
|
||||
keys: Final = _request_count_keys(model_group, healthy_deployments)
|
||||
deployment_ids: Final = _deployment_ids(healthy_deployments)
|
||||
keys: Final = _request_count_keys(model_group, deployment_ids)
|
||||
redis_cache: Final = self.router_cache.redis_cache
|
||||
if redis_cache is not None:
|
||||
try:
|
||||
|
|
@ -265,14 +304,17 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
except Exception as e:
|
||||
_warn_unreadable(model_group, e)
|
||||
else:
|
||||
return self._select(healthy_deployments, keys, self._at_least_this_workers(keys, shared), reserve)
|
||||
return self._select(healthy_deployments, keys, self._this_workers_counts(keys), reserve)
|
||||
return self._select(
|
||||
healthy_deployments, deployment_ids, self._at_least_this_workers(keys, shared), reserve
|
||||
)
|
||||
return self._select(healthy_deployments, deployment_ids, self._this_workers_counts(keys), reserve)
|
||||
|
||||
@with_service_target("router_usage")
|
||||
async def async_get_available_deployments(
|
||||
self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]], reserve: bool = False
|
||||
) -> Mapping[str, object] | None:
|
||||
keys: Final = _request_count_keys(model_group, healthy_deployments)
|
||||
deployment_ids: Final = _deployment_ids(healthy_deployments)
|
||||
keys: Final = _request_count_keys(model_group, deployment_ids)
|
||||
redis_cache: Final = self.router_cache.redis_cache
|
||||
if redis_cache is not None:
|
||||
try:
|
||||
|
|
@ -280,20 +322,22 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
except Exception as e:
|
||||
_warn_unreadable(model_group, e)
|
||||
else:
|
||||
return self._select(healthy_deployments, keys, self._at_least_this_workers(keys, shared), reserve)
|
||||
return self._select(healthy_deployments, keys, self._this_workers_counts(keys), reserve)
|
||||
return self._select(
|
||||
healthy_deployments, deployment_ids, self._at_least_this_workers(keys, shared), reserve
|
||||
)
|
||||
return self._select(healthy_deployments, deployment_ids, self._this_workers_counts(keys), reserve)
|
||||
|
||||
def _settle(self, kwargs: Mapping[str, object]) -> None:
|
||||
if not _counted_in_flight(kwargs):
|
||||
self._release_reservation(kwargs)
|
||||
ref: Final = _deployment_ref(kwargs)
|
||||
if ref is None or not self._end_attempt(ref, kwargs.get("litellm_call_id")):
|
||||
return
|
||||
self._increment(kwargs, -1)
|
||||
self._increment(_request_count_key(*ref), -1)
|
||||
|
||||
async def _async_settle(self, kwargs: Mapping[str, object]) -> None:
|
||||
if not _counted_in_flight(kwargs):
|
||||
self._release_reservation(kwargs)
|
||||
ref: Final = _deployment_ref(kwargs)
|
||||
if ref is None or not self._end_attempt(ref, kwargs.get("litellm_call_id")):
|
||||
return
|
||||
await self._async_increment(kwargs, -1)
|
||||
await self._async_increment(_request_count_key(*ref), -1)
|
||||
|
||||
def _this_workers_counts(self, keys: tuple[str, ...]) -> tuple[int, ...]:
|
||||
return _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys)
|
||||
|
|
@ -301,11 +345,7 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
def _at_least_this_workers(self, keys: tuple[str, ...], shared: tuple[int, ...]) -> tuple[int, ...]:
|
||||
return tuple(max(counts) for counts in zip(shared, self._this_workers_counts(keys), strict=True))
|
||||
|
||||
def _increment(self, kwargs: Mapping[str, object], delta: int) -> None:
|
||||
ref: Final = _deployment_ref(kwargs)
|
||||
if ref is None:
|
||||
return
|
||||
key: Final = _request_count_key(*ref)
|
||||
def _increment(self, key: str, delta: int) -> None:
|
||||
redis_cache: Final = self.router_cache.redis_cache
|
||||
try:
|
||||
local: Final = self.router_cache.increment_cache(
|
||||
|
|
@ -319,11 +359,7 @@ class LeastBusyLoggingHandler(CustomLogger):
|
|||
except Exception as e:
|
||||
_warn_unwritable(key, e)
|
||||
|
||||
async def _async_increment(self, kwargs: Mapping[str, object], delta: int) -> None:
|
||||
ref: Final = _deployment_ref(kwargs)
|
||||
if ref is None:
|
||||
return
|
||||
key: Final = _request_count_key(*ref)
|
||||
async def _async_increment(self, key: str, delta: int) -> None:
|
||||
redis_cache: Final = self.router_cache.redis_cache
|
||||
try:
|
||||
local: Final = await self.router_cache.async_increment_cache(
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import datetime
|
|||
import json
|
||||
import logging
|
||||
import threading
|
||||
from collections.abc import Callable, Iterator
|
||||
from collections.abc import Callable, Generator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -25,22 +25,21 @@ from litellm.router_strategy.least_busy import (
|
|||
IN_FLIGHT_COUNT_TTL_SECONDS,
|
||||
SELECTION_RESERVATION_TTL_SECONDS,
|
||||
LeastBusyLoggingHandler,
|
||||
picks_reserve,
|
||||
)
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
GROUP: Final = "least-busy-group"
|
||||
CALL_STARTED_AT: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc)
|
||||
DEPLOYMENT_A: Final[dict[str, object]] = {"model_info": {"id": "dep-a"}}
|
||||
DEPLOYMENT_B: Final[dict[str, object]] = {"model_info": {"id": "dep-b"}}
|
||||
HEALTHY: Final = [DEPLOYMENT_A, DEPLOYMENT_B]
|
||||
|
||||
|
||||
def _routed_call_kwargs(deployment_id: str) -> dict[str, object]:
|
||||
return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}}
|
||||
|
||||
|
||||
def _call_kwargs(deployment_id: str) -> dict[str, object]:
|
||||
return {**_routed_call_kwargs(deployment_id), "api_call_start_time": CALL_STARTED_AT}
|
||||
def _call_kwargs(deployment_id: str, call_id: str = "call-1") -> dict[str, object]:
|
||||
return {
|
||||
"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}},
|
||||
"litellm_call_id": call_id,
|
||||
}
|
||||
|
||||
|
||||
class SharedRedisCounters:
|
||||
|
|
@ -145,7 +144,7 @@ def _routed_through_litellm_metadata(deployment_id: str) -> dict[str, object]:
|
|||
"litellm_metadata": {"model_group": GROUP},
|
||||
"model_info": {"id": deployment_id},
|
||||
},
|
||||
"api_call_start_time": CALL_STARTED_AT,
|
||||
"litellm_call_id": "call-1",
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -561,7 +560,7 @@ class _StoreWhereThreadsReadBeforeWriting:
|
|||
self._meeting: Final = threading.Event()
|
||||
|
||||
@contextlib.contextmanager
|
||||
def first_reads_meet(self) -> Iterator[None]:
|
||||
def first_reads_meet(self) -> Generator[None, None, None]:
|
||||
self._meeting.set()
|
||||
yield
|
||||
self._meeting.clear()
|
||||
|
|
@ -616,11 +615,25 @@ def _router_pick(router: Router) -> str:
|
|||
def test_a_pick_stays_reserved_while_other_router_keys_fill_its_cache() -> None:
|
||||
router: Final = _least_busy_router()
|
||||
|
||||
first: Final = _router_pick(router)
|
||||
for index in range(DEFAULT_MAX_SIZE_IN_MEMORY):
|
||||
router.health().store.set_cache(f"deployment-usage-{index}", 1, local_only=True, ttl=60)
|
||||
with picks_reserve(True):
|
||||
first: Final = _router_pick(router)
|
||||
for index in range(DEFAULT_MAX_SIZE_IN_MEMORY):
|
||||
router.health().store.set_cache(f"deployment-usage-{index}", 1, local_only=True, ttl=60)
|
||||
|
||||
assert _router_pick(router) != first
|
||||
assert _router_pick(router) != first
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_looking_up_a_deployment_does_not_hold_it_for_later_picks() -> None:
|
||||
router: Final = _least_busy_router()
|
||||
|
||||
sync_picks: Final = [_router_pick(router) for _ in range(3)]
|
||||
async_picks: Final = [
|
||||
(await router.async_get_available_deployment(ROUTER_GROUP, request_kwargs={"metadata": {}}))["model_info"]["id"]
|
||||
for _ in range(3)
|
||||
]
|
||||
|
||||
assert sync_picks + async_picks == [ROUTER_DEPLOYMENT_IDS[0]] * 6
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -638,7 +651,7 @@ def test_a_cache_hit_leaves_requests_in_flight_on_its_deployment_counted() -> No
|
|||
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)
|
||||
worker.log_success_event({**_call_kwargs("dep-a", "cached-call"), "cache_hit": True}, None, None, None)
|
||||
|
||||
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
|
||||
|
||||
|
|
@ -650,14 +663,14 @@ def test_a_call_that_ends_before_reaching_the_provider_frees_its_pick(
|
|||
worker: Final = _worker(None)
|
||||
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
|
||||
|
||||
outcome(worker, _routed_call_kwargs("dep-a"), None, None, None)
|
||||
outcome(worker, _call_kwargs("dep-a"), None, None, None)
|
||||
|
||||
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
|
||||
|
||||
|
||||
def test_a_call_the_proxy_rejected_after_logging_its_start_leaves_no_count_behind() -> None:
|
||||
worker: Final = _worker(None)
|
||||
rejected: Final = {**_routed_call_kwargs("dep-a"), LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True}
|
||||
rejected: Final = {**_call_kwargs("dep-a"), LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True}
|
||||
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=rejected)
|
||||
worker.log_failure_event(rejected, None, None, None)
|
||||
|
|
@ -665,6 +678,44 @@ def test_a_call_the_proxy_rejected_after_logging_its_start_leaves_no_count_behin
|
|||
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
|
||||
|
||||
|
||||
def test_a_call_logged_under_its_callers_group_name_frees_the_pick_made_for_another_group() -> None:
|
||||
worker: Final = _worker(None)
|
||||
assert worker.get_available_deployments("tier-group", HEALTHY, reserve=True) is DEPLOYMENT_A
|
||||
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a"))
|
||||
worker.log_success_event(_call_kwargs("dep-a"), None, None, None)
|
||||
|
||||
assert worker.get_available_deployments("tier-group", HEALTHY) is DEPLOYMENT_A
|
||||
|
||||
|
||||
def test_a_fallback_attempt_that_fails_before_reaching_its_provider_frees_its_pick() -> None:
|
||||
worker: Final = _worker(None)
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "busy-call"))
|
||||
first_attempt_started_at: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc)
|
||||
first_attempt: Final = {**_call_kwargs("dep-a", "fallback-call"), "api_call_start_time": first_attempt_started_at}
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=first_attempt)
|
||||
worker.log_failure_event(first_attempt, None, None, None)
|
||||
|
||||
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
|
||||
second_attempt: Final = {**_call_kwargs("dep-b", "fallback-call"), "api_call_start_time": first_attempt_started_at}
|
||||
worker.log_failure_event(second_attempt, None, None, None)
|
||||
|
||||
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
|
||||
|
||||
|
||||
def test_a_retry_that_starts_before_the_failed_attempt_is_logged_counts_each_attempt_once() -> None:
|
||||
worker: Final = _worker(None)
|
||||
attempt: Final = _call_kwargs("dep-a", "retried-call")
|
||||
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=attempt)
|
||||
worker.log_pre_api_call(model="m", messages=[], kwargs=attempt)
|
||||
worker.log_failure_event(attempt, None, None, None)
|
||||
worker.log_success_event(attempt, None, None, None)
|
||||
|
||||
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
|
||||
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
|
||||
|
||||
|
||||
async def _drain_logging() -> None:
|
||||
for _ in range(10):
|
||||
await asyncio.sleep(0)
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ import litellm
|
|||
from litellm import Router
|
||||
from litellm.caching.redis_cache import RedisPipelineIncrementOperation
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.router_strategy.least_busy import picks_reserve
|
||||
from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup, RoutingStrategy
|
||||
from litellm.utils import Rules, function_setup
|
||||
|
||||
|
|
@ -56,9 +57,6 @@ def _model_list():
|
|||
]
|
||||
|
||||
|
||||
CALL_STARTED_AT: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc)
|
||||
|
||||
|
||||
def _build_router(routing_strategy="simple-shuffle", routing_groups=None):
|
||||
return Router(
|
||||
model_list=_model_list(),
|
||||
|
|
@ -471,8 +469,7 @@ def test_two_least_busy_groups_count_a_request_once(monkeypatch):
|
|||
"litellm_params": {
|
||||
"metadata": {"model_group": "filtered-model"},
|
||||
"model_info": {"id": "deploy-1"},
|
||||
},
|
||||
"api_call_start_time": CALL_STARTED_AT,
|
||||
}
|
||||
}
|
||||
|
||||
for callback in litellm.input_callback:
|
||||
|
|
@ -485,7 +482,6 @@ 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")
|
||||
|
|
@ -501,10 +497,11 @@ def test_a_least_busy_group_pick_stops_counting_once_its_request_finishes(monkey
|
|||
routing_groups=[{"group_name": "fast", "models": ["filtered-model"], "routing_strategy": "least-busy"}],
|
||||
)
|
||||
|
||||
first: Final = _filtered_model_pick(router)
|
||||
with picks_reserve(True):
|
||||
first: Final = _filtered_model_pick(router)
|
||||
kwargs: Final = {
|
||||
"litellm_params": {"metadata": {"model_group": "filtered-model"}, "model_info": {"id": first}},
|
||||
"api_call_start_time": CALL_STARTED_AT,
|
||||
"litellm_call_id": "call-1",
|
||||
}
|
||||
for callback in litellm.input_callback:
|
||||
if isinstance(callback, CustomLogger):
|
||||
|
|
@ -515,6 +512,7 @@ def test_a_least_busy_group_pick_stops_counting_once_its_request_finishes(monkey
|
|||
|
||||
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
|
||||
|
|
@ -532,8 +530,7 @@ def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch):
|
|||
"litellm_params": {
|
||||
"metadata": {"model_group": "filtered-model"},
|
||||
"model_info": {"id": "deploy-1"},
|
||||
},
|
||||
"api_call_start_time": CALL_STARTED_AT,
|
||||
}
|
||||
}
|
||||
|
||||
for callback in litellm.input_callback:
|
||||
|
|
@ -1241,12 +1238,12 @@ async def test_least_busy_override_sees_the_overriding_request_in_flight():
|
|||
stream = await router.acompletion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy", stream=True))
|
||||
busy = stream._hidden_params["model_id"]
|
||||
idle = "d2" if busy == "d1" else "d1"
|
||||
assert await _async_override_pick(router, "least-busy") == idle
|
||||
assert [await _async_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle]
|
||||
|
||||
async for _ in stream:
|
||||
pass
|
||||
await _async_wait_until(lambda: _in_flight(router, busy) == 0)
|
||||
assert await _async_override_pick(router, "least-busy") == busy
|
||||
assert await _async_override_pick(router, "least-busy") == "d1"
|
||||
assert _selector_is_not_global(router._override_selectors["least-busy"])
|
||||
|
||||
|
||||
|
|
@ -1256,12 +1253,12 @@ def test_sync_least_busy_override_sees_the_overriding_request_in_flight():
|
|||
stream = router.completion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy", stream=True))
|
||||
busy = stream._hidden_params["model_id"]
|
||||
idle = "d2" if busy == "d1" else "d1"
|
||||
assert _sync_override_pick(router, "least-busy") == idle
|
||||
assert [_sync_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle]
|
||||
|
||||
for _ in stream:
|
||||
pass
|
||||
_sync_wait_until(lambda: _in_flight(router, busy) == 0)
|
||||
assert _sync_override_pick(router, "least-busy") == busy
|
||||
assert _sync_override_pick(router, "least-busy") == "d1"
|
||||
assert _selector_is_not_global(router._override_selectors["least-busy"])
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -29096,11 +29096,21 @@ def test_reset_custom_routing_strategy():
|
|||
|
||||
assert "get_available_deployment" not in router.__dict__
|
||||
assert "async_get_available_deployment" not in router.__dict__
|
||||
assert router.async_get_available_deployment.__func__ is Router.async_get_available_deployment
|
||||
assert router.async_get_available_deployment.__wrapped__.__func__ is Router.async_get_available_deployment
|
||||
|
||||
router._reset_custom_routing_strategy()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_select_picks_through_a_custom_routing_strategy() -> None:
|
||||
router: Final = _create_custom_routing_router()
|
||||
router.set_custom_routing_strategy(SpecialEndpointRoutingStrategy(router))
|
||||
|
||||
picked: Final = await router.select("azure-model", {"metadata": {}})
|
||||
|
||||
assert picked["model_info"]["id"] == "very-special-endpoint"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_router_acompletion_passes_deployment_credentials_for_plain_and_stream(
|
||||
monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
|
|
|
|||
|
|
@ -616,7 +616,7 @@ async def test_responses_websocket_turns_do_not_make_their_deployment_look_busy_
|
|||
await admit_turn()
|
||||
|
||||
picks: Final = [
|
||||
(await router.async_get_available_deployment(model="ws-least-busy", request_kwargs={"input": "hi"}))[
|
||||
(await router._core_async_get_available_deployment(model="ws-least-busy", request_kwargs={"input": "hi"}))[
|
||||
"model_info"
|
||||
]["id"]
|
||||
for _ in range(2)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue