diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index 415d5d546e9..3b73f65c917 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -557,7 +557,6 @@ def prepare_call( litellm_params={ "litellm_call_id": kwargs.get("litellm_call_id"), "api_base": provider_config.get_complete_url(resolved_api_base, canonical_model), - **({} if (model_info := kwargs.get("model_info")) is None else {"model_info": model_info}), }, custom_llm_provider=provider, ) diff --git a/litellm/router.py b/litellm/router.py index 62024bc9e74..b34b621c189 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -200,7 +200,7 @@ from litellm.router_strategy.complexity_router.context_compaction import ( reject_recursive_compactor, surface_for_call, ) -from litellm.router_strategy.least_busy import LeastBusyLoggingHandler +from litellm.router_strategy.least_busy import SELECTION_RESERVATION_MAX_KEYS, LeastBusyLoggingHandler 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 @@ -1463,6 +1463,7 @@ class Router: self.cache = DualCache( redis_cache=redis_cache, in_memory_cache=InMemoryCache() ) # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc. + self._least_busy_reservations: InMemoryCache = InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS) self._claude_code_session_router_cache: DualCache = DualCache( redis_cache=redis_cache, in_memory_cache=InMemoryCache(), @@ -1967,7 +1968,9 @@ class Router: selector: RouterStrategySelector | None = None match self._normalize_strategy(strategy): case RoutingStrategy.LEAST_BUSY.value: - selector = LeastBusyLoggingHandler(router_cache=self.cache) + selector = LeastBusyLoggingHandler( + router_cache=self.cache, reservation_store=self._least_busy_reservations + ) case RoutingStrategy.USAGE_BASED_ROUTING.value: selector = LowestTPMLoggingHandler( router_cache=self.cache, diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 29fb99ff63d..b69620aea23 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -9,12 +9,14 @@ from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target 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.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0 +SELECTION_RESERVATION_MAX_KEYS: Final = 10_000 class _ModelInfo(TypedDict, total=False): @@ -23,6 +25,7 @@ class _ModelInfo(TypedDict, total=False): class _Metadata(TypedDict, total=False): model_group: ReadOnly[str | None] + model_info: ReadOnly[_ModelInfo | None] class _LitellmParams(TypedDict, total=False): @@ -57,6 +60,14 @@ def _model_group(metadata: _Metadata | None) -> str | None: return metadata.get("model_group") if metadata else None +def _deployment_id(model_info: _ModelInfo | None) -> str | int | None: + return model_info.get("id") if model_info else None + + +def _routed_deployment_id(metadata: _Metadata | None) -> str | int | None: + return _deployment_id(metadata.get("model_info")) 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 @@ -65,11 +76,19 @@ def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: except ValidationError: return 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") + routing_metadata: Final = (litellm_params.get("litellm_metadata"), litellm_params.get("metadata")) + model_group: Final = next((group for group in map(_model_group, routing_metadata) if group is not None), None) + deployment_id: Final = next( + ( + deployment_id + for deployment_id in ( + _deployment_id(litellm_params.get("model_info")), + *map(_routed_deployment_id, routing_metadata), + ) + if deployment_id is not None + ), + None, ) - deployment_id: Final = model_info.get("id") if model_info else None if model_group is None or deployment_id is None: return None return model_group, str(deployment_id) @@ -137,24 +156,30 @@ class LeastBusyLoggingHandler(CustomLogger): logged_success: int = 0 logged_failure: int = 0 - def __init__(self, router_cache: DualCache, clock: Callable[[], float] = time.monotonic): + def __init__( + self, + router_cache: DualCache, + clock: Callable[[], float] = time.monotonic, + reservation_store: _ReservationStore | None = None, + ): self.router_cache = router_cache self.router_cache_id = str(id(router_cache)) self._clock = clock - - def _reservation_store(self) -> _ReservationStore: - return self.router_cache.in_memory_cache + self._reservations: Final[_ReservationStore] = ( + reservation_store + if reservation_store is not None + 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._reservation_store().get_cache(_reservation_key(key)) + raw: Final = self._reservations.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 - store: Final = self._reservation_store() - store.delete_cache(_reservation_key(key)) + self._reservations.delete_cache(_reservation_key(key)) if times: - store.set_cache(_reservation_key(key), times, ttl=SELECTION_RESERVATION_TTL_SECONDS) + self._reservations.set_cache(_reservation_key(key), times, ttl=SELECTION_RESERVATION_TTL_SECONDS) def _select( self, @@ -235,9 +260,8 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unreadable(model_group, e) else: - return self._select(healthy_deployments, keys, shared, reserve) - local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) - return self._select(healthy_deployments, keys, local, reserve) + 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) @with_service_target("router_usage") async def async_get_available_deployments( @@ -251,9 +275,14 @@ class LeastBusyLoggingHandler(CustomLogger): except Exception as e: _warn_unreadable(model_group, e) else: - return self._select(healthy_deployments, keys, shared, reserve) - local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys) - return self._select(healthy_deployments, keys, local, reserve) + 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) + + 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) + + 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) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 3dbd90b814b..177a3968f76 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -1,6 +1,7 @@ import asyncio import json import logging +from collections.abc import Callable from typing import Final import httpx @@ -11,7 +12,7 @@ import litellm from litellm import Router from litellm.caching.caching import DualCache -from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.router_strategy.least_busy import ( @@ -62,6 +63,20 @@ class SharedRedisCounters: return self.batch_get_counts(key_list) +class SlowRedisRead(SharedRedisCounters): + """Answers a count read with the counts from when the read started, after letting other work run.""" + + def __init__(self) -> None: + super().__init__() + self.while_reading: Callable[[], None] = lambda: None + + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + counts_when_asked: Final = self.batch_get_counts(key_list) + await asyncio.sleep(0) + self.while_reading() + return counts_when_asked + + def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler: cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # duck-typed Redis double return LeastBusyLoggingHandler(router_cache=cache) @@ -85,6 +100,16 @@ async def test_worker_routes_around_a_request_another_worker_started() -> None: assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A +@pytest.mark.asyncio +async def test_a_request_that_starts_while_redis_is_read_still_counts_on_its_deployment() -> None: + redis: Final = SlowRedisRead() + worker: Final = _worker(redis) + assert await worker.async_get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + redis.while_reading = lambda: 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 + + def test_sync_pick_reads_the_shared_counts() -> None: shared: Final = SharedRedisCounters() streaming_worker: Final = _worker(shared) @@ -503,10 +528,11 @@ def test_a_started_request_moves_its_reservation_into_the_in_flight_count() -> N assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A -def test_a_request_started_through_another_selector_on_the_same_cache_releases_the_reservation() -> None: +def test_a_request_started_through_another_selector_of_the_router_releases_the_reservation() -> None: cache: Final = DualCache() - picker: Final = LeastBusyLoggingHandler(router_cache=cache) - registered: Final = LeastBusyLoggingHandler(router_cache=cache) + reservations: Final = InMemoryCache() + picker: Final = LeastBusyLoggingHandler(router_cache=cache, reservation_store=reservations) + registered: Final = LeastBusyLoggingHandler(router_cache=cache, reservation_store=reservations) assert picker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A registered.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) @@ -515,6 +541,21 @@ 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_A +def _router_pick(router: Router) -> str: + deployment: Final = router.get_available_deployment(ROUTER_GROUP, request_kwargs={"metadata": {}}) + return deployment["model_info"]["id"] + + +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.cache.set_cache(f"deployment-usage-{index}", 1, local_only=True, ttl=60) + + assert _router_pick(router) != first + + @pytest.mark.asyncio async def test_a_cache_hit_releases_the_pick_it_reserved() -> None: worker: Final = _worker(None)