mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
fix(router): free least-busy picks of calls made without a call id
This commit is contained in:
mateo-berri
parent
252012a469
commit
59924d3b89
3 changed files with 161 additions and 102 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue