[MLI-8206] Set an explicit, configurable timeout on the sync forwarder

model-engine/model_engine_server/inference/forwarding/forwarding.py

@@ -1,5 +1,6 @@
 import ast
 import json
+import math
 import os
 import time
 from dataclasses import dataclass
@@ -42,6 +43,8 @@
 
 DEFAULT_PORT: int = 5005
 
+DEFAULT_SYNC_TIMEOUT_SECONDS: float = 3600
+
 
 class ModelEngineSerializationMixin:
     """Mixin class for optionally wrapping Model Engine requests."""
@@ -169,6 +172,10 @@ class Forwarder(ModelEngineSerializationMixin):
     # We do this to avoid having to put this data in any sync response and only do it for async responses
     forward_http_status_in_body: bool
     post_inference_hooks_handler: Optional[PostInferenceHooksHandler] = None
+    # Cap on the full round-trip to the user-defined service. Must be explicit: without
+    # one, aiohttp applies its default total=300s and long-running non-streaming
+    # generations get cut off with a 500 while the model server keeps computing.
+    timeout_seconds: float = DEFAULT_SYNC_TIMEOUT_SECONDS
 
     async def forward(self, json_payload: Any, trace_config: Optional[str] = None) -> Any:
         json_payload, using_serialize_results_as_string = self.unwrap_json_payload(json_payload)
@@ -185,6 +192,7 @@ async def forward(self, json_payload: Any, trace_config: Optional[str] = None) -
                     self.predict_endpoint,
                     json=json_payload,
                     headers=headers,
+                    timeout=aiohttp.ClientTimeout(total=self.timeout_seconds),
                 )
                 response = await response_raw.json(
                     content_type=None
@@ -297,8 +305,17 @@ class LoadForwarder:
     wrap_response: bool = True
     forward_http_status: bool = False
     forward_http_status_in_body: bool = False
+    timeout_seconds: float = DEFAULT_SYNC_TIMEOUT_SECONDS
 
     def load(self, resources: Optional[Path], cache: Any) -> Forwarder:
+        if (
+            not isinstance(self.timeout_seconds, (int, float))
+            or isinstance(self.timeout_seconds, bool)
+            or not math.isfinite(self.timeout_seconds)
+            or self.timeout_seconds <= 0
+        ):
+            raise ValueError(f"timeout_seconds must be a positive number: {self.timeout_seconds=}")
+
         if self.use_grpc:
             raise NotImplementedError(
                 "User-defined service **MUST** use HTTP at the moment. "
@@ -405,6 +422,7 @@ def endpoint(route: str) -> str:
             wrap_response=self.wrap_response,
             forward_http_status=self.forward_http_status,
             forward_http_status_in_body=self.forward_http_status_in_body,
+            timeout_seconds=self.timeout_seconds,
         )
 
 

model-engine/tests/unit/inference/test_forwarding.py

@@ -10,6 +10,7 @@
 from model_engine_server.core.utils.env import environment
 from model_engine_server.domain.entities import ModelEndpointConfig
 from model_engine_server.inference.forwarding.forwarding import (
+    DEFAULT_SYNC_TIMEOUT_SECONDS,
     ENV_SERIALIZE_RESULTS_AS_STRING,
     KEY_SERIALIZE_RESULTS_AS_STRING,
     Forwarder,
@@ -480,6 +481,40 @@ def test_forwarder_loader():
     _check_responses_not_wrapped(json_response)
 
 
+@mock.patch("requests.post", mocked_post)
+@mock.patch("requests.get", mocked_get)
+@mock.patch(
+    "model_engine_server.inference.forwarding.forwarding.get_endpoint_config",
+    mocked_get_endpoint_config,
+)
+@pytest.mark.parametrize(
+    "loader_kwargs, expected_timeout",
+    [
+        pytest.param({}, DEFAULT_SYNC_TIMEOUT_SECONDS, id="default"),
+        pytest.param({"timeout_seconds": 123.0}, 123.0, id="override"),
+    ],
+)
+def test_forwarder_loader_timeout(loader_kwargs, expected_timeout):
+    fwd = LoadForwarder(**loader_kwargs).load(None, None)  # type: ignore
+    assert fwd.timeout_seconds == expected_timeout
+
+
+@pytest.mark.parametrize(
+    "invalid_timeout",
+    [
+        pytest.param(0, id="zero"),
+        pytest.param(-1, id="negative"),
+        pytest.param(None, id="null"),
+        pytest.param(float("inf"), id="non-finite"),
+        pytest.param("60", id="string"),
+        pytest.param(True, id="bool"),
+    ],
+)
+def test_forwarder_loader_invalid_timeout(invalid_timeout):
+    with pytest.raises(ValueError, match="timeout_seconds"):
+        LoadForwarder(timeout_seconds=invalid_timeout).load(None, None)  # type: ignore
+
+
 @mock.patch("requests.post", mocked_post)
 @mock.patch("requests.get", mocked_get)
 @mock.patch(