Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion litellm/llms/azure/azure.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
1 change: 1 addition & 0 deletions tests/e2e/coverage_registry/reliability.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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"}
Expand Down
29 changes: 29 additions & 0 deletions tests/e2e/e2e_http.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
1 change: 1 addition & 0 deletions tests/e2e/gateway/stage_mirror_ci_config.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
18 changes: 18 additions & 0 deletions tests/e2e/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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


Expand Down
15 changes: 15 additions & 0 deletions tests/e2e/proxy_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,8 @@
AnthropicMessagesResponse,
ChatBody,
ChatResponse,
ConfigFieldList,
ConfigListParams,
CostMap,
CostMapEntry,
CountTokensBody,
Expand Down Expand Up @@ -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:
Expand Down
25 changes: 23 additions & 2 deletions tests/e2e/router/reliability_support.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -53,6 +53,7 @@
)

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
Expand Down Expand Up @@ -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,
Expand All @@ -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))
Expand Down
139 changes: 139 additions & 0 deletions tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py
Original file line number Diff line number Diff line change
@@ -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.
"""

Comment thread
greptile-apps[bot] marked this conversation as resolved.
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"
)
2 changes: 1 addition & 1 deletion tests/e2e/router/test_reliability_cooldowns_e2e.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
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,
Expand All @@ -57,7 +58,6 @@
pytestmark = pytest.mark.e2e

RECOVERY_GRACE_SECONDS = 10
REPLICA_PROPAGATION_SECONDS = 15.0
PROPAGATION_POLL_SECONDS = 0.25
BENCH_MARGIN_SECONDS = 4.0

Expand Down
15 changes: 15 additions & 0 deletions tests/e2e/transport.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
import e2e_http
from e2e_http import (
URL,
AbandonedRequest,
AuthHeaders,
BinaryStream,
NetworkError,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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)

Expand Down
Loading
Loading