Skip to content
This repository was archived by the owner on May 23, 2024. It is now read-only.

Commit 1d28892

Browse files
author
Jinpeng Qi
committed
refactor naming of some variables
1 parent 98b7b6b commit 1d28892

File tree

2 files changed

+17
-18
lines changed

2 files changed

+17
-18
lines changed

docker/build_artifacts/sagemaker/python_service.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,8 @@
3131

3232
SAGEMAKER_BATCHING_ENABLED = os.environ.get("SAGEMAKER_TFS_ENABLE_BATCHING", "false").lower()
3333
MODEL_CONFIG_FILE_PATH = "/sagemaker/model-config.cfg"
34-
TFS_GRPC_PORT_RANGE = os.environ.get("TFS_GRPC_PORT_RANGE")
35-
TFS_REST_PORT_RANGE = os.environ.get("TFS_REST_PORT_RANGE")
34+
TFS_GRPC_SELECTED_PORTS = os.environ.get("TFS_GRPC_SELECTED_PORTS")
35+
TFS_REST_SELECTED_PORTS = os.environ.get("TFS_REST_SELECTED_PORTS")
3636
SAGEMAKER_TFS_PORT_RANGE = os.environ.get("SAGEMAKER_SAFE_PORT_RANGE")
3737
TFS_INSTANCE_COUNT = int(os.environ.get("SAGEMAKER_TFS_INSTANCE_COUNT", "1"))
3838

@@ -69,8 +69,8 @@ def __init__(self):
6969
# during the _handle_load_model_post()
7070
self.model_handlers = {}
7171
else:
72-
self._tfs_grpc_ports = self._parse_sagemaker_port_range(TFS_GRPC_PORT_RANGE)
73-
self._tfs_rest_ports = self._parse_sagemaker_port_range(TFS_REST_PORT_RANGE)
72+
self._tfs_grpc_ports = self._reconstruct_sagemaker_selected_ports(TFS_GRPC_SELECTED_PORTS)
73+
self._tfs_rest_ports = self._reconstruct_sagemaker_selected_ports(TFS_REST_SELECTED_PORTS)
7474

7575
self._channels = {}
7676
for grpc_port in self._tfs_grpc_ports:
@@ -98,8 +98,8 @@ def on_post(self, req, res, model_name=None):
9898
data = json.loads(req.stream.read().decode("utf-8"))
9999
self._handle_load_model_post(res, data)
100100

101-
def _parse_sagemaker_port_range(self, port_range):
102-
lower, upper = port_range.split('-')
101+
def _reconstruct_sagemaker_selected_ports(self, selected_ports):
102+
lower, upper = selected_ports.split('-')
103103
lower = int(lower)
104104
upper = int(upper)
105105
if lower == upper:

docker/build_artifacts/sagemaker/serve.py

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -94,25 +94,24 @@ def __init__(self):
9494
if low + 2 * self._tfs_instance_count > hi:
9595
raise ValueError("not enough ports available in SAGEMAKER_SAFE_PORT_RANGE ({})"
9696
.format(self._sagemaker_port_range))
97-
self._tfs_grpc_port_range = "{}-{}".format(low,
97+
# select non-overlapping grpc and rest ports based on tfs instance count
98+
# record the init and end value in order to reconstruct them in python service
99+
self._tfs_grpc_selected_ports = "{}-{}".format(low,
98100
low + 2 * self._tfs_instance_count)
99-
self._tfs_rest_port_range = "{}-{}".format(low + 1,
101+
self._tfs_rest_selected_ports = "{}-{}".format(low + 1,
100102
low + 2 * self._tfs_instance_count + 1)
101103
for i in range(self._tfs_instance_count):
102104
self._tfs_grpc_port.append(str(low + 2 * i))
103105
self._tfs_rest_port.append(str(low + 2 * i + 1))
104-
# set environment variable for python service
105-
os.environ["TFS_GRPC_PORT_RANGE"] = self._tfs_grpc_port_range
106-
os.environ["TFS_REST_PORT_RANGE"] = self._tfs_rest_port_range
107106
else:
108107
# just use the standard default ports
109108
self._tfs_grpc_port = ["9000"]
110109
self._tfs_rest_port = ["8501"]
111-
self._tfs_grpc_port_range = "9000-9000"
112-
self._tfs_rest_port_range = "8501-8501"
113-
# set environment variable for python service
114-
os.environ["TFS_GRPC_PORT_RANGE"] = self._tfs_grpc_port_range
115-
os.environ["TFS_REST_PORT_RANGE"] = self._tfs_rest_port_range
110+
self._tfs_grpc_selected_ports = "9000-9000"
111+
self._tfs_rest_selected_ports = "8501-8501"
112+
# set environment variable for python service
113+
os.environ["TFS_GRPC_SELECTED_PORTS"] = self._tfs_grpc_selected_ports
114+
os.environ["TFS_REST_SELECTED_PORTS"] = self._tfs_rest_selected_ports
116115

117116
def _need_python_service(self):
118117
if os.path.exists(INFERENCE_PATH):
@@ -194,13 +193,13 @@ def _setup_gunicorn(self):
194193
gunicorn_command = (
195194
"gunicorn -b unix:/tmp/gunicorn.sock -k {} --chdir /sagemaker "
196195
"--workers {} --threads {} "
197-
"{}{} -e TFS_GRPC_PORT_RANGE={} -e TFS_REST_PORT_RANGE={} "
196+
"{}{} -e TFS_GRPC_SELECTED_PORTS={} -e TFS_REST_SELECTED_PORTS={} "
198197
"-e SAGEMAKER_MULTI_MODEL={} -e SAGEMAKER_SAFE_PORT_RANGE={} "
199198
"-e SAGEMAKER_TFS_WAIT_TIME_SECONDS={} "
200199
"python_service:app").format(self._gunicorn_worker_class,
201200
self._gunicorn_workers, self._gunicorn_threads,
202201
python_path_option, ",".join(python_path_content),
203-
self._tfs_grpc_port_range, self._tfs_rest_port_range,
202+
self._tfs_grpc_selected_ports, self._tfs_rest_selected_ports,
204203
self._tfs_enable_multi_model_endpoint,
205204
self._sagemaker_port_range,
206205
self._tfs_wait_time_seconds)

0 commit comments

Comments
 (0)