scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

refactor(http): shrink retry client stream ownership fix

Browse source 
Generated by Mistral Vibe.
Co-Authored-By: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Simon Sorg 2026-11-05 19:38:50 +01:00
parent 23614f5f54
commit deb53fe784
No known key found for this signature in database
1 changed files with 76 additions and 74 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

150
litellm/llms/custom_httpx/http_handler.py
View file

@ -256,11 +256,10 @@ def _anchor_handler_to(response: httpx.Response, handler: object) -> None:
weakref.finalize(response, _drop_streaming_anchor, handler)
class _SingleUseClientStream(httpx.AsyncByteStream):
class _RetryClientStream(httpx.AsyncByteStream):
def __init__(self, stream: httpx.AsyncByteStream, client: httpx.AsyncClient) -> None:
self._stream = stream
self._client = client
self._closed = False
async def __aiter__(self) -> AsyncIterator[bytes]:
try:
@ -270,13 +269,20 @@ class _SingleUseClientStream(httpx.AsyncByteStream):
await self.aclose()
async def aclose(self) -> None:
if self._closed:
return
self._closed = True
try:
await self._stream.aclose()
finally:
await self._client.aclose()
if not self._client.is_closed:
try:
await self._stream.aclose()
finally:
await self._client.aclose()
async def _hand_retry_client_to(response: httpx.Response, client: httpx.AsyncClient) -> httpx.Response:
if response.is_closed:
await client.aclose()
else:
assert isinstance(response.stream, httpx.AsyncByteStream), "An async retry must return an async stream"
response.stream = _RetryClientStream(response.stream, client)
return response
def blocked_cookie_jar() -> CookieJar:
@ -994,17 +1000,21 @@ class AsyncHTTPHandler:
except httpx.ConnectError:
# Retry the request with a new session if there is a connection error
new_client: Final = self.create_client(timeout=timeout, event_hooks=self.event_hooks)
return await self.single_connection_request(
url=url,
method="POST",
client=new_client,
data=data,
json=json,
params=params,
headers=headers,
stream=stream,
content=content,
)
try:
return await self.single_connection_request(
url=url,
method="POST",
client=new_client,
data=data,
json=json,
params=params,
headers=headers,
stream=stream,
content=content,
)
except BaseException:
await new_client.aclose()
raise
except httpx.TimeoutException as e:
_raise_timeout(e, start_time)
except httpx.HTTPStatusError as e:
@ -1066,18 +1076,10 @@ class AsyncHTTPHandler:
retried.raise_for_status()
except httpx.HTTPStatusError as retried_error:
await _raise_masked_async_error(retried_error, stream)
return await _hand_retry_client_to(retried, new_client)
except BaseException:
await new_client.aclose()
raise
else:
if stream and not retried.is_closed:
assert isinstance(retried.stream, httpx.AsyncByteStream), (
"An async retry must return an async stream"
)
retried.stream = _SingleUseClientStream(retried.stream, new_client)
else:
await new_client.aclose()
return retried
except httpx.TimeoutException as e:
_raise_timeout(e, start_time)
except httpx.HTTPStatusError as e:
@ -1120,17 +1122,21 @@ class AsyncHTTPHandler:
except httpx.ConnectError:
# Retry the request with a new session if there is a connection error
new_client: Final = self.create_client(timeout=timeout, event_hooks=self.event_hooks)
return await self.single_connection_request(
url=url,
method="PATCH",
client=new_client,
data=data,
json=json,
params=params,
headers=headers,
stream=stream,
content=content,
)
try:
return await self.single_connection_request(
url=url,
method="PATCH",
client=new_client,
data=data,
json=json,
params=params,
headers=headers,
stream=stream,
content=content,
)
except BaseException:
await new_client.aclose()
raise
except httpx.TimeoutException as e:
_raise_timeout(e, start_time)
except httpx.HTTPStatusError as e:
@ -1174,17 +1180,21 @@ class AsyncHTTPHandler:
except (httpx.RemoteProtocolError, httpx.ConnectError):
# Retry the request with a new session if there is a connection error
new_client: Final = self.create_client(timeout=timeout, event_hooks=self.event_hooks)
return await self.single_connection_request(
url=url,
method="DELETE",
client=new_client,
data=data,
json=json,
params=params,
headers=headers,
stream=stream,
content=content,
)
try:
return await self.single_connection_request(
url=url,
method="DELETE",
client=new_client,
data=data,
json=json,
params=params,
headers=headers,
stream=stream,
content=content,
)
except BaseException:
await new_client.aclose()
raise
except httpx.HTTPStatusError as e:
await _raise_masked_async_error(e, stream)
except Exception as e:
@ -1207,29 +1217,21 @@ class AsyncHTTPHandler:
Used for retrying connection client errors.
"""
try:
request_data, request_content = _prepare_request_data_and_content(data, content)
req: Final = client.build_request(
method,
url,
data=request_data,
json=json,
params=params,
headers=headers,
content=request_content,
)
response: Final = await client.send(req, stream=stream)
response.raise_for_status()
except BaseException:
await client.aclose()
raise
else:
if stream and not response.is_closed:
assert isinstance(response.stream, httpx.AsyncByteStream), "An async retry must return an async stream"
response.stream = _SingleUseClientStream(response.stream, client)
else:
await client.aclose()
return response
# Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix)
request_data, request_content = _prepare_request_data_and_content(data, content)
req: Final = client.build_request(
method,
url,
data=request_data,
json=json,
params=params,
headers=headers,
content=request_content,
)
response: Final = await client.send(req, stream=stream)
response.raise_for_status()
return await _hand_retry_client_to(response, client)
# Strong references to finalizer-scheduled client-close tasks. A bare
# create_task() result may be garbage-collected before it runs, leaving