Read port range from container support for TFS port (#106)

Author: ChoiByungWook

src/tf_container/serve.py

@@ -32,8 +32,7 @@
 from tf_container.run import logger
 import time
 
-
-TF_SERVING_PORT = 9000
+DEFAULT_TF_SERVING_PORT = 9000
 GENERIC_MODEL_NAME = "generic_model"
 TF_SERVING_MAXIMUM_LOAD_MODEL_TIME_IN_SECONDS = 60 * 15
 
@@ -97,17 +96,23 @@ def _recursive_copy(src, dst):
 
 
 def transformer(user_module):
-    grpc_proxy_client = proxy_client.GRPCProxyClient(TF_SERVING_PORT)
+    env = cs.HostingEnvironment()
+
+    port = int(cs.Server.next_safe_port(env.port_range)) if env.port_range else DEFAULT_TF_SERVING_PORT
+
+    grpc_proxy_client = proxy_client.GRPCProxyClient(port)
     _wait_model_to_load(grpc_proxy_client, TF_SERVING_MAXIMUM_LOAD_MODEL_TIME_IN_SECONDS)
 
     return Transformer.from_module(user_module, grpc_proxy_client)
 
 
 def load_dependencies():
     env = cs.HostingEnvironment()
+
+    port = cs.Server.next_safe_port(env.port_range) if env.port_range else DEFAULT_TF_SERVING_PORT
     saved_model_path = os.path.join(env.model_dir, 'export/Servo')
     subprocess.Popen(['tensorflow_model_server',
-                      '--port={}'.format(TF_SERVING_PORT),
+                      '--port={}'.format(port),
                       '--model_name={}'.format(GENERIC_MODEL_NAME),
                       '--model_base_path={}'.format(saved_model_path)])
 

test/unit/test_serve.py

@@ -21,6 +21,9 @@
 from container_support.serving import UnsupportedAcceptTypeError, UnsupportedContentTypeError
 
 JSON_CONTENT_TYPE = "application/json"
+FIRST_PORT = '1111'
+LAST_PORT = '2222'
+SAFE_PORT_RANGE = '{}-{}'.format(FIRST_PORT, LAST_PORT)
 
 
 @pytest.fixture(scope="module")
@@ -271,11 +274,15 @@ def test_transformer_method(proxy_client, serve):
 
 
 @patch('subprocess.Popen')
-def test_load_dependencies(popen, serve):
+@patch('container_support.HostingEnvironment')
+def test_load_dependencies_with_default_port(hosting_env, popen, serve):
     with patch('os.environ') as env:
         env['SAGEMAKER_PROGRAM'] = 'script.py'
         env['SAGEMAKER_SUBMIT_DIRECTORY'] = 's3://what/ever'
 
+        hosting_env.return_value.port_range = None
+        hosting_env.return_value.model_dir = '/opt/ml/model'
+
         serve.Transformer.from_module = Mock()
         serve.load_dependencies()
 
@@ -285,6 +292,25 @@ def test_load_dependencies(popen, serve):
                                   '--model_base_path=/opt/ml/model/export/Servo'])
 
 
+@patch('subprocess.Popen')
+@patch('container_support.HostingEnvironment')
+def test_load_dependencies_with_safe_port(hosting_env, popen, serve):
+    with patch('os.environ') as env:
+        env['SAGEMAKER_PROGRAM'] = 'script.py'
+        env['SAGEMAKER_SUBMIT_DIRECTORY'] = 's3://what/ever'
+
+        hosting_env.return_value.port_range = SAFE_PORT_RANGE
+        hosting_env.return_value.model_dir = '/opt/ml/model'
+
+        serve.Transformer.from_module = Mock()
+        serve.load_dependencies()
+
+        popen.assert_called_with(['tensorflow_model_server',
+                                  '--port={}'.format(FIRST_PORT),
+                                  '--model_name=generic_model',
+                                  '--model_base_path=/opt/ml/model/export/Servo'])
+
+
 @patch('tf_container.proxy_client.GRPCProxyClient')
 def test_wait_model_to_load(proxy_client, serve):
     client = proxy_client()