using plural ports instead of switching to singular

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_PORT = os.environ.get("TFS_GRPC_PORT")
-TFS_REST_PORT = os.environ.get("TFS_REST_PORT")
+TFS_GRPC_PORTS = os.environ.get("TFS_GRPC_PORTS")
+TFS_REST_PORTS = os.environ.get("TFS_REST_PORTS")
 SAGEMAKER_TFS_PORT_RANGE = os.environ.get("SAGEMAKER_SAFE_PORT_RANGE")
 TFS_INSTANCE_COUNT = int(os.environ.get("SAGEMAKER_TFS_INSTANCE_COUNT", "1"))
 
@@ -69,11 +69,11 @@ def __init__(self):
             # during the _handle_load_model_post()
             self.model_handlers = {}
         else:
-            self._tfs_grpc_port = self._parse_concat_port(TFS_GRPC_PORT)
-            self._tfs_rest_port = self._parse_concat_port(TFS_REST_PORT)
+            self._tfs_grpc_ports = self._parse_concat_ports(TFS_GRPC_PORTS)
+            self._tfs_rest_ports = self._parse_concat_ports(TFS_REST_PORTS)
 
             self._channels = {}
-            for grpc_port in self._tfs_grpc_port:
+            for grpc_port in self._tfs_grpc_ports:
                 # Initialize grpc channel here so gunicorn worker could have mapping
                 # between each grpc port and channel
                 self._setup_channel(grpc_port)
@@ -98,8 +98,8 @@ 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 _parse_concat_port(self, concat_port):
-        return concat_port.split(",")
+    def _parse_concat_ports(self, concat_ports):
+        return concat_ports.split(",")
 
     def _pick_port(self, ports):
         return str(random.choice(ports))
@@ -245,8 +245,8 @@ def _handle_invocation_post(self, req, res, model_name=None):
                 })
         else:
             # Randomly pick port used for routing incoming request.
-            grpc_port = self._pick_port(self._tfs_grpc_port)
-            rest_port = self._pick_port(self._tfs_rest_port)
+            grpc_port = self._pick_port(self._tfs_grpc_ports)
+            rest_port = self._pick_port(self._tfs_rest_ports)
             data, context = tfs_utils.parse_request(req, rest_port, grpc_port,
                                                     self._tfs_default_model_name,
                                                     channel=self._channels[int(grpc_port)])

docker/build_artifacts/sagemaker/serve.py

@@ -89,29 +89,29 @@ def __init__(self):
             parts = self._sagemaker_port_range.split("-")
             low = int(parts[0])
             hi = int(parts[1])
-            self._tfs_grpc_port = []
-            self._tfs_rest_port = []
+            self._tfs_grpc_ports = []
+            self._tfs_rest_ports = []
             if low + 2 * self._tfs_instance_count > hi:
                 raise ValueError("not enough ports available in SAGEMAKER_SAFE_PORT_RANGE ({})"
                                  .format(self._sagemaker_port_range))
             # select non-overlapping grpc and rest ports based on 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_ports.append(str(low + 2 * i))
+                self._tfs_rest_ports.append(str(low + 2 * i + 1))
             # concat selected ports respectively in order to pass them to python service
-            self._tfs_grpc_concat_port = self._concat_port(self._tfs_grpc_port)
-            self._tfs_rest_concat_port = self._concat_port(self._tfs_rest_port)
+            self._tfs_grpc_concat_ports = self._concat_ports(self._tfs_grpc_ports)
+            self._tfs_rest_concat_ports = self._concat_ports(self._tfs_rest_ports)
         else:
             # just use the standard default ports
-            self._tfs_grpc_port = ["9000"]
-            self._tfs_rest_port = ["8501"]
+            self._tfs_grpc_ports = ["9000"]
+            self._tfs_rest_ports = ["8501"]
             # provide single concat port here for default case
-            self._tfs_grpc_concat_port = "9000"
-            self._tfs_rest_concat_port = "8501"
+            self._tfs_grpc_concat_ports = "9000"
+            self._tfs_rest_concat_ports = "8501"
 
         # set environment variable for python service
-        os.environ["TFS_GRPC_PORT"] = self._tfs_grpc_concat_port
-        os.environ["TFS_REST_PORT"] = self._tfs_rest_concat_port
+        os.environ["TFS_GRPC_PORTS"] = self._tfs_grpc_concat_ports
+        os.environ["TFS_REST_PORTS"] = self._tfs_rest_concat_ports
 
     def _need_python_service(self):
         if os.path.exists(INFERENCE_PATH):
@@ -120,7 +120,7 @@ def _need_python_service(self):
                 and os.environ.get("SAGEMAKER_MULTI_MODEL_UNIVERSAL_PREFIX"):
             self._enable_python_service = True
 
-    def _concat_port(self, ports):
+    def _concat_ports(self, ports):
         str_ports = [str(port) for port in ports]
         concat_str_ports = ",".join(str_ports)
         return concat_str_ports
@@ -198,13 +198,13 @@ def _setup_gunicorn(self):
         gunicorn_command = (
             "gunicorn -b unix:/tmp/gunicorn.sock -k {} --chdir /sagemaker "
             "--workers {} --threads {} "
-            "{}{} -e TFS_GRPC_PORT={} -e TFS_REST_PORT={} "
+            "{}{} -e TFS_GRPC_PORTS={} -e TFS_REST_PORTS={} "
             "-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_concat_port, self._tfs_rest_concat_port,
+                                         self._tfs_grpc_concat_ports, self._tfs_rest_concat_ports,
                                          self._tfs_enable_multi_model_endpoint,
                                          self._sagemaker_port_range,
                                          self._tfs_wait_time_seconds)
@@ -234,7 +234,7 @@ def _download_scripts(self, bucket, prefix):
     def _create_nginx_tfs_upstream(self):
         indentation = "    "
         tfs_upstream = ""
-        for port in self._tfs_rest_port:
+        for port in self._tfs_rest_ports:
             tfs_upstream += "{}server localhost:{};\n".format(indentation, port)
         tfs_upstream = tfs_upstream[len(indentation):-2]
 
@@ -338,7 +338,7 @@ def _wait_for_gunicorn(self):
 
     def _wait_for_tfs(self):
         for i in range(self._tfs_instance_count):
-            tfs_utils.wait_for_model(self._tfs_rest_port[i],
+            tfs_utils.wait_for_model(self._tfs_rest_ports[i],
                                      self._tfs_default_model_name, self._tfs_wait_time_seconds)
 
     @contextmanager
@@ -374,8 +374,8 @@ def _restart_single_tfs(self, pid):
 
     def _start_single_tfs(self, instance_id):
         cmd = tfs_utils.tfs_command(
-            self._tfs_grpc_port[instance_id],
-            self._tfs_rest_port[instance_id],
+            self._tfs_grpc_ports[instance_id],
+            self._tfs_rest_ports[instance_id],
             self._tfs_config_path,
             self._tfs_enable_batching,
             self._tfs_batching_config_path,