diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 449319c1b95e..2e7e7bb0c9d0 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -481,7 +481,7 @@ async def acompletion( additional_args={"complete_input_dict": data}, original_response=str(e), ) - raise AzureOpenAIError(status_code=500, message=str(e)) + raise except Exception as e: message: Final = getattr(e, "message", str(e)) body: Final = getattr(e, "body", None) diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 65354100f58b..334780eda533 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -12,6 +12,7 @@ - {id: reliability.cooldown.429.trips_then_recovers, module: reliability, tier: P0, behavior: cooldown, variant: "429", assertions: [trips_then_recovers], exercised_on: [chat_completions], source: "cooldown_handlers.py:69", rationale: "Cools on 429, avoids hammering exhausted provider"} - {id: reliability.cooldown.auth.trips_then_recovers, module: reliability, tier: P1, behavior: cooldown, variant: auth, assertions: [trips_then_recovers], exercised_on: [chat_completions], source: "cooldown_handlers.py:74", rationale: "Cools on 401 auth error"} - {id: reliability.cooldown.timeout.trips_then_recovers, module: reliability, tier: P1, behavior: cooldown, variant: timeout, assertions: [trips_then_recovers], exercised_on: [chat_completions], source: "cooldown_handlers.py:77", rationale: "Cools on 408 timeout"} +- {id: reliability.cooldown.client_disconnect.stays_healthy, module: reliability, tier: P0, behavior: cooldown, variant: client_disconnect, assertions: [stays_healthy], exercised_on: [chat_completions], source: "llms/azure/azure.py:484", fail_before_fix: proven, rationale: "A client hanging up mid-request under cancel_on_disconnect never benches the Azure deployment it was talking to: the cancellation used to surface as a fake 500 that tripped the cooldown and sent every caller behind it to billed fallbacks (GitHub issues #35329 and #42222)"} - {id: reliability.routing.simple_shuffle.picks_healthy_deployment, module: reliability, tier: P1, behavior: routing, variant: simple_shuffle, assertions: [picks_healthy_deployment], exercised_on: [chat_completions], source: "router_strategy/simple_shuffle.py", rationale: "Baseline weighted/uniform pick"} - {id: reliability.routing.latency_based.picks_lowest_latency, module: reliability, tier: P1, behavior: routing, variant: latency_based, assertions: [picks_lowest_latency], exercised_on: [chat_completions], source: "router_strategy/lowest_latency.py", rationale: "Routes to lowest-latency deployment"} - {id: reliability.routing.cost_based.picks_lowest_cost, module: reliability, tier: P1, behavior: routing, variant: cost_based, assertions: [picks_lowest_cost], exercised_on: [chat_completions], source: "router_strategy/lowest_cost.py", rationale: "Spend-aware routing"} diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index d4978601b200..58817d399a86 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -676,6 +676,38 @@ def send( return streaming_outcome(resp, stream, sent_at=sent_at) +class AbandonedRequest(BaseModel): + """A non-streaming request the client walked away from: the socket was closed + ``after`` seconds in, before the proxy had answered, so the proxy saw a client + disconnect with the upstream call still in flight.""" + + kind: Literal["abandoned"] = "abandoned" + after: float + + +def abandon( + url: URL, *, headers: BaseModel, json: BaseModel, after: float, connect_timeout: float = 10.0 +) -> AbandonedRequest | StreamingResponse: + """POST and hang up ``after`` seconds if no response head has arrived by then, + closing the connection so the proxy observes the disconnect. Returns the + response instead when the proxy answered first, so a test can tell a real + disconnect from a generation that finished too fast to be cancelled.""" + sent_at: Final = time.monotonic() + session: Final = requests.Session() + try: + resp = session.post( + str(url), + headers=_headers(headers), + json=wire_body(json), + timeout=(connect_timeout, after), + ) + except requests.exceptions.ReadTimeout: + return AbandonedRequest(after=after) + finally: + session.close() + return streaming_outcome(resp, False, sent_at=sent_at) + + def stream(url: URL, *, headers: BaseModel, json: BaseModel, timeout: float = 60.0) -> StreamingResponse: """Streaming (SSE) call: consumes the stream counting events, and captures the x-litellm-call-id + content-type headers. Body is elided.""" diff --git a/tests/e2e/gateway/stage_mirror_ci_config.yml b/tests/e2e/gateway/stage_mirror_ci_config.yml index 352caddf5882..2d02fedddaeb 100644 --- a/tests/e2e/gateway/stage_mirror_ci_config.yml +++ b/tests/e2e/gateway/stage_mirror_ci_config.yml @@ -10,6 +10,7 @@ general_settings: store_prompts_in_spend_logs: true database_connection_pool_limit: 10 forward_client_headers_to_llm_api: false + cancel_on_disconnect: true maximum_spend_logs_retention_period: "60d" maximum_spend_logs_cleanup_cron: "0 1 * * *" proxy_budget_rescheduler_min_time: 15 diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 47ef672ebec4..a0f1fa47c67b 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -916,6 +916,23 @@ class RouterSettingsResponse(BaseModel): current_values: RouterCurrentValues +class ConfigListParams(BaseModel): + config_type: Literal["general_settings"] + + +class ConfigField(BaseModel): + """One row of GET /config/list: a general_settings field and the value the + proxy is running with, the two fields a test preconditions on.""" + + model_config = ConfigDict(extra="ignore") + field_name: str + field_value: JsonValue = None + + +class ConfigFieldList(RootModel[tuple[ConfigField, ...]]): + """GET /config/list answers with a bare array of general_settings fields.""" + + class CostMapEntry(BaseModel): model_config = ConfigDict(extra="ignore") litellm_provider: str | None = None @@ -1030,6 +1047,7 @@ class ModelInfoBody(BaseModel): mode: ModelMode | None = None access_groups: list[str] | None = None team_id: str | None = None + allowed_fails: int | None = None allowed_fails_policy: dict[str, int] | None = None diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index c6ede240c3b5..ffde24d9da54 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -47,6 +47,8 @@ AnthropicMessagesResponse, ChatBody, ChatResponse, + ConfigFieldList, + ConfigListParams, CostMap, CostMapEntry, CountTokensBody, @@ -628,6 +630,19 @@ def create_model( provider_live=provider_live, ) + def general_setting_enabled(self, field_name: str) -> bool: + """Whether the proxy is running with the named general_settings flag on, for + a test whose behavior only exists under a config flag the stack has to carry.""" + fields = unwrap( + self.transport.get( + "/config/list", + headers=self.transport.master, + params=ConfigListParams(config_type="general_settings"), + response_type=ConfigFieldList, + ) + ).root + return any(entry.field_name == field_name and entry.field_value is True for entry in fields) + def register_model( self, body: ModelNewBody, listed_for: str | None = None, *, provider_live: bool = False ) -> str: diff --git a/tests/e2e/router/reliability_support.py b/tests/e2e/router/reliability_support.py index 976c05ffceb7..887954bc7fd2 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -42,7 +42,7 @@ CACHING_MODEL = "anthropic/claude-haiku-4-5" CACHING_KEY = "os.environ/ANTHROPIC_API_KEY" -CONTENT_FILTERED_MODEL = "azure/gpt-5.4-nano" +AZURE_MODEL = "azure/gpt-5.4-nano" AZURE_KEY = "os.environ/AZURE_API_KEY" AZURE_BASE = "os.environ/AZURE_API_BASE" AZURE_API_VERSION = "2024-10-21" @@ -111,7 +111,7 @@ def create_content_filtered_deployment(proxy: ProxyClient, name: str) -> str: return proxy.create_model( name, LiteLLMParamsBody( - model=CONTENT_FILTERED_MODEL, + model=AZURE_MODEL, api_key=AZURE_KEY, api_base=AZURE_BASE, api_version=AZURE_API_VERSION, @@ -120,6 +120,30 @@ def create_content_filtered_deployment(proxy: ProxyClient, name: str) -> str: ) +def create_azure_benched_on_first_failure_deployment(proxy: ProxyClient, name: str, cooldown_time: float) -> str: + """The live Azure OpenAI deployment holding all of the group's shuffle weight, + benched on its first failure of any class, with the client's own retries off. + The 500 the proxy used to book against a call the client hung up on carries no + provider body, so litellm maps it to a bare APIError that no named + allowed_fails_policy class covers; the deployment-wide allowed_fails=0 is the + knob that makes that undeserved bench show on the very next call.""" + return proxy.register_model( + ModelNewBody( + model_name=name, + litellm_params=LiteLLMParamsBody( + model=AZURE_MODEL, + api_key=AZURE_KEY, + api_base=AZURE_BASE, + api_version=AZURE_API_VERSION, + max_retries=0, + weight=1, + cooldown_time=cooldown_time, + ), + model_info=ModelInfoBody(allowed_fails=0), + ) + ) + + def create_caching_deployment(proxy: ProxyClient, name: str) -> str: """Register the Anthropic deployment whose prompt cache the affinity check pins to.""" return proxy.create_model(name, LiteLLMParamsBody(model=CACHING_MODEL, api_key=CACHING_KEY, weight=1)) diff --git a/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py b/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py new file mode 100644 index 000000000000..7d46a9800349 --- /dev/null +++ b/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py @@ -0,0 +1,131 @@ +"""Live e2e: a client hanging up mid-request under cancel_on_disconnect never +benches the deployment it was talking to. + +With `general_settings.cancel_on_disconnect: true` the proxy cancels the in-flight +provider call the moment the client's socket closes. The Azure handler used to +turn that cancellation into a fake 500, which the router booked as a deployment +failure: one impatient client benched a healthy deployment and every caller +behind it paid for fallbacks (GitHub issues #35329 and #42222). This cell pins the +fix at the seam a customer sees. The group is the cooldown suite's pair: the live +Azure deployment holding all of the shuffle weight, benched on its first failure +of any class (the fake 500 carried no provider body, so litellm mapped it to a +bare APIError no named policy class covers) with a cooldown long enough to +outlast the test, plus a healthy backup at weight 0 the shuffle can only reach +once the Azure deployment is benched. One cheap call first proves the Azure +deployment answers the key and leaves the key's auth path warm. The test then +asks for a long answer, retries off, and hangs up a few seconds in: the client's +read timeout closes the socket well after the proxy has handed the call to Azure +(a cold virtual-key auth can take a couple of seconds on its own, and a hang-up +that lands before the provider call is in flight cancels nothing the router could +bench, so a shorter window passes vacuously) and well before the answer is done. +After a settle window wide enough for a sibling replica to have read any bench +from Redis, every one of the next calls has to come back 200 from the Azure +deployment itself, named in x-litellm-model-id; a single answer from the backup +means the hang-up was booked as a failure. + +The test reads `cancel_on_disconnect` back from the proxy first: without the flag +the hang-up cancels nothing and the cell would pass vacuously. +""" + +from __future__ import annotations + +import time + +import pytest +from complexity_router_client import ComplexityRouterClient +from e2e_config import unique_marker +from e2e_http import AbandonedRequest, StreamingResponse +from lifecycle import ResourceManager +from models import ChatMessage, ReliabilityChatBody, RouterSettingsOverride +from reliability_support import ( + chat_override, + create_azure_benched_on_first_failure_deployment, + create_zero_weight_backup_deployment, + model_id_of, +) + +pytestmark = pytest.mark.e2e + +CLIENT_HANGS_UP_AFTER_SECONDS = 8.0 +LONG_ANSWER_MAX_TOKENS = 4096 +BENCH_OUTLASTS_TEST_SECONDS = 300.0 +SETTLE_AFTER_HANGUP_SECONDS = 3.0 +CALLS_AFTER_HANGUP = 6 + + +def _say_hi(client: ComplexityRouterClient, key: str, group: str) -> StreamingResponse: + return chat_override( + client.proxy, + key, + group, + f"say hi {unique_marker()}", + override=RouterSettingsOverride(num_retries=0), + ) + + +def _hang_up_mid_answer(client: ComplexityRouterClient, key: str, group: str) -> None: + """Send a request whose answer takes far longer than the client waits, so the + client closes the socket while the provider is still generating.""" + outcome = client.proxy.transport.abandon( + "/chat/completions", + headers=client.proxy.transport.bearer(key), + json=ReliabilityChatBody( + model=group, + messages=[ + ChatMessage( + role="user", + content=f"Write a 3000 word essay on the history of the telegraph. {unique_marker()}", + ) + ], + max_tokens=LONG_ANSWER_MAX_TOKENS, + router_settings_override=RouterSettingsOverride(num_retries=0), + ), + after=CLIENT_HANGS_UP_AFTER_SECONDS, + ) + match outcome: + case AbandonedRequest(): + return + case StreamingResponse(status_code=status_code, body=body): + pytest.fail( + f"the client should have hung up {CLIENT_HANGS_UP_AFTER_SECONDS:.0f}s into a long answer with the " + f"call still in flight, but the proxy answered first with {status_code}: {body[:300]}" + ) + + +class TestReliabilityCancelOnDisconnect: + @pytest.mark.covers("reliability.cooldown.client_disconnect.stays_healthy") + def test_client_hanging_up_never_benches_the_deployment( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + assert client.proxy.general_setting_enabled("cancel_on_disconnect"), ( + "this cell needs general_settings.cancel_on_disconnect: true in the proxy config; without it the " + "hang-up cancels nothing and the bench it guards against can never happen" + ) + + group = f"reliability-cooldown-disconnect-{unique_marker()}" + azure = create_azure_benched_on_first_failure_deployment( + client.proxy, group, cooldown_time=BENCH_OUTLASTS_TEST_SECONDS + ) + resources.defer(lambda: client.proxy.delete_model(azure)) + backup = create_zero_weight_backup_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(backup)) + + warm_up = _say_hi(client, scoped_key, group) + assert warm_up.status_code == 200 and model_id_of(warm_up) == azure, ( + f"before any hang-up the Azure deployment {azure} should answer the group, got {warm_up.status_code} " + f"from {model_id_of(warm_up)!r}: {warm_up.body[:300]}" + ) + + _hang_up_mid_answer(client, scoped_key, group) + time.sleep(SETTLE_AFTER_HANGUP_SECONDS) + + for call in range(1, CALLS_AFTER_HANGUP + 1): + resp = _say_hi(client, scoped_key, group) + assert resp.status_code == 200, ( + f"call {call} after the hang-up should have been a plain 200 from the group, got " + f"{resp.status_code}: {resp.body[:300]}" + ) + assert model_id_of(resp) == azure, ( + f"call {call} after the hang-up should have been served by the Azure deployment {azure}, the proxy " + f"named {model_id_of(resp)!r}: the cancelled call was booked as a failure and benched it" + ) diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 0022c0c4355a..a3eec815441c 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -13,6 +13,7 @@ import e2e_http from e2e_http import ( URL, + AbandonedRequest, AuthHeaders, BinaryStream, NetworkError, @@ -58,6 +59,10 @@ def send( stream: bool = False, ) -> StreamingResponse: ... + def abandon( + self, path: str, *, headers: BaseModel, json: BaseModel, after: float + ) -> AbandonedRequest | StreamingResponse: ... + def get[R: BaseModel]( self, path: str, @@ -243,6 +248,11 @@ def send( timeout=self.request_timeout, ) + def abandon( + self, path: str, *, headers: BaseModel, json: BaseModel, after: float + ) -> AbandonedRequest | StreamingResponse: + return e2e_http.abandon(self._url(path), headers=headers, json=json, after=after) + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: return e2e_http.probe( self._url(path), @@ -420,6 +430,11 @@ def send( ) -> StreamingResponse: return self._route(path).send(path, headers=headers, json=json, params=params, stream=stream) + def abandon( + self, path: str, *, headers: BaseModel, json: BaseModel, after: float + ) -> AbandonedRequest | StreamingResponse: + return self._route(path).abandon(path, headers=headers, json=json, after=after) + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: return self._route(path).probe(path, params=params, headers=headers) diff --git a/tests/test_litellm/llms/azure/test_azure.py b/tests/test_litellm/llms/azure/test_azure.py index 6b6832f623c2..86065c7adc69 100644 --- a/tests/test_litellm/llms/azure/test_azure.py +++ b/tests/test_litellm/llms/azure/test_azure.py @@ -1,10 +1,13 @@ """Tests for litellm/llms/azure/azure.py AzureChatCompletion handler behaviour.""" +import asyncio import time from typing import Final -from openai import AzureOpenAI +import pytest +from openai import AsyncAzureOpenAI, AzureOpenAI +import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.azure import AzureChatCompletion @@ -52,3 +55,25 @@ def test_sync_streaming_stamps_response_headers_on_the_logging_obj() -> None: ) assert logging_obj.model_call_details["response_headers"] == {"x-ms-is-spilled-over": "true"} + + +class _CancelledRawCompletions: + async def create(self, **kwargs): + raise asyncio.CancelledError() + + +@pytest.mark.asyncio +async def test_acompletion_propagates_cancelled_error() -> None: + client = AsyncAzureOpenAI( + api_key="fake-key", + api_version="2024-02-01", + azure_endpoint="https://fake-resource.openai.azure.com", + ) + client.chat.completions.with_raw_response = _CancelledRawCompletions() + + with pytest.raises(asyncio.CancelledError): + await litellm.acompletion( + model="azure/fake-deployment", + messages=[{"role": "user", "content": "hi"}], + client=client, + )