diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 30b05456987..bf7e6fd2417 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -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): diff --git a/tests/unit/test_router_weighted_failover.py b/tests/unit/test_router_weighted_failover.py index c00670473af..c878b79b8e6 100644 --- a/tests/unit/test_router_weighted_failover.py +++ b/tests/unit/test_router_weighted_failover.py @@ -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