Default GRPC timeout for EI & Allow timeout to be configurable (#135)

Author: ChoiByungWook

setup.py

@@ -37,6 +37,6 @@ def read(fname):
     install_requires=['sagemaker-container-support'],
     extras_require={
         'test': ['tox', 'flake8', 'pytest', 'pytest-cov', 'pytest-xdist', 'mock',
-                 'sagemaker', 'tensorflow']
+                 'sagemaker', 'tensorflow', 'tensorflow-serving-api']
     },
 )

src/tf_container/proxy_client.py

@@ -10,6 +10,7 @@
 #  on an "AS IS" BASIS, 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 os
 
 import numpy as np
 from google.protobuf import json_format
@@ -25,6 +26,11 @@
 
 from tf_container.run import logger as _logger
 
+INFERENCE_ACCELERATOR_PRESENT_ENV = 'SAGEMAKER_INFERENCE_ACCELERATOR_PRESENT'
+TF_SERVING_GRPC_REQUEST_TIMEOUT_ENV = 'SAGEMAKER_TFS_GRPC_REQUEST_TIMEOUT'
+
+DEFAULT_GRPC_REQUEST_TIMEOUT_FOR_INFERENCE_ACCELERATOR = 30.0
+
 REGRESSION = 'tensorflow/serving/regression'
 CLASSIFY = 'tensorflow/serving/classify'
 INFERENCE = 'tensorflow/serving/inference'
@@ -37,6 +43,11 @@ def __init__(self, tf_serving_port, host='localhost', request_timeout=10.0,
                  model_name=GENERIC_MODEL_NAME,
                  input_tensor_name=PREDICT_INPUTS,
                  signature_name=DEFAULT_SERVING_SIGNATURE_DEF_KEY):
+        if os.environ.get(TF_SERVING_GRPC_REQUEST_TIMEOUT_ENV):
+            request_timeout = float(os.environ.get(TF_SERVING_GRPC_REQUEST_TIMEOUT_ENV))
+        elif os.environ.get(INFERENCE_ACCELERATOR_PRESENT_ENV) == 'true':
+            request_timeout = DEFAULT_GRPC_REQUEST_TIMEOUT_FOR_INFERENCE_ACCELERATOR
+
         self.tf_serving_port = tf_serving_port
         self.host = host
         self.request_timeout = request_timeout

test/unit/test_experiment_trainer.py

@@ -1,14 +1,14 @@
 #  Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
-#  
+#
 #  Licensed under the Apache License, Version 2.0 (the "License").
 #  You may not use this file except in compliance with the License.
 #  A copy of the License is located at
-#  
+#
 #      http://www.apache.org/licenses/LICENSE-2.0
-#  
-#  or in the "license" file accompanying this file. This file is distributed 
-#  on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either 
-#  express or implied. See the License for the specific language governing 
+#
+#  or in the "license" file accompanying this file. This file is distributed
+#  on an "AS IS" BASIS, 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 pytest