Skip to content

Commit 9037772

Browse files
shriharsha98shriharsha98
committed
fix(aiohttp): enable TCP keepalive on connector to prevent NAT-dropped connection hangs
1 parent 3919ce1 commit 9037772

3 files changed

Lines changed: 168 additions & 0 deletions

File tree

‎litellm/constants.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,20 @@
203203
)
204204
AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120))
205205
AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300))
206+
# TCP keepalive probes on the aiohttp connector sockets.
207+
# keepalive_timeout above only governs *idle pooled* connections. A connection
208+
# actively waiting on a response (long prefill / high TTFT) is not idle, so a
209+
# NAT/gateway between the proxy and the backend can silently drop the flow
210+
# mid-request -> the backend's response is blackholed and the request hangs until
211+
# the read timeout (e.g. 540s). SO_KEEPALIVE probes keep the NAT mapping warm and
212+
# tear a genuinely-dead socket down in ~IDLE + INTVL*CNT seconds instead.
213+
AIOHTTP_TCP_KEEPALIVE = os.getenv("AIOHTTP_TCP_KEEPALIVE", "True").lower() in (
214+
"true",
215+
"1",
216+
)
217+
AIOHTTP_TCP_KEEPALIVE_IDLE = int(os.getenv("AIOHTTP_TCP_KEEPALIVE_IDLE", 30))
218+
AIOHTTP_TCP_KEEPALIVE_INTVL = int(os.getenv("AIOHTTP_TCP_KEEPALIVE_INTVL", 10))
219+
AIOHTTP_TCP_KEEPALIVE_CNT = int(os.getenv("AIOHTTP_TCP_KEEPALIVE_CNT", 3))
206220
# enable_cleanup_closed is only needed for Python versions with the SSL leak bug
207221
# Fixed in Python 3.12.7+ and 3.13.1+ (see https://github.com/python/cpython/pull/118960)
208222
# Reference: https://github.com/aio-libs/aiohttp/blob/master/aiohttp/connector.py#L74-L78

‎litellm/llms/custom_httpx/http_handler.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
import asyncio
2+
import inspect
23
import os
4+
import socket
35
import ssl
46
import sys
57
import time
@@ -29,6 +31,10 @@
2931
AIOHTTP_CONNECTOR_LIMIT_PER_HOST,
3032
AIOHTTP_KEEPALIVE_TIMEOUT,
3133
AIOHTTP_NEEDS_CLEANUP_CLOSED,
34+
AIOHTTP_TCP_KEEPALIVE,
35+
AIOHTTP_TCP_KEEPALIVE_CNT,
36+
AIOHTTP_TCP_KEEPALIVE_IDLE,
37+
AIOHTTP_TCP_KEEPALIVE_INTVL,
3238
AIOHTTP_TTL_DNS_CACHE,
3339
DEFAULT_SSL_CIPHERS,
3440
)
@@ -72,6 +78,43 @@ def get_default_headers() -> dict:
7278
# https://www.python-httpx.org/advanced/timeouts
7379
_DEFAULT_TIMEOUT = httpx.Timeout(timeout=5.0, connect=5.0)
7480

81+
# aiohttp gained the public ``socket_factory`` TCPConnector arg in 3.11. Detect it
82+
# once so the keepalive wiring is a no-op (rather than a TypeError) on older pins.
83+
_AIOHTTP_SUPPORTS_SOCKET_FACTORY = (
84+
"socket_factory" in inspect.signature(TCPConnector.__init__).parameters
85+
)
86+
87+
88+
def _tcp_keepalive_socket_factory(
89+
addr_info: Tuple[Any, ...],
90+
) -> socket.socket:
91+
"""
92+
socket_factory for aiohttp's TCPConnector that enables TCP keepalive probes.
93+
94+
aiohttp's ``keepalive_timeout`` only governs *idle pooled* connections. A
95+
connection blocked on a slow response (long prefill / high TTFT) is not idle,
96+
so a NAT/gateway can drop the flow mid-request and the response is blackholed
97+
until the read timeout. SO_KEEPALIVE probes keep the NAT mapping warm for a
98+
live-but-slow backend (its kernel ACKs the probes regardless of load) and tear
99+
a genuinely-dead socket down in ~IDLE + INTVL*CNT seconds.
100+
"""
101+
family, type_, proto = addr_info[0], addr_info[1], addr_info[2]
102+
sock = socket.socket(family, type_, proto)
103+
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
104+
if hasattr(socket, "TCP_KEEPIDLE"):
105+
sock.setsockopt(
106+
socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, AIOHTTP_TCP_KEEPALIVE_IDLE
107+
)
108+
if hasattr(socket, "TCP_KEEPINTVL"):
109+
sock.setsockopt(
110+
socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, AIOHTTP_TCP_KEEPALIVE_INTVL
111+
)
112+
if hasattr(socket, "TCP_KEEPCNT"):
113+
sock.setsockopt(
114+
socket.IPPROTO_TCP, socket.TCP_KEEPCNT, AIOHTTP_TCP_KEEPALIVE_CNT
115+
)
116+
return sock
117+
75118

76119
def _prepare_request_data_and_content(
77120
data: Optional[Union[dict, str, bytes]] = None,
@@ -881,6 +924,17 @@ def _create_aiohttp_transport(
881924
"ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE,
882925
**connector_kwargs,
883926
}
927+
# Enable TCP keepalive probes so half-open connections across a NAT are
928+
# detected (and torn down) fast instead of blackholing until the read
929+
# timeout. Requires aiohttp >= 3.11 (socket_factory) and platform support.
930+
if (
931+
AIOHTTP_TCP_KEEPALIVE
932+
and _AIOHTTP_SUPPORTS_SOCKET_FACTORY
933+
and hasattr(socket, "TCP_KEEPIDLE")
934+
):
935+
transport_connector_kwargs[
936+
"socket_factory"
937+
] = _tcp_keepalive_socket_factory
884938
if AIOHTTP_NEEDS_CLEANUP_CLOSED:
885939
transport_connector_kwargs["enable_cleanup_closed"] = True
886940
if AIOHTTP_CONNECTOR_LIMIT > 0:

‎tests/test_litellm/llms/custom_httpx/test_aiohttp_tcp_keepalive.py‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
import socket
2+
from unittest.mock import MagicMock, patch
3+
4+
5+
def test_create_aiohttp_transport_sets_socket_factory_when_enabled(monkeypatch):
6+
from litellm.llms.custom_httpx import http_handler as http_handler_module
7+
8+
connector_mock = MagicMock(name="connector")
9+
session_mock = MagicMock(name="session")
10+
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPALIVE", True)
11+
monkeypatch.setattr(
12+
http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True
13+
)
14+
15+
with patch.object(
16+
http_handler_module, "TCPConnector", return_value=connector_mock
17+
) as mock_tcp_connector:
18+
with patch.object(
19+
http_handler_module, "ClientSession", return_value=session_mock
20+
):
21+
transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(
22+
shared_session=None
23+
)
24+
transport._get_valid_client_session()
25+
26+
assert (
27+
mock_tcp_connector.call_args.kwargs["socket_factory"]
28+
is http_handler_module._tcp_keepalive_socket_factory
29+
)
30+
31+
32+
def test_create_aiohttp_transport_omits_socket_factory_when_disabled(monkeypatch):
33+
from litellm.llms.custom_httpx import http_handler as http_handler_module
34+
35+
connector_mock = MagicMock(name="connector")
36+
session_mock = MagicMock(name="session")
37+
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPALIVE", False)
38+
monkeypatch.setattr(
39+
http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True
40+
)
41+
42+
with patch.object(
43+
http_handler_module, "TCPConnector", return_value=connector_mock
44+
) as mock_tcp_connector:
45+
with patch.object(
46+
http_handler_module, "ClientSession", return_value=session_mock
47+
):
48+
transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(
49+
shared_session=None
50+
)
51+
transport._get_valid_client_session()
52+
53+
assert "socket_factory" not in mock_tcp_connector.call_args.kwargs
54+
55+
56+
def test_create_aiohttp_transport_omits_socket_factory_when_unsupported(monkeypatch):
57+
"""Older aiohttp without socket_factory must not receive the kwarg (would TypeError)."""
58+
from litellm.llms.custom_httpx import http_handler as http_handler_module
59+
60+
connector_mock = MagicMock(name="connector")
61+
session_mock = MagicMock(name="session")
62+
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPALIVE", True)
63+
monkeypatch.setattr(
64+
http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", False
65+
)
66+
67+
with patch.object(
68+
http_handler_module, "TCPConnector", return_value=connector_mock
69+
) as mock_tcp_connector:
70+
with patch.object(
71+
http_handler_module, "ClientSession", return_value=session_mock
72+
):
73+
transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(
74+
shared_session=None
75+
)
76+
transport._get_valid_client_session()
77+
78+
assert "socket_factory" not in mock_tcp_connector.call_args.kwargs
79+
80+
81+
def test_tcp_keepalive_socket_factory_sets_socket_options():
82+
from litellm.llms.custom_httpx import http_handler as http_handler_module
83+
84+
addr_info = socket.getaddrinfo(
85+
"127.0.0.1", 80, socket.AF_INET, socket.SOCK_STREAM
86+
)[0]
87+
sock = http_handler_module._tcp_keepalive_socket_factory(addr_info)
88+
try:
89+
assert sock.getsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE) == 1
90+
assert sock.getsockopt(
91+
socket.IPPROTO_TCP, socket.TCP_KEEPIDLE
92+
) == http_handler_module.AIOHTTP_TCP_KEEPALIVE_IDLE
93+
assert sock.getsockopt(
94+
socket.IPPROTO_TCP, socket.TCP_KEEPINTVL
95+
) == http_handler_module.AIOHTTP_TCP_KEEPALIVE_INTVL
96+
assert sock.getsockopt(
97+
socket.IPPROTO_TCP, socket.TCP_KEEPCNT
98+
) == http_handler_module.AIOHTTP_TCP_KEEPALIVE_CNT
99+
finally:
100+
sock.close()

0 commit comments

Comments
 (0)