From 04892c814df428757e16f4b508b7d15e0981333d Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 27 Jul 2026 16:55:06 -0700 Subject: [PATCH 1/5] test(e2e): harden harness and tests against data-plane pod churn A stage autoscaler scale-down produced a 2s window of ALB 502s that killed six budget tests on their first management call, and a freshly scaled-up pod that had not run its 30s DB object sync yet failed two MCP tests and one prometheus cardinality test. Retry transient gateway errors (502/503/504, connection errors) once at the shared e2e_http dispatch seam, poll MCP server registration to the poll deadline instead of asserting a single-shot listing, anchor the MCP guardrail full-sync wait to the later of the guardrail and server writes, and turn the prometheus alias poll into a drive-and-scrape convergence loop that re-sends traffic for missing aliases and unions results across scrapes --- tests/e2e/e2e_http.py | 295 ++++++++++++------ .../test_prometheus_cardinality_e2e.py | 27 +- tests/e2e/mcp/mcp_client.py | 21 ++ tests/e2e/mcp/test_mcp_access_group_e2e.py | 7 +- tests/e2e/mcp/test_mcp_guardrail_e2e.py | 59 +++- tests/e2e/mcp/test_mcp_key_access_e2e.py | 9 +- tests/e2e/test_e2e_http.py | 165 ++++++++++ 7 files changed, 454 insertions(+), 129 deletions(-) create mode 100644 tests/e2e/test_e2e_http.py diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 74e57f86b885..315759798092 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -10,6 +10,9 @@ from __future__ import annotations +import time +from collections.abc import Callable +from dataclasses import dataclass from typing import Generic, Iterator, Literal, NewType, TypeVar, cast import pytest @@ -229,6 +232,90 @@ def _params(params: BaseModel | None) -> dict[str, str]: return {key: str(value) for key, value in dumped.items()} +_TRANSIENT_STATUSES: frozenset[int] = frozenset({502, 503, 504}) + + +@dataclass(frozen=True, slots=True) +class _RetryPolicy: + budget_seconds: float + initial_backoff_seconds: float + + +_GATEWAY_RETRY = _RetryPolicy(budget_seconds=15.0, initial_backoff_seconds=0.5) + + +@dataclass(frozen=True, slots=True) +class _Delivered: + response: requests.Response + + +@dataclass(frozen=True, slots=True) +class _Undelivered: + error: requests.RequestException + + +type _Attempt = _Delivered | _Undelivered + + +def _attempt(call: Callable[[], requests.Response]) -> _Attempt: + try: + return _Delivered(response=call()) + except requests.RequestException as exc: + return _Undelivered(error=exc) + + +def _is_transient(attempt: _Attempt) -> bool: + """True only when the gateway dropped the call instead of the app answering it: a + 502/503/504 from the load balancer while pods roll or scale down, or that same + event one layer lower as a refused/reset connection. A 429, a 500 and every 4xx + are the app's own answer and must reach the caller untouched, because tests assert + on them (rate limits, budget blocks, error mapping).""" + match attempt: + case _Delivered(response=response): + return response.status_code in _TRANSIENT_STATUSES + case _Undelivered(error=error): + return isinstance(error, requests.ConnectionError) + + +def _retry( + call: Callable[[], requests.Response], + *, + deadline: float, + backoff: float, + last: _Attempt, +) -> _Attempt: + remaining = deadline - time.monotonic() + if remaining <= 0: + return last + time.sleep(min(backoff, remaining)) + attempt = _attempt(call) + if not _is_transient(attempt): + return attempt + return _retry(call, deadline=deadline, backoff=backoff * 2, last=attempt) + + +def _dispatch( + call: Callable[[], requests.Response], + *, + policy: _RetryPolicy | None = _GATEWAY_RETRY, +) -> _Attempt: + """The single seam every non-streaming verb sends through, so one bounded retry + covers all of them. A transient gateway failure is retried with exponential + backoff until the policy's budget is spent, after which the last attempt is + handed back unchanged - a permanent failure therefore looks exactly like it did + before, just later. `policy=None` opts a call out (the streaming paths, whose + body is consumed live).""" + first = _attempt(call) + if policy is None or not _is_transient(first): + return first + return _retry( + call, + deadline=time.monotonic() + policy.budget_seconds, + backoff=policy.initial_backoff_seconds, + last=first, + ) + + def _classify[R: BaseModel]( resp: requests.Response, response_type: type[R] ) -> Result[R]: @@ -244,6 +331,14 @@ def _classify[R: BaseModel]( return ValidationError(message=str(exc)) +def _result[R: BaseModel](attempt: _Attempt, response_type: type[R]) -> Result[R]: + match attempt: + case _Delivered(response=response): + return _classify(response, response_type) + case _Undelivered(error=error): + return NetworkError(message=str(error)) + + def post[R: BaseModel]( url: URL, *, @@ -252,16 +347,17 @@ def post[R: BaseModel]( response_type: type[R], timeout: float = 30.0, ) -> Result[R]: - try: - resp = requests.post( - str(url), - headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), - timeout=timeout, - ) - except requests.RequestException as exc: - return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return _result( + _dispatch( + lambda: requests.post( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + ), + response_type, + ) def get[R: BaseModel]( @@ -272,16 +368,17 @@ def get[R: BaseModel]( response_type: type[R], timeout: float = 30.0, ) -> Result[R]: - try: - resp = requests.get( - str(url), - headers=_headers(headers), - params=params.model_dump(by_alias=True, exclude_none=True), - timeout=timeout, - ) - except requests.RequestException as exc: - return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return _result( + _dispatch( + lambda: requests.get( + str(url), + headers=_headers(headers), + params=params.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + ), + response_type, + ) def get_external[R: BaseModel]( @@ -293,15 +390,16 @@ def get_external[R: BaseModel]( """GET an absolute URL outside the proxy (e.g. a public /.well-known document). Unlike the transport wrappers there is no proxy base url and no proxy auth; the response still gets the same tagged-union classification as every other call.""" - try: - resp = requests.get( - url, - headers={"Accept": "application/json"}, - timeout=timeout, - ) - except requests.RequestException as exc: - return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return _result( + _dispatch( + lambda: requests.get( + url, + headers={"Accept": "application/json"}, + timeout=timeout, + ) + ), + response_type, + ) def delete[R: BaseModel]( @@ -313,17 +411,18 @@ def delete[R: BaseModel]( params: BaseModel | None = None, timeout: float = 30.0, ) -> Result[R]: - try: - resp = requests.delete( - str(url), - headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), - params=_params(params), - timeout=timeout, - ) - except requests.RequestException as exc: - return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return _result( + _dispatch( + lambda: requests.delete( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + params=_params(params), + timeout=timeout, + ) + ), + response_type, + ) def patch[R: BaseModel]( @@ -334,16 +433,17 @@ def patch[R: BaseModel]( response_type: type[R], timeout: float = 30.0, ) -> Result[R]: - try: - resp = requests.patch( - str(url), - headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), - timeout=timeout, - ) - except requests.RequestException as exc: - return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return _result( + _dispatch( + lambda: requests.patch( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + ), + response_type, + ) def put[R: BaseModel]( @@ -354,31 +454,34 @@ def put[R: BaseModel]( response_type: type[R], timeout: float = 30.0, ) -> Result[R]: - try: - resp = requests.put( - str(url), - headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), - timeout=timeout, - ) - except requests.RequestException as exc: - return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return _result( + _dispatch( + lambda: requests.put( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + ), + response_type, + ) def probe( url: URL, *, headers: BaseModel, params: BaseModel, timeout: float = 30.0 ) -> ProbeResult: - try: - resp = requests.get( + match _dispatch( + lambda: requests.get( str(url), headers=_headers(headers), params=params.model_dump(by_alias=True, exclude_none=True), timeout=timeout, ) - except requests.RequestException as exc: - return ProbeResult(status_code=-1, body=str(exc)) - return ProbeResult(status_code=resp.status_code, body=resp.text) + ): + case _Delivered(response=response): + return ProbeResult(status_code=response.status_code, body=response.text) + case _Undelivered(error=error): + return ProbeResult(status_code=-1, body=str(error)) def _parse_response_cost(resp: requests.Response) -> float | None: @@ -455,18 +558,21 @@ def send( x-litellm-call-id header. For native/passthrough bodies and for calls judged by status rather than a typed JSON model (e.g. a budget block is a non-2xx). With ``stream=True`` the SSE body is consumed and its events counted instead.""" - try: - resp = requests.post( + match _dispatch( + lambda: requests.post( str(url), headers=_headers(headers), params=_params(params), json=json.model_dump(by_alias=True, exclude_none=True), stream=stream, timeout=timeout, - ) - except requests.RequestException as exc: - return StreamingResponse(status_code=-1, body=str(exc)) - return _streaming_outcome(resp, stream) + ), + policy=None if stream else _GATEWAY_RETRY, + ): + case _Delivered(response=response): + return _streaming_outcome(response, stream) + case _Undelivered(error=error): + return StreamingResponse(status_code=-1, body=str(error)) def stream( @@ -496,18 +602,19 @@ def upload[R: BaseModel]( routing (e.g. ?model=). requests sets the multipart Content-Type itself.""" dumped: dict[str, object] = form.model_dump(by_alias=True, exclude_none=True) data = {key: str(value) for key, value in dumped.items()} - try: - resp = requests.post( - str(url), - headers=_headers(headers), - params=_params(params), - data=data, - files={file_field: (filename, content, file_content_type)}, - timeout=timeout, - ) - except requests.RequestException as exc: - return NetworkError(message=str(exc)) - return _classify(resp, response_type) + return _result( + _dispatch( + lambda: requests.post( + str(url), + headers=_headers(headers), + params=_params(params), + data=data, + files={file_field: (filename, content, file_content_type)}, + timeout=timeout, + ) + ), + response_type, + ) def stream_binary( @@ -563,13 +670,15 @@ def download( ) -> StreamingResponse: """Raw GET for file content (/v1/files/{id}/content): provider-native bytes, no schema. Returns the decoded body and the x-litellm-call-id header.""" - try: - resp = requests.get(str(url), headers=_headers(headers), timeout=timeout) - except requests.RequestException as exc: - return StreamingResponse(status_code=-1, body=str(exc)) - return StreamingResponse( - status_code=resp.status_code, - call_id=_hdr(resp, "x-litellm-call-id"), - content_type=_hdr(resp, "content-type"), - body=resp.text, - ) + match _dispatch( + lambda: requests.get(str(url), headers=_headers(headers), timeout=timeout) + ): + case _Delivered(response=response): + return StreamingResponse( + status_code=response.status_code, + call_id=_hdr(response, "x-litellm-call-id"), + content_type=_hdr(response, "content-type"), + body=response.text, + ) + case _Undelivered(error=error): + return StreamingResponse(status_code=-1, body=str(error)) diff --git a/tests/e2e/logging/test_prometheus_cardinality_e2e.py b/tests/e2e/logging/test_prometheus_cardinality_e2e.py index 44e3d93c07b5..72747a6f82a1 100644 --- a/tests/e2e/logging/test_prometheus_cardinality_e2e.py +++ b/tests/e2e/logging/test_prometheus_cardinality_e2e.py @@ -9,8 +9,10 @@ the aliases and fail here. Scraping goes through ``transport.probe`` (raw text) and is parsed with -prometheus_client; the metric is eventually consistent (it increments on the -success-logging callback), so the scrape polls to a deadline. +prometheus_client. ``/metrics`` is per-pod behind a round-robin LB and the metric +is eventually consistent (it increments on the success-logging callback), so the +poll re-drives each still-missing alias with fresh traffic and unions the aliases +seen across scrapes until the deadline. """ from __future__ import annotations @@ -48,23 +50,34 @@ def test_distinct_key_aliases_produce_distinct_series( self, client: LoggingClient, resources: ResourceManager ) -> None: aliases = tuple(f"e2e-prom-{unique_marker()}" for _ in range(DISTINCT_KEYS)) - for alias in aliases: + + def provisioned_key(alias: str) -> str: key = client.key_with_alias(alias, models=[DRIVER_MODEL]) - resources.defer(lambda k=key: client.delete_key(k)) - response = client.chat(key, DRIVER_MODEL, f"reply with one word {alias}") + resources.defer(lambda: client.delete_key(key)) + return key + + def drive(alias: str, key: str) -> None: + response = client.chat(key, DRIVER_MODEL, f"reply with one word {alias} {unique_marker()}") assert response.model, f"driver call for {alias} returned no model: {response}" + keys_by_alias = {alias: provisioned_key(alias) for alias in aliases} + for alias, key in keys_by_alias.items(): + drive(alias, key) + wanted = frozenset(aliases) deadline = time.monotonic() + client.proxy.poll_timeout seen: frozenset[str] = frozenset() while time.monotonic() < deadline: - seen = _aliases_in_metric(client.scrape_metrics(), REQUESTS_METRIC, ALIAS_LABEL) + seen = seen | _aliases_in_metric(client.scrape_metrics(), REQUESTS_METRIC, ALIAS_LABEL) if wanted <= seen: break + for alias in sorted(wanted - seen): + drive(alias, keys_by_alias[alias]) time.sleep(client.proxy.poll_interval) missing = wanted - seen assert not missing, ( - f"{REQUESTS_METRIC} is missing a per-key series for aliases {sorted(missing)}; " + f"{REQUESTS_METRIC} never exposed a per-key series for aliases {sorted(missing)} " + f"on any scraped pod within the deadline despite repeated driver calls; " f"each distinct {ALIAS_LABEL} must grow its own series" ) diff --git a/tests/e2e/mcp/mcp_client.py b/tests/e2e/mcp/mcp_client.py index 4b1725bb2058..6d0f6ddc7602 100644 --- a/tests/e2e/mcp/mcp_client.py +++ b/tests/e2e/mcp/mcp_client.py @@ -195,6 +195,27 @@ def registered_servers(self) -> list[McpServerRow]: ) ).root + def await_registered(self, server_id: str) -> None: + """Poll /v1/mcp/server until `server_id` is listed. Fails at poll_timeout. + + The DB row exists the moment registration returns, but a data-plane pod + answers the listing from a registry it refreshes on a periodic DB sync, so a + pod that joined the load balancer after the write reports the server as + absent until its first sync. + """ + deadline = time.monotonic() + self.proxy.poll_timeout + while True: + registered = frozenset(row.server_id for row in self.registered_servers()) + if server_id in registered: + return + if time.monotonic() >= deadline: + raise AssertionError( + f"registered server {server_id} still absent from /v1/mcp/server " + f"{self.proxy.poll_timeout}s after registration (the data plane never synced " + f"the row): {registered}" + ) + time.sleep(self.proxy.poll_interval) + def generate_key( self, *, diff --git a/tests/e2e/mcp/test_mcp_access_group_e2e.py b/tests/e2e/mcp/test_mcp_access_group_e2e.py index d7ff9736896c..1b53d1ca0b4e 100644 --- a/tests/e2e/mcp/test_mcp_access_group_e2e.py +++ b/tests/e2e/mcp/test_mcp_access_group_e2e.py @@ -44,12 +44,7 @@ def test_access_group_scopes_tool_selection( ) resources.defer(lambda: client.proxy.delete_key(other)) - granted_tools = unwrap(client.list_tools(granted)) - assert granted_tools.tool_name_containing(server_id, SEARCH_LOGS_TOOL) is not None, ( - f"key granted access group {group} did not see the tagged server's tool " - f"(upstream dead or access-group grant not applied): " - f"{granted_tools.tool_names_for_server(server_id)}" - ) + _ = client.await_tool(granted, server_id, SEARCH_LOGS_TOOL) other_tools = unwrap(client.list_tools(other)).tool_names_for_server(server_id) assert other_tools == frozenset(), ( diff --git a/tests/e2e/mcp/test_mcp_guardrail_e2e.py b/tests/e2e/mcp/test_mcp_guardrail_e2e.py index dcab235465f4..9e3a8c483956 100644 --- a/tests/e2e/mcp/test_mcp_guardrail_e2e.py +++ b/tests/e2e/mcp/test_mcp_guardrail_e2e.py @@ -29,11 +29,11 @@ pytestmark = pytest.mark.e2e # Stage runs several data-plane pods behind the shared key, and each picks up a -# newly registered guardrail only on its next periodic DB sync (~30s in +# newly registered guardrail or MCP server only on its next periodic DB sync (~30s in # proxy_server.py). Every pod is guaranteed to have refreshed only once a full sync -# interval has elapsed since the create; before then a banned call routed to a -# lagging pod passes through as legitimate in-flight propagation, not a leak. -GUARDRAIL_FULL_SYNC_SECONDS = 40.0 +# interval has elapsed since the later of those two writes; before then a banned call +# routed to a lagging pod passes through as legitimate in-flight propagation, not a leak. +FULL_SYNC_SECONDS = 40.0 POST_SYNC_VERIFICATION_CALLS = 4 @@ -53,6 +53,30 @@ def _poll_until_blocked( return last +def _pod_lacks_mcp_server(result: Result[McpCallToolResponse]) -> bool: + """True when the pod that served the call answered as though the MCP server or its + tool does not exist (500 "Tool ... not found"), i.e. its MCP registry has not synced + yet and the request never reached the guardrail at all.""" + if not isinstance(result, UnknownApiError) or result.status_code != 500: + return False + body = result.body.lower() + return "not found" in body and ("tool" in body or "server" in body) + + +def _search_on_synced_pod( + search: Callable[[str], Result[McpCallToolResponse]], query: str, client: McpClient +) -> Result[McpCallToolResponse]: + """Issue `query`, retrying to the poll deadline only while the serving pod does not + know the MCP server yet. Every other outcome, guardrail block or pass-through, comes + back untouched so the caller's assertion still decides it.""" + deadline = time.monotonic() + client.proxy.poll_timeout + last = search(query) + while _pod_lacks_mcp_server(last) and time.monotonic() < deadline: + time.sleep(client.proxy.poll_interval) + last = search(query) + return last + + class TestMcpToolCallGuardrail: @pytest.mark.covers( "guardrail.litellm_content_filter.pre_mcp_call.blocks", @@ -72,6 +96,7 @@ def test_content_filter_blocks_banned_keyword_in_tool_args( resources.defer(lambda: client.delete_guardrail(guardrail_id)) server_id = register_datadog_mcp(client, resources) + server_registered_at = time.monotonic() key = client.generate_key(user_id=f"e2e-mcp-guard-{marker}", mcp_servers=[server_id]) resources.defer(lambda: client.proxy.delete_key(key)) @@ -105,31 +130,33 @@ def search(query: str) -> Result[McpCallToolResponse]: case _: pytest.fail( "content_filter never blocked the banned keyword on the MCP tool call within " - f"{client.proxy.poll_timeout}s (guardrail sync to the data plane never landed); " - f"last result: {blocked}" + f"{client.proxy.poll_timeout}s (the guardrail or the MCP server never synced to " + f"the data plane); last result: {blocked}" ) # The block above only proves the one pod that served it has synced; another # pod could still lack the guardrail and let the banned call reach Datadog. - # Wait out the full sync interval from the create so every pod has refreshed - # from the DB, then require the banned call to stay blocked across several - # attempts. A pass-through now is a genuine partial-propagation leak, not a - # race. Client load balancing still can't guarantee every pod is hit, so this + # Wait out the full sync interval from the later of the guardrail create and the + # MCP server registration (each syncs on its own clock, so the earlier write's + # deadline can elapse while a pod still lacks the other) so every pod has + # refreshed from the DB, then require the banned call to stay blocked across + # several attempts. A pass-through now is a genuine partial-propagation leak, not + # a race. Client load balancing still can't guarantee every pod is hit, so this # samples several worker selections rather than proving all pods synced. - sync_remaining = guardrail_created_at + GUARDRAIL_FULL_SYNC_SECONDS - time.monotonic() + sync_remaining = max(guardrail_created_at, server_registered_at) + FULL_SYNC_SECONDS - time.monotonic() if sync_remaining > 0: time.sleep(sync_remaining) for attempt in range(1, POST_SYNC_VERIFICATION_CALLS + 1): - reblocked = search(f"still about {banned_keyword} #{attempt}") + reblocked = _search_on_synced_pod(search, f"still about {banned_keyword} #{attempt}", client) assert isinstance(reblocked, UnknownApiError) and reblocked.status_code == 400, ( - "after the guardrail sync interval every data-plane pod must block the banned " - f"keyword, but attempt {attempt} of {POST_SYNC_VERIFICATION_CALLS} was allowed " - f"through (a pod still lacks the guardrail): {reblocked}" + "after the sync interval every data-plane pod must block the banned keyword, but " + f"attempt {attempt} of {POST_SYNC_VERIFICATION_CALLS} was not blocked (a pod still " + f"lacks the guardrail, or never synced the MCP server): {reblocked}" ) if attempt < POST_SYNC_VERIFICATION_CALLS: time.sleep(client.proxy.poll_interval) - allowed = search(f"e2e-clean-{marker}") + allowed = _search_on_synced_pod(search, f"e2e-clean-{marker}", client) match allowed: case Success(data=result): assert result.is_error is not True, ( diff --git a/tests/e2e/mcp/test_mcp_key_access_e2e.py b/tests/e2e/mcp/test_mcp_key_access_e2e.py index 4aeb811a64f4..35c864c07d8d 100644 --- a/tests/e2e/mcp/test_mcp_key_access_e2e.py +++ b/tests/e2e/mcp/test_mcp_key_access_e2e.py @@ -30,11 +30,6 @@ def _key(client: McpClient, resources: ResourceManager, *, mcp_servers: list[str return key -def _assert_registered(client: McpClient, server_id: str) -> None: - registered = {row.server_id for row in client.registered_servers()} - assert server_id in registered, f"registered server {server_id} absent from /v1/mcp/server: {registered}" - - class TestMcpKeyWithoutAccessIsDenied: @pytest.mark.covers("mcp.list_tools.api_key.denied_without_permission") def test_list_tools_denied_without_permission( @@ -43,7 +38,7 @@ def test_list_tools_denied_without_permission( resources: ResourceManager, ) -> None: server_id = register_datadog_mcp(client, resources) - _assert_registered(client, server_id) + client.await_registered(server_id) permitted_key = _key(client, resources, mcp_servers=[server_id]) denied_key = _key(client, resources, mcp_servers=None) @@ -63,7 +58,7 @@ def test_call_tool_denied_without_permission( resources: ResourceManager, ) -> None: server_id = register_datadog_mcp(client, resources) - _assert_registered(client, server_id) + client.await_registered(server_id) permitted_key = _key(client, resources, mcp_servers=[server_id]) denied_key = _key(client, resources, mcp_servers=None) diff --git a/tests/e2e/test_e2e_http.py b/tests/e2e/test_e2e_http.py new file mode 100644 index 000000000000..8e2f46b794b6 --- /dev/null +++ b/tests/e2e/test_e2e_http.py @@ -0,0 +1,165 @@ +"""Harness coverage for the transport's transient-failure retry (no proxy needed). + +A pod scale-down makes the load balancer answer 502 for a second or two, which used +to kill whatever test was mid-call. These pin which failures are retried and, just +as importantly, which are not: a 429 or a 500 is the app's own answer that +rate-limit and error-mapping tests assert on, so it must arrive unretried. + +The scripted server serves real HTTP on an ephemeral port and records every hit, so +"did it retry" is a request count rather than a timing guess. +""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Generator, Sequence +from contextlib import contextmanager +from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +from pydantic import BaseModel + +from e2e_http import URL, NoBody, Result, post, probe, send + +DROP_CONNECTION = 0 +SLOW_RESPONSE_SECONDS = 2.0 +RETRY_BUDGET_CEILING_SECONDS = 40.0 + + +class OkBody(BaseModel): + ok: bool + + +class QuietServer(ThreadingHTTPServer): + """A client that times out or walks away mid-response is the point of some of + these tests, so the resulting broken pipe is expected, not a server crash.""" + + daemon_threads = True + + def handle_error(self, request: object, client_address: object) -> None: + return + + +@dataclass(frozen=True, slots=True) +class ScriptedServer: + url: URL + hits: list[int] = field(default_factory=list) + + +@contextmanager +def scripted_server(statuses: Sequence[int], *, slow: bool = False) -> Generator[ScriptedServer]: + """Answer with `statuses` in order, repeating the last one forever. A status of + DROP_CONNECTION closes the socket without answering (what a load balancer does + when it drops a backend mid-request); `slow` stalls every response so the client + can hit its own read timeout.""" + scripted = iter(statuses) + hits: list[int] = [] + + class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def _answer(self) -> None: + status = next(scripted, statuses[-1]) + hits.append(status) + if slow: + time.sleep(SLOW_RESPONSE_SECONDS) + if status == DROP_CONNECTION: + self.close_connection = True + return + body = b'{"ok": true}' if status == 200 else b'{"error": "transient"}' + self.send_response(status) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(body))) + self.end_headers() + _ = self.wfile.write(body) + + do_GET = _answer + do_POST = _answer + + def log_message(self, format: str, *args: object) -> None: + return + + server = QuietServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield ScriptedServer(url=URL(f"http://127.0.0.1:{server.server_port}/anything"), hits=hits) + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +def _post(url: URL, timeout: float = 30.0) -> Result[OkBody]: + return post(url, headers=NoBody(), json=NoBody(), response_type=OkBody, timeout=timeout) + + +class TestTransientRetry: + def test_post_retries_a_bad_gateway_until_it_succeeds(self) -> None: + with scripted_server((502, 503, 200)) as server: + result = _post(server.url) + assert result.kind == "success", result + assert server.hits == [502, 503, 200] + + def test_post_retries_a_dropped_connection_until_it_succeeds(self) -> None: + with scripted_server((DROP_CONNECTION, 200)) as server: + result = _post(server.url) + assert result.kind == "success", result + assert server.hits == [DROP_CONNECTION, 200] + + def test_post_gives_up_on_a_permanent_bad_gateway_with_the_same_error_as_before(self) -> None: + started = time.monotonic() + with scripted_server((502,)) as server: + result = _post(server.url) + elapsed = time.monotonic() - started + assert result.kind == "unknown", result + assert result.status_code == 502 + assert "transient" in result.body + assert len(server.hits) >= 4, server.hits + assert elapsed < RETRY_BUDGET_CEILING_SECONDS, elapsed + + def test_post_does_not_retry_a_rate_limit(self) -> None: + with scripted_server((429, 200)) as server: + result = _post(server.url) + assert result.kind == "rate_limited", result + assert server.hits == [429] + + def test_post_does_not_retry_a_server_error(self) -> None: + with scripted_server((500, 200)) as server: + result = _post(server.url) + assert result.kind == "unknown", result + assert result.status_code == 500 + assert server.hits == [500] + + def test_post_does_not_retry_a_client_error(self) -> None: + with scripted_server((400, 200)) as server: + result = _post(server.url) + assert result.kind == "unknown", result + assert result.status_code == 400 + assert server.hits == [400] + + def test_post_does_not_retry_a_read_timeout(self) -> None: + with scripted_server((200,), slow=True) as server: + result = _post(server.url, timeout=0.25) + assert result.kind == "network", result + assert server.hits == [200] + + def test_probe_retries_a_bad_gateway_until_it_succeeds(self) -> None: + with scripted_server((503, 200)) as server: + result = probe(server.url, headers=NoBody(), params=NoBody()) + assert result.status_code == 200, result + assert result.healthy + assert server.hits == [503, 200] + + def test_send_retries_a_bad_gateway_until_it_succeeds(self) -> None: + with scripted_server((504, 200)) as server: + result = send(server.url, headers=NoBody(), json=NoBody()) + assert result.ok, result + assert server.hits == [504, 200] + + def test_send_leaves_a_streaming_call_alone(self) -> None: + with scripted_server((502, 200)) as server: + result = send(server.url, headers=NoBody(), json=NoBody(), stream=True) + assert result.status_code == 502, result + assert server.hits == [502] From 3c4d5902a208a6888063a16e59e16c3023d90679 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 27 Jul 2026 17:04:09 -0700 Subject: [PATCH 2/5] test(e2e): drain request body in retry stub handler so keep-alive reuse cannot misparse leftovers as requests --- tests/e2e/test_e2e_http.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/e2e/test_e2e_http.py b/tests/e2e/test_e2e_http.py index 8e2f46b794b6..5a0cc601dfe5 100644 --- a/tests/e2e/test_e2e_http.py +++ b/tests/e2e/test_e2e_http.py @@ -60,6 +60,8 @@ class Handler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" def _answer(self) -> None: + length = int(self.headers.get("content-length") or 0) + _ = self.rfile.read(length) status = next(scripted, statuses[-1]) hits.append(status) if slow: From 793fa583f319952c3ce3e2b35b08b3be60770271 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 27 Jul 2026 17:04:38 -0700 Subject: [PATCH 3/5] revert(e2e): drop the transient-502 retry seam A raw 502 during a pod scale-down is what a real client sees, so the suite retrying past it hides an availability gap instead of flagging it. The gateway-side fix is graceful drain on the deployment; until then the failures are signal --- tests/e2e/e2e_http.py | 295 ++++++++++++------------------------- tests/e2e/test_e2e_http.py | 167 --------------------- 2 files changed, 93 insertions(+), 369 deletions(-) delete mode 100644 tests/e2e/test_e2e_http.py diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 315759798092..74e57f86b885 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -10,9 +10,6 @@ from __future__ import annotations -import time -from collections.abc import Callable -from dataclasses import dataclass from typing import Generic, Iterator, Literal, NewType, TypeVar, cast import pytest @@ -232,90 +229,6 @@ def _params(params: BaseModel | None) -> dict[str, str]: return {key: str(value) for key, value in dumped.items()} -_TRANSIENT_STATUSES: frozenset[int] = frozenset({502, 503, 504}) - - -@dataclass(frozen=True, slots=True) -class _RetryPolicy: - budget_seconds: float - initial_backoff_seconds: float - - -_GATEWAY_RETRY = _RetryPolicy(budget_seconds=15.0, initial_backoff_seconds=0.5) - - -@dataclass(frozen=True, slots=True) -class _Delivered: - response: requests.Response - - -@dataclass(frozen=True, slots=True) -class _Undelivered: - error: requests.RequestException - - -type _Attempt = _Delivered | _Undelivered - - -def _attempt(call: Callable[[], requests.Response]) -> _Attempt: - try: - return _Delivered(response=call()) - except requests.RequestException as exc: - return _Undelivered(error=exc) - - -def _is_transient(attempt: _Attempt) -> bool: - """True only when the gateway dropped the call instead of the app answering it: a - 502/503/504 from the load balancer while pods roll or scale down, or that same - event one layer lower as a refused/reset connection. A 429, a 500 and every 4xx - are the app's own answer and must reach the caller untouched, because tests assert - on them (rate limits, budget blocks, error mapping).""" - match attempt: - case _Delivered(response=response): - return response.status_code in _TRANSIENT_STATUSES - case _Undelivered(error=error): - return isinstance(error, requests.ConnectionError) - - -def _retry( - call: Callable[[], requests.Response], - *, - deadline: float, - backoff: float, - last: _Attempt, -) -> _Attempt: - remaining = deadline - time.monotonic() - if remaining <= 0: - return last - time.sleep(min(backoff, remaining)) - attempt = _attempt(call) - if not _is_transient(attempt): - return attempt - return _retry(call, deadline=deadline, backoff=backoff * 2, last=attempt) - - -def _dispatch( - call: Callable[[], requests.Response], - *, - policy: _RetryPolicy | None = _GATEWAY_RETRY, -) -> _Attempt: - """The single seam every non-streaming verb sends through, so one bounded retry - covers all of them. A transient gateway failure is retried with exponential - backoff until the policy's budget is spent, after which the last attempt is - handed back unchanged - a permanent failure therefore looks exactly like it did - before, just later. `policy=None` opts a call out (the streaming paths, whose - body is consumed live).""" - first = _attempt(call) - if policy is None or not _is_transient(first): - return first - return _retry( - call, - deadline=time.monotonic() + policy.budget_seconds, - backoff=policy.initial_backoff_seconds, - last=first, - ) - - def _classify[R: BaseModel]( resp: requests.Response, response_type: type[R] ) -> Result[R]: @@ -331,14 +244,6 @@ def _classify[R: BaseModel]( return ValidationError(message=str(exc)) -def _result[R: BaseModel](attempt: _Attempt, response_type: type[R]) -> Result[R]: - match attempt: - case _Delivered(response=response): - return _classify(response, response_type) - case _Undelivered(error=error): - return NetworkError(message=str(error)) - - def post[R: BaseModel]( url: URL, *, @@ -347,17 +252,16 @@ def post[R: BaseModel]( response_type: type[R], timeout: float = 30.0, ) -> Result[R]: - return _result( - _dispatch( - lambda: requests.post( - str(url), - headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), - timeout=timeout, - ) - ), - response_type, - ) + try: + resp = requests.post( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) def get[R: BaseModel]( @@ -368,17 +272,16 @@ def get[R: BaseModel]( response_type: type[R], timeout: float = 30.0, ) -> Result[R]: - return _result( - _dispatch( - lambda: requests.get( - str(url), - headers=_headers(headers), - params=params.model_dump(by_alias=True, exclude_none=True), - timeout=timeout, - ) - ), - response_type, - ) + try: + resp = requests.get( + str(url), + headers=_headers(headers), + params=params.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) def get_external[R: BaseModel]( @@ -390,16 +293,15 @@ def get_external[R: BaseModel]( """GET an absolute URL outside the proxy (e.g. a public /.well-known document). Unlike the transport wrappers there is no proxy base url and no proxy auth; the response still gets the same tagged-union classification as every other call.""" - return _result( - _dispatch( - lambda: requests.get( - url, - headers={"Accept": "application/json"}, - timeout=timeout, - ) - ), - response_type, - ) + try: + resp = requests.get( + url, + headers={"Accept": "application/json"}, + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) def delete[R: BaseModel]( @@ -411,18 +313,17 @@ def delete[R: BaseModel]( params: BaseModel | None = None, timeout: float = 30.0, ) -> Result[R]: - return _result( - _dispatch( - lambda: requests.delete( - str(url), - headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), - params=_params(params), - timeout=timeout, - ) - ), - response_type, - ) + try: + resp = requests.delete( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + params=_params(params), + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) def patch[R: BaseModel]( @@ -433,17 +334,16 @@ def patch[R: BaseModel]( response_type: type[R], timeout: float = 30.0, ) -> Result[R]: - return _result( - _dispatch( - lambda: requests.patch( - str(url), - headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), - timeout=timeout, - ) - ), - response_type, - ) + try: + resp = requests.patch( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) def put[R: BaseModel]( @@ -454,34 +354,31 @@ def put[R: BaseModel]( response_type: type[R], timeout: float = 30.0, ) -> Result[R]: - return _result( - _dispatch( - lambda: requests.put( - str(url), - headers=_headers(headers), - json=json.model_dump(by_alias=True, exclude_none=True), - timeout=timeout, - ) - ), - response_type, - ) + try: + resp = requests.put( + str(url), + headers=_headers(headers), + json=json.model_dump(by_alias=True, exclude_none=True), + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) def probe( url: URL, *, headers: BaseModel, params: BaseModel, timeout: float = 30.0 ) -> ProbeResult: - match _dispatch( - lambda: requests.get( + try: + resp = requests.get( str(url), headers=_headers(headers), params=params.model_dump(by_alias=True, exclude_none=True), timeout=timeout, ) - ): - case _Delivered(response=response): - return ProbeResult(status_code=response.status_code, body=response.text) - case _Undelivered(error=error): - return ProbeResult(status_code=-1, body=str(error)) + except requests.RequestException as exc: + return ProbeResult(status_code=-1, body=str(exc)) + return ProbeResult(status_code=resp.status_code, body=resp.text) def _parse_response_cost(resp: requests.Response) -> float | None: @@ -558,21 +455,18 @@ def send( x-litellm-call-id header. For native/passthrough bodies and for calls judged by status rather than a typed JSON model (e.g. a budget block is a non-2xx). With ``stream=True`` the SSE body is consumed and its events counted instead.""" - match _dispatch( - lambda: requests.post( + try: + resp = requests.post( str(url), headers=_headers(headers), params=_params(params), json=json.model_dump(by_alias=True, exclude_none=True), stream=stream, timeout=timeout, - ), - policy=None if stream else _GATEWAY_RETRY, - ): - case _Delivered(response=response): - return _streaming_outcome(response, stream) - case _Undelivered(error=error): - return StreamingResponse(status_code=-1, body=str(error)) + ) + except requests.RequestException as exc: + return StreamingResponse(status_code=-1, body=str(exc)) + return _streaming_outcome(resp, stream) def stream( @@ -602,19 +496,18 @@ def upload[R: BaseModel]( routing (e.g. ?model=). requests sets the multipart Content-Type itself.""" dumped: dict[str, object] = form.model_dump(by_alias=True, exclude_none=True) data = {key: str(value) for key, value in dumped.items()} - return _result( - _dispatch( - lambda: requests.post( - str(url), - headers=_headers(headers), - params=_params(params), - data=data, - files={file_field: (filename, content, file_content_type)}, - timeout=timeout, - ) - ), - response_type, - ) + try: + resp = requests.post( + str(url), + headers=_headers(headers), + params=_params(params), + data=data, + files={file_field: (filename, content, file_content_type)}, + timeout=timeout, + ) + except requests.RequestException as exc: + return NetworkError(message=str(exc)) + return _classify(resp, response_type) def stream_binary( @@ -670,15 +563,13 @@ def download( ) -> StreamingResponse: """Raw GET for file content (/v1/files/{id}/content): provider-native bytes, no schema. Returns the decoded body and the x-litellm-call-id header.""" - match _dispatch( - lambda: requests.get(str(url), headers=_headers(headers), timeout=timeout) - ): - case _Delivered(response=response): - return StreamingResponse( - status_code=response.status_code, - call_id=_hdr(response, "x-litellm-call-id"), - content_type=_hdr(response, "content-type"), - body=response.text, - ) - case _Undelivered(error=error): - return StreamingResponse(status_code=-1, body=str(error)) + try: + resp = requests.get(str(url), headers=_headers(headers), timeout=timeout) + except requests.RequestException as exc: + return StreamingResponse(status_code=-1, body=str(exc)) + return StreamingResponse( + status_code=resp.status_code, + call_id=_hdr(resp, "x-litellm-call-id"), + content_type=_hdr(resp, "content-type"), + body=resp.text, + ) diff --git a/tests/e2e/test_e2e_http.py b/tests/e2e/test_e2e_http.py deleted file mode 100644 index 5a0cc601dfe5..000000000000 --- a/tests/e2e/test_e2e_http.py +++ /dev/null @@ -1,167 +0,0 @@ -"""Harness coverage for the transport's transient-failure retry (no proxy needed). - -A pod scale-down makes the load balancer answer 502 for a second or two, which used -to kill whatever test was mid-call. These pin which failures are retried and, just -as importantly, which are not: a 429 or a 500 is the app's own answer that -rate-limit and error-mapping tests assert on, so it must arrive unretried. - -The scripted server serves real HTTP on an ephemeral port and records every hit, so -"did it retry" is a request count rather than a timing guess. -""" - -from __future__ import annotations - -import threading -import time -from collections.abc import Generator, Sequence -from contextlib import contextmanager -from dataclasses import dataclass, field -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer - -from pydantic import BaseModel - -from e2e_http import URL, NoBody, Result, post, probe, send - -DROP_CONNECTION = 0 -SLOW_RESPONSE_SECONDS = 2.0 -RETRY_BUDGET_CEILING_SECONDS = 40.0 - - -class OkBody(BaseModel): - ok: bool - - -class QuietServer(ThreadingHTTPServer): - """A client that times out or walks away mid-response is the point of some of - these tests, so the resulting broken pipe is expected, not a server crash.""" - - daemon_threads = True - - def handle_error(self, request: object, client_address: object) -> None: - return - - -@dataclass(frozen=True, slots=True) -class ScriptedServer: - url: URL - hits: list[int] = field(default_factory=list) - - -@contextmanager -def scripted_server(statuses: Sequence[int], *, slow: bool = False) -> Generator[ScriptedServer]: - """Answer with `statuses` in order, repeating the last one forever. A status of - DROP_CONNECTION closes the socket without answering (what a load balancer does - when it drops a backend mid-request); `slow` stalls every response so the client - can hit its own read timeout.""" - scripted = iter(statuses) - hits: list[int] = [] - - class Handler(BaseHTTPRequestHandler): - protocol_version = "HTTP/1.1" - - def _answer(self) -> None: - length = int(self.headers.get("content-length") or 0) - _ = self.rfile.read(length) - status = next(scripted, statuses[-1]) - hits.append(status) - if slow: - time.sleep(SLOW_RESPONSE_SECONDS) - if status == DROP_CONNECTION: - self.close_connection = True - return - body = b'{"ok": true}' if status == 200 else b'{"error": "transient"}' - self.send_response(status) - self.send_header("content-type", "application/json") - self.send_header("content-length", str(len(body))) - self.end_headers() - _ = self.wfile.write(body) - - do_GET = _answer - do_POST = _answer - - def log_message(self, format: str, *args: object) -> None: - return - - server = QuietServer(("127.0.0.1", 0), Handler) - thread = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - try: - yield ScriptedServer(url=URL(f"http://127.0.0.1:{server.server_port}/anything"), hits=hits) - finally: - server.shutdown() - server.server_close() - thread.join(timeout=5) - - -def _post(url: URL, timeout: float = 30.0) -> Result[OkBody]: - return post(url, headers=NoBody(), json=NoBody(), response_type=OkBody, timeout=timeout) - - -class TestTransientRetry: - def test_post_retries_a_bad_gateway_until_it_succeeds(self) -> None: - with scripted_server((502, 503, 200)) as server: - result = _post(server.url) - assert result.kind == "success", result - assert server.hits == [502, 503, 200] - - def test_post_retries_a_dropped_connection_until_it_succeeds(self) -> None: - with scripted_server((DROP_CONNECTION, 200)) as server: - result = _post(server.url) - assert result.kind == "success", result - assert server.hits == [DROP_CONNECTION, 200] - - def test_post_gives_up_on_a_permanent_bad_gateway_with_the_same_error_as_before(self) -> None: - started = time.monotonic() - with scripted_server((502,)) as server: - result = _post(server.url) - elapsed = time.monotonic() - started - assert result.kind == "unknown", result - assert result.status_code == 502 - assert "transient" in result.body - assert len(server.hits) >= 4, server.hits - assert elapsed < RETRY_BUDGET_CEILING_SECONDS, elapsed - - def test_post_does_not_retry_a_rate_limit(self) -> None: - with scripted_server((429, 200)) as server: - result = _post(server.url) - assert result.kind == "rate_limited", result - assert server.hits == [429] - - def test_post_does_not_retry_a_server_error(self) -> None: - with scripted_server((500, 200)) as server: - result = _post(server.url) - assert result.kind == "unknown", result - assert result.status_code == 500 - assert server.hits == [500] - - def test_post_does_not_retry_a_client_error(self) -> None: - with scripted_server((400, 200)) as server: - result = _post(server.url) - assert result.kind == "unknown", result - assert result.status_code == 400 - assert server.hits == [400] - - def test_post_does_not_retry_a_read_timeout(self) -> None: - with scripted_server((200,), slow=True) as server: - result = _post(server.url, timeout=0.25) - assert result.kind == "network", result - assert server.hits == [200] - - def test_probe_retries_a_bad_gateway_until_it_succeeds(self) -> None: - with scripted_server((503, 200)) as server: - result = probe(server.url, headers=NoBody(), params=NoBody()) - assert result.status_code == 200, result - assert result.healthy - assert server.hits == [503, 200] - - def test_send_retries_a_bad_gateway_until_it_succeeds(self) -> None: - with scripted_server((504, 200)) as server: - result = send(server.url, headers=NoBody(), json=NoBody()) - assert result.ok, result - assert server.hits == [504, 200] - - def test_send_leaves_a_streaming_call_alone(self) -> None: - with scripted_server((502, 200)) as server: - result = send(server.url, headers=NoBody(), json=NoBody(), stream=True) - assert result.status_code == 502, result - assert server.hits == [502] From 043ff03e3078da8ded5ad59910bf5c589c3de6f7 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 27 Jul 2026 17:41:19 -0700 Subject: [PATCH 4/5] test(e2e): cap per-alias driver re-drives in the prometheus cardinality poll Bounds worst-case provider spend to 4 completions per alias while scrapes keep polling to the deadline; counters persist on whichever pod served them, so the cap costs no convergence unless that pod dies --- .../e2e/logging/test_prometheus_cardinality_e2e.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/e2e/logging/test_prometheus_cardinality_e2e.py b/tests/e2e/logging/test_prometheus_cardinality_e2e.py index 72747a6f82a1..9eb57a5b1f39 100644 --- a/tests/e2e/logging/test_prometheus_cardinality_e2e.py +++ b/tests/e2e/logging/test_prometheus_cardinality_e2e.py @@ -11,8 +11,9 @@ Scraping goes through ``transport.probe`` (raw text) and is parsed with prometheus_client. ``/metrics`` is per-pod behind a round-robin LB and the metric is eventually consistent (it increments on the success-logging callback), so the -poll re-drives each still-missing alias with fresh traffic and unions the aliases -seen across scrapes until the deadline. +poll re-drives each still-missing alias with fresh traffic (bounded per alias to +cap provider spend; scrapes continue to the deadline regardless since counters +persist on whichever pod served them) and unions the aliases seen across scrapes. """ from __future__ import annotations @@ -32,6 +33,7 @@ REQUESTS_METRIC = "litellm_requests_metric_total" ALIAS_LABEL = "api_key_alias" DISTINCT_KEYS = 3 +MAX_DRIVER_CALLS_PER_ALIAS = 4 def _aliases_in_metric(exposition: str, metric: str, label: str) -> frozenset[str]: @@ -67,12 +69,17 @@ def drive(alias: str, key: str) -> None: wanted = frozenset(aliases) deadline = time.monotonic() + client.proxy.poll_timeout seen: frozenset[str] = frozenset() + drive_counts = dict.fromkeys(aliases, 1) while time.monotonic() < deadline: seen = seen | _aliases_in_metric(client.scrape_metrics(), REQUESTS_METRIC, ALIAS_LABEL) if wanted <= seen: break - for alias in sorted(wanted - seen): + redriven = tuple( + alias for alias in sorted(wanted - seen) if drive_counts[alias] < MAX_DRIVER_CALLS_PER_ALIAS + ) + for alias in redriven: drive(alias, keys_by_alias[alias]) + drive_counts = {alias: count + (alias in redriven) for alias, count in drive_counts.items()} time.sleep(client.proxy.poll_interval) missing = wanted - seen From 59908d6746d511e52755f1080cb13d27c63bd1b9 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Mon, 27 Jul 2026 17:59:01 -0700 Subject: [PATCH 5/5] test(e2e): drop driver re-drives from the prometheus cardinality poll The per-key cardinality contract is process-local and counters persist on whichever pod served the driver call, so unioning aliases across free scrape polls converges without re-sending billable traffic. The residual gap, a pod dying inside the poll window, is deferred to direct per-pod scraping --- .../test_prometheus_cardinality_e2e.py | 30 +++++-------------- 1 file changed, 7 insertions(+), 23 deletions(-) diff --git a/tests/e2e/logging/test_prometheus_cardinality_e2e.py b/tests/e2e/logging/test_prometheus_cardinality_e2e.py index 9eb57a5b1f39..e2d164d8b4cc 100644 --- a/tests/e2e/logging/test_prometheus_cardinality_e2e.py +++ b/tests/e2e/logging/test_prometheus_cardinality_e2e.py @@ -11,9 +11,9 @@ Scraping goes through ``transport.probe`` (raw text) and is parsed with prometheus_client. ``/metrics`` is per-pod behind a round-robin LB and the metric is eventually consistent (it increments on the success-logging callback), so the -poll re-drives each still-missing alias with fresh traffic (bounded per alias to -cap provider spend; scrapes continue to the deadline regardless since counters -persist on whichever pod served them) and unions the aliases seen across scrapes. +poll unions the aliases seen across scrapes until the deadline: counters persist +on whichever pod served the driver call, so repeated scrapes converge without +re-sending any billable traffic. """ from __future__ import annotations @@ -33,7 +33,6 @@ REQUESTS_METRIC = "litellm_requests_metric_total" ALIAS_LABEL = "api_key_alias" DISTINCT_KEYS = 3 -MAX_DRIVER_CALLS_PER_ALIAS = 4 def _aliases_in_metric(exposition: str, metric: str, label: str) -> frozenset[str]: @@ -52,39 +51,24 @@ def test_distinct_key_aliases_produce_distinct_series( self, client: LoggingClient, resources: ResourceManager ) -> None: aliases = tuple(f"e2e-prom-{unique_marker()}" for _ in range(DISTINCT_KEYS)) - - def provisioned_key(alias: str) -> str: + for alias in aliases: key = client.key_with_alias(alias, models=[DRIVER_MODEL]) - resources.defer(lambda: client.delete_key(key)) - return key - - def drive(alias: str, key: str) -> None: - response = client.chat(key, DRIVER_MODEL, f"reply with one word {alias} {unique_marker()}") + resources.defer(lambda k=key: client.delete_key(k)) + response = client.chat(key, DRIVER_MODEL, f"reply with one word {alias}") assert response.model, f"driver call for {alias} returned no model: {response}" - keys_by_alias = {alias: provisioned_key(alias) for alias in aliases} - for alias, key in keys_by_alias.items(): - drive(alias, key) - wanted = frozenset(aliases) deadline = time.monotonic() + client.proxy.poll_timeout seen: frozenset[str] = frozenset() - drive_counts = dict.fromkeys(aliases, 1) while time.monotonic() < deadline: seen = seen | _aliases_in_metric(client.scrape_metrics(), REQUESTS_METRIC, ALIAS_LABEL) if wanted <= seen: break - redriven = tuple( - alias for alias in sorted(wanted - seen) if drive_counts[alias] < MAX_DRIVER_CALLS_PER_ALIAS - ) - for alias in redriven: - drive(alias, keys_by_alias[alias]) - drive_counts = {alias: count + (alias in redriven) for alias, count in drive_counts.items()} time.sleep(client.proxy.poll_interval) missing = wanted - seen assert not missing, ( f"{REQUESTS_METRIC} never exposed a per-key series for aliases {sorted(missing)} " - f"on any scraped pod within the deadline despite repeated driver calls; " + f"on any scraped pod within the deadline; " f"each distinct {ALIAS_LABEL} must grow its own series" )