diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 65354100f58..002d231745b 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -11,6 +11,7 @@ - {id: reliability.cooldown.5xx.trips_then_recovers, module: reliability, tier: P0, behavior: cooldown, variant: "5xx", assertions: [trips_then_recovers], exercised_on: [chat_completions], source: "cooldown_handlers.py:40", rationale: "Deployment cools after repeated 5xx, recovers after cooldown_time"} - {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.client_disconnect.stays_healthy, module: reliability, tier: P0, behavior: cooldown, variant: client_disconnect, assertions: [stays_healthy], exercised_on: [chat_completions], source: "litellm/llms/azure/azure.py", fail_before_fix: proven, rationale: "With cancel_on_disconnect on, a client hanging up mid-request cancels the upstream call; that cancellation must not be recorded as a deployment 500 that benches a healthy Azure deployment"} - {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.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"} diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index d4978601b20..7f0940ced05 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.""" @@ -877,7 +909,10 @@ class PreparedForward: def prepare_forward( - method: str, url: str, headers: dict[str, str], body: bytes | None, + method: str, + url: str, + headers: dict[str, str], + body: bytes | None, ) -> PreparedForward | NetworkError: try: with requests.Session() as session: @@ -896,7 +931,8 @@ def forward_prepared_stream(prepared: PreparedForward, timeout: float) -> Stream except requests.RequestException as exc: return NetworkError(message=str(exc)) return StreamHead( - resp.status_code, {name.lower(): value for name, value in resp.headers.items()}, + resp.status_code, + {name.lower(): value for name, value in resp.headers.items()}, primed_steps(_stream_steps(resp)), ) diff --git a/tests/e2e/gateway/stage_mirror_ci_config.yml b/tests/e2e/gateway/stage_mirror_ci_config.yml index 8c8e64443cb..0ce9a4c0be8 100644 --- a/tests/e2e/gateway/stage_mirror_ci_config.yml +++ b/tests/e2e/gateway/stage_mirror_ci_config.yml @@ -1,5 +1,6 @@ general_settings: proxy_batch_write_at: 5 + cancel_on_disconnect: true enable_jwt_auth: true litellm_jwtauth: user_id_jwt_field: sub diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 47ef672ebec..dd01e3ad301 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -1030,6 +1030,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/router/reliability_support.py b/tests/e2e/router/reliability_support.py index 976c05ffceb..790d8c7aee0 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -17,9 +17,6 @@ from __future__ import annotations from collections.abc import Sequence -from pydantic import ValidationError - -from proxy_client import ProxyClient from e2e_config import CHEAP_OPENAI_MODEL, PROXY_BASE_URL, unique_marker from e2e_http import NetworkError, StreamHead, StreamingResponse from models import ( @@ -35,6 +32,8 @@ from models import ( TextContentPart, Usage, ) +from proxy_client import ProxyClient +from pydantic import ValidationError REAL_MODEL = "openai/gpt-5.5" REAL_KEY = "os.environ/OPENAI_API_KEY" @@ -120,6 +119,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: + """A healthy real Azure deployment benched on its first failure of any kind, so a cancellation the proxy + wrongly records as a 500 shows up as the next call landing on the backup.""" + return proxy.register_model( + ModelNewBody( + model_name=name, + litellm_params=LiteLLMParamsBody( + model=CONTENT_FILTERED_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..557103bdfab --- /dev/null +++ b/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py @@ -0,0 +1,99 @@ +"""Live e2e: a client hanging up on an in-flight request must not bench the healthy +deployment that was serving it. + +The proxy runs with `cancel_on_disconnect: true`, so when the client closes the +socket before the answer arrives the proxy cancels the upstream call. That +cancellation is the client's doing, so it must never count as a failure of the +deployment: a deployment that benches on its very first failure of any kind has +to keep serving the next request, and a request served right after the hang-up +has to come from that same deployment rather than its zero-weight backup. + +The disconnect is real: a non-streaming /chat/completions asking the real Azure +OpenAI deployment for a long generation, with the client closing the connection +ABANDON_AFTER_SECONDS in, well before any answer. If Azure ever answers within +that window the test fails loudly rather than passing without a disconnect. +""" + +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 +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 + +ABANDON_AFTER_SECONDS = 2.0 +LONG_GENERATION_MAX_TOKENS = 4000 +FOLLOW_UP_CALLS = 3 +FOLLOW_UP_SPACING_SECONDS = 1.0 +COOLDOWN_SECONDS = 60.0 + + +def _long_generation_prompt(marker: str) -> str: + return ( + f"Write a detailed, multi-chapter short story of at least 3000 words about {marker}. " + "Do not stop early and do not summarize." + ) + + +class TestReliabilityCancelOnDisconnect: + @pytest.mark.covers("reliability.cooldown.client_disconnect.stays_healthy") + def test_client_disconnect_does_not_bench_healthy_deployment( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + group = f"reliability-cancel-on-disconnect-{unique_marker()}" + azure_deployment = create_azure_benched_on_first_failure_deployment( + client.proxy, group, cooldown_time=COOLDOWN_SECONDS + ) + resources.defer(lambda: client.proxy.delete_model(azure_deployment)) + backup_deployment = create_zero_weight_backup_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(backup_deployment)) + + abandoned = client.proxy.transport.abandon( + "/chat/completions", + headers=client.proxy.transport.bearer(scoped_key), + json=ReliabilityChatBody( + model=group, + messages=[ChatMessage(role="user", content=_long_generation_prompt(unique_marker()))], + max_tokens=LONG_GENERATION_MAX_TOKENS, + stream=False, + router_settings_override=RouterSettingsOverride(num_retries=0), + cache={"no-cache": True}, + ), + after=ABANDON_AFTER_SECONDS, + ) + assert isinstance(abandoned, AbandonedRequest), ( + f"the proxy answered within {ABANDON_AFTER_SECONDS}s so the client never disconnected mid-request, " + f"got {abandoned.status_code}: {abandoned.body[:300]}" + ) + + for attempt in range(1, FOLLOW_UP_CALLS + 1): + time.sleep(FOLLOW_UP_SPACING_SECONDS) + resp = chat_override( + client.proxy, + scoped_key, + group, + f"say hi {unique_marker()}", + override=RouterSettingsOverride(num_retries=0), + ) + assert resp.status_code == 200, ( + f"follow-up {attempt}/{FOLLOW_UP_CALLS} should still land on the Azure deployment the client hung up on; " + f"landing on the backup means the cancellation was recorded as a deployment failure and benched it, " + f"got {resp.status_code}: {resp.body[:300]}" + ) + assert model_id_of(resp) == azure_deployment, ( + f"follow-up {attempt}/{FOLLOW_UP_CALLS} should still land on the Azure deployment the client hung up on; " + f"landing on the backup means the cancellation was recorded as a deployment failure and benched it, " + f"got {model_id_of(resp)!r}" + ) 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)