Add support for SAGEMAKER_MODEL_SERVER_TIMEOUT_SECONDS variable (#129)

Author: davidthomas426

setup.cfg

@@ -2,4 +2,4 @@
 universal=1
 
 [metadata]
-description-file = README.md
+description_file = README.md

src/sagemaker_inference/environment.py

@@ -16,6 +16,7 @@
 from __future__ import absolute_import
 
 import os
+from typing import Optional
 
 from sagemaker_inference import content_types, logging, parameters
 
@@ -51,7 +52,9 @@ class Environment(object):
 
     Attributes:
         module_name (str): The name of the user-provided module. Default is inference.py.
-        model_server_timeout (int): Timeout in seconds for the model server. Default is 60.
+        model_server_timeout (int): Timeout for the model server. Default is 60.
+        model_server_timeout_seconds (Optional[int]): Timeout in seconds for the model server.
+            Default is None.
         model_server_workers (str): Number of worker processes the model server will use.
 
         default_accept (str): The desired default MIME type of the inference in the response
@@ -71,7 +74,13 @@ def __init__(self):
         self._model_server_timeout = int(
             os.environ.get(parameters.MODEL_SERVER_TIMEOUT_ENV, DEFAULT_MODEL_SERVER_TIMEOUT)
         )
+        timeout_seconds_var = os.environ.get(parameters.MODEL_SERVER_TIMEOUT_SECONDS_ENV)
+        self._model_server_timeout_seconds = (
+            int(timeout_seconds_var) if timeout_seconds_var is not None else None
+        )
+
         self._model_server_workers = os.environ.get(parameters.MODEL_SERVER_WORKERS_ENV)
+
         self._startup_timeout = int(
             os.environ.get(parameters.STARTUP_TIMEOUT_ENV, DEFAULT_STARTUP_TIMEOUT)
         )
@@ -107,53 +116,61 @@ def module_name(self):  # type: () -> str
         return self._parse_module_name(self._module_name)
 
     @property
-    def model_server_timeout(self):  # type: () -> int
+    def model_server_timeout(self) -> int:
+        """int: Timeout used for model server's backend workers before they are
+        deemed unresponsive and rebooted.
+
+        """
+        return self._model_server_timeout
+
+    @property
+    def model_server_timeout_seconds(self) -> Optional[int]:
         """int: Timeout, in seconds, used for model server's backend workers before
         they are deemed unresponsive and rebooted.
         """
-        return self._model_server_timeout
+        return self._model_server_timeout_seconds
 
     @property
-    def model_server_workers(self):  # type: () -> str
+    def model_server_workers(self) -> Optional[str]:
         """str: Number of worker processes the model server is configured to use."""
         return self._model_server_workers
 
     @property
-    def startup_timeout(self):  # type () -> int
+    def startup_timeout(self) -> int:
         """int: Timeout, in seconds, used for starting up the model server and fetching
         its process id, before giving up and throwing error.
         """
         return self._startup_timeout
 
     @property
-    def default_accept(self):  # type: () -> str
+    def default_accept(self) -> str:
         """str: The desired default MIME type of the inference in the response."""
         return self._default_accept
 
     @property
-    def inference_http_port(self):  # type: () -> str
+    def inference_http_port(self) -> str:
         """str: HTTP port that SageMaker uses to handle invocations and pings."""
         return self._inference_http_port
 
     @property
-    def management_http_port(self):  # type: () -> str
+    def management_http_port(self) -> str:
         """str: HTTP port that SageMaker uses to handle model management requests."""
         return self._management_http_port
 
     @property
-    def safe_port_range(self):  # type: () -> str
+    def safe_port_range(self) -> Optional[str]:
         """str: HTTP port range that can be used by users to avoid collisions with the HTTP port
         specified by SageMaker for handling pings and invocations.
         """
         return self._safe_port_range
 
     @property
-    def vmargs(self):  # type: () -> str
+    def vmargs(self) -> str:
         """str: vmargs can be provided for the JVM, to be overriden"""
         return self._vmargs
 
     @property
-    def max_request_size(self):  # type: () -> str
+    def max_request_size(self) -> Optional[int]:
         """str: max request size set by Sagemaker platform in bytes"""
         if self._max_request_size_in_mb is not None:
             return int(self._max_request_size_in_mb) * 1024 * 1024

src/sagemaker_inference/model_server.py

@@ -166,6 +166,11 @@ def _generate_mms_config_properties(env, handler_service=None):
     if handler_service:
         user_defined_configuration["default_service_handler"] = handler_service
 
+    if env.model_server_timeout_seconds:
+        user_defined_configuration[
+            "default_response_timeout_seconds"
+        ] = env.model_server_timeout_seconds
+
     custom_configuration = str()
 
     for key in user_defined_configuration:

src/sagemaker_inference/parameters.py

@@ -20,6 +20,7 @@
 DEFAULT_INVOCATIONS_ACCEPT_ENV = "SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT"  # type: str
 MODEL_SERVER_WORKERS_ENV = "SAGEMAKER_MODEL_SERVER_WORKERS"  # type: str
 MODEL_SERVER_TIMEOUT_ENV = "SAGEMAKER_MODEL_SERVER_TIMEOUT"  # type: str
+MODEL_SERVER_TIMEOUT_SECONDS_ENV = "SAGEMAKER_MODEL_SERVER_TIMEOUT_SECONDS"  # type: str
 MODEL_SERVER_VMARGS = "SAGEMAKER_MODEL_SERVER_VMARGS"  # type: str
 STARTUP_TIMEOUT_ENV = "SAGEMAKER_STARTUP_TIMEOUT"  # type: str
 BIND_TO_PORT_ENV = "SAGEMAKER_BIND_TO_PORT"  # type: str

test/unit/test_environment.py

@@ -23,6 +23,7 @@
     {
         parameters.USER_PROGRAM_ENV: "main.py",
         parameters.MODEL_SERVER_TIMEOUT_ENV: "20",
+        parameters.MODEL_SERVER_TIMEOUT_SECONDS_ENV: "30",
         parameters.MODEL_SERVER_WORKERS_ENV: "8",
         parameters.STARTUP_TIMEOUT_ENV: "50",
         parameters.DEFAULT_INVOCATIONS_ACCEPT_ENV: "text/html",
@@ -41,6 +42,7 @@ def test_env():
     assert environment.code_dir.endswith("opt/ml/model/code")
     assert env.module_name == "main"
     assert env.model_server_timeout == 20
+    assert env.model_server_timeout_seconds == 30
     assert env.startup_timeout == 50
     assert env.model_server_workers == "8"
     assert env.default_accept == "text/html"