feature: MME support (#42)

* Set model_dir based on SAGEMAKER_MULTI_MODEL env variable

* Adapt to MMS format

* Update Dockerfile

* Update build artifacts

* Format with black and fix pylint violation

* Fix filenames for build artifacts

* Download models in Dockerfile

* Fix directory for build artifacts

* Update directories for downloaded models in Dockerfile

* Set WORKDIR back to root after downloading models in Dockerfile

* Update directories for downloaded models in MXNet multi-model integration tests

* Update Dockerfile with sagemaker_inference tarball and remove SAGEMAKER_HANDLER

* Add sagemaker_inference tarball

* Replace SAGEMAKER_HANDLER

* Rename model files to default model name

* Create model archive if MME not enabled

* Update sagemaker_inference tarball

* Get the model_dir from the context

* Update default handler service unit tests

* Update sagemaker_mxnet_inference tarball (modified HandlerService)

* Revert change to Transformer.transform

* Update call to validate_and_initialize in Transformer.transform and update unit tests

* Update build artifact tarballs

* Update sagemaker_mxnet_inference tarball

* Temporarily skip test_mxnet_multi_model.py::test_load_model_multiple_times to isolate bug

* Fix model names in integ test

* Temporarily skip test_load_model_multiple_times and test_invocation

* Remove skip on test_load_model_multiple_times

* Remove test_invocation since default_inference_handler does not implement input_fn to deserialize image

* Fix flake8 violation

Author: ajaykarpur

src/sagemaker_inference/default_handler_service.py

@@ -43,10 +43,14 @@ def handle(self, data, context):
                 with the context set appropriately.
 
         """
+        self.initialize(context)
+
         return self._service.transform(data, context)
 
-    def initialize(self):
+    def initialize(self, context):
         """Calls the Transformer method that validates the user module against
         the SageMaker inference contract.
         """
-        self._service.validate_and_initialize()
+        properties = context.system_properties
+        model_dir = properties.get("model_dir")
+        self._service.validate_and_initialize(model_dir=model_dir)

src/sagemaker_inference/environment.py

@@ -29,8 +29,11 @@
 
 base_dir = os.environ.get(parameters.BASE_PATH_ENV, SAGEMAKER_BASE_PATH)  # type: str
 
-model_dir = os.path.join(base_dir, "model")  # type: str
-"""str: the directory where models should be saved, e.g., /opt/ml/model/"""
+if os.environ.get(parameters.MULTI_MODEL_ENV) == "true":
+    model_dir = os.path.join(base_dir, "models")  # type: str
+else:
+    model_dir = os.path.join(base_dir, "model")  # type: str
+# str: the directory where models should be saved, e.g., /opt/ml/model/
 
 code_dir = os.path.join(model_dir, "code")  # type: str
 """str: the path of the user's code directory, e.g., /opt/ml/model/code/"""

src/sagemaker_inference/model_server.py

@@ -70,6 +70,7 @@ def start_model_server(handler_service=DEFAULT_HANDLER_SERVICE):
     if ENABLE_MULTI_MODEL:
         if not os.getenv("SAGEMAKER_HANDLER"):
             os.environ["SAGEMAKER_HANDLER"] = handler_service
+        _set_python_path()
     else:
         _adapt_to_mms_format(handler_service)
 

src/sagemaker_inference/parameters.py

@@ -22,3 +22,4 @@
 MODEL_SERVER_TIMEOUT_ENV = "SAGEMAKER_MODEL_SERVER_TIMEOUT"  # type: str
 BIND_TO_PORT_ENV = "SAGEMAKER_BIND_TO_PORT"  # type: str
 SAFE_PORT_RANGE_ENV = "SAGEMAKER_SAFE_PORT_RANGE"  # type: str
+MULTI_MODEL_ENV = "SAGEMAKER_MULTI_MODEL"  # type: str

src/sagemaker_inference/transformer.py

@@ -105,7 +105,9 @@ def transform(self, data, context):
                 with the context set appropriately.
         """
         try:
-            self.validate_and_initialize()
+            properties = context.system_properties
+            model_dir = properties.get("model_dir")
+            self.validate_and_initialize(model_dir=model_dir)
 
             input_data = data[0].get("body")
 
@@ -141,7 +143,7 @@ def transform(self, data, context):
                     context, GenericInferenceToolkitError(http_client.INTERNAL_SERVER_ERROR, str(e))
                 )
 
-    def validate_and_initialize(self):  # type: () -> None
+    def validate_and_initialize(self, model_dir=environment.model_dir):  # type: () -> None
         """Validates the user module against the SageMaker inference contract.
 
         Load the model as defined by the ``model_fn`` to prepare handling predictions.
@@ -150,7 +152,7 @@ def validate_and_initialize(self):  # type: () -> None
         if not self._initialized:
             self._environment = environment.Environment()
             self._validate_user_module_and_set_functions()
-            self._model = self._model_fn(environment.model_dir)
+            self._model = self._model_fn(model_dir)
             self._initialized = True
 
     def _validate_user_module_and_set_functions(self):

test/container/mxnet/Dockerfile

@@ -10,12 +10,11 @@ LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
 LABEL com.amazonaws.sagemaker.capabilities.multi-models=true
 
 ARG MMS_VERSION=1.0.8
-ARG MX_URL=https://aws-mxnet-pypi.s3-us-west-2.amazonaws.com/1.6.0/aws_mxnet_mkl-1.6.0rc0-py2.py3-none-manylinux1_x86_64.whl
+ARG MX_URL=https://aws-mxnet-pypi.s3-us-west-2.amazonaws.com/1.6.0/aws_mxnet_mkl-1.6.0-py2.py3-none-manylinux1_x86_64.whl
 ARG PYTHON=python3
 ARG PYTHON_PIP=python3-pip
 ARG PIP=pip3
 ARG PYTHON_VERSION=3.6.8
-ARG PIP_VERSION=19.3.1
 
 ENV PYTHONDONTWRITEBYTECODE=1 \
     PYTHONUNBUFFERED=1 \
@@ -63,21 +62,24 @@ RUN wget https://www.python.org/ftp/python/$PYTHON_VERSION/Python-$PYTHON_VERSIO
 RUN ln -s $(which ${PYTHON}) /usr/local/bin/python
 
 RUN ${PIP} --no-cache-dir install --upgrade \
-    pip==${PIP_VERSION} \
+    pip \
     setuptools
 
 WORKDIR /
 
-COPY mxnet/sagemaker_mxnet_serving_container.tar.gz /sagemaker_mxnet_serving_container.tar.gz
+COPY mxnet/sagemaker_mxnet_inference.tar.gz /sagemaker_mxnet_inference.tar.gz
+COPY mxnet/sagemaker_inference.tar.gz /sagemaker_inference.tar.gz
 
 RUN ${PIP} install --no-cache-dir \
     ${MX_URL} \
+    git+git://github.com/dmlc/[email protected] \
     mxnet-model-server==$MMS_VERSION \
     keras-mxnet==2.2.4.1 \
     numpy==1.17.4 \
     onnx==1.4.1 \
-    /sagemaker_mxnet_serving_container.tar.gz \
- && rm /sagemaker_mxnet_serving_container.tar.gz
+    /sagemaker_mxnet_inference.tar.gz \
+    /sagemaker_inference.tar.gz \
+ && rm /sagemaker_mxnet_inference.tar.gz /sagemaker_inference.tar.gz
 
 # This is here to make our installed version of OpenCV work.
 # https://stackoverflow.com/questions/29274638/opencv-libdc1394-error-failed-to-initialize-libdc1394
@@ -87,45 +89,34 @@ RUN ln -s /dev/null /dev/raw1394
 ##################################
 # download models to model_store
 ##################################
-RUN mkdir resnet_152 \
- && cd resnet_152 \
- && wget -O resnet-152-0000.params http://data.mxnet.io/models/imagenet/resnet/152-layers/resnet-152-0000.params \
- && wget -O resnet-152-symbol.json http://data.mxnet.io/models/imagenet/resnet/152-layers/resnet-152-symbol.json \
+WORKDIR /opt/ml/models/resnet_152/model
+RUN wget -O model-0000.params http://data.mxnet.io/models/imagenet/resnet/152-layers/resnet-152-0000.params \
+ && wget -O model-symbol.json http://data.mxnet.io/models/imagenet/resnet/152-layers/resnet-152-symbol.json \
  && wget -O synset.txt http://data.mxnet.io/models/imagenet/synset.txt \
- && echo '[{"shape": [1, 3, 224, 224], "name": "data"}]' > resnet-152-shapes.json \
- && cd ..
+ && echo '[{"shape": [1, 3, 224, 224], "name": "data"}]' > model-shapes.json \
+ && cd /
 
-RUN mkdir resnet_18 \
- && cd resnet_18 \
- && wget -O resnet-18-0000.params http://data.mxnet.io/models/imagenet/resnet/18-layers/resnet-18-0000.params \
- && wget -O resnet-18-symbol.json http://data.mxnet.io/models/imagenet/resnet/18-layers/resnet-18-symbol.json \
+WORKDIR /opt/ml/models/resnet_18/model
+RUN wget -O model-0000.params http://data.mxnet.io/models/imagenet/resnet/18-layers/resnet-18-0000.params \
+ && wget -O model-symbol.json http://data.mxnet.io/models/imagenet/resnet/18-layers/resnet-18-symbol.json \
  && wget -O synset.txt http://data.mxnet.io/models/imagenet/synset.txt \
- && echo '[{"shape": [1, 3, 224, 224], "name": "data"}]' > resnet-18-shapes.json \
- && cd ..
+ && echo '[{"shape": [1, 3, 224, 224], "name": "data"}]' > model-shapes.json \
+ && cd /
 
-ENV PYTHONDONTWRITEBYTECODE=1 \
-    PYTHONUNBUFFERED=1 \
-    LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/lib" \
-    PYTHONIOENCODING=UTF-8 \
-    LANG=C.UTF-8 \
-    LC_ALL=C.UTF-8
+WORKDIR /
 
 RUN useradd -m model-server \
-    && mkdir -p /home/model-server/tmp \
-    && chown -R model-server /home/model-server
-
-#################################################
-# custom_entrypoint uses custom_handler_service
-#################################################
-COPY mxnet/custom_handler.py /usr/local/bin/custom_handler.py
-COPY mxnet/custom_entrypoint.py /usr/local/bin/dockerd-entrypoint.py
+ && mkdir -p /home/model-server/tmp \
+ && chown -R model-server /home/model-server
+
+COPY mxnet/mms_entrypoint.py /usr/local/bin/dockerd_entrypoint.py
 COPY mxnet/config.properties /home/model-server
-COPY mxnet/deep_learning_container.py /usr/local/bin/deep_learning_container.py
+COPY mxnet/inference.py /opt/ml/models/code/inference.py
+
+RUN chmod +x /usr/local/bin/dockerd_entrypoint.py
 
-RUN chmod +x /usr/local/bin/dockerd-entrypoint.py && \
-    chmod +x /usr/local/bin/deep_learning_container.py
+RUN curl https://aws-dlc-licenses.s3.amazonaws.com/aws-mxnet-1.6.0/license.txt -o /license.txt
 
 EXPOSE 8080 8081
-ENV TEMP=/home/model-server/tmp
-ENTRYPOINT ["python", "/usr/local/bin/dockerd-entrypoint.py"]
+ENTRYPOINT ["python", "/usr/local/bin/dockerd_entrypoint.py"]
 CMD ["mxnet-model-server", "--start", "--mms-config", "/home/model-server/config.properties"]