From 59924d3b89de8c802406131417ae2fd19260caa0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 05:46:25 -0700 Subject: [PATCH] fix(router): free least-busy picks of calls made without a call id --- litellm/router.py | 188 +++++++++--------- litellm/router_strategy/least_busy.py | 22 +- tests/unit/router_strategy/test_least_busy.py | 53 ++++- 3 files changed, 161 insertions(+), 102 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index b21dc5cc212..845c8eba37b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -224,9 +224,9 @@ from litellm.router_strategy.least_busy import ( SELECTION_RESERVATION_MAX_KEYS, LeastBusyLoggingHandler, call_id_of, + owning_picks, pick_reserves, picks_reserve, - release_unstarted_picks, ) from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler @@ -7766,72 +7766,71 @@ class Router(DeprecatedRouterMembers): Returns: VectorStoreCreateResponse """ - try: - # If model is None, use the factory function approach (direct SDK call) - if model is None: - from litellm.vector_stores.main import acreate + with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): + try: + # If model is None, use the factory function approach (direct SDK call) + if model is None: + from litellm.vector_stores.main import acreate - # Use the factory function to handle the call - factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create") - return await factory_fn(**kwargs) + # Use the factory function to handle the call + factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create") + return await factory_fn(**kwargs) - from litellm.vector_stores import acreate as avector_store_create_sdk + from litellm.vector_stores import acreate as avector_store_create_sdk - 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": "vector-store-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="avector_store_create", - ) - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - # Get custom provider from deployment params - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await avector_store_create_sdk( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } + 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": "vector-store-api-fake-text"}], + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) + data: Final = deployment["litellm_params"].copy() + model_name: Final = data["model"] + self._update_kwargs_with_deployment( + deployment=deployment, + kwargs=kwargs, + function_name="avector_store_create", ) - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) + model_client: Final = self._get_async_openai_model_client( + deployment=deployment, + kwargs=kwargs, + ) + self.total_calls[model_name] += 1 - return response - except Exception as e: - verbose_router_logger.exception( - "litellm.avector_store_create(model=%s)\x1b[31m Exception %s\x1b[0m", model, e - ) - if model is not None: - self.fail_calls[model] += 1 - raise e - finally: - release_unstarted_picks( - self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs)) - ) + # Get custom provider from deployment params + custom_llm_provider = data.get("custom_llm_provider") + _, inferred_custom_llm_provider, _, _ = get_llm_provider( + model=data["model"], + custom_llm_provider=custom_llm_provider, + ) + custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider + + async with self._deployment_slot( + deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span + ): + response = await avector_store_create_sdk( + **{ + **data, + "custom_llm_provider": custom_llm_provider, + "caching": self._core_cache_responses, + "client": model_client, + **kwargs, + } + ) + + self.success_calls[model_name] += 1 + verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) + + return response + except Exception as e: + verbose_router_logger.exception( + "litellm.avector_store_create(model=%s)\x1b[31m Exception %s\x1b[0m", model, e + ) + if model is not None: + self.fail_calls[model] += 1 + raise e def _override_vector_store_methods_for_router(self): """ @@ -9219,42 +9218,39 @@ class Router(DeprecatedRouterMembers): mock_timeout: Final = kwargs.pop("mock_timeout", None) - try: - self._handle_mock_testing_fallbacks( - kwargs=kwargs, - model_group=model_group, - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - ) + with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): + try: + self._handle_mock_testing_fallbacks( + kwargs=kwargs, + model_group=model_group, + fallbacks=fallbacks, + context_window_fallbacks=context_window_fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + ) - if mock_timeout is not None: - response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout) - else: - response = await self.async_function_with_retries(*args, **kwargs) - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("Async Response: %s", response) - response = add_fallback_headers_to_response( - response=response, - attempted_fallbacks=0, - ) - return response - except Exception as e: - return await self.async_function_with_fallbacks_common_utils( - e, - disable_fallbacks, - fallbacks, - context_window_fallbacks, - content_policy_fallbacks, - model_group, - args, - kwargs, - include_fallback_errors=include_fallback_errors, - ) - finally: - release_unstarted_picks( - self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs)) - ) + if mock_timeout is not None: + response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout) + else: + response = await self.async_function_with_retries(*args, **kwargs) + if verbose_router_logger.isEnabledFor(logging.DEBUG): + verbose_router_logger.debug("Async Response: %s", response) + response = add_fallback_headers_to_response( + response=response, + attempted_fallbacks=0, + ) + return response + except Exception as e: + return await self.async_function_with_fallbacks_common_utils( + e, + disable_fallbacks, + fallbacks, + context_window_fallbacks, + content_policy_fallbacks, + model_group, + args, + kwargs, + include_fallback_errors=include_fallback_errors, + ) def _handle_mock_testing_fallbacks( self, diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index c01693586a3..1250dcd19c4 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -2,6 +2,7 @@ import contextlib import logging import threading import time +import uuid from collections.abc import Callable, Generator, Mapping, Sequence from contextvars import ContextVar from typing import Final, Protocol, TypeAlias @@ -22,6 +23,7 @@ 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") +_PICK_OWNER: Final[ContextVar[str | None]] = ContextVar("least_busy_pick_owner", default=None) @contextlib.contextmanager @@ -164,7 +166,7 @@ def _unexpired(picks: Sequence[_Pick], now: float) -> tuple[_Pick, ...]: def _without_one_pick(picks: tuple[_Pick, ...], call_id: str | None) -> tuple[_Pick, ...]: tags: Final = tuple(tag for _, tag in picks) - owner: Final = call_id if call_id in tags else None + owner: Final = next((tag for tag in (call_id, _PICK_OWNER.get()) if tag is not None and tag in tags), None) if owner not in tags: return picks index: Final = tags.index(owner) @@ -237,6 +239,17 @@ def release_unstarted_picks(store: _ReservationStore, call_id: str | None) -> No store.delete_cache(_picked_deployments_key(call_id)) +@contextlib.contextmanager +def owning_picks(store: _ReservationStore, call_id: str | None) -> Generator[None, None, None]: + owner: Final = call_id or _PICK_OWNER.get() or uuid.uuid4().hex + token: Final = _PICK_OWNER.set(owner) + try: + yield + finally: + _PICK_OWNER.reset(token) + release_unstarted_picks(store, owner) + + class LeastBusyLoggingHandler(CustomLogger): test_flag: bool = False logged_success: int = 0 @@ -285,9 +298,10 @@ class LeastBusyLoggingHandler(CustomLogger): 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(deployment_ids[index], (*live[index], (now, call_id))) - if call_id is not None: - self._record_pick(call_id, deployment_ids[index]) + tag: Final = call_id if call_id is not None else _PICK_OWNER.get() + self._store_reservations(deployment_ids[index], (*live[index], (now, tag))) + if tag is not None: + self._record_pick(tag, deployment_ids[index]) return healthy_deployments[index] def _start_attempt(self, ref: tuple[str, str], call_id: str | None) -> None: diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index 12e220e6294..7d84d6ca997 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -1,5 +1,6 @@ import asyncio import contextlib +import contextvars import datetime import json import logging @@ -25,6 +26,7 @@ from litellm.router_strategy.least_busy import ( IN_FLIGHT_COUNT_TTL_SECONDS, SELECTION_RESERVATION_TTL_SECONDS, LeastBusyLoggingHandler, + owning_picks, picks_reserve, release_unstarted_picks, ) @@ -719,6 +721,31 @@ def test_ending_one_request_frees_only_the_picks_it_never_started() -> None: ] +def test_a_call_without_a_call_id_turns_its_own_pick_into_its_in_flight_count() -> None: + reservations: Final = InMemoryCache() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=reservations) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "busy-call")) + + with owning_picks(reservations, None): + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a", "attempt-call")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_another_call_starting_leaves_the_pick_of_a_call_without_a_call_id_reserved() -> None: + reservations: Final = InMemoryCache() + worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=reservations) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-b", "busy-call")) + other_call: Final = contextvars.copy_context() + + with owning_picks(reservations, None): + assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A + other_call.run(worker.log_pre_api_call, model="m", messages=[], kwargs=_call_kwargs("dep-a", "other-call")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + def test_a_call_started_without_a_pick_of_its_own_leaves_another_calls_pick_reserved() -> None: worker: Final = _worker(None) assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True, call_id="picked-call") is DEPLOYMENT_A @@ -842,21 +869,43 @@ _VECTOR_STORE: Final = { "status": "completed", } _REJECTED_VECTOR_STORE: Final = {"error": {"message": "Invalid name", "type": "invalid_request_error"}} +_CALL_IDENTITIES: Final = ({"litellm_call_id": "least-busy-call"}, {}) @pytest.mark.asyncio +@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) @pytest.mark.parametrize( "provider_response", [httpx.Response(200, json=_VECTOR_STORE), httpx.Response(400, json=_REJECTED_VECTOR_STORE)] ) async def test_a_vector_store_create_does_not_keep_its_deployment_busy( - provider_response: httpx.Response, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + provider_response: httpx.Response, + call_identity: dict[str, str], + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) route: Final = respx_mock.post("https://api.openai.com/v1/vector_stores").mock(return_value=provider_response) router: Final = _least_busy_router() with contextlib.suppress(openai.BadRequestError): - await router.avector_store_create(model=ROUTER_GROUP, name="least-busy", litellm_call_id="vector-store-call") + await router.avector_store_create(model=ROUTER_GROUP, name="least-busy", **call_identity) + await _drain_logging() + + assert _router_pick(router) == _called_deployment(route) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES) +async def test_listing_fine_tuning_jobs_does_not_keep_its_deployment_busy( + call_identity: dict[str, str], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.get(url__startswith="https://api.openai.com/v1/fine_tuning/jobs").mock( + return_value=httpx.Response(200, json={"object": "list", "data": [], "has_more": False}) + ) + router: Final = _least_busy_router() + + await router.alist_fine_tuning_jobs(model=ROUTER_GROUP, **call_identity) await _drain_logging() assert _router_pick(router) == _called_deployment(route)