diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 449319c1b95..2e7e7bb0c9d 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -481,7 +481,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): 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 65354100f58..334780eda53 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 d4978601b20..97f1e1671f8 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -676,6 +676,35 @@ def send( return streaming_outcome(resp, stream, sent_at=sent_at) +class AbandonedRequest(BaseModel): + """A non-streaming request whose socket the client closed ``after`` seconds in, + before the proxy had answered.""" + + 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 close the connection ``after`` seconds if no response head has arrived + by then; returns the response instead when the proxy answered first.""" + 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 352caddf588..2d02fedddae 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 3eb18cdf899..2608da2e4fc 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -951,6 +951,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 @@ -1066,6 +1083,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 2f32361e083..76c57452c69 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -47,6 +47,8 @@ from models import ( AnthropicMessagesResponse, ChatBody, ChatResponse, + ConfigFieldList, + ConfigListParams, CostMap, CostMapEntry, CountTokensBody, @@ -630,6 +632,19 @@ class ProxyClient: 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 976c05ffceb..3d5b76f6408 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -42,7 +42,7 @@ REAL_KEY = "os.environ/OPENAI_API_KEY" 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" @@ -53,6 +53,7 @@ CONTENT_POLICY_PROMPT = ( ) COOLDOWN_SECONDS = 30.0 +REPLICA_PROPAGATION_SECONDS = 15.0 # The smallest-context chat model OpenAI still serves (16385 tokens). A prompt # past that limit comes back as a real `context_length_exceeded` 400, which is @@ -111,7 +112,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 +121,26 @@ 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.""" + 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 00000000000..06174e97d20 --- /dev/null +++ b/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py @@ -0,0 +1,139 @@ +"""Live e2e: a client hanging up mid-request under cancel_on_disconnect never +benches the deployment it was talking to. + +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 with a cooldown that +outlasts the test, plus a healthy backup at weight 0 the shuffle only reaches once +the Azure deployment is benched. A cheap call first proves the Azure deployment +answers the key and warms its auth path. The test then asks for an answer far +longer than CLIENT_HANGS_UP_AFTER_SECONDS of generation, retries off, and hangs up +that many seconds in: late enough that the proxy has handed the call to Azure (a +hang-up before the provider call is in flight cancels nothing the router could +bench, so the cell would pass vacuously). An answer that comes back inside the +window proves nothing and benches nothing either, since a success never counts +against the deployment, so the cell asks again up to HANG_UP_ATTEMPTS times and +fails out loud naming the window only when every ask came back early. After the +cooldown suite's replica propagation window, 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 ( + REPLICA_PROPAGATION_SECONDS, + 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 = 5.0 +HANG_UP_ATTEMPTS = 3 +LONG_ANSWER_MAX_TOKENS = 16384 +BENCH_OUTLASTS_TEST_SECONDS = 300.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 _ask_for_a_long_answer_then_hang_up( + client: ComplexityRouterClient, key: str, group: str +) -> AbandonedRequest | StreamingResponse: + return client.proxy.transport.abandon( + "/chat/completions", + headers=client.proxy.transport.bearer(key), + json=ReliabilityChatBody( + model=group, + messages=[ + ChatMessage( + role="user", + content=( + "Write an essay on the history of the telegraph with one section per decade from the 1830s " + f"to the 2020s, each section at least 300 words. {unique_marker()}" + ), + ) + ], + max_tokens=LONG_ANSWER_MAX_TOKENS, + router_settings_override=RouterSettingsOverride(num_retries=0), + ), + after=CLIENT_HANGS_UP_AFTER_SECONDS, + ) + + +def _hang_up_mid_answer(client: ComplexityRouterClient, key: str, group: str) -> None: + for attempt in range(1, HANG_UP_ATTEMPTS + 1): + match _ask_for_a_long_answer_then_hang_up(client, key, group): + case AbandonedRequest(): + return + case StreamingResponse(status_code=200): + continue + case StreamingResponse(status_code=status_code, body=body): + pytest.fail( + f"hang-up attempt {attempt} should have found the long answer still in flight after " + f"{CLIENT_HANGS_UP_AFTER_SECONDS:.0f}s, but the proxy answered {status_code}: {body[:300]}" + ) + pytest.fail( + f"the proxy answered all {HANG_UP_ATTEMPTS} long asks within {CLIENT_HANGS_UP_AFTER_SECONDS:.0f}s, so the " + "client never hung up with a call still in flight and the bench this cell guards against could not happen" + ) + + +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(REPLICA_PROPAGATION_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/router/test_reliability_cooldowns_e2e.py b/tests/e2e/router/test_reliability_cooldowns_e2e.py index 5b5cec09f06..769971e1533 100644 --- a/tests/e2e/router/test_reliability_cooldowns_e2e.py +++ b/tests/e2e/router/test_reliability_cooldowns_e2e.py @@ -43,6 +43,7 @@ from lifecycle import ResourceManager from models import KeyGenerateBody, RouterSettingsOverride from reliability_support import ( COOLDOWN_SECONDS, + REPLICA_PROPAGATION_SECONDS, chat_override, create_always_5xx_deployment, create_always_rate_limited_deployment, @@ -57,7 +58,6 @@ from reliability_support import ( pytestmark = pytest.mark.e2e RECOVERY_GRACE_SECONDS = 10 -REPLICA_PROPAGATION_SECONDS = 15.0 PROPAGATION_POLL_SECONDS = 0.25 BENCH_MARGIN_SECONDS = 4.0 diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 0022c0c4355..a3eec815441 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -13,6 +13,7 @@ from typing import Protocol import e2e_http from e2e_http import ( URL, + AbandonedRequest, AuthHeaders, BinaryStream, NetworkError, @@ -58,6 +59,10 @@ class Transport(Protocol): 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 @@ class HttpTransport: 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 @@ class SplitTransport: ) -> 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 6b6832f623c..86065c7adc6 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, + )