From 4ebef8a884954c3820ad3d27e716df377863323d Mon Sep 17 00:00:00 2001 From: "Martina G. Vilas" Date: Wed, 8 Jul 2026 16:32:15 +0200 Subject: [PATCH 1/4] feat(server_utils): enable TCP keepalive on global aiohttp connector Detect silently dropped upstream sockets that otherwise cause long-running rollouts to stall forever. _TCPKeepAliveConnector sets SO_KEEPALIVE + TCP_KEEPIDLE/INTVL/CNT via _wrap_create_connection (aiohttp<3.11 compat); detection budget ~90s. Signed-off-by: Martina G. Vilas --- nemo_gym/server_utils.py | 25 ++++++++++- tests/unit_tests/test_server_utils.py | 65 +++++++++++++++++++++++++++ 2 files changed, 89 insertions(+), 1 deletion(-) diff --git a/nemo_gym/server_utils.py b/nemo_gym/server_utils.py index 607e477e41..b250e5a9db 100644 --- a/nemo_gym/server_utils.py +++ b/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 @@ -100,6 +101,28 @@ def get_global_aiohttp_client( return set_global_aiohttp_client(cfg) +class _TCPKeepAliveConnector(TCPConnector): + _KEEPALIVE_IDLE_SECONDS: int = 60 + _KEEPALIVE_INTERVAL_SECONDS: int = 10 + _KEEPALIVE_PROBES: int = 3 + + async def _wrap_create_connection(self, *args, **kwargs): + transport, protocol = await super()._wrap_create_connection(*args, **kwargs) + sock = transport.get_extra_info("socket") + if sock is None: + raise RuntimeError("_TCPKeepAliveConnector: transport has no underlying socket") + sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + for opt_name, opt_value in ( + ("TCP_KEEPIDLE", self._KEEPALIVE_IDLE_SECONDS), + ("TCP_KEEPINTVL", self._KEEPALIVE_INTERVAL_SECONDS), + ("TCP_KEEPCNT", self._KEEPALIVE_PROBES), + ): + opt = getattr(socket, opt_name, None) + if opt is not None: + sock.setsockopt(socket.IPPROTO_TCP, opt, opt_value) + return transport, protocol + + 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!" @@ -107,7 +130,7 @@ def set_global_aiohttp_client(cfg: GlobalAIOHTTPAsyncClientConfig) -> ClientSess num_workers = get_nemo_gym_fastapi_num_workers() client_session = ClientSession( - connector=TCPConnector( + connector=_TCPKeepAliveConnector( limit=cfg.global_aiohttp_connector_limit // num_workers, limit_per_host=cfg.global_aiohttp_connector_limit_per_host // num_workers, keepalive_timeout=15.0, diff --git a/tests/unit_tests/test_server_utils.py b/tests/unit_tests/test_server_utils.py index cae89e1933..10f7eeb1ef 100644 --- a/tests/unit_tests/test_server_utils.py +++ b/tests/unit_tests/test_server_utils.py @@ -12,8 +12,10 @@ # 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 aiohttp import TCPConnector from pytest import MonkeyPatch, raises import nemo_gym.global_config @@ -29,6 +31,7 @@ HeadServer, ServerClient, SimpleServer, + _TCPKeepAliveConnector, initialize_ray, ) @@ -228,6 +231,68 @@ 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() + async def test_TCPKeepAliveConnector_sets_keepalive_sockopts(self, monkeypatch: MonkeyPatch) -> None: + mock_socket = MagicMock() + mock_transport = MagicMock() + mock_transport.get_extra_info.return_value = mock_socket + mock_protocol = MagicMock() + + super_wrap_mock = AsyncMock(return_value=(mock_transport, mock_protocol)) + monkeypatch.setattr(TCPConnector, "_wrap_create_connection", super_wrap_mock) + + connector = _TCPKeepAliveConnector() + try: + transport, protocol = await connector._wrap_create_connection() + finally: + await connector.close() + + assert transport is mock_transport + assert protocol is mock_protocol + mock_transport.get_extra_info.assert_called_once_with("socket") + mock_socket.setsockopt.assert_any_call(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + for opt_name, opt_value in ( + ("TCP_KEEPIDLE", _TCPKeepAliveConnector._KEEPALIVE_IDLE_SECONDS), + ("TCP_KEEPINTVL", _TCPKeepAliveConnector._KEEPALIVE_INTERVAL_SECONDS), + ("TCP_KEEPCNT", _TCPKeepAliveConnector._KEEPALIVE_PROBES), + ): + opt = getattr(socket, opt_name, None) + if opt is not None: + mock_socket.setsockopt.assert_any_call(socket.IPPROTO_TCP, opt, opt_value) + + async def test_TCPKeepAliveConnector_skips_missing_platform_sockopts(self, monkeypatch: MonkeyPatch) -> None: + mock_socket = MagicMock() + mock_transport = MagicMock() + mock_transport.get_extra_info.return_value = mock_socket + mock_protocol = MagicMock() + + super_wrap_mock = AsyncMock(return_value=(mock_transport, mock_protocol)) + monkeypatch.setattr(TCPConnector, "_wrap_create_connection", super_wrap_mock) + for opt_name in ("TCP_KEEPIDLE", "TCP_KEEPINTVL", "TCP_KEEPCNT"): + monkeypatch.delattr(socket, opt_name, raising=False) + + connector = _TCPKeepAliveConnector() + try: + await connector._wrap_create_connection() + finally: + await connector.close() + + mock_socket.setsockopt.assert_called_once_with(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + + async def test_TCPKeepAliveConnector_raises_when_no_socket(self, monkeypatch: MonkeyPatch) -> None: + mock_transport = MagicMock() + mock_transport.get_extra_info.return_value = None + mock_protocol = MagicMock() + + super_wrap_mock = AsyncMock(return_value=(mock_transport, mock_protocol)) + monkeypatch.setattr(TCPConnector, "_wrap_create_connection", super_wrap_mock) + + connector = _TCPKeepAliveConnector() + try: + with raises(RuntimeError): + await connector._wrap_create_connection() + finally: + await connector.close() + def test_dry_run_skips_webserver_spinup(self, monkeypatch: MonkeyPatch) -> None: self._mock_ray_return_value(monkeypatch, True) From 4248314367cc62dae09df8e03dfa757e4944fcf7 Mon Sep 17 00:00:00 2001 From: "Martina G. Vilas" Date: Mon, 13 Jul 2026 10:55:21 +0200 Subject: [PATCH 2/4] refactor(server_utils): make TCP keepalive knobs configurable Move idle/interval/probe from hardcoded connector attrs to GlobalAIOHTTPAsyncClientConfig fields. Defaults unchanged (60s / 10s / 3). Signed-off-by: Martina G. Vilas --- nemo_gym/server_utils.py | 39 ++++++++++++--- tests/unit_tests/test_server_utils.py | 69 ++++++++++++++++++++++++--- 2 files changed, 95 insertions(+), 13 deletions(-) diff --git a/nemo_gym/server_utils.py b/nemo_gym/server_utils.py index b250e5a9db..60c9e75dfc 100644 --- a/nemo_gym/server_utils.py +++ b/nemo_gym/server_utils.py @@ -50,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 @@ -82,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, @@ -102,9 +115,18 @@ def get_global_aiohttp_client( class _TCPKeepAliveConnector(TCPConnector): - _KEEPALIVE_IDLE_SECONDS: int = 60 - _KEEPALIVE_INTERVAL_SECONDS: int = 10 - _KEEPALIVE_PROBES: int = 3 + def __init__( + self, + *args, + tcp_keepalive_idle_seconds: int, + tcp_keepalive_interval_seconds: int, + tcp_keepalive_probes: int, + **kwargs, + ): + super().__init__(*args, **kwargs) + self._tcp_keepalive_idle_seconds = tcp_keepalive_idle_seconds + self._tcp_keepalive_interval_seconds = tcp_keepalive_interval_seconds + self._tcp_keepalive_probes = tcp_keepalive_probes async def _wrap_create_connection(self, *args, **kwargs): transport, protocol = await super()._wrap_create_connection(*args, **kwargs) @@ -113,9 +135,9 @@ async def _wrap_create_connection(self, *args, **kwargs): raise RuntimeError("_TCPKeepAliveConnector: transport has no underlying socket") sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) for opt_name, opt_value in ( - ("TCP_KEEPIDLE", self._KEEPALIVE_IDLE_SECONDS), - ("TCP_KEEPINTVL", self._KEEPALIVE_INTERVAL_SECONDS), - ("TCP_KEEPCNT", self._KEEPALIVE_PROBES), + ("TCP_KEEPIDLE", self._tcp_keepalive_idle_seconds), + ("TCP_KEEPINTVL", self._tcp_keepalive_interval_seconds), + ("TCP_KEEPCNT", self._tcp_keepalive_probes), ): opt = getattr(socket, opt_name, None) if opt is not None: @@ -134,6 +156,9 @@ 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, + tcp_keepalive_idle_seconds=cfg.global_aiohttp_tcp_keepalive_idle_seconds, + tcp_keepalive_interval_seconds=cfg.global_aiohttp_tcp_keepalive_interval_seconds, + tcp_keepalive_probes=cfg.global_aiohttp_tcp_keepalive_probes, ), timeout=ClientTimeout(), cookie_jar=DummyCookieJar(), diff --git a/tests/unit_tests/test_server_utils.py b/tests/unit_tests/test_server_utils.py index 10f7eeb1ef..ed7c06d920 100644 --- a/tests/unit_tests/test_server_utils.py +++ b/tests/unit_tests/test_server_utils.py @@ -28,6 +28,7 @@ BaseServerConfig, ConnectionError, DictConfig, + GlobalAIOHTTPAsyncClientConfig, HeadServer, ServerClient, SimpleServer, @@ -36,6 +37,11 @@ ) +_TCP_KEEPALIVE_TEST_IDLE = 42 +_TCP_KEEPALIVE_TEST_INTERVAL = 7 +_TCP_KEEPALIVE_TEST_PROBES = 2 + + 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) @@ -240,7 +246,11 @@ async def test_TCPKeepAliveConnector_sets_keepalive_sockopts(self, monkeypatch: super_wrap_mock = AsyncMock(return_value=(mock_transport, mock_protocol)) monkeypatch.setattr(TCPConnector, "_wrap_create_connection", super_wrap_mock) - connector = _TCPKeepAliveConnector() + connector = _TCPKeepAliveConnector( + tcp_keepalive_idle_seconds=_TCP_KEEPALIVE_TEST_IDLE, + tcp_keepalive_interval_seconds=_TCP_KEEPALIVE_TEST_INTERVAL, + tcp_keepalive_probes=_TCP_KEEPALIVE_TEST_PROBES, + ) try: transport, protocol = await connector._wrap_create_connection() finally: @@ -251,9 +261,9 @@ async def test_TCPKeepAliveConnector_sets_keepalive_sockopts(self, monkeypatch: mock_transport.get_extra_info.assert_called_once_with("socket") mock_socket.setsockopt.assert_any_call(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) for opt_name, opt_value in ( - ("TCP_KEEPIDLE", _TCPKeepAliveConnector._KEEPALIVE_IDLE_SECONDS), - ("TCP_KEEPINTVL", _TCPKeepAliveConnector._KEEPALIVE_INTERVAL_SECONDS), - ("TCP_KEEPCNT", _TCPKeepAliveConnector._KEEPALIVE_PROBES), + ("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: @@ -270,7 +280,11 @@ async def test_TCPKeepAliveConnector_skips_missing_platform_sockopts(self, monke for opt_name in ("TCP_KEEPIDLE", "TCP_KEEPINTVL", "TCP_KEEPCNT"): monkeypatch.delattr(socket, opt_name, raising=False) - connector = _TCPKeepAliveConnector() + connector = _TCPKeepAliveConnector( + tcp_keepalive_idle_seconds=_TCP_KEEPALIVE_TEST_IDLE, + tcp_keepalive_interval_seconds=_TCP_KEEPALIVE_TEST_INTERVAL, + tcp_keepalive_probes=_TCP_KEEPALIVE_TEST_PROBES, + ) try: await connector._wrap_create_connection() finally: @@ -286,13 +300,56 @@ async def test_TCPKeepAliveConnector_raises_when_no_socket(self, monkeypatch: Mo super_wrap_mock = AsyncMock(return_value=(mock_transport, mock_protocol)) monkeypatch.setattr(TCPConnector, "_wrap_create_connection", super_wrap_mock) - connector = _TCPKeepAliveConnector() + connector = _TCPKeepAliveConnector( + tcp_keepalive_idle_seconds=_TCP_KEEPALIVE_TEST_IDLE, + tcp_keepalive_interval_seconds=_TCP_KEEPALIVE_TEST_INTERVAL, + tcp_keepalive_probes=_TCP_KEEPALIVE_TEST_PROBES, + ) try: with raises(RuntimeError): await connector._wrap_create_connection() finally: await connector.close() + 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 + + async def test_TCPKeepAliveConnector_uses_configured_values(self, monkeypatch: MonkeyPatch) -> None: + mock_socket = MagicMock() + mock_transport = MagicMock() + mock_transport.get_extra_info.return_value = mock_socket + mock_protocol = MagicMock() + + super_wrap_mock = AsyncMock(return_value=(mock_transport, mock_protocol)) + monkeypatch.setattr(TCPConnector, "_wrap_create_connection", super_wrap_mock) + + cfg = GlobalAIOHTTPAsyncClientConfig( + global_aiohttp_tcp_keepalive_idle_seconds=123, + global_aiohttp_tcp_keepalive_interval_seconds=45, + global_aiohttp_tcp_keepalive_probes=6, + ) + connector = _TCPKeepAliveConnector( + tcp_keepalive_idle_seconds=cfg.global_aiohttp_tcp_keepalive_idle_seconds, + tcp_keepalive_interval_seconds=cfg.global_aiohttp_tcp_keepalive_interval_seconds, + tcp_keepalive_probes=cfg.global_aiohttp_tcp_keepalive_probes, + ) + try: + await connector._wrap_create_connection() + finally: + await connector.close() + + 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_socket.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) From 10906a8a96fa2385ba17a09c29318dc96e66ea21 Mon Sep 17 00:00:00 2001 From: "Martina G. Vilas" Date: Mon, 13 Jul 2026 11:11:45 +0200 Subject: [PATCH 3/4] fix(server_utils): degrade instead of raising when transport exposes no socket Raising RuntimeError from _wrap_create_connection is caught by client.request()'s retry loop and would fail every outbound request on a transport type that doesn't expose a socket. Log once and skip keepalive for that transport instead, so the worst case is losing the optimization, not the client. Signed-off-by: Martina G. Vilas --- nemo_gym/server_utils.py | 11 ++++++++++- tests/unit_tests/test_server_utils.py | 17 +++++++++++++---- 2 files changed, 23 insertions(+), 5 deletions(-) diff --git a/nemo_gym/server_utils.py b/nemo_gym/server_utils.py index 60c9e75dfc..277d8dabe6 100644 --- a/nemo_gym/server_utils.py +++ b/nemo_gym/server_utils.py @@ -75,6 +75,8 @@ _GLOBAL_AIOHTTP_CLIENT: Union[None, ClientSession] = None _GLOBAL_AIOHTTP_CLIENT_REQUEST_DEBUG: bool = False +_LOGGER = getLogger(__name__) + class GlobalAIOHTTPAsyncClientConfig(BaseModel): global_aiohttp_connector_limit: int = 100 * 1024 @@ -127,12 +129,19 @@ def __init__( self._tcp_keepalive_idle_seconds = tcp_keepalive_idle_seconds self._tcp_keepalive_interval_seconds = tcp_keepalive_interval_seconds self._tcp_keepalive_probes = tcp_keepalive_probes + self._warned_missing_socket: bool = False async def _wrap_create_connection(self, *args, **kwargs): transport, protocol = await super()._wrap_create_connection(*args, **kwargs) sock = transport.get_extra_info("socket") if sock is None: - raise RuntimeError("_TCPKeepAliveConnector: transport has no underlying socket") + if not self._warned_missing_socket: + _LOGGER.warning( + "_TCPKeepAliveConnector: transport exposes no underlying socket; " + "TCP keepalive will not be applied for connections on this transport." + ) + self._warned_missing_socket = True + return transport, protocol sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) for opt_name, opt_value in ( ("TCP_KEEPIDLE", self._tcp_keepalive_idle_seconds), diff --git a/tests/unit_tests/test_server_utils.py b/tests/unit_tests/test_server_utils.py index ed7c06d920..b7ceb7935c 100644 --- a/tests/unit_tests/test_server_utils.py +++ b/tests/unit_tests/test_server_utils.py @@ -16,7 +16,7 @@ from unittest.mock import AsyncMock, MagicMock from aiohttp import TCPConnector -from pytest import MonkeyPatch, raises +from pytest import LogCaptureFixture, MonkeyPatch, raises import nemo_gym.global_config import nemo_gym.server_utils @@ -292,7 +292,9 @@ async def test_TCPKeepAliveConnector_skips_missing_platform_sockopts(self, monke mock_socket.setsockopt.assert_called_once_with(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) - async def test_TCPKeepAliveConnector_raises_when_no_socket(self, monkeypatch: MonkeyPatch) -> None: + async def test_TCPKeepAliveConnector_degrades_when_no_socket( + self, monkeypatch: MonkeyPatch, caplog: LogCaptureFixture + ) -> None: mock_transport = MagicMock() mock_transport.get_extra_info.return_value = None mock_protocol = MagicMock() @@ -306,11 +308,18 @@ async def test_TCPKeepAliveConnector_raises_when_no_socket(self, monkeypatch: Mo tcp_keepalive_probes=_TCP_KEEPALIVE_TEST_PROBES, ) try: - with raises(RuntimeError): - await connector._wrap_create_connection() + with caplog.at_level("WARNING", logger=nemo_gym.server_utils.__name__): + transport1, protocol1 = await connector._wrap_create_connection() + transport2, protocol2 = await connector._wrap_create_connection() finally: await connector.close() + assert (transport1, protocol1) == (mock_transport, mock_protocol) + assert (transport2, protocol2) == (mock_transport, mock_protocol) + warnings = [r for r in caplog.records if r.levelname == "WARNING"] + assert len(warnings) == 1 + assert "TCP keepalive will not be applied" in warnings[0].message + def test_GlobalAIOHTTPAsyncClientConfig_keepalive_defaults(self) -> None: cfg = GlobalAIOHTTPAsyncClientConfig() assert cfg.global_aiohttp_tcp_keepalive_idle_seconds == 60 From 9d34aa715f7d90fa5c366cfa051eaa6d5c761a39 Mon Sep 17 00:00:00 2001 From: "Martina G. Vilas" Date: Mon, 13 Jul 2026 15:24:41 +0200 Subject: [PATCH 4/4] refactor(server_utils): use aiohttp socket_factory for TCP keepalive Drop the _TCPKeepAliveConnector subclass and its _wrap_create_connection override in favor of TCPConnector's public socket_factory= parameter (aiohttp >= 3.11). Same sockopts applied with the same values; stops depending on a private aiohttp method. Signed-off-by: Martina G. Vilas --- nemo_gym/server_utils.py | 56 ++++------- tests/unit_tests/test_server_utils.py | 132 +++++++++----------------- 2 files changed, 65 insertions(+), 123 deletions(-) diff --git a/nemo_gym/server_utils.py b/nemo_gym/server_utils.py index 277d8dabe6..89cd31b395 100644 --- a/nemo_gym/server_utils.py +++ b/nemo_gym/server_utils.py @@ -75,8 +75,6 @@ _GLOBAL_AIOHTTP_CLIENT: Union[None, ClientSession] = None _GLOBAL_AIOHTTP_CLIENT_REQUEST_DEBUG: bool = False -_LOGGER = getLogger(__name__) - class GlobalAIOHTTPAsyncClientConfig(BaseModel): global_aiohttp_connector_limit: int = 100 * 1024 @@ -116,42 +114,26 @@ def get_global_aiohttp_client( return set_global_aiohttp_client(cfg) -class _TCPKeepAliveConnector(TCPConnector): - def __init__( - self, - *args, - tcp_keepalive_idle_seconds: int, - tcp_keepalive_interval_seconds: int, - tcp_keepalive_probes: int, - **kwargs, - ): - super().__init__(*args, **kwargs) - self._tcp_keepalive_idle_seconds = tcp_keepalive_idle_seconds - self._tcp_keepalive_interval_seconds = tcp_keepalive_interval_seconds - self._tcp_keepalive_probes = tcp_keepalive_probes - self._warned_missing_socket: bool = False - - async def _wrap_create_connection(self, *args, **kwargs): - transport, protocol = await super()._wrap_create_connection(*args, **kwargs) - sock = transport.get_extra_info("socket") - if sock is None: - if not self._warned_missing_socket: - _LOGGER.warning( - "_TCPKeepAliveConnector: transport exposes no underlying socket; " - "TCP keepalive will not be applied for connections on this transport." - ) - self._warned_missing_socket = True - return transport, protocol +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", self._tcp_keepalive_idle_seconds), - ("TCP_KEEPINTVL", self._tcp_keepalive_interval_seconds), - ("TCP_KEEPCNT", self._tcp_keepalive_probes), + ("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 transport, protocol + return sock + + return factory def set_global_aiohttp_client(cfg: GlobalAIOHTTPAsyncClientConfig) -> ClientSession: # pragma: no cover @@ -161,13 +143,15 @@ def set_global_aiohttp_client(cfg: GlobalAIOHTTPAsyncClientConfig) -> ClientSess num_workers = get_nemo_gym_fastapi_num_workers() client_session = ClientSession( - connector=_TCPKeepAliveConnector( + connector=TCPConnector( limit=cfg.global_aiohttp_connector_limit // num_workers, limit_per_host=cfg.global_aiohttp_connector_limit_per_host // num_workers, keepalive_timeout=15.0, - tcp_keepalive_idle_seconds=cfg.global_aiohttp_tcp_keepalive_idle_seconds, - tcp_keepalive_interval_seconds=cfg.global_aiohttp_tcp_keepalive_interval_seconds, - tcp_keepalive_probes=cfg.global_aiohttp_tcp_keepalive_probes, + 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(), diff --git a/tests/unit_tests/test_server_utils.py b/tests/unit_tests/test_server_utils.py index b7ceb7935c..22f75f9b4d 100644 --- a/tests/unit_tests/test_server_utils.py +++ b/tests/unit_tests/test_server_utils.py @@ -15,8 +15,7 @@ import socket from unittest.mock import AsyncMock, MagicMock -from aiohttp import TCPConnector -from pytest import LogCaptureFixture, MonkeyPatch, raises +from pytest import MonkeyPatch, raises import nemo_gym.global_config import nemo_gym.server_utils @@ -32,7 +31,7 @@ HeadServer, ServerClient, SimpleServer, - _TCPKeepAliveConnector, + _make_keepalive_socket_factory, initialize_ray, ) @@ -40,6 +39,13 @@ _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: @@ -237,29 +243,23 @@ 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() - async def test_TCPKeepAliveConnector_sets_keepalive_sockopts(self, monkeypatch: MonkeyPatch) -> None: - mock_socket = MagicMock() - mock_transport = MagicMock() - mock_transport.get_extra_info.return_value = mock_socket - mock_protocol = MagicMock() + 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) - super_wrap_mock = AsyncMock(return_value=(mock_transport, mock_protocol)) - monkeypatch.setattr(TCPConnector, "_wrap_create_connection", super_wrap_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) - connector = _TCPKeepAliveConnector( - tcp_keepalive_idle_seconds=_TCP_KEEPALIVE_TEST_IDLE, - tcp_keepalive_interval_seconds=_TCP_KEEPALIVE_TEST_INTERVAL, - tcp_keepalive_probes=_TCP_KEEPALIVE_TEST_PROBES, + 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] ) - try: - transport, protocol = await connector._wrap_create_connection() - finally: - await connector.close() - - assert transport is mock_transport - assert protocol is mock_protocol - mock_transport.get_extra_info.assert_called_once_with("socket") - mock_socket.setsockopt.assert_any_call(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) + 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), @@ -267,58 +267,23 @@ async def test_TCPKeepAliveConnector_sets_keepalive_sockopts(self, monkeypatch: ): opt = getattr(socket, opt_name, None) if opt is not None: - mock_socket.setsockopt.assert_any_call(socket.IPPROTO_TCP, opt, opt_value) - - async def test_TCPKeepAliveConnector_skips_missing_platform_sockopts(self, monkeypatch: MonkeyPatch) -> None: - mock_socket = MagicMock() - mock_transport = MagicMock() - mock_transport.get_extra_info.return_value = mock_socket - mock_protocol = MagicMock() + mock_sock.setsockopt.assert_any_call(socket.IPPROTO_TCP, opt, opt_value) - super_wrap_mock = AsyncMock(return_value=(mock_transport, mock_protocol)) - monkeypatch.setattr(TCPConnector, "_wrap_create_connection", super_wrap_mock) + 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) - connector = _TCPKeepAliveConnector( - tcp_keepalive_idle_seconds=_TCP_KEEPALIVE_TEST_IDLE, - tcp_keepalive_interval_seconds=_TCP_KEEPALIVE_TEST_INTERVAL, - tcp_keepalive_probes=_TCP_KEEPALIVE_TEST_PROBES, + factory = _make_keepalive_socket_factory( + idle_seconds=_TCP_KEEPALIVE_TEST_IDLE, + interval_seconds=_TCP_KEEPALIVE_TEST_INTERVAL, + probes=_TCP_KEEPALIVE_TEST_PROBES, ) - try: - await connector._wrap_create_connection() - finally: - await connector.close() - - mock_socket.setsockopt.assert_called_once_with(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) - - async def test_TCPKeepAliveConnector_degrades_when_no_socket( - self, monkeypatch: MonkeyPatch, caplog: LogCaptureFixture - ) -> None: - mock_transport = MagicMock() - mock_transport.get_extra_info.return_value = None - mock_protocol = MagicMock() - - super_wrap_mock = AsyncMock(return_value=(mock_transport, mock_protocol)) - monkeypatch.setattr(TCPConnector, "_wrap_create_connection", super_wrap_mock) - - connector = _TCPKeepAliveConnector( - tcp_keepalive_idle_seconds=_TCP_KEEPALIVE_TEST_IDLE, - tcp_keepalive_interval_seconds=_TCP_KEEPALIVE_TEST_INTERVAL, - tcp_keepalive_probes=_TCP_KEEPALIVE_TEST_PROBES, - ) - try: - with caplog.at_level("WARNING", logger=nemo_gym.server_utils.__name__): - transport1, protocol1 = await connector._wrap_create_connection() - transport2, protocol2 = await connector._wrap_create_connection() - finally: - await connector.close() - - assert (transport1, protocol1) == (mock_transport, mock_protocol) - assert (transport2, protocol2) == (mock_transport, mock_protocol) - warnings = [r for r in caplog.records if r.levelname == "WARNING"] - assert len(warnings) == 1 - assert "TCP keepalive will not be applied" in warnings[0].message + 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() @@ -326,29 +291,22 @@ def test_GlobalAIOHTTPAsyncClientConfig_keepalive_defaults(self) -> None: assert cfg.global_aiohttp_tcp_keepalive_interval_seconds == 10 assert cfg.global_aiohttp_tcp_keepalive_probes == 3 - async def test_TCPKeepAliveConnector_uses_configured_values(self, monkeypatch: MonkeyPatch) -> None: - mock_socket = MagicMock() - mock_transport = MagicMock() - mock_transport.get_extra_info.return_value = mock_socket - mock_protocol = MagicMock() - - super_wrap_mock = AsyncMock(return_value=(mock_transport, mock_protocol)) - monkeypatch.setattr(TCPConnector, "_wrap_create_connection", super_wrap_mock) + 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, ) - connector = _TCPKeepAliveConnector( - tcp_keepalive_idle_seconds=cfg.global_aiohttp_tcp_keepalive_idle_seconds, - tcp_keepalive_interval_seconds=cfg.global_aiohttp_tcp_keepalive_interval_seconds, - tcp_keepalive_probes=cfg.global_aiohttp_tcp_keepalive_probes, + 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, ) - try: - await connector._wrap_create_connection() - finally: - await connector.close() + factory(_TEST_ADDR_INFO) for opt_name, opt_value in ( ("TCP_KEEPIDLE", 123), @@ -357,7 +315,7 @@ async def test_TCPKeepAliveConnector_uses_configured_values(self, monkeypatch: M ): opt = getattr(socket, opt_name, None) if opt is not None: - mock_socket.setsockopt.assert_any_call(socket.IPPROTO_TCP, opt, opt_value) + 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)