diff --git a/src/_bentoml_impl/client/proxy2.py b/src/_bentoml_impl/client/proxy2.py index 602d10b9344..1fbfcb9ff00 100644 --- a/src/_bentoml_impl/client/proxy2.py +++ b/src/_bentoml_impl/client/proxy2.py @@ -100,6 +100,7 @@ def _make_client(self) -> aiohttp.ClientSession: base_url = "http://127.0.0.1:3000" elif self._parsed_url.scheme == "tcp": base_url = f"http://{self._parsed_url.netloc}" + connector = aiohttp.TCPConnector(keepalive_timeout=4.0) else: base_url = self._parsed_url.geturl() @@ -353,7 +354,14 @@ async def _call( try: headers = CIMultiDict({"Content-Type": self.media_type, **(headers or {})}) body = self._build_payload(endpoint, args, kwargs, headers) - resp = await client.post(endpoint.route, data=body, headers=headers) + try: + resp = await client.post(endpoint.route, data=body, headers=headers) + except (aiohttp.ClientConnectionError, aiohttp.ServerDisconnectedError): + client = await self._session_manager.get_session() + retry_body = self._build_payload(endpoint, args, kwargs, headers) + resp = await client.post( + endpoint.route, data=retry_body, headers=headers + ) if not resp.ok: raise await map_exception(resp) if endpoint.stream_output: diff --git a/tests/unit/_internal/client/test_session_manager.py b/tests/unit/_internal/client/test_session_manager.py index 29a97f1d53a..cdef2176611 100644 --- a/tests/unit/_internal/client/test_session_manager.py +++ b/tests/unit/_internal/client/test_session_manager.py @@ -1,3 +1,4 @@ +import aiohttp import pytest from starlette.applications import Starlette from starlette.responses import JSONResponse @@ -28,3 +29,24 @@ async def test_session_refresh_creates_new_connector() -> None: assert resp.status == 200 finally: await manager.close() + + +@pytest.mark.asyncio +async def test_tcp_connector_has_explicit_keepalive_timeout() -> None: + # Regression test for #4971: the tcp-scheme branch of _make_client + # previously left connector=None, so aiohttp fell back to its default + # keepalive_timeout (15s), which could exceed the server's uvicorn + # timeout_keep_alive default (5s) and cause intermittent + # ServerDisconnectedError/RemoteProtocolError under load. + manager = SessionManager( + url="tcp://127.0.0.1:3000", + timeout=30.0, + headers={}, + ) + try: + session = await manager.get_session() + connector = session.connector + assert isinstance(connector, aiohttp.TCPConnector) + assert connector._keepalive_timeout == 4.0 + finally: + await manager.close()