modify to pick ports in linear fashion

Author: Jinpeng Qi

docker/build_artifacts/sagemaker/python_service.py

@@ -31,8 +31,8 @@
 
 SAGEMAKER_BATCHING_ENABLED = os.environ.get("SAGEMAKER_TFS_ENABLE_BATCHING", "false").lower()
 MODEL_CONFIG_FILE_PATH = "/sagemaker/model-config.cfg"
-TFS_GRPC_SELECTED_PORTS = os.environ.get("TFS_GRPC_SELECTED_PORTS")
-TFS_REST_SELECTED_PORTS = os.environ.get("TFS_REST_SELECTED_PORTS")
+TFS_GRPC_PORT_RANGE = os.environ.get("TFS_GRPC_PORT_RANGE")
+TFS_REST_PORT_RANGE = os.environ.get("TFS_REST_PORT_RANGE")
 SAGEMAKER_TFS_PORT_RANGE = os.environ.get("SAGEMAKER_SAFE_PORT_RANGE")
 TFS_INSTANCE_COUNT = int(os.environ.get("SAGEMAKER_TFS_INSTANCE_COUNT", "1"))
 
@@ -69,10 +69,8 @@ def __init__(self):
             # during the _handle_load_model_post()
             self.model_handlers = {}
         else:
-            self._tfs_grpc_ports = self._reconstruct_sagemaker_selected_ports(
-                TFS_GRPC_SELECTED_PORTS)
-            self._tfs_rest_ports = self._reconstruct_sagemaker_selected_ports(
-                TFS_REST_SELECTED_PORTS)
+            self._tfs_grpc_ports = self._parse_sagemaker_port_range(TFS_GRPC_PORT_RANGE)
+            self._tfs_rest_ports = self._parse_sagemaker_port_range(TFS_REST_PORT_RANGE)
 
             self._channels = {}
             for grpc_port in self._tfs_grpc_ports:
@@ -100,13 +98,13 @@ def on_post(self, req, res, model_name=None):
             data = json.loads(req.stream.read().decode("utf-8"))
             self._handle_load_model_post(res, data)
 
-    def _reconstruct_sagemaker_selected_ports(self, selected_ports):
-        lower, upper = selected_ports.split('-')
+    def _parse_sagemaker_port_range(self, port_range):
+        lower, upper = port_range.split('-')
         lower = int(lower)
         upper = int(upper)
         if lower == upper:
             return [lower]
-        return [lower + 2 * i for i in range(TFS_INSTANCE_COUNT)]
+        return [lower + i for i in range(TFS_INSTANCE_COUNT)]
 
     def _pick_port(self, ports):
         return str(random.choice(ports))

docker/build_artifacts/sagemaker/serve.py

@@ -96,22 +96,23 @@ def __init__(self):
                                  .format(self._sagemaker_port_range))
             # select non-overlapping grpc and rest ports based on tfs instance count
             # record the init and end value in order to reconstruct them in python service
-            self._tfs_grpc_selected_ports = "{}-{}".format(low,
-                                                           low + 2 * self._tfs_instance_count)
-            self._tfs_rest_selected_ports = "{}-{}".format(low + 1,
-                                                           low + 2 * self._tfs_instance_count + 1)
+            self._tfs_grpc_port_range = "{}-{}".format(low,
+                                                       low + self._tfs_instance_count)
+            self._tfs_rest_port_range = "{}-{}".format(low + self._tfs_instance_count,
+                                                       low + 2 * self._tfs_instance_count)
             for i in range(self._tfs_instance_count):
-                self._tfs_grpc_port.append(str(low + 2 * i))
-                self._tfs_rest_port.append(str(low + 2 * i + 1))
+                self._tfs_grpc_port.append(str(low + i))
+                self._tfs_rest_port.append(str(low + self._tfs_instance_count + i))
         else:
             # just use the standard default ports
             self._tfs_grpc_port = ["9000"]
             self._tfs_rest_port = ["8501"]
-            self._tfs_grpc_selected_ports = "9000-9000"
-            self._tfs_rest_selected_ports = "8501-8501"
+            self._tfs_grpc_port_range = "9000-9000"
+            self._tfs_rest_port_range = "8501-8501"
+
         # set environment variable for python service
-        os.environ["TFS_GRPC_SELECTED_PORTS"] = self._tfs_grpc_selected_ports
-        os.environ["TFS_REST_SELECTED_PORTS"] = self._tfs_rest_selected_ports
+        os.environ["TFS_GRPC_PORT_RANGE"] = self._tfs_grpc_port_range
+        os.environ["TFS_REST_PORT_RANGE"] = self._tfs_rest_port_range
 
     def _need_python_service(self):
         if os.path.exists(INFERENCE_PATH):
@@ -193,14 +194,13 @@ def _setup_gunicorn(self):
         gunicorn_command = (
             "gunicorn -b unix:/tmp/gunicorn.sock -k {} --chdir /sagemaker "
             "--workers {} --threads {} "
-            "{}{} -e TFS_GRPC_SELECTED_PORTS={} -e TFS_REST_SELECTED_PORTS={} "
+            "{}{} -e TFS_GRPC_PORT_RANGE={} -e TFS_REST_PORT_RANGE={} "
             "-e SAGEMAKER_MULTI_MODEL={} -e SAGEMAKER_SAFE_PORT_RANGE={} "
             "-e SAGEMAKER_TFS_WAIT_TIME_SECONDS={} "
             "python_service:app").format(self._gunicorn_worker_class,
                                          self._gunicorn_workers, self._gunicorn_threads,
                                          python_path_option, ",".join(python_path_content),
-                                         self._tfs_grpc_selected_ports,
-                                         self._tfs_rest_selected_ports,
+                                         self._tfs_grpc_port_range, self._tfs_rest_port_range,
                                          self._tfs_enable_multi_model_endpoint,
                                          self._sagemaker_port_range,
                                          self._tfs_wait_time_seconds)