diff --git a/model-engine/model_engine_server/inference/forwarding/forwarding.py b/model-engine/model_engine_server/inference/forwarding/forwarding.py index 5183955b..e42f982b 100644 --- a/model-engine/model_engine_server/inference/forwarding/forwarding.py +++ b/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, ) diff --git a/model-engine/tests/unit/inference/test_forwarding.py b/model-engine/tests/unit/inference/test_forwarding.py index 3e0141c8..0e085b54 100644 --- a/model-engine/tests/unit/inference/test_forwarding.py +++ b/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(