Spaces:
Build error
Build error
File size: 2,731 Bytes
51ff9e5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
"""
LiteLLM currently have an issue where HttpHandlers are being created but not
closed. We have submitted a PR to them, (https://github.com/BerriAI/litellm/pull/8711)
and their dev team say they are in the process of a refactor that will fix this, but
in the meantime, we need to manage the lifecycle of the httpx.Client manually.
We can't simply pass in our own client object, because all the different implementations use
different types of client object.
So we monkey patch the httpx.Client class to track newly created instances and close these
when the operations complete. (Since some paths create a single shared client and reuse these,
we actually need to create a proxy object that allows these clients to be reusable.)
Hopefully, this will be fixed soon and we can remove this abomination.
"""
import contextlib
from typing import Callable
import httpx
@contextlib.contextmanager
def ensure_httpx_close():
wrapped_class = httpx.Client
proxys = []
class ClientProxy:
"""
Sometimes LiteLLM opens a new httpx client for each connection, and does not close them.
Sometimes it does close them. Sometimes, it reuses a client between connections. For cases
where a client is reused, we need to be able to reuse the client even after closing it.
"""
client_constructor: Callable
args: tuple
kwargs: dict
client: httpx.Client
def __init__(self, *args, **kwargs):
self.args = args
self.kwargs = kwargs
self.client = wrapped_class(*self.args, **self.kwargs)
proxys.append(self)
def __getattr__(self, name):
# Invoke a method on the proxied client - create one if required
if self.client is None:
self.client = wrapped_class(*self.args, **self.kwargs)
return getattr(self.client, name)
def close(self):
# Close the client if it is open
if self.client:
self.client.close()
self.client = None
def __iter__(self, *args, **kwargs):
# We have to override this as debuggers invoke it causing the client to reopen
if self.client:
return self.client.iter(*args, **kwargs)
return object.__getattribute__(self, 'iter')(*args, **kwargs)
@property
def is_closed(self):
# Check if closed
if self.client is None:
return True
return self.client.is_closed
httpx.Client = ClientProxy
try:
yield
finally:
httpx.Client = wrapped_class
while proxys:
proxy = proxys.pop()
proxy.close()
|