diff --git a/litellm/constants.py b/litellm/constants.py index 05ccfb5993..5ab56bcaa2 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -100,8 +100,10 @@ RUNWAYML_POLLING_TIMEOUT = int( ########## Networking constants ############################################################## _DEFAULT_TTL_FOR_HTTPX_CLIENTS = 3600 # 1 hour, re-use the same httpx client for 1 hour -# Aiohttp connection pooling constants -AIOHTTP_CONNECTOR_LIMIT = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 0)) +# Aiohttp connection pooling - prevents memory leaks from unbounded connection growth +# Set to 0 for unlimited (not recommended for production) +AIOHTTP_CONNECTOR_LIMIT = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 300)) +AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 50)) AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120)) AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300)) # enable_cleanup_closed is only needed for Python versions with the SSL leak bug diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index b06e8463ab..d9a9d4f9dc 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -16,6 +16,7 @@ from litellm._logging import verbose_logger from litellm.constants import ( _DEFAULT_TTL_FOR_HTTPX_CLIENTS, AIOHTTP_CONNECTOR_LIMIT, + AIOHTTP_CONNECTOR_LIMIT_PER_HOST, AIOHTTP_KEEPALIVE_TIMEOUT, AIOHTTP_NEEDS_CLEANUP_CLOSED, AIOHTTP_TTL_DNS_CACHE, @@ -793,15 +794,20 @@ class AsyncHTTPHandler: verbose_logger.debug( "NEW SESSION: Creating new ClientSession (no shared session provided)" ) + transport_connector_kwargs = { + "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, + "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, + "enable_cleanup_closed": True, + **connector_kwargs, + } + if AIOHTTP_CONNECTOR_LIMIT > 0: + transport_connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT + if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0: + transport_connector_kwargs["limit_per_host"] = AIOHTTP_CONNECTOR_LIMIT_PER_HOST + return LiteLLMAiohttpTransport( client=lambda: ClientSession( - connector=TCPConnector( - limit=AIOHTTP_CONNECTOR_LIMIT, - keepalive_timeout=AIOHTTP_KEEPALIVE_TIMEOUT, - ttl_dns_cache=AIOHTTP_TTL_DNS_CACHE, - enable_cleanup_closed=AIOHTTP_NEEDS_CLEANUP_CLOSED, - **connector_kwargs, - ), + connector=TCPConnector(**transport_connector_kwargs), trust_env=trust_env, ), ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a1e01caddc..f2c7c361c4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -31,6 +31,7 @@ from typing import ( from litellm._uuid import uuid from litellm.constants import ( AIOHTTP_CONNECTOR_LIMIT, + AIOHTTP_CONNECTOR_LIMIT_PER_HOST, AIOHTTP_KEEPALIVE_TIMEOUT, AIOHTTP_NEEDS_CLEANUP_CLOSED, AIOHTTP_TTL_DNS_CACHE, @@ -627,21 +628,26 @@ async def proxy_shutdown_event(): async def _initialize_shared_aiohttp_session(): - """Initialize shared aiohttp session for connection reuse.""" + """Initialize shared aiohttp session for connection reuse with connection limits.""" try: from aiohttp import ClientSession, TCPConnector - # Create connector with connection pooling settings optimized for long-lived connections - connector = TCPConnector( - limit=AIOHTTP_CONNECTOR_LIMIT, - keepalive_timeout=AIOHTTP_KEEPALIVE_TIMEOUT, - ttl_dns_cache=AIOHTTP_TTL_DNS_CACHE, - enable_cleanup_closed=AIOHTTP_NEEDS_CLEANUP_CLOSED, - ) - + connector_kwargs = { + "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, + "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, + "enable_cleanup_closed": True, + } + if AIOHTTP_CONNECTOR_LIMIT > 0: + connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT + if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0: + connector_kwargs["limit_per_host"] = AIOHTTP_CONNECTOR_LIMIT_PER_HOST + + connector = TCPConnector(**connector_kwargs) session = ClientSession(connector=connector) + verbose_proxy_logger.info( - f"SESSION REUSE: Created shared aiohttp session for connection pooling (ID: {id(session)})" + f"SESSION REUSE: Created shared aiohttp session for connection pooling (ID: {id(session)}, " + f"limit={AIOHTTP_CONNECTOR_LIMIT}, limit_per_host={AIOHTTP_CONNECTOR_LIMIT_PER_HOST})" ) return session except Exception as e: