fix(server_utils): enable TCP keepalive on global aiohttp connector

nemo_gym/server_utils.py

@@ -16,6 +16,7 @@
 import atexit
 import json
 import resource
+import socket
 import sys
 import time
 from abc import abstractmethod
@@ -49,7 +50,7 @@
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse
 from omegaconf import DictConfig, OmegaConf, open_dict
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, Field
 from requests.exceptions import ConnectionError
 from starlette.middleware.sessions import SessionMiddleware
 
@@ -81,6 +82,19 @@ class GlobalAIOHTTPAsyncClientConfig(BaseModel):
 
     global_aiohttp_client_request_debug: bool = False
 
+    global_aiohttp_tcp_keepalive_idle_seconds: int = Field(
+        default=60,
+        description=("TCP_KEEPIDLE: seconds a socket must be idle before the kernel starts sending keepalive probes."),
+    )
+    global_aiohttp_tcp_keepalive_interval_seconds: int = Field(
+        default=10,
+        description=("TCP_KEEPINTVL: seconds between successive keepalive probes."),
+    )
+    global_aiohttp_tcp_keepalive_probes: int = Field(
+        default=3,
+        description=("TCP_KEEPCNT: number of unanswered probes before the kernel drops the connection."),
+    )
+
 
 def get_global_aiohttp_client(
     global_config_dict_parser_config: Optional[GlobalConfigDictParserConfig] = None,
@@ -100,6 +114,28 @@ def get_global_aiohttp_client(
     return set_global_aiohttp_client(cfg)
 
 
+def _make_keepalive_socket_factory(
+    idle_seconds: int,
+    interval_seconds: int,
+    probes: int,
+):
+    def factory(addr_info) -> socket.socket:
+        family, type_, proto, _canonname, _sockaddr = addr_info
+        sock = socket.socket(family=family, type=type_, proto=proto)
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+        for opt_name, opt_value in (
+            ("TCP_KEEPIDLE", idle_seconds),
+            ("TCP_KEEPINTVL", interval_seconds),
+            ("TCP_KEEPCNT", probes),
+        ):
+            opt = getattr(socket, opt_name, None)
+            if opt is not None:
+                sock.setsockopt(socket.IPPROTO_TCP, opt, opt_value)
+        return sock
+
+    return factory
+
+
 def set_global_aiohttp_client(cfg: GlobalAIOHTTPAsyncClientConfig) -> ClientSession:  # pragma: no cover
     assert not is_global_aiohttp_client_setup(), (
         "There is already a global aiohttp client setup. Please refactor your code or call `global_aiohttp_client_exit` if you want to explicitly re-make the client!"
@@ -111,6 +147,11 @@ def set_global_aiohttp_client(cfg: GlobalAIOHTTPAsyncClientConfig) -> ClientSess
             limit=cfg.global_aiohttp_connector_limit // num_workers,
             limit_per_host=cfg.global_aiohttp_connector_limit_per_host // num_workers,
             keepalive_timeout=15.0,
+            socket_factory=_make_keepalive_socket_factory(
+                idle_seconds=cfg.global_aiohttp_tcp_keepalive_idle_seconds,
+                interval_seconds=cfg.global_aiohttp_tcp_keepalive_interval_seconds,
+                probes=cfg.global_aiohttp_tcp_keepalive_probes,
+            ),
         ),
         timeout=ClientTimeout(),
         cookie_jar=DummyCookieJar(),

tests/unit_tests/test_server_utils.py

@@ -12,6 +12,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import socket
 from unittest.mock import AsyncMock, MagicMock
 
 from pytest import MonkeyPatch, raises
@@ -26,13 +27,27 @@
     BaseServerConfig,
     ConnectionError,
     DictConfig,
+    GlobalAIOHTTPAsyncClientConfig,
     HeadServer,
     ServerClient,
     SimpleServer,
+    _make_keepalive_socket_factory,
     initialize_ray,
 )
 
 
+_TCP_KEEPALIVE_TEST_IDLE = 42
+_TCP_KEEPALIVE_TEST_INTERVAL = 7
+_TCP_KEEPALIVE_TEST_PROBES = 2
+_TEST_ADDR_INFO = (
+    socket.AF_INET,
+    socket.SOCK_STREAM,
+    socket.IPPROTO_TCP,
+    "",
+    ("203.0.113.1", 443),
+)
+
+
 class TestServerUtils:
     def test_global_aiohttp_client_request_debug_enabled(self, monkeypatch: MonkeyPatch) -> None:
         monkeypatch.setattr(nemo_gym.server_utils, "_GLOBAL_AIOHTTP_CLIENT_REQUEST_DEBUG", False)
@@ -228,6 +243,80 @@ def test_initialize_ray_without_address(self, monkeypatch: MonkeyPatch) -> None:
         ray_init_mock.assert_called_once_with(ignore_reinit_error=True)
         ray_get_runtime_context_mock.assert_called_once()
 
+    def test_keepalive_socket_factory_sets_keepalive_sockopts(self, monkeypatch: MonkeyPatch) -> None:
+        mock_sock = MagicMock()
+        socket_ctor_mock = MagicMock(return_value=mock_sock)
+        monkeypatch.setattr(socket, "socket", socket_ctor_mock)
+
+        factory = _make_keepalive_socket_factory(
+            idle_seconds=_TCP_KEEPALIVE_TEST_IDLE,
+            interval_seconds=_TCP_KEEPALIVE_TEST_INTERVAL,
+            probes=_TCP_KEEPALIVE_TEST_PROBES,
+        )
+        result = factory(_TEST_ADDR_INFO)
+
+        assert result is mock_sock
+        socket_ctor_mock.assert_called_once_with(
+            family=_TEST_ADDR_INFO[0], type=_TEST_ADDR_INFO[1], proto=_TEST_ADDR_INFO[2]
+        )
+        mock_sock.setsockopt.assert_any_call(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+        for opt_name, opt_value in (
+            ("TCP_KEEPIDLE", _TCP_KEEPALIVE_TEST_IDLE),
+            ("TCP_KEEPINTVL", _TCP_KEEPALIVE_TEST_INTERVAL),
+            ("TCP_KEEPCNT", _TCP_KEEPALIVE_TEST_PROBES),
+        ):
+            opt = getattr(socket, opt_name, None)
+            if opt is not None:
+                mock_sock.setsockopt.assert_any_call(socket.IPPROTO_TCP, opt, opt_value)
+
+    def test_keepalive_socket_factory_skips_missing_platform_sockopts(self, monkeypatch: MonkeyPatch) -> None:
+        mock_sock = MagicMock()
+        socket_ctor_mock = MagicMock(return_value=mock_sock)
+        monkeypatch.setattr(socket, "socket", socket_ctor_mock)
+        for opt_name in ("TCP_KEEPIDLE", "TCP_KEEPINTVL", "TCP_KEEPCNT"):
+            monkeypatch.delattr(socket, opt_name, raising=False)
+
+        factory = _make_keepalive_socket_factory(
+            idle_seconds=_TCP_KEEPALIVE_TEST_IDLE,
+            interval_seconds=_TCP_KEEPALIVE_TEST_INTERVAL,
+            probes=_TCP_KEEPALIVE_TEST_PROBES,
+        )
+        factory(_TEST_ADDR_INFO)
+
+        mock_sock.setsockopt.assert_called_once_with(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+
+    def test_GlobalAIOHTTPAsyncClientConfig_keepalive_defaults(self) -> None:
+        cfg = GlobalAIOHTTPAsyncClientConfig()
+        assert cfg.global_aiohttp_tcp_keepalive_idle_seconds == 60
+        assert cfg.global_aiohttp_tcp_keepalive_interval_seconds == 10
+        assert cfg.global_aiohttp_tcp_keepalive_probes == 3
+
+    def test_keepalive_socket_factory_uses_configured_values(self, monkeypatch: MonkeyPatch) -> None:
+        mock_sock = MagicMock()
+        socket_ctor_mock = MagicMock(return_value=mock_sock)
+        monkeypatch.setattr(socket, "socket", socket_ctor_mock)
+
+        cfg = GlobalAIOHTTPAsyncClientConfig(
+            global_aiohttp_tcp_keepalive_idle_seconds=123,
+            global_aiohttp_tcp_keepalive_interval_seconds=45,
+            global_aiohttp_tcp_keepalive_probes=6,
+        )
+        factory = _make_keepalive_socket_factory(
+            idle_seconds=cfg.global_aiohttp_tcp_keepalive_idle_seconds,
+            interval_seconds=cfg.global_aiohttp_tcp_keepalive_interval_seconds,
+            probes=cfg.global_aiohttp_tcp_keepalive_probes,
+        )
+        factory(_TEST_ADDR_INFO)
+
+        for opt_name, opt_value in (
+            ("TCP_KEEPIDLE", 123),
+            ("TCP_KEEPINTVL", 45),
+            ("TCP_KEEPCNT", 6),
+        ):
+            opt = getattr(socket, opt_name, None)
+            if opt is not None:
+                mock_sock.setsockopt.assert_any_call(socket.IPPROTO_TCP, opt, opt_value)
+
     def test_dry_run_skips_webserver_spinup(self, monkeypatch: MonkeyPatch) -> None:
         self._mock_ray_return_value(monkeypatch, True)