Skip to content

Commit 1f6cfe1

Browse files
marcuswood-oaimarcuswood-oai
authored
fix(streaming): restore async HTTP/1.1 connection reuse (#4041)
## Changes being requested Restore **async HTTP/1.1 streaming connection reuse** after `[DONE]` with a 100 ms drain budget. When the HTTP body finishes within that budget, the connection returns to the pool. A timeout or transport failure closes it without failing completed output. The drain reads raw transport bytes without decoding trailing SSE or content. An AnyIO cancellation deadline handles stalled and trickling endings, including when the request read timeout is disabled. Response cleanup runs outside that deadline and is shielded from caller cancellation; cancellation still propagates. This applies only to async streams receiving `[DONE]` over HTTP/1.1. Synchronous streaming, HTTP/2, Responses streams without `[DONE]`, early exits, and pre-completion errors keep their existing behavior. No watchdog threads, dependency changes, or public API changes. Builds on the fresh transport-iterator approach in [Michał Rembalski's proposal](mrembalski@5d615ad). ## Why async only The synchronous approach we evaluated needed a temporary watchdog thread per stream to enforce a total drain deadline. A per-read timeout alone does not bound a response that keeps trickling bytes. In three synthetic runs of 1,024 requests at concurrency 128, sync HTTPX2 median batch time increased from **2.70 s to 3.23 s**. A diagnostic variant that retained reuse but removed the watchdog took **2.73 s**, isolating thread creation/joining as the main added cost. We therefore leave synchronous streaming unchanged rather than introduce that overhead in a performance fix. Async cleanup can enforce the deadline through cancellation without watchdog threads. A bounded sync implementation without the measured regression is separate work; HTTP/2 remains a connection-reuse workaround where supported. ## Validation Three focused regression tests cover: - Async connection reuse and recovery from silent, trickling, and truncated endings with read timeouts disabled, across HTTPX2, legacy HTTPX, and aiohttp. - Draining only async HTTP/1.1 streams, leaving sync and HTTP/2 unchanged. - Underlying-stream cleanup on caller cancellation and deadline expiry. After narrowing the scope: - Streaming/chat/agent suites: 217 passed, 8 skipped. - Pydantic v1 streaming/chat suites: 81 passed, 34 skipped. - Ruff, repository Pyright, and mypy passed. - Structural comparison confirms sync `Stream` matches main and the async implementation matches the version exercised by the TCP/TLS stress tests and 120 asyncio/Trio cancellation/recovery cycles. ## Performance and tradeoffs In three synthetic TLS comparisons of 1,024 measured requests at concurrency 128, async HTTPX2 median batch time improved from **3.16 s to 2.95 s** with connection reuse. Legacy HTTPX async has a separate pooling cost at high concurrency: a synthetic batch of 512 requests at concurrency 64 increased from 2.16 s to 5.06 s. Consuming ordinary EOF with an unpatched SDK reproduced that pooling cost. This PR does not change the transport's pool implementation. There is no fixed 100 ms sleep. Cleanup returns as soon as the HTTP body ends; a delayed ending can add up to the drain budget, plus scheduling and close overhead. These synthetic results are not production latency estimates or a guarantee that every workload becomes faster. ## Additional context & links Addresses the async streaming case of #4040. Related to #3440.
1 parent c5452af commit 1f6cfe1

3 files changed

Lines changed: 203 additions & 1 deletion

File tree

‎src/openai/_streaming.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from ._utils import is_mapping, extract_type_var_from_base
1414
from ._httpx2 import request_exceptions, timeout_exceptions
1515
from ._exceptions import APIError, APITimeoutError, APIConnectionError
16+
from ._streaming_cleanup import drain_response_async
1617

1718
if TYPE_CHECKING:
1819
from ._client import OpenAI, AsyncOpenAI
@@ -183,6 +184,11 @@ async def __stream__(self) -> AsyncIterator[_T]:
183184
try:
184185
async for sse in iterator:
185186
if sse.data.startswith("[DONE]"):
187+
# HTTP/1.1 needs the remaining body consumed for connection reuse.
188+
# HTTP/2 can close this stream and keep the connection, so draining
189+
# it would add unnecessary latency after [DONE].
190+
if response.http_version == "HTTP/1.1":
191+
await drain_response_async(response)
186192
break
187193

188194
# we have to special case the Assistants `thread.` events since we won't have an "event" key in the data

‎src/openai/_streaming_cleanup.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
from __future__ import annotations
2+
3+
from typing import AsyncIterator, cast
4+
5+
import anyio
6+
import httpx2
7+
from anyio.lowlevel import checkpoint
8+
9+
from ._httpx2 import request_exceptions, stream_consumed_exceptions
10+
11+
_DRAIN_TIMEOUT = 0.1
12+
13+
14+
async def drain_response_async(response: httpx2.Response) -> None:
15+
if response.is_closed:
16+
return
17+
try:
18+
with anyio.move_on_after(_DRAIN_TIMEOUT):
19+
# Discard raw transport bytes without decoding trailing SSE or content.
20+
async for _ in cast(AsyncIterator[bytes], response.stream):
21+
# Also enforce the deadline on transports yielding buffered data.
22+
await checkpoint()
23+
except request_exceptions() + stream_consumed_exceptions():
24+
pass # Cleanup must not turn completed output into a request error.
25+
finally:
26+
# Close outside our timeout scope and shield cleanup from caller
27+
# cancellation. Cancellation still propagates after the response closes.
28+
with anyio.CancelScope(shield=True):
29+
await response.aclose()

‎tests/test_streaming.py‎

Lines changed: 168 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,21 @@
22

33
import os
44
import ssl
5+
import time
6+
import socket
7+
import asyncio
58
import importlib
9+
import threading
610
from typing import Any, Iterator, AsyncIterator
711
from contextlib import aclosing, nullcontext
12+
from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler
13+
from typing_extensions import override
814

915
import anyio
1016
import httpx2
1117
import pytest
1218

13-
from openai import OpenAI, AsyncOpenAI, APITimeoutError, APIConnectionError
19+
from openai import OpenAI, AsyncOpenAI, APITimeoutError, APIConnectionError, DefaultAioHttpClient
1420
from openai._streaming import Stream, AsyncStream, ServerSentEvent
1521

1622

@@ -29,6 +35,167 @@ def http_module(request: pytest.FixtureRequest) -> Any:
2935
return importlib.import_module(request.param)
3036

3137

38+
@pytest.mark.parametrize("mode", ["async", "aiohttp"])
39+
@pytest.mark.parametrize("ending", ["complete", "truncate", "stall", "trickle"])
40+
async def test_done_http11_connection_reuse(mode: str, ending: str, http_module: Any) -> None:
41+
if mode == "aiohttp" and http_module.__name__ == "httpx":
42+
pytest.skip("aiohttp uses HTTPX2")
43+
connections = 0
44+
requests = 0
45+
release = threading.Event()
46+
47+
class Handler(BaseHTTPRequestHandler):
48+
protocol_version = "HTTP/1.1"
49+
50+
@override
51+
def setup(self) -> None:
52+
nonlocal connections
53+
super().setup()
54+
connections += 1
55+
self.connection.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
56+
57+
def do_POST(self) -> None:
58+
nonlocal requests
59+
self.rfile.read(int(self.headers["Content-Length"]))
60+
requests += 1
61+
self.send_response(200)
62+
self.send_header("Content-Type", "text/event-stream")
63+
self.send_header("Transfer-Encoding", "chunked")
64+
self.end_headers()
65+
payload = (
66+
b'data: {"id":"synthetic","object":"chat.completion.chunk","created":0,"model":"synthetic",'
67+
b'"choices":[{"index":0,"delta":{"content":"hello"}}]}\n\ndata: [DONE]\n\n'
68+
)
69+
try:
70+
self.wfile.write(f"{len(payload):x}\r\n".encode() + payload + b"\r\n")
71+
if requests == 1:
72+
if ending == "truncate":
73+
self.close_connection = True
74+
return
75+
if ending == "stall":
76+
release.wait(timeout=0.6)
77+
if ending == "trickle":
78+
# Trailer bytes do not yield body chunks, but must still
79+
# obey the total cleanup deadline, not a per-read timeout.
80+
self.wfile.write(b"0\r\nX-Synthetic: ")
81+
for _ in range(15):
82+
if release.wait(timeout=0.04):
83+
return
84+
self.wfile.write(b"a")
85+
self.wfile.write(b"\r\n\r\n")
86+
return
87+
self.wfile.write(b"0\r\n\r\n")
88+
except ConnectionError:
89+
self.close_connection = True
90+
91+
@override
92+
def log_message(self, format: str, *args: object) -> None: # noqa: A002
93+
pass
94+
95+
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
96+
thread = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01})
97+
thread.start()
98+
options: dict[str, Any] = {
99+
"api_key": "synthetic",
100+
"base_url": f"http://127.0.0.1:{server.server_port}/v1",
101+
"max_retries": 0,
102+
"timeout": http_module.Timeout(5, read=None),
103+
}
104+
transport_options: dict[str, Any] = {"trust_env": False, "limits": http_module.Limits(max_connections=1)}
105+
try:
106+
http_client = (
107+
DefaultAioHttpClient(**transport_options)
108+
if mode == "aiohttp"
109+
else http_module.AsyncClient(**transport_options)
110+
)
111+
async with AsyncOpenAI(**options, http_client=http_client) as async_client:
112+
for _ in range(2):
113+
async with await async_client.chat.completions.create(
114+
model="synthetic", messages=[], stream=True
115+
) as async_stream:
116+
assert (await async_stream.__anext__()).choices[0].delta.content == "hello"
117+
start = time.monotonic()
118+
assert [chunk async for chunk in async_stream] == []
119+
assert time.monotonic() - start < 0.4
120+
finally:
121+
release.set()
122+
server.shutdown()
123+
thread.join(timeout=5)
124+
server.server_close()
125+
126+
assert not thread.is_alive()
127+
assert requests == 2
128+
# A failed drain must discard the connection and release the only pool slot.
129+
assert connections == (1 if ending == "complete" else 2)
130+
131+
132+
@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
133+
@pytest.mark.parametrize("http_version", [b"HTTP/1.1", b"HTTP/2"])
134+
async def test_done_discards_bytes_only_for_async_http11(
135+
sync: bool, http_version: bytes, http_module: Any, client: OpenAI, async_client: AsyncOpenAI
136+
) -> None:
137+
read_tail = False
138+
139+
def body() -> Iterator[bytes]:
140+
nonlocal read_tail
141+
yield b"data: [DONE]\n\n"
142+
read_tail = True
143+
yield b"\xff\n\n" # Trailing data must not be decoded as SSE.
144+
145+
iterator = body()
146+
async_iterator = to_aiter(iterator)
147+
148+
class Body(http_module.SyncByteStream, http_module.AsyncByteStream):
149+
def __iter__(self) -> Iterator[bytes]:
150+
return iterator
151+
152+
def __aiter__(self) -> AsyncIterator[bytes]:
153+
return async_iterator
154+
155+
response = http_module.Response(200, stream=Body(), extensions={"http_version": http_version})
156+
if sync:
157+
assert list(Stream(cast_to=object, response=response, client=client)) == []
158+
else:
159+
assert [event async for event in AsyncStream(cast_to=object, response=response, client=async_client)] == []
160+
assert response.is_closed
161+
assert read_tail == (not sync and http_version == b"HTTP/1.1")
162+
163+
164+
@pytest.mark.parametrize("cancel", [True, False], ids=["caller-cancel", "deadline"])
165+
async def test_cancellation_during_done_drain_closes_response(
166+
cancel: bool, http_module: Any, async_client: AsyncOpenAI
167+
) -> None:
168+
draining = asyncio.Event()
169+
170+
class Body(http_module.AsyncByteStream):
171+
closed = False
172+
173+
async def __aiter__(self) -> AsyncIterator[bytes]:
174+
yield b"data: [DONE]\n\n"
175+
draining.set()
176+
await asyncio.Event().wait()
177+
178+
async def aclose(self) -> None:
179+
await anyio.sleep(0)
180+
self.closed = True
181+
182+
body = Body()
183+
response = http_module.Response(200, stream=body)
184+
stream = AsyncStream(cast_to=object, response=response, client=async_client)
185+
task = asyncio.create_task(stream.__anext__())
186+
try:
187+
await asyncio.wait_for(draining.wait(), timeout=5)
188+
if cancel:
189+
task.cancel()
190+
with pytest.raises(asyncio.CancelledError if cancel else StopAsyncIteration):
191+
await task
192+
assert response.is_closed
193+
assert body.closed
194+
finally:
195+
task.cancel()
196+
await asyncio.gather(task, return_exceptions=True)
197+
198+
32199
@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
33200
@pytest.mark.parametrize("delivered", [False, True], ids=["before-first-event", "after-first-event"])
34201
@pytest.mark.parametrize(

0 commit comments

Comments
 (0)