From 6b7c77b7d309daca2cebecf2871752ee1623d6e1 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 10:50:11 -0700 Subject: [PATCH] fix(http): restore the 15s aiohttp keep-alive default in the Python and Rust clients (#46355) * fix: lower aiohttp keepalive timeout Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(http): restore the 15s keep-alive default and apply it to the Rust client 15s is aiohttp's own default and what every release before v1.78.0 used. The Rust HTTP client reads the same AIOHTTP_KEEPALIVE_TIMEOUT variable, so its default moves with it. * test(integration): cover upstream connections dropped after the keep-alive idle window --------- Co-authored-by: nate Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/http/src/settings.rs | 4 +- litellm/constants.py | 2 +- tests/integration/_support/idle_gateway.py | 207 +++++ .../providers/_idle_upstream_support.py | 753 ++++++++++++++++++ ...test_idle_upstream_keepalive_edges_wire.py | 89 +++ ...t_idle_upstream_keepalive_override_wire.py | 56 ++ ...test_idle_upstream_reconnect_chaos_wire.py | 135 ++++ ...est_idle_upstream_reconnect_routes_wire.py | 60 ++ ...e_upstream_reconnect_rust_messages_wire.py | 56 ++ .../test_idle_upstream_reconnect_wire.py | 62 ++ .../sdk/test_idle_upstream_reconnect_sdk.py | 94 +++ 11 files changed, 1515 insertions(+), 3 deletions(-) create mode 100644 tests/integration/_support/idle_gateway.py create mode 100644 tests/integration/providers/_idle_upstream_support.py create mode 100644 tests/integration/providers/test_idle_upstream_keepalive_edges_wire.py create mode 100644 tests/integration/providers/test_idle_upstream_keepalive_override_wire.py create mode 100644 tests/integration/providers/test_idle_upstream_reconnect_chaos_wire.py create mode 100644 tests/integration/providers/test_idle_upstream_reconnect_routes_wire.py create mode 100644 tests/integration/providers/test_idle_upstream_reconnect_rust_messages_wire.py create mode 100644 tests/integration/providers/test_idle_upstream_reconnect_wire.py create mode 100644 tests/integration/sdk/test_idle_upstream_reconnect_sdk.py diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs index e1edc6d37e1..1f437df225c 100644 --- a/litellm-rust/crates/http/src/settings.rs +++ b/litellm-rust/crates/http/src/settings.rs @@ -140,7 +140,7 @@ impl Default for HttpSettings { proxies: EnvironmentProxies::default(), connect_timeout: Duration::from_secs(10), tcp_keepalive: None, - pool_idle_timeout: Duration::from_secs(120), + pool_idle_timeout: Duration::from_secs(15), } } } @@ -330,7 +330,7 @@ mod tests { &[("AIOHTTP_TCP_KEEPIDLE", "45")], ))]); assert_eq!(settings.tcp_keepalive, None); - assert_eq!(settings.pool_idle_timeout, Duration::from_secs(120)); + assert_eq!(settings.pool_idle_timeout, Duration::from_secs(15)); } fn proxy_flags( diff --git a/litellm/constants.py b/litellm/constants.py index a7ea4a2975b..aed2b04b16b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -345,7 +345,7 @@ EVICTED_LLM_CLIENT_CLOSE_MAX_PENDING: Final = 10_000 # Set to 0 for unlimited (not recommended for production) AIOHTTP_CONNECTOR_LIMIT: Final = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 1000)) AIOHTTP_CONNECTOR_LIMIT_PER_HOST: Final = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 500)) -AIOHTTP_KEEPALIVE_TIMEOUT: Final = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120)) +AIOHTTP_KEEPALIVE_TIMEOUT: Final = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 15)) AIOHTTP_TTL_DNS_CACHE: Final = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300)) # TCP keep-alive (SO_KEEPALIVE) — opt-in. Required when running behind NAT/LBs # whose idle timeout is shorter than provider response timeouts (e.g. AWS NAT diff --git a/tests/integration/_support/idle_gateway.py b/tests/integration/_support/idle_gateway.py new file mode 100644 index 00000000000..5a86d7e0a35 --- /dev/null +++ b/tests/integration/_support/idle_gateway.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +import itertools +import re +import socket +import struct +import threading +import time +from collections.abc import Callable, Generator, Mapping +from contextlib import contextmanager, suppress +from dataclasses import dataclass +from enum import Enum +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from io import BufferedIOBase +from queue import SimpleQueue +from socketserver import BaseServer +from typing import Final + +from integration._support.wire import Reply, Request + +MARKER: Final = re.compile(rb"idle-[0-9a-f]{16}") +_LOOPBACK: Final = "127.0.0.1" + + +def _exactly(stream: BufferedIOBase, size: int) -> bytes | None: + data: Final = stream.read(size) + return data if len(data) == size else None + + +def _body(headers: Mapping[str, str], stream: BufferedIOBase) -> bytes | None: + if headers.get("transfer-encoding", "").lower() == "chunked": + return None + return _exactly(stream, int(headers.get("content-length", "0"))) + + +class Drop(Enum): + CLOSE = "close" + RESET = "reset" + + +@dataclass(frozen=True, slots=True) +class Arrival: + marker: str + target: str + connection: int + dropped: bool + + +@dataclass(frozen=True, slots=True) +class IdleGateway: + """A keep-alive HTTP/1.1 peer that never closes an idle connection itself and, like a load balancer whose + idle timer already fired, swallows a request that reuses a connection idle for `drop_after` seconds.""" + + url: str + port: int + arrivals: SimpleQueue[Arrival] + + def pending(self) -> int: + return self.arrivals.qsize() + + def drain(self) -> tuple[Arrival, ...]: + return tuple(self.arrivals.get_nowait() for _ in range(self.arrivals.qsize())) + + +class _Server(ThreadingHTTPServer): + daemon_threads = True + block_on_close = False + request_queue_size = 128 + + def __init__( + self, + port: int, + respond: Callable[[Request], Reply], + drop_after: float | None, + drop: Drop, + arrivals: SimpleQueue[Arrival], + ) -> None: + super().__init__((_LOOPBACK, port), _Handler) + self.respond = respond + self.drop_after = drop_after + self.drop = drop + self.arrivals = arrivals + self.connections = itertools.count(1) + self.sockets: SimpleQueue[socket.socket] = SimpleQueue() + self.failures: SimpleQueue[Exception] = SimpleQueue() + + def close_connections(self) -> None: + for client in tuple(self.sockets.get_nowait() for _ in range(self.sockets.qsize())): + with suppress(OSError): + client.shutdown(socket.SHUT_RDWR) + + def handle_error(self, request: object, client_address: object) -> None: + return + + +class _Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def __init__( + self, request: socket.socket | tuple[bytes, socket.socket], client_address: object, server: BaseServer + ) -> None: + assert isinstance(server, _Server) + assert isinstance(request, socket.socket) + self.gateway = server + self.client = request + super().__init__(request, client_address, server) + + def setup(self) -> None: + super().setup() + self.connection_id = next(self.gateway.connections) + self.idle_since: float | None = None + self.arrived_at = time.monotonic() + self.gateway.sockets.put(self.client) + + def parse_request(self) -> bool: + self.arrived_at = time.monotonic() + return super().parse_request() + + def _stale(self) -> bool: + return ( + self.gateway.drop_after is not None + and self.idle_since is not None + and self.arrived_at - self.idle_since >= self.gateway.drop_after + ) + + def answer(self) -> None: + headers: Final = {name.lower(): value for name, value in self.headers.items()} + body: Final = _body(headers, self.rfile) + if body is None: + self.gateway.failures.put(ValueError(f"unreadable body for {self.path}: {headers}")) + self.close_connection = True + return + found: Final = MARKER.search(body) or MARKER.search(self.path.encode()) + stale: Final = self._stale() + self.gateway.arrivals.put( + Arrival(found.group().decode() if found else "", self.path, self.connection_id, stale) + ) + if stale: + self._drop() + return + try: + reply: Final = self.gateway.respond(Request(self.command, self.path, headers, body)) + except (ValueError, KeyError, TypeError) as error: + self.gateway.failures.put(error) + self._write(Reply(status=500)) + return + self._write(reply) + self.idle_since = time.monotonic() + + do_POST = answer + do_GET = answer + + def _drop(self) -> None: + self.close_connection = True + match self.gateway.drop: + case Drop.CLOSE: + with suppress(OSError): + self.client.shutdown(socket.SHUT_RDWR) + case Drop.RESET: + self.client.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0)) + self.rfile.close() + self.client.close() + + def _write(self, reply: Reply) -> None: + self.send_response(reply.status) + self.send_header("content-type", reply.content_type) + for name, value in reply.headers.items(): + self.send_header(name, value) + if reply.chunks is None: + self.send_header("content-length", str(len(reply.body))) + else: + self.send_header("transfer-encoding", "chunked") + self.end_headers() + if reply.chunks is None: + self.wfile.write(reply.body) + return + for index, chunk in enumerate(reply.chunks): + if index and reply.pause_between_chunks: + time.sleep(reply.pause_between_chunks) + self.wfile.write(b"%x\r\n%s\r\n" % (len(chunk), chunk)) + self.wfile.write(b"0\r\n\r\n") + + def log_message(self, format: str, *args: object) -> None: + return + + +@contextmanager +def idle_gateway( + respond: Callable[[Request], Reply], + *, + drop_after: float | None, + drop: Drop = Drop.CLOSE, + port: int = 0, +) -> Generator[IdleGateway]: + arrivals: Final[SimpleQueue[Arrival]] = SimpleQueue() + with _Server(port, respond, drop_after, drop, arrivals) as server: + thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}) + thread.start() + try: + yield IdleGateway(f"http://{_LOOPBACK}:{server.server_port}", server.server_port, arrivals) + finally: + server.shutdown() + thread.join(timeout=6) + server.close_connections() + assert not thread.is_alive(), "Idle gateway survived cleanup" + failure: Final = None if server.failures.empty() else server.failures.get_nowait() + assert failure is None, f"Idle gateway failed to answer: {failure!r}" diff --git a/tests/integration/providers/_idle_upstream_support.py b/tests/integration/providers/_idle_upstream_support.py new file mode 100644 index 00000000000..ca808abc008 --- /dev/null +++ b/tests/integration/providers/_idle_upstream_support.py @@ -0,0 +1,753 @@ +from __future__ import annotations + +import asyncio +import io +import json +import uuid +import wave +from collections.abc import Callable, Generator, Mapping +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import anthropic +import httpx +import openai +import yaml +from anthropic.types import RawContentBlockDeltaEvent, RawMessageStartEvent, TextBlock, TextDelta +from integration._support.anthropic_sse import message_json, message_stream, stream_reply +from integration._support.client import Gateway, eventually, held, list_value, object_value, string_value +from integration._support.database import read_rows +from integration._support.idle_gateway import MARKER, Arrival, Drop, IdleGateway, idle_gateway +from integration._support.openai_wire import chat_reply, chat_stream, responses_reply +from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_proxy_process +from integration._support.wire import Reply, Request +from openai.types.chat import ChatCompletionChunk +from openai.types.responses import ResponseCompletedEvent, ResponseTextDeltaEvent +from pydantic import JsonValue, TypeAdapter + +IDLE_SECONDS: Final = 17.0 +DROP_AFTER_SECONDS: Final = 16.0 +REUSE_IDLE_SECONDS: Final = 5.0 +QUIET_STREAM_SECONDS: Final = 16.0 +PROVIDER_KEY: Final = "sk-idle-upstream-provider" +CALL_ID: Final = "x-litellm-call-id" +RUST_SERVED: Final = "x-litellm-rust" +GUARDRAIL: Final = "idle-upstream-guardrail" +PASS_THROUGH: Final = "/idle-upstream-pass-through" +EMBEDDING: Final = (0.25, 0.5) +FRESH_CONNECTIONS: Final = httpx.Limits(max_keepalive_connections=0) +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" + + +def new_marker() -> str: + return f"idle-{uuid.uuid4().hex[:16]}" + + +def expected_text(marker: str) -> str: + return f"answer {marker}" + + +def expected_embedding(_marker: str) -> str: + return json.dumps(list(EMBEDDING)) + + +def _marker_of(request: Request) -> str: + found: Final = MARKER.search(request.body) + return found.group().decode() if found else "" + + +def _converse(text: str) -> bytes: + return json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": text}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 5, "outputTokens": 3, "totalTokens": 8}, + "metrics": {"latencyMs": 1}, + } + ).encode() + + +def _ocr(text: str) -> bytes: + return json.dumps( + { + "model": "mistral-ocr-latest", + "pages": [{"index": 0, "markdown": text}], + "usage_info": {"pages_processed": 1, "doc_size_bytes": 1024}, + } + ).encode() + + +def _embeddings() -> bytes: + return json.dumps( + { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": list(EMBEDDING)}], + "model": "text-embedding-idle-upstream", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } + ).encode() + + +def _model_reply(target: str, marker: str, body: Mapping[str, JsonValue]) -> Reply: + model: Final = str(body.get("model", "idle-upstream")) + stream: Final = body.get("stream") is True + if target.endswith("/v1/messages"): + if stream: + return stream_reply(message_stream(f"msg_{marker}", model, expected_text(marker))) + return Reply(body=message_json(f"msg_{marker}", model, expected_text(marker))) + if target.endswith("/responses"): + return responses_reply(f"resp_{marker}", model, expected_text(marker), stream=stream) + if target.endswith("/chat/completions"): + return chat_reply(f"chatcmpl-{marker}", model, expected_text(marker), stream=stream) + return Reply(status=404, body=b'{"error": "unscripted route"}') + + +def answer(request: Request) -> Reply: + marker: Final = _marker_of(request) + target: Final = request.target.split("?", 1)[0] + if target.endswith("/beta/litellm_basic_guardrail_api"): + return Reply(body=json.dumps({"action": "NONE"}).encode()) + if target.startswith("/passthrough"): + return Reply(body=json.dumps({"echo": expected_text(marker)}).encode()) + if target.endswith("/converse"): + return Reply(body=_converse(expected_text(marker))) + if target.endswith("/ocr"): + return Reply(body=_ocr(expected_text(marker))) + if target.endswith("/embeddings"): + return Reply(body=_embeddings()) + return _model_reply(target, marker, _JSON.validate_json(request.body)) + + +def quiet_stream_answer(request: Request) -> Reply: + marker: Final = _marker_of(request) + body: Final = _JSON.validate_json(request.body) + model: Final = str(body.get("model", "idle-upstream")) + frames: Final = ( + message_stream(f"msg_{marker}", model, expected_text(marker)) + if request.target.endswith("/v1/messages") + else chat_stream(f"chatcmpl-{marker}", model, expected_text(marker)) + ) + return Reply( + content_type="text/event-stream", + chunks=(frames[0], b"".join(frames[1:])), + pause_between_chunks=QUIET_STREAM_SECONDS, + ) + + +@dataclass(frozen=True, slots=True) +class ProxyTarget: + url: str + key: str + + +@dataclass(frozen=True, slots=True) +class Answer: + status: int + text: str + call_id: str + response_id: str + rust: bool = False + + +@dataclass(frozen=True, slots=True) +class Sent: + marker: str + answer: Answer + + +def _call_id(headers: httpx.Headers) -> str: + return next(iter(headers.get_list(CALL_ID)), "") + + +def _rust(headers: httpx.Headers) -> bool: + return "true" in headers.get_list(RUST_SERVED) + + +def _refused(status: int, response: httpx.Response) -> Answer: + return Answer(status, response.text, _call_id(response.headers), "") + + +def _auth(target: ProxyTarget) -> Mapping[str, str]: + return {"Authorization": f"Bearer {target.key}"} + + +def _openai(target: ProxyTarget) -> openai.OpenAI: + return openai.OpenAI( + base_url=f"{target.url}/v1", + api_key=target.key, + max_retries=0, + http_client=httpx.Client(limits=FRESH_CONNECTIONS, timeout=60, trust_env=False), + ) + + +def _async_openai(target: ProxyTarget) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI( + base_url=f"{target.url}/v1", + api_key=target.key, + max_retries=0, + http_client=httpx.AsyncClient(limits=FRESH_CONNECTIONS, timeout=60, trust_env=False), + ) + + +def _anthropic(target: ProxyTarget) -> anthropic.Anthropic: + return anthropic.Anthropic( + base_url=target.url, + api_key=target.key, + max_retries=0, + http_client=httpx.Client(limits=FRESH_CONNECTIONS, timeout=60, trust_env=False), + ) + + +def _async_anthropic(target: ProxyTarget) -> anthropic.AsyncAnthropic: + return anthropic.AsyncAnthropic( + base_url=target.url, + api_key=target.key, + max_retries=0, + http_client=httpx.AsyncClient(limits=FRESH_CONNECTIONS, timeout=60, trust_env=False), + ) + + +def _chat(target: ProxyTarget, model: str, marker: str, extra: Mapping[str, JsonValue]) -> Answer: + with _openai(target) as client: + try: + raw: Final = client.chat.completions.with_raw_response.create( + model=model, messages=[{"role": "user", "content": marker}], extra_body=dict(extra) + ) + except openai.APIStatusError as error: + return _refused(error.status_code, error.response) + completion: Final = raw.parse() + return Answer(200, completion.choices[0].message.content or "", _call_id(raw.headers), completion.id) + + +def chat_sdk(target: ProxyTarget, model: str, marker: str) -> Answer: + return _chat(target, model, marker, {}) + + +def chat_sdk_with_default_retries(target: ProxyTarget, model: str, marker: str) -> Answer: + return _chat(target, model, marker, {"num_retries": 2}) + + +def _chunk_text(chunk: ChatCompletionChunk) -> str: + return "".join(choice.delta.content or "" for choice in chunk.choices) + + +async def _chat_stream(target: ProxyTarget, model: str, marker: str) -> Answer: + async with _async_openai(target) as client: + try: + raw: Final = await client.chat.completions.with_raw_response.create( + model=model, messages=[{"role": "user", "content": marker}], stream=True + ) + except openai.APIStatusError as error: + return _refused(error.status_code, error.response) + try: + chunks: Final = tuple([chunk async for chunk in raw.parse()]) + except openai.APIError as error: + return Answer(200, f"stream broke: {error}", _call_id(raw.headers), "") + return Answer( + 200, + "".join(_chunk_text(chunk) for chunk in chunks), + _call_id(raw.headers), + chunks[0].id if chunks else "", + ) + + +def chat_sdk_stream(target: ProxyTarget, model: str, marker: str) -> Answer: + return asyncio.run(_chat_stream(target, model, marker)) + + +def _post( + target: ProxyTarget, + path: str, + body: Mapping[str, JsonValue], + read: Callable[[Mapping[str, JsonValue]], tuple[str, str]], +) -> Answer: + with httpx.Client(limits=FRESH_CONNECTIONS, timeout=60, trust_env=False) as client: + response: Final = client.post(f"{target.url}{path}", json=dict(body), headers=dict(_auth(target))) + if response.status_code != 200: + return _refused(response.status_code, response) + text, identity = read(_JSON.validate_json(response.content)) + return Answer(200, text, _call_id(response.headers), identity, _rust(response.headers)) + + +def _chat_body(body: Mapping[str, JsonValue]) -> tuple[str, str]: + message: Final = object_value(object_value(list_value(body["choices"])[0])["message"]) + return string_value(message["content"]), string_value(body["id"]) + + +def _echo_body(body: Mapping[str, JsonValue]) -> tuple[str, str]: + return string_value(body["echo"]), "" + + +def _ocr_body(body: Mapping[str, JsonValue]) -> tuple[str, str]: + return string_value(object_value(list_value(body["pages"])[0])["markdown"]), "" + + +def chat_httpx(target: ProxyTarget, model: str, marker: str) -> Answer: + return _post( + target, "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}]}, _chat_body + ) + + +def chat_with_guardrail(target: ProxyTarget, model: str, marker: str) -> Answer: + return _post( + target, + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": marker}], "guardrails": [GUARDRAIL]}, + _chat_body, + ) + + +def pass_through(target: ProxyTarget, _model: str, marker: str) -> Answer: + return _post(target, PASS_THROUGH, {"input": marker}, _echo_body) + + +def ocr(target: ProxyTarget, model: str, marker: str) -> Answer: + return _post( + target, + "/v1/ocr", + { + "model": model, + "document": {"type": "document_url", "document_url": f"https://documents.invalid/{marker}.pdf"}, + }, + _ocr_body, + ) + + +def _wav() -> bytes: + buffer: Final = io.BytesIO() + with wave.open(buffer, "wb") as writer: + writer.setnchannels(1) + writer.setsampwidth(2) + writer.setframerate(8000) + writer.writeframes(b"\x00\x00" * 800) + return buffer.getvalue() + + +def transcription(target: ProxyTarget, model: str, marker: str) -> Answer: + with httpx.Client(limits=FRESH_CONNECTIONS, timeout=60, trust_env=False) as client: + response: Final = client.post( + f"{target.url}/v1/audio/transcriptions", + files={"file": ("idle.wav", _wav(), "audio/wav")}, + data={"model": model, "prompt": marker}, + headers=dict(_auth(target)), + ) + if response.status_code != 200: + return _refused(response.status_code, response) + return Answer( + 200, + string_value(_JSON.validate_json(response.content)["text"]), + _call_id(response.headers), + "", + _rust(response.headers), + ) + + +def embeddings_sdk(target: ProxyTarget, model: str, marker: str) -> Answer: + with _openai(target) as client: + try: + raw: Final = client.embeddings.with_raw_response.create(model=model, input=marker, encoding_format="float") + except openai.APIStatusError as error: + return _refused(error.status_code, error.response) + result: Final = raw.parse() + return Answer(200, json.dumps(result.data[0].embedding), _call_id(raw.headers), "") + + +def responses_sdk(target: ProxyTarget, model: str, marker: str) -> Answer: + with _openai(target) as client: + try: + raw: Final = client.responses.with_raw_response.create(model=model, input=marker) + except openai.APIStatusError as error: + return _refused(error.status_code, error.response) + response: Final = raw.parse() + return Answer(200, response.output_text, _call_id(raw.headers), response.id) + + +async def _responses_stream(target: ProxyTarget, model: str, marker: str) -> Answer: + async with _async_openai(target) as client: + try: + raw: Final = await client.responses.with_raw_response.create(model=model, input=marker, stream=True) + except openai.APIStatusError as error: + return _refused(error.status_code, error.response) + try: + events: Final = tuple([event async for event in raw.parse()]) + except openai.APIError as error: + return Answer(200, f"stream broke: {error}", _call_id(raw.headers), "") + return Answer( + 200, + "".join(event.delta for event in events if isinstance(event, ResponseTextDeltaEvent)), + _call_id(raw.headers), + next((event.response.id for event in events if isinstance(event, ResponseCompletedEvent)), ""), + ) + + +def responses_sdk_stream(target: ProxyTarget, model: str, marker: str) -> Answer: + return asyncio.run(_responses_stream(target, model, marker)) + + +def messages_sdk(target: ProxyTarget, model: str, marker: str) -> Answer: + with _anthropic(target) as client: + try: + raw: Final = client.messages.with_raw_response.create( + model=model, max_tokens=32, messages=[{"role": "user", "content": marker}] + ) + except anthropic.APIStatusError as error: + return _refused(error.status_code, error.response) + message: Final = raw.parse() + return Answer( + 200, + "".join(block.text for block in message.content if isinstance(block, TextBlock)), + _call_id(raw.headers), + message.id, + _rust(raw.headers), + ) + + +def _delta_text(event: object) -> str: + if isinstance(event, RawContentBlockDeltaEvent) and isinstance(event.delta, TextDelta): + return event.delta.text + return "" + + +async def _messages_stream(target: ProxyTarget, model: str, marker: str) -> Answer: + async with _async_anthropic(target) as client: + try: + raw: Final = await client.messages.with_raw_response.create( + model=model, max_tokens=32, messages=[{"role": "user", "content": marker}], stream=True + ) + except anthropic.APIStatusError as error: + return _refused(error.status_code, error.response) + try: + events: Final = tuple([event async for event in raw.parse()]) + except anthropic.APIError as error: + return Answer(200, f"stream broke: {error}", _call_id(raw.headers), "") + return Answer( + 200, + "".join(_delta_text(event) for event in events), + _call_id(raw.headers), + next((event.message.id for event in events if isinstance(event, RawMessageStartEvent)), ""), + _rust(raw.headers), + ) + + +def messages_sdk_stream(target: ProxyTarget, model: str, marker: str) -> Answer: + return asyncio.run(_messages_stream(target, model, marker)) + + +@dataclass(frozen=True, slots=True) +class ConfigPart: + models: tuple[Mapping[str, JsonValue], ...] = () + pass_through: tuple[Mapping[str, JsonValue], ...] = () + guardrails: tuple[Mapping[str, JsonValue], ...] = () + + +def openai_deployment(model: str, url: str, _steady: str) -> ConfigPart: + return ConfigPart( + models=( + { + "model_name": model, + "litellm_params": {"model": "openai/idle-upstream", "api_base": url, "api_key": PROVIDER_KEY}, + }, + ) + ) + + +def embedding_deployment(model: str, url: str, _steady: str) -> ConfigPart: + return ConfigPart( + models=( + { + "model_name": model, + "litellm_params": { + "model": "openai/text-embedding-idle-upstream", + "api_base": url, + "api_key": PROVIDER_KEY, + }, + "model_info": {"mode": "embedding"}, + }, + ) + ) + + +def anthropic_deployment(model: str, url: str, _steady: str) -> ConfigPart: + return ConfigPart( + models=( + { + "model_name": model, + "litellm_params": {"model": "anthropic/claude-idle-upstream", "api_base": url, "api_key": PROVIDER_KEY}, + }, + ) + ) + + +def ocr_deployment(model: str, url: str, _steady: str) -> ConfigPart: + return ConfigPart( + models=( + { + "model_name": model, + "litellm_params": {"model": "mistral/mistral-ocr-latest", "api_base": url, "api_key": PROVIDER_KEY}, + }, + ) + ) + + +def transcription_deployment(model: str, url: str, _steady: str) -> ConfigPart: + return ConfigPart( + models=( + { + "model_name": model, + "litellm_params": { + "model": "bedrock/mistral.voxtral-mini-3b-2507", + "api_base": url, + "aws_access_key_id": "AKIAIDLEUPSTREAM0000", + "aws_secret_access_key": "idle-upstream-synthetic-secret", + "aws_region_name": "us-east-1", + }, + "model_info": {"mode": "audio_transcription"}, + }, + ) + ) + + +def pass_through_deployment(_model: str, url: str, _steady: str) -> ConfigPart: + return ConfigPart(pass_through=({"path": PASS_THROUGH, "target": f"{url}/passthrough"},)) + + +def guardrail_deployment(model: str, url: str, steady: str) -> ConfigPart: + return ConfigPart( + models=openai_deployment(model, steady, steady).models, + guardrails=( + { + "guardrail_name": GUARDRAIL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": url, + "api_key": PROVIDER_KEY, + }, + }, + ), + ) + + +@dataclass(frozen=True, slots=True) +class Row: + name: str + deploy: Callable[[str, str, str], ConfigPart] + send: Callable[[ProxyTarget, str, str], Answer] + expected: Callable[[str], str] = expected_text + drop: Drop = Drop.CLOSE + rust: bool = False + respond: Callable[[Request], Reply] = answer + + @property + def model(self) -> str: + return f"idle-{self.name}" + + +def send(row: Row, target: ProxyTarget) -> Sent: + marker: Final = new_marker() + return Sent(marker, row.send(target, row.model, marker)) + + +def proxy_config(directory: Path, parts: tuple[ConfigPart, ...]) -> Path: + base: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text())) + config: Final = { + **base, + "model_list": list(chain.from_iterable(part.models for part in parts)), + "guardrails": list(chain.from_iterable(part.guardrails for part in parts)), + "general_settings": { + **object_value(base["general_settings"]), + "pass_through_endpoints": list(chain.from_iterable(part.pass_through for part in parts)), + }, + "router_settings": {**object_value(base["router_settings"]), "num_retries": 0}, + } + path: Final = directory / f"idle-upstream-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(json.loads(json.dumps(config)))) + return path + + +def proxy_target(owned: OwnedProxy) -> ProxyTarget: + return ProxyTarget(str(owned.gateway.client.base_url).rstrip("/"), owned.gateway.key) + + +def two_workers(owned: OwnedProxy) -> tuple[int, int]: + first, second = sorted(eventually(owned.worker_pids, lambda pids: len(pids) == 2, seconds=120)) + return first, second + + +def on_each_worker( + owned: OwnedProxy, workers: tuple[int, int], call: Callable[[], tuple[Sent, ...]] +) -> tuple[Sent, ...]: + first, second = workers + with owned.suspended(second): + on_first: Final = call() + with owned.suspended(first): + on_second: Final = call() + return (*on_first, *on_second) + + +def stay_idle(gateways: tuple[IdleGateway, ...], seconds: float) -> None: + held( + lambda: sum(gateway.pending() for gateway in gateways), + lambda arrived: arrived == 0, + holding=seconds, + seconds=seconds + 30, + ) + + +def answered(sents: tuple[Sent, ...]) -> tuple[tuple[int, str], ...]: + return tuple((sent.answer.status, sent.answer.text) for sent in sents) + + +def expected_answers(row: Row, sents: tuple[Sent, ...]) -> tuple[tuple[int, str], ...]: + return tuple((200, row.expected(sent.marker)) for sent in sents) + + +def served_by_rust(sents: tuple[Sent, ...]) -> tuple[bool, ...]: + return tuple(sent.answer.rust for sent in sents) + + +def connections(arrivals: tuple[Arrival, ...]) -> tuple[int, ...]: + return tuple(sorted(arrival.connection for arrival in arrivals)) + + +def landed(arrivals: tuple[Arrival, ...]) -> tuple[tuple[str, bool], ...]: + return tuple(sorted((arrival.marker, arrival.dropped) for arrival in arrivals)) + + +def landed_once_undropped(sents: tuple[Sent, ...]) -> tuple[tuple[str, bool], ...]: + return tuple(sorted((sent.marker, False) for sent in sents)) + + +def spend_statuses(sents: tuple[Sent, ...]) -> tuple[str, ...]: + call_ids: Final = ",".join(sorted(sent.answer.call_id for sent in sents if sent.answer.call_id)) + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE litellm_call_id = ANY(string_to_array(%s, %s))', + (call_ids, ","), + ), + lambda found: len(found) >= len(sents), + seconds=90, + return_last_on_timeout=True, + ) + return tuple(sorted(string_value(row["status"]) for row in rows)) + + +@dataclass(frozen=True, slots=True) +class UpstreamProxy: + owned: OwnedProxy + workers: tuple[int, int] + rows: Mapping[str, Row] + gateways: Mapping[str, IdleGateway] + + @property + def target(self) -> ProxyTarget: + return proxy_target(self.owned) + + def probe(self, name: str) -> tuple[tuple[Sent, ...], tuple[Arrival, ...]]: + row: Final = self.rows[name] + target: Final = self.target + sents: Final = on_each_worker(self.owned, self.workers, lambda: (send(row, target),)) + return sents, self.gateways[name].drain() + + def burst(self, names: tuple[str, ...], each: int) -> tuple[tuple[Row, Sent], ...]: + rows: Final = tuple(chain.from_iterable((self.rows[name],) * each for name in names)) + target: Final = self.target + with ThreadPoolExecutor(max_workers=len(rows)) as pool: + return tuple(zip(rows, pool.map(send, rows, (target,) * len(rows)), strict=True)) + + +@dataclass(frozen=True, slots=True) +class StalePools: + proxy: UpstreamProxy + warmed: Mapping[str, tuple[Arrival, ...]] + + @property + def rows(self) -> Mapping[str, Row]: + return self.proxy.rows + + def probe(self, name: str) -> tuple[tuple[Sent, ...], tuple[Arrival, ...]]: + return self.proxy.probe(name) + + +def burst_answers(sent: tuple[tuple[Row, Sent], ...]) -> tuple[tuple[int, str], ...]: + return tuple((one.answer.status, one.answer.text) for _, one in sent) + + +def burst_expected(sent: tuple[tuple[Row, Sent], ...]) -> tuple[tuple[int, str], ...]: + return tuple((200, row.expected(one.marker)) for row, one in sent) + + +def drain_all(gateways: tuple[IdleGateway, ...]) -> tuple[Arrival, ...]: + return tuple(chain.from_iterable(gateway.drain() for gateway in gateways)) + + +@contextmanager +def proxy_for( + gateway: Gateway, + directory: Path, + rows: tuple[Row, ...], + gateways: Mapping[str, IdleGateway], + steady: str, + overrides: Mapping[str, str], +) -> Generator[UpstreamProxy]: + config: Final = proxy_config( + directory, tuple(row.deploy(row.model, gateways[row.name].url, steady) for row in rows) + ) + with owned_proxy_process( + gateway, directory, overrides, config=config, workers=2, healthcheck=WorkerHealthcheck.LENIENT + ) as owned: + yield UpstreamProxy(owned, two_workers(owned), MappingProxyType({row.name: row for row in rows}), gateways) + + +@contextmanager +def upstream_proxy( + gateway: Gateway, + directory: Path, + rows: tuple[Row, ...], + overrides: Mapping[str, str], + drop_after: float | None = DROP_AFTER_SECONDS, +) -> Generator[UpstreamProxy]: + with ExitStack() as stack: + steady: Final = stack.enter_context(idle_gateway(answer, drop_after=None)) + gateways: Final = MappingProxyType( + { + row.name: stack.enter_context(idle_gateway(row.respond, drop_after=drop_after, drop=row.drop)) + for row in rows + } + ) + yield stack.enter_context(proxy_for(gateway, directory, rows, gateways, steady.url, overrides)) + + +def _warm(rows: tuple[Row, ...], target: ProxyTarget) -> tuple[Sent, ...]: + return tuple(send(row, target) for row in rows) + + +def warm_and_idle(proxy: UpstreamProxy) -> StalePools: + """Gives every row one pooled upstream connection per worker, then leaves each idle for IDLE_SECONDS: past the + 15 s default keep-alive, inside the 120 s one it replaced.""" + target: Final = proxy.target + warm: Final = on_each_worker(proxy.owned, proxy.workers, lambda: _warm(tuple(proxy.rows.values()), target)) + refused: Final = tuple(sent for sent in warm if sent.answer.status != 200) + assert refused == (), refused + warmed: Final = MappingProxyType({name: upstream.drain() for name, upstream in proxy.gateways.items()}) + stay_idle(tuple(proxy.gateways.values()), IDLE_SECONDS) + return StalePools(proxy, warmed) + + +@contextmanager +def stale_pools( + gateway: Gateway, + directory: Path, + rows: tuple[Row, ...], + overrides: Mapping[str, str], + drop_after: float | None = DROP_AFTER_SECONDS, +) -> Generator[StalePools]: + with upstream_proxy(gateway, directory, rows, overrides, drop_after) as proxy: + yield warm_and_idle(proxy) diff --git a/tests/integration/providers/test_idle_upstream_keepalive_edges_wire.py b/tests/integration/providers/test_idle_upstream_keepalive_edges_wire.py new file mode 100644 index 00000000000..59a3d492be7 --- /dev/null +++ b/tests/integration/providers/test_idle_upstream_keepalive_edges_wire.py @@ -0,0 +1,89 @@ +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.process import graceful_stop_seconds, refused_boot_log +from integration.providers._idle_upstream_support import ( + QUIET_STREAM_SECONDS, + REUSE_IDLE_SECONDS, + Row, + UpstreamProxy, + answered, + anthropic_deployment, + chat_sdk, + chat_sdk_stream, + connections, + expected_answers, + landed, + landed_once_undropped, + messages_sdk, + messages_sdk_stream, + openai_deployment, + quiet_stream_answer, + served_by_rust, + spend_statuses, + stay_idle, + upstream_proxy, +) + +pytestmark: Final = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + 2 * QUIET_STREAM_SECONDS)) + +REUSE_ROWS: Final = ( + Row("reuse-chat-openai-sdk-sync", openai_deployment, chat_sdk), + Row("reuse-messages-rust-anthropic-sdk-sync", anthropic_deployment, messages_sdk, rust=True), +) +QUIET_ROWS: Final = ( + Row("quiet-chat-openai-sdk-async-stream", openai_deployment, chat_sdk_stream, respond=quiet_stream_answer), + Row( + "quiet-messages-rust-anthropic-sdk-async-stream", + anthropic_deployment, + messages_sdk_stream, + rust=True, + respond=quiet_stream_answer, + ), +) + + +@pytest.fixture(scope="module") +def edges(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[UpstreamProxy]: + with upstream_proxy( + session_gateway, + tmp_path_factory.mktemp("idle-upstream-edges"), + (*REUSE_ROWS, *QUIET_ROWS), + {"LITELLM_RUST": "1"}, + drop_after=None, + ) as proxy: + yield proxy + + +@pytest.mark.parametrize("name", [row.name for row in REUSE_ROWS]) +def test_request_idle_under_the_keepalive_reuses_the_pooled_connection_on_each_worker( + edges: UpstreamProxy, name: str +) -> None: + row: Final = edges.rows[name] + warm, warmed = edges.probe(name) + stay_idle((edges.gateways[name],), REUSE_IDLE_SECONDS) + probe, arrivals = edges.probe(name) + sents: Final = (*warm, *probe) + assert answered(sents) == expected_answers(row, sents), sents + assert landed((*warmed, *arrivals)) == landed_once_undropped(sents), (warmed, arrivals) + assert connections(arrivals) == connections(warmed), (warmed, arrivals) + assert served_by_rust(sents) == (row.rust,) * len(sents), sents + + +@pytest.mark.parametrize("name", [row.name for row in QUIET_ROWS]) +def test_stream_quiet_for_longer_than_the_keepalive_completes_on_each_worker(edges: UpstreamProxy, name: str) -> None: + row: Final = edges.rows[name] + sents, arrivals = edges.probe(name) + assert answered(sents) == expected_answers(row, sents), sents + assert landed(arrivals) == landed_once_undropped(sents), arrivals + assert served_by_rust(sents) == (row.rust,) * len(sents), sents + assert spend_statuses(sents) == ("success",) * len(sents), sents + + +@pytest.mark.parametrize("value", ["fifteen", ""], ids=["word", "empty"]) +def test_malformed_keepalive_timeout_refuses_to_boot(gateway: Gateway, tmp_path: Path, value: str) -> None: + log: Final = refused_boot_log(gateway, tmp_path, {"AIOHTTP_KEEPALIVE_TIMEOUT": value}) + assert f"invalid literal for int() with base 10: {value!r}" in log, log diff --git a/tests/integration/providers/test_idle_upstream_keepalive_override_wire.py b/tests/integration/providers/test_idle_upstream_keepalive_override_wire.py new file mode 100644 index 00000000000..8b5a69d729e --- /dev/null +++ b/tests/integration/providers/test_idle_upstream_keepalive_override_wire.py @@ -0,0 +1,56 @@ +from collections.abc import Iterator +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.process import graceful_stop_seconds +from integration.providers._idle_upstream_support import ( + IDLE_SECONDS, + Row, + StalePools, + answered, + anthropic_deployment, + chat_sdk, + connections, + expected_answers, + landed, + landed_once_undropped, + messages_sdk, + ocr, + ocr_deployment, + openai_deployment, + served_by_rust, + stale_pools, +) + +pytestmark: Final = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + IDLE_SECONDS)) + +ROWS: Final = ( + Row("override-chat-openai-sdk-sync", openai_deployment, chat_sdk), + Row("override-messages-rust-anthropic-sdk-sync", anthropic_deployment, messages_sdk, rust=True), + Row("override-ocr-mistral-rust-bridge", ocr_deployment, ocr, rust=True), +) + + +@pytest.fixture(scope="module") +def pools(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[StalePools]: + with stale_pools( + session_gateway, + tmp_path_factory.mktemp("idle-upstream-override"), + ROWS, + {"AIOHTTP_KEEPALIVE_TIMEOUT": "120", "LITELLM_RUST": "1"}, + drop_after=None, + ) as value: + yield value + + +@pytest.mark.parametrize("name", [row.name for row in ROWS]) +def test_explicit_keepalive_timeout_keeps_reusing_the_pooled_connection_past_the_default( + pools: StalePools, name: str +) -> None: + row: Final = pools.rows[name] + sents, arrivals = pools.probe(name) + assert answered(sents) == expected_answers(row, sents), sents + assert landed(arrivals) == landed_once_undropped(sents), arrivals + assert connections(arrivals) == connections(pools.warmed[name]), (arrivals, pools.warmed[name]) + assert served_by_rust(sents) == (row.rust,) * len(sents), sents diff --git a/tests/integration/providers/test_idle_upstream_reconnect_chaos_wire.py b/tests/integration/providers/test_idle_upstream_reconnect_chaos_wire.py new file mode 100644 index 00000000000..727b7ff7f5b --- /dev/null +++ b/tests/integration/providers/test_idle_upstream_reconnect_chaos_wire.py @@ -0,0 +1,135 @@ +from contextlib import ExitStack +from dataclasses import replace +from itertools import chain +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.idle_gateway import idle_gateway +from integration._support.process import graceful_stop_seconds +from integration.providers._idle_upstream_support import ( + DROP_AFTER_SECONDS, + IDLE_SECONDS, + REUSE_IDLE_SECONDS, + Row, + answer, + answered, + anthropic_deployment, + burst_answers, + burst_expected, + chat_sdk, + chat_sdk_stream, + drain_all, + expected_answers, + landed, + landed_once_undropped, + messages_sdk, + messages_sdk_stream, + on_each_worker, + openai_deployment, + proxy_for, + responses_sdk, + responses_sdk_stream, + send, + spend_statuses, + stale_pools, + stay_idle, + two_workers, + upstream_proxy, + warm_and_idle, +) + +BURST_EACH: Final = 4 +BURST_ROWS: Final = ( + Row("burst-chat-openai-sdk-sync", openai_deployment, chat_sdk), + Row("burst-chat-openai-sdk-async-stream", openai_deployment, chat_sdk_stream), + Row("burst-messages-anthropic-sdk-sync", anthropic_deployment, messages_sdk), + Row("burst-messages-anthropic-sdk-async-stream", anthropic_deployment, messages_sdk_stream), + Row("burst-responses-openai-sdk-sync", openai_deployment, responses_sdk), + Row("burst-responses-openai-sdk-async-stream", openai_deployment, responses_sdk_stream), +) +KILL_ROWS: Final = ( + Row("kill-chat-openai-sdk-sync", openai_deployment, chat_sdk), + Row("kill-messages-anthropic-sdk-sync", anthropic_deployment, messages_sdk), +) +RESTART_ROWS: Final = ( + Row("restart-chat-openai-sdk-sync", openai_deployment, chat_sdk), + Row("restart-messages-anthropic-sdk-sync", anthropic_deployment, messages_sdk), +) + + +@pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + IDLE_SECONDS)) +def test_concurrent_mixed_burst_after_upstream_idle_past_keepalive_lands_every_request_once( + gateway: Gateway, tmp_path: Path +) -> None: + with stale_pools(gateway, tmp_path, BURST_ROWS, {}) as pools: + sent: Final = pools.proxy.burst(tuple(row.name for row in BURST_ROWS), BURST_EACH) + arrivals: Final = drain_all(tuple(pools.proxy.gateways.values())) + sents: Final = tuple(one for _, one in sent) + statuses: Final = spend_statuses(sents) + assert burst_answers(sent) == burst_expected(sent), sent + assert landed(arrivals) == landed_once_undropped(sents), arrivals + assert statuses == ("success",) * len(sents), sents + + +@pytest.mark.timeout(int(3 * graceful_stop_seconds() + 180 + IDLE_SECONDS)) +def test_worker_killed_while_the_other_serves_then_both_open_fresh_connections_after_upstream_idle( + gateway: Gateway, tmp_path: Path +) -> None: + with upstream_proxy(gateway, tmp_path, KILL_ROWS, {}) as proxy: + victim, _ = proxy.workers + with proxy.owned.replacing_worker(victim) as kill: + kill() + during: Final = send(KILL_ROWS[0], proxy.target) + stale: Final = warm_and_idle(replace(proxy, workers=two_workers(proxy.owned))) + probes: Final = tuple(stale.probe(row.name) for row in KILL_ROWS) + sents: Final = tuple(chain.from_iterable(probe_sents for probe_sents, _ in probes)) + arrivals: Final = tuple(chain.from_iterable(probe_arrivals for _, probe_arrivals in probes)) + statuses: Final = spend_statuses(sents) + expected: Final = tuple( + chain.from_iterable( + expected_answers(row, probe_sents) for row, (probe_sents, _) in zip(KILL_ROWS, probes, strict=True) + ) + ) + assert answered((during,)) == expected_answers(KILL_ROWS[0], (during,)), during + assert answered(sents) == expected, sents + assert landed(arrivals) == landed_once_undropped(sents), arrivals + assert statuses == ("success",) * len(sents), sents + + +@pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + IDLE_SECONDS)) +def test_upstream_restarted_on_the_same_port_with_connections_pooled_serves_a_burst_on_fresh_connections( + gateway: Gateway, tmp_path: Path +) -> None: + with ExitStack() as stack: + steady: Final = stack.enter_context(idle_gateway(answer, drop_after=None)) + before_restart: Final = stack.enter_context(ExitStack()) + first: Final = MappingProxyType( + { + row.name: before_restart.enter_context(idle_gateway(answer, drop_after=DROP_AFTER_SECONDS)) + for row in RESTART_ROWS + } + ) + proxy: Final = stack.enter_context(proxy_for(gateway, tmp_path, RESTART_ROWS, first, steady.url, {})) + target: Final = proxy.target + warm: Final = on_each_worker( + proxy.owned, proxy.workers, lambda: tuple(send(row, target) for row in RESTART_ROWS) + ) + before_restart.close() + second: Final = MappingProxyType( + { + name: stack.enter_context(idle_gateway(answer, drop_after=DROP_AFTER_SECONDS, port=upstream.port)) + for name, upstream in first.items() + } + ) + stay_idle(tuple(second.values()), REUSE_IDLE_SECONDS) + sent: Final = replace(proxy, gateways=second).burst(tuple(row.name for row in RESTART_ROWS), BURST_EACH) + arrivals: Final = drain_all(tuple(second.values())) + sents: Final = tuple(one for _, one in sent) + statuses: Final = spend_statuses(sents) + assert tuple(one.answer.status for one in warm) == (200,) * 2 * len(RESTART_ROWS), warm + assert burst_answers(sent) == burst_expected(sent), sent + assert landed(arrivals) == landed_once_undropped(sents), arrivals + assert statuses == ("success",) * len(sents), sents diff --git a/tests/integration/providers/test_idle_upstream_reconnect_routes_wire.py b/tests/integration/providers/test_idle_upstream_reconnect_routes_wire.py new file mode 100644 index 00000000000..2a8b640f654 --- /dev/null +++ b/tests/integration/providers/test_idle_upstream_reconnect_routes_wire.py @@ -0,0 +1,60 @@ +from collections.abc import Iterator +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.idle_gateway import Drop +from integration._support.process import graceful_stop_seconds +from integration.providers._idle_upstream_support import ( + IDLE_SECONDS, + Row, + StalePools, + answered, + anthropic_deployment, + chat_with_guardrail, + expected_answers, + guardrail_deployment, + landed, + landed_once_undropped, + messages_sdk, + messages_sdk_stream, + ocr, + ocr_deployment, + pass_through, + pass_through_deployment, + served_by_rust, + spend_statuses, + stale_pools, + transcription, + transcription_deployment, +) + +pytestmark: Final = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + IDLE_SECONDS)) + +ROWS: Final = ( + Row("messages-anthropic-sdk-sync", anthropic_deployment, messages_sdk), + Row("messages-anthropic-sdk-async-stream", anthropic_deployment, messages_sdk_stream), + Row("messages-upstream-resets-stale-connection", anthropic_deployment, messages_sdk, drop=Drop.RESET), + Row("custom-pass-through-raw-httpx", pass_through_deployment, pass_through), + Row("generic-guardrail-api-pre-call", guardrail_deployment, chat_with_guardrail), + Row("ocr-mistral-rust-bridge", ocr_deployment, ocr, rust=True), + Row("audio-transcription-bedrock-rust-bridge", transcription_deployment, transcription, rust=True), +) + + +@pytest.fixture(scope="module") +def pools(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[StalePools]: + with stale_pools(session_gateway, tmp_path_factory.mktemp("idle-upstream-routes"), ROWS, {}) as value: + yield value + + +@pytest.mark.parametrize("name", [row.name for row in ROWS]) +def test_route_after_upstream_idle_past_keepalive_opens_a_fresh_connection_on_each_worker( + pools: StalePools, name: str +) -> None: + row: Final = pools.rows[name] + sents, arrivals = pools.probe(name) + assert answered(sents) == expected_answers(row, sents), sents + assert landed(arrivals) == landed_once_undropped(sents), arrivals + assert served_by_rust(sents) == (row.rust,) * len(sents), sents + assert spend_statuses(sents) == ("success",) * len(sents), sents diff --git a/tests/integration/providers/test_idle_upstream_reconnect_rust_messages_wire.py b/tests/integration/providers/test_idle_upstream_reconnect_rust_messages_wire.py new file mode 100644 index 00000000000..8d92fb609e9 --- /dev/null +++ b/tests/integration/providers/test_idle_upstream_reconnect_rust_messages_wire.py @@ -0,0 +1,56 @@ +from collections.abc import Iterator +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.idle_gateway import Drop +from integration._support.process import graceful_stop_seconds +from integration.providers._idle_upstream_support import ( + IDLE_SECONDS, + Row, + StalePools, + answered, + anthropic_deployment, + expected_answers, + landed, + landed_once_undropped, + messages_sdk, + messages_sdk_stream, + served_by_rust, + spend_statuses, + stale_pools, +) + +pytestmark: Final = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + IDLE_SECONDS)) + +ROWS: Final = ( + Row("messages-rust-anthropic-sdk-sync", anthropic_deployment, messages_sdk, rust=True), + Row("messages-rust-anthropic-sdk-async-stream", anthropic_deployment, messages_sdk_stream, rust=True), + Row( + "messages-rust-upstream-resets-stale-connection", + anthropic_deployment, + messages_sdk, + drop=Drop.RESET, + rust=True, + ), +) + + +@pytest.fixture(scope="module") +def pools(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[StalePools]: + with stale_pools( + session_gateway, tmp_path_factory.mktemp("idle-upstream-rust"), ROWS, {"LITELLM_RUST": "1"} + ) as value: + yield value + + +@pytest.mark.parametrize("name", [row.name for row in ROWS]) +def test_rust_messages_after_upstream_idle_past_keepalive_open_a_fresh_connection_on_each_worker( + pools: StalePools, name: str +) -> None: + row: Final = pools.rows[name] + sents, arrivals = pools.probe(name) + assert answered(sents) == expected_answers(row, sents), sents + assert landed(arrivals) == landed_once_undropped(sents), arrivals + assert served_by_rust(sents) == (row.rust,) * len(sents), sents + assert spend_statuses(sents) == ("success",) * len(sents), sents diff --git a/tests/integration/providers/test_idle_upstream_reconnect_wire.py b/tests/integration/providers/test_idle_upstream_reconnect_wire.py new file mode 100644 index 00000000000..820eff4546f --- /dev/null +++ b/tests/integration/providers/test_idle_upstream_reconnect_wire.py @@ -0,0 +1,62 @@ +from collections.abc import Iterator +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.idle_gateway import Drop +from integration._support.process import graceful_stop_seconds +from integration.providers._idle_upstream_support import ( + IDLE_SECONDS, + Row, + StalePools, + answered, + anthropic_deployment, + chat_httpx, + chat_sdk, + chat_sdk_stream, + chat_sdk_with_default_retries, + embedding_deployment, + embeddings_sdk, + expected_answers, + expected_embedding, + landed, + landed_once_undropped, + openai_deployment, + responses_sdk, + responses_sdk_stream, + served_by_rust, + spend_statuses, + stale_pools, +) + +pytestmark: Final = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120 + IDLE_SECONDS)) + +ROWS: Final = ( + Row("chat-openai-sdk-sync", openai_deployment, chat_sdk), + Row("chat-openai-sdk-async-stream", openai_deployment, chat_sdk_stream), + Row("chat-raw-httpx", openai_deployment, chat_httpx), + Row("chat-to-anthropic-deployment", anthropic_deployment, chat_sdk), + Row("embeddings-openai-sdk", embedding_deployment, embeddings_sdk, expected=expected_embedding), + Row("responses-openai-sdk-sync", openai_deployment, responses_sdk), + Row("responses-openai-sdk-async-stream", openai_deployment, responses_sdk_stream), + Row("chat-upstream-resets-stale-connection", openai_deployment, chat_sdk, drop=Drop.RESET), + Row("chat-with-default-router-retries", openai_deployment, chat_sdk_with_default_retries), +) + + +@pytest.fixture(scope="module") +def pools(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[StalePools]: + with stale_pools(session_gateway, tmp_path_factory.mktemp("idle-upstream"), ROWS, {}) as value: + yield value + + +@pytest.mark.parametrize("name", [row.name for row in ROWS]) +def test_request_after_upstream_idle_past_keepalive_opens_a_fresh_connection_on_each_worker( + pools: StalePools, name: str +) -> None: + row: Final = pools.rows[name] + sents, arrivals = pools.probe(name) + assert answered(sents) == expected_answers(row, sents), sents + assert landed(arrivals) == landed_once_undropped(sents), arrivals + assert served_by_rust(sents) == (row.rust,) * len(sents), sents + assert spend_statuses(sents) == ("success",) * len(sents), sents diff --git a/tests/integration/sdk/test_idle_upstream_reconnect_sdk.py b/tests/integration/sdk/test_idle_upstream_reconnect_sdk.py new file mode 100644 index 00000000000..617e756426c --- /dev/null +++ b/tests/integration/sdk/test_idle_upstream_reconnect_sdk.py @@ -0,0 +1,94 @@ +import asyncio +from collections.abc import Awaitable, Callable +from types import MappingProxyType +from typing import Final + +import openai +import pytest +from integration._support.idle_gateway import idle_gateway +from integration.providers._idle_upstream_support import ( + DROP_AFTER_SECONDS, + IDLE_SECONDS, + PROVIDER_KEY, + answer, + expected_text, + landed, + new_marker, + stay_idle, +) + +import litellm +from litellm.types.utils import Choices, ModelResponse + + +def _text(response: object) -> str: + assert isinstance(response, ModelResponse), response + choice: Final = response.choices[0] + assert isinstance(choice, Choices), choice + return choice.message.content or "" + + +async def _chat(model: str, api_base: str, marker: str) -> str: + return _text( + await litellm.acompletion( + model=model, + api_base=api_base, + api_key=PROVIDER_KEY, + messages=[{"role": "user", "content": marker}], + num_retries=0, + max_retries=0, + ) + ) + + +async def openai_chat(api_base: str, marker: str) -> str: + return await _chat("openai/idle-upstream", api_base, marker) + + +async def anthropic_chat(api_base: str, marker: str) -> str: + return await _chat("anthropic/claude-idle-upstream", api_base, marker) + + +async def mistral_ocr_rust_bridge(api_base: str, marker: str) -> str: + response: Final = await litellm.aocr( + model="mistral/mistral-ocr-latest", + document={"type": "document_url", "document_url": f"https://documents.invalid/{marker}.pdf"}, + api_base=api_base, + api_key=PROVIDER_KEY, + ) + return response.pages[0].markdown + + +CALLS: Final[MappingProxyType[str, Callable[[str, str], Awaitable[str]]]] = MappingProxyType( + { + "acompletion-openai": openai_chat, + "acompletion-anthropic": anthropic_chat, + "aocr-mistral-rust-bridge": mistral_ocr_rust_bridge, + } +) + + +async def _outcome(call: Callable[[str, str], Awaitable[str]], api_base: str, marker: str) -> str: + try: + return await call(api_base, marker) + except openai.APIError as error: + return f"raised {type(error).__name__}: {error}" + + +@pytest.mark.timeout(int(IDLE_SECONDS + 120)) +@pytest.mark.parametrize("name", list(CALLS)) +async def test_sdk_call_after_upstream_idle_past_keepalive_opens_a_fresh_connection(name: str) -> None: + call: Final = CALLS[name] + warm_marker: Final = new_marker() + probe_marker: Final = new_marker() + with idle_gateway(answer, drop_after=DROP_AFTER_SECONDS) as upstream: + warm: Final = await _outcome(call, upstream.url, warm_marker) + warmed: Final = upstream.drain() + await asyncio.to_thread(stay_idle, (upstream,), IDLE_SECONDS) + probe: Final = await _outcome(call, upstream.url, probe_marker) + arrivals: Final = upstream.drain() + assert (warm, probe) == (expected_text(warm_marker), expected_text(probe_marker)), (warm, probe) + assert landed((*warmed, *arrivals)) == tuple(sorted(((warm_marker, False), (probe_marker, False)))), ( + warmed, + arrivals, + )