diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 3789f546d7..3dfef07d42 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -28,6 +28,7 @@ from litellm.constants import ( AIOHTTP_CONNECTOR_LIMIT, AIOHTTP_CONNECTOR_LIMIT_PER_HOST, AIOHTTP_KEEPALIVE_TIMEOUT, + AIOHTTP_NEEDS_CLEANUP_CLOSED, AIOHTTP_TTL_DNS_CACHE, DEFAULT_SSL_CIPHERS, ) @@ -876,9 +877,10 @@ class AsyncHTTPHandler: transport_connector_kwargs = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, - "enable_cleanup_closed": True, **connector_kwargs, } + if AIOHTTP_NEEDS_CLEANUP_CLOSED: + transport_connector_kwargs["enable_cleanup_closed"] = True if AIOHTTP_CONNECTOR_LIMIT > 0: transport_connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5d0c7a89d8..462e865354 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -714,8 +714,9 @@ async def _initialize_shared_aiohttp_session(): connector_kwargs = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, - "enable_cleanup_closed": True, } + if AIOHTTP_NEEDS_CLEANUP_CLOSED: + connector_kwargs["enable_cleanup_closed"] = True if AIOHTTP_CONNECTOR_LIMIT > 0: connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0: diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py new file mode 100644 index 0000000000..c8c0e09c08 --- /dev/null +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_cleanup_closed.py @@ -0,0 +1,31 @@ +from unittest.mock import MagicMock, patch + + +def test_create_aiohttp_transport_sets_enable_cleanup_closed_when_needed(monkeypatch): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + monkeypatch.setattr(http_handler_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", True) + + with patch.object(http_handler_module, "TCPConnector", return_value=connector_mock) as mock_tcp_connector: + with patch.object(http_handler_module, "ClientSession", return_value=session_mock): + transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(shared_session=None) + transport._get_valid_client_session() + + assert mock_tcp_connector.call_args.kwargs["enable_cleanup_closed"] is True + + +def test_create_aiohttp_transport_omits_enable_cleanup_closed_when_not_needed(monkeypatch): + from litellm.llms.custom_httpx import http_handler as http_handler_module + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + monkeypatch.setattr(http_handler_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", False) + + with patch.object(http_handler_module, "TCPConnector", return_value=connector_mock) as mock_tcp_connector: + with patch.object(http_handler_module, "ClientSession", return_value=session_mock): + transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(shared_session=None) + transport._get_valid_client_session() + + assert "enable_cleanup_closed" not in mock_tcp_connector.call_args.kwargs diff --git a/tests/test_litellm/proxy/test_aiohttp_cleanup_closed.py b/tests/test_litellm/proxy/test_aiohttp_cleanup_closed.py new file mode 100644 index 0000000000..f16b687a24 --- /dev/null +++ b/tests/test_litellm/proxy/test_aiohttp_cleanup_closed.py @@ -0,0 +1,34 @@ +import asyncio +from unittest.mock import MagicMock, patch + + +def test_initialize_shared_aiohttp_session_sets_enable_cleanup_closed_when_needed( + monkeypatch, +): + from litellm.proxy import proxy_server as proxy_server_module + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + monkeypatch.setattr(proxy_server_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", True) + + with patch("aiohttp.TCPConnector", return_value=connector_mock) as mock_tcp_connector: + with patch("aiohttp.ClientSession", return_value=session_mock): + asyncio.run(proxy_server_module._initialize_shared_aiohttp_session()) + + assert mock_tcp_connector.call_args.kwargs["enable_cleanup_closed"] is True + + +def test_initialize_shared_aiohttp_session_omits_enable_cleanup_closed_when_not_needed( + monkeypatch, +): + from litellm.proxy import proxy_server as proxy_server_module + + connector_mock = MagicMock(name="connector") + session_mock = MagicMock(name="session") + monkeypatch.setattr(proxy_server_module, "AIOHTTP_NEEDS_CLEANUP_CLOSED", False) + + with patch("aiohttp.TCPConnector", return_value=connector_mock) as mock_tcp_connector: + with patch("aiohttp.ClientSession", return_value=session_mock): + asyncio.run(proxy_server_module._initialize_shared_aiohttp_session()) + + assert "enable_cleanup_closed" not in mock_tcp_connector.call_args.kwargs