mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(router): stamp the failed deployment on every endpoint's attempt
Weighted failover (#27980) re-picks within a model group by excluding the deployment named in the error's failed_deployment_id. #27980 stamped that id in only three attempt functions: _completion, _acompletion and _ageneric_api_call_with_fallbacks_helper. Embedding, image generation, transcription, speech, text completion, rerank, adapter completion, the sync generic call and vector store create never stamped it, so with enable_weighted_failover on, a failure on one of their deployments skipped the in-group re-pick and went straight to the cross-group fallbacks, or raised. Each of those attempts now stamps the deployment it picked, read from the model_info that _update_kwargs_with_deployment put on the request (so a client-side-credential call names its own dynamic deployment, as on completions). A failed pick stamps nothing. File and batch creation and batch cancellation are left alone: they are bound to a resource on the deployment they picked, so re-picking another one cannot help.
This commit is contained in:
narmstrong-berri
parent
60fc121ce8
commit
8affaa20ef
2 changed files with 173 additions and 11 deletions
|
|
@ -5172,6 +5172,13 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
effective_model_info: Final = kwargs.get("model_info") or deployment.get("model_info") or MappingProxyType({})
|
||||
self._set_failed_deployment_id_on_exception(exception, MappingProxyType({"model_info": effective_model_info}))
|
||||
|
||||
def _stamp_failed_attempt(self, exception: Exception, kwargs: Mapping[str, object], *, picked: bool) -> None:
|
||||
"""An attempt's exception handler: once it picked a deployment, stamp that deployment (the
|
||||
``model_info`` ``_update_kwargs_with_deployment`` put on ``kwargs``) as ``failed_deployment_id``,
|
||||
which weighted failover reads. A failed pick stamps nothing."""
|
||||
if picked:
|
||||
self._stamp_failed_deployment_id_with_effective_model_info(exception, MappingProxyType({}), kwargs)
|
||||
|
||||
@staticmethod
|
||||
def _stamp_retry_skip_deployment_id(exception: Exception, kwargs: Mapping[str, object]) -> None:
|
||||
effective_model_info: Final = kwargs.get("model_info")
|
||||
|
|
@ -6047,9 +6054,10 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
|
||||
def _image_generation(self, prompt: str, model: str, **kwargs):
|
||||
model_name: Final = ""
|
||||
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
|
||||
try:
|
||||
verbose_router_logger.debug("Inside _image_generation()- model: %s", model)
|
||||
deployment: Final = self.get_available_deployment(
|
||||
deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "prompt"}],
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
|
|
@ -6085,6 +6093,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
if model_name is not None:
|
||||
self.fail_calls[model_name] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
self._stamp_failed_attempt(
|
||||
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
|
||||
)
|
||||
raise e
|
||||
|
||||
async def aimage_generation(self, prompt: str, model: str, **kwargs):
|
||||
|
|
@ -6110,11 +6121,12 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
async def _aimage_generation(self, prompt: str, model: str, **kwargs):
|
||||
model_name = model
|
||||
start_time: Final = datetime.now() # noqa: DTZ005 # SDK logging uses naive local timestamps
|
||||
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
|
||||
try:
|
||||
verbose_router_logger.debug("Inside _image_generation()- model: %s", model)
|
||||
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
|
||||
try:
|
||||
deployment: Final = await self._core_async_get_available_deployment(
|
||||
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "prompt"}],
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
|
|
@ -6175,6 +6187,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
if model_name is not None:
|
||||
self.fail_calls[model_name] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
self._stamp_failed_attempt(
|
||||
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
|
||||
)
|
||||
raise e
|
||||
|
||||
async def atranscription(self, file: FileTypes, model: str, **kwargs):
|
||||
|
|
@ -6221,10 +6236,11 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
|
||||
async def _atranscription(self, file: FileTypes, model: str, **kwargs):
|
||||
model_name: Final = model
|
||||
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
|
||||
try:
|
||||
verbose_router_logger.debug("Inside _atranscription()- model: %s", model)
|
||||
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
|
||||
deployment: Final = await self._core_async_get_available_deployment(
|
||||
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "prompt"}],
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
|
|
@ -6258,6 +6274,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
if model_name is not None:
|
||||
self.fail_calls[model_name] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
self._stamp_failed_attempt(
|
||||
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
|
||||
)
|
||||
raise e
|
||||
|
||||
async def aspeech(self, model: str, input: str, voice: str | None = None, **kwargs):
|
||||
|
|
@ -6314,10 +6333,11 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
|
||||
async def _aspeech(self, model: str, input: str, voice: str | None = None, **kwargs):
|
||||
model_name: Final = model
|
||||
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
|
||||
try:
|
||||
verbose_router_logger.debug("Inside _aspeech()- model: %s", model)
|
||||
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
|
||||
deployment: Final = await self._core_async_get_available_deployment(
|
||||
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "prompt"}],
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
|
|
@ -6351,6 +6371,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
if model_name is not None:
|
||||
self.fail_calls[model_name] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
self._stamp_failed_attempt(
|
||||
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
|
||||
)
|
||||
raise e
|
||||
|
||||
async def arerank(self, model: str, **kwargs):
|
||||
|
|
@ -6376,9 +6399,10 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
|
||||
async def _arerank(self, model: str, **kwargs):
|
||||
model_name = None
|
||||
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
|
||||
try:
|
||||
verbose_router_logger.debug("Inside _rerank()- model: %s", model)
|
||||
deployment: Final = await self._core_async_get_available_deployment(
|
||||
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
|
||||
model=model,
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
request_kwargs=kwargs,
|
||||
|
|
@ -6410,6 +6434,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
if model_name is not None:
|
||||
self.fail_calls[model_name] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
self._stamp_failed_attempt(
|
||||
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
|
||||
)
|
||||
raise e
|
||||
|
||||
def text_completion(
|
||||
|
|
@ -6491,10 +6518,11 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
raise e
|
||||
|
||||
async def _atext_completion(self, model: str, prompt: str, **kwargs):
|
||||
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
|
||||
try:
|
||||
verbose_router_logger.debug("Inside _atext_completion()- model: %s", model)
|
||||
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
|
||||
deployment: Final = await self._core_async_get_available_deployment(
|
||||
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": prompt}],
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
|
|
@ -6530,6 +6558,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
if model is not None:
|
||||
self.fail_calls[model] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
self._stamp_failed_attempt(
|
||||
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
|
||||
)
|
||||
raise e
|
||||
|
||||
async def aadapter_completion(
|
||||
|
|
@ -6561,10 +6592,11 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
raise e
|
||||
|
||||
async def _aadapter_completion(self, adapter_id: str, model: str, **kwargs):
|
||||
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
|
||||
try:
|
||||
verbose_router_logger.debug("Inside _aadapter_completion()- model: %s", model)
|
||||
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
|
||||
deployment: Final = await self._core_async_get_available_deployment(
|
||||
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "default text"}],
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
|
|
@ -6600,6 +6632,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
if model is not None:
|
||||
self.fail_calls[model] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
self._stamp_failed_attempt(
|
||||
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
|
||||
)
|
||||
raise e
|
||||
|
||||
async def _asearch_with_fallbacks(self, original_function: Callable, **kwargs) -> object:
|
||||
|
|
@ -7570,6 +7605,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
"""
|
||||
handler_name: Final = original_function.__name__
|
||||
metadata_variable_name: Final = get_router_metadata_variable_name(function_name="generic_api_call")
|
||||
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
|
||||
try:
|
||||
verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model)
|
||||
self._update_kwargs_before_fallbacks(
|
||||
|
|
@ -7578,7 +7614,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
metadata_variable_name=metadata_variable_name,
|
||||
)
|
||||
with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name):
|
||||
deployment: Final = self.get_available_deployment(
|
||||
deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init
|
||||
model=model,
|
||||
messages=kwargs.get("messages", None),
|
||||
input=kwargs.get("input", None),
|
||||
|
|
@ -7624,6 +7660,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e)
|
||||
if model is not None:
|
||||
self.fail_calls[model] += 1
|
||||
self._stamp_failed_attempt(
|
||||
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
|
||||
)
|
||||
raise e
|
||||
|
||||
def embedding(
|
||||
|
|
@ -7645,10 +7684,11 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
|
||||
def _embedding(self, input: str | list[object], model: str, **kwargs):
|
||||
model_name = None
|
||||
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
|
||||
try:
|
||||
verbose_router_logger.debug("Inside embedding()- model: %s", model)
|
||||
with embedding_deployment_selection():
|
||||
deployment: Final = self.get_available_deployment(
|
||||
deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init
|
||||
model=model,
|
||||
input=input,
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
|
|
@ -7705,6 +7745,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
if model_name is not None:
|
||||
self.fail_calls[model_name] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
self._stamp_failed_attempt(
|
||||
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
|
||||
)
|
||||
raise e
|
||||
|
||||
async def aembedding(
|
||||
|
|
@ -7734,11 +7777,12 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
|
||||
async def _aembedding(self, input: str | list[object], model: str, **kwargs):
|
||||
model_name = None
|
||||
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
|
||||
try:
|
||||
verbose_router_logger.debug("Inside _aembedding()- model: %s", model)
|
||||
parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs)
|
||||
with embedding_deployment_selection():
|
||||
deployment: Final = await self._core_async_get_available_deployment(
|
||||
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
|
||||
model=model,
|
||||
input=input,
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
|
|
@ -7785,6 +7829,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
if model_name is not None:
|
||||
self.fail_calls[model_name] += 1
|
||||
self._stamp_retry_skip_deployment_id(e, kwargs)
|
||||
self._stamp_failed_attempt(
|
||||
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
|
||||
)
|
||||
raise e
|
||||
|
||||
#### FILES API ####
|
||||
|
|
@ -7949,6 +7996,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
Returns:
|
||||
VectorStoreCreateResponse
|
||||
"""
|
||||
deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked
|
||||
try:
|
||||
# If model is None, use the factory function approach (direct SDK call)
|
||||
if model is None:
|
||||
|
|
@ -7961,7 +8009,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
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(
|
||||
deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": "vector-store-api-fake-text"}],
|
||||
specific_deployment=kwargs.pop("specific_deployment", None),
|
||||
|
|
@ -8010,6 +8058,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
)
|
||||
if model is not None:
|
||||
self.fail_calls[model] += 1
|
||||
self._stamp_failed_attempt(
|
||||
e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None
|
||||
)
|
||||
raise e
|
||||
|
||||
def _override_vector_store_methods_for_router(self):
|
||||
|
|
|
|||
|
|
@ -1019,3 +1019,114 @@ async def test_failover_reaches_healthy_sibling_when_affinity_pins_failed_deploy
|
|||
assert while_a_is_down._hidden_params["model_id"] == "B"
|
||||
assert while_a_is_down._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1
|
||||
assert after_a_recovers._hidden_params["model_id"] == "A"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Every endpoint's attempt stamps the failed deployment, so weighted failover
|
||||
# covers more than chat completions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _two_deployment_router(model: str) -> Router:
|
||||
"""Deployment A (always picked first) fails, deployment B answers."""
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": model, "api_key": "bad", "weight": 1},
|
||||
"model_info": {"id": "A"},
|
||||
},
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": model, "api_key": "good", "weight": 0},
|
||||
"model_info": {"id": "B"},
|
||||
},
|
||||
],
|
||||
routing_strategy="simple-shuffle",
|
||||
num_retries=0,
|
||||
disable_cooldowns=True,
|
||||
enable_weighted_failover=True,
|
||||
)
|
||||
|
||||
|
||||
_ENDPOINTS: Final = [
|
||||
("aembedding", "openai/text-embedding-3-small", {"input": ["hi"]}),
|
||||
("aimage_generation", "openai/dall-e-3", {"prompt": "a cat"}),
|
||||
("atranscription", "openai/whisper-1", {"file": b"audio"}),
|
||||
("aspeech", "openai/tts-1", {"input": "hi", "voice": "alloy"}),
|
||||
("atext_completion", "openai/gpt-3.5-turbo-instruct", {"prompt": "hi"}),
|
||||
("arerank", "cohere/rerank-english-v3.0", {"query": "q", "documents": ["d"]}),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("endpoint", "model", "arguments"), _ENDPOINTS, ids=[e[0] for e in _ENDPOINTS])
|
||||
async def test_weighted_failover_moves_every_endpoint_to_another_deployment(
|
||||
monkeypatch: pytest.MonkeyPatch, endpoint: str, model: str, arguments: dict[str, object]
|
||||
) -> None:
|
||||
"""A failed attempt on any endpoint names its deployment, so weighted failover re-picks
|
||||
within the group instead of raising (#27980 only stamped completions and generic calls)."""
|
||||
keys_called: list[object] = []
|
||||
answer: Final = object()
|
||||
|
||||
async def provider(**kwargs: object) -> object:
|
||||
keys_called.append(kwargs.get("api_key"))
|
||||
if kwargs.get("api_key") == "bad":
|
||||
raise litellm.InternalServerError(message="A is down", llm_provider="openai", model=model)
|
||||
return answer
|
||||
|
||||
monkeypatch.setattr(litellm, endpoint, provider)
|
||||
router: Final = _two_deployment_router(model)
|
||||
|
||||
response: Final = await getattr(router, endpoint)(model="test-model", **arguments)
|
||||
|
||||
assert response is answer
|
||||
assert keys_called == ["bad", "good"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("endpoint", "attempt", "arguments"),
|
||||
[
|
||||
("embedding", "_embedding", {"input": ["hi"]}),
|
||||
("image_generation", "_image_generation", {"prompt": "a cat"}),
|
||||
],
|
||||
)
|
||||
def test_a_sync_attempt_stamps_the_failed_deployment(
|
||||
monkeypatch: pytest.MonkeyPatch, endpoint: str, attempt: str, arguments: dict[str, object]
|
||||
) -> None:
|
||||
def provider(**kwargs: object) -> object:
|
||||
raise litellm.InternalServerError(message="down", llm_provider="openai", model="m")
|
||||
|
||||
monkeypatch.setattr(litellm, endpoint, provider)
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "openai/m", "api_key": "k"},
|
||||
"model_info": {"id": "A"},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.InternalServerError) as raised:
|
||||
getattr(router, attempt)(model="test-model", **arguments)
|
||||
|
||||
assert getattr(raised.value, "failed_deployment_id", None) == "A"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_pick_stamps_no_deployment() -> None:
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "test-model",
|
||||
"litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "k"},
|
||||
"model_info": {"id": "A"},
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
with pytest.raises(Exception) as raised:
|
||||
await router._aembedding(input=["hi"], model="no-such-model")
|
||||
|
||||
assert getattr(raised.value, "failed_deployment_id", None) is None
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue