fix: add model_dir to python path at service initialization

src/sagemaker_inference/default_handler_service.py

@@ -13,8 +13,12 @@
 """This module contains functionality for the default handler service."""
 from __future__ import absolute_import
 
+import os
+
 from sagemaker_inference.transformer import Transformer
 
+PYTHON_PATH_ENV = "PYTHONPATH"
+
 
 class DefaultHandlerService(object):
     """Default handler service that is executed by the model server.
@@ -51,4 +55,12 @@ def initialize(self, context):
         """
         properties = context.system_properties
         model_dir = properties.get("model_dir")
+
+        # add model_dir/code to python path
+        code_dir_path = "{}:".format(model_dir + "/code")
+        if PYTHON_PATH_ENV in os.environ:
+            os.environ[PYTHON_PATH_ENV] = code_dir_path + os.environ[PYTHON_PATH_ENV]
+        else:
+            os.environ[PYTHON_PATH_ENV] = code_dir_path
+
         self._service.validate_and_initialize(model_dir=model_dir)

test/container/mxnet/Dockerfile

@@ -111,7 +111,10 @@ RUN useradd -m model-server \
 
 COPY mxnet/mms_entrypoint.py /usr/local/bin/dockerd_entrypoint.py
 COPY mxnet/config.properties /home/model-server
-COPY mxnet/inference.py /opt/ml/models/code/inference.py
+
+# TEST: model_dir -> /opt/ml/models/<model_name>/inference.py
+COPY mxnet/inference.py /opt/ml/models/resnet_18/code/inference.py
+COPY mxnet/inference.py /opt/ml/models/resnet_152/code/inference.py
 
 RUN chmod +x /usr/local/bin/dockerd_entrypoint.py
 

test/unit/test_default_handler_service.py

@@ -10,7 +10,7 @@
 # distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-from mock import Mock, patch
+from mock import MagicMock, Mock, patch
 
 from sagemaker_inference.default_handler_service import DefaultHandlerService
 from sagemaker_inference.transformer import Transformer
@@ -48,7 +48,13 @@ def test_handle():
 
 def test_initialize():
     transformer = Mock()
+    properties = {"model_dir": "/opt/ml/models/model-name"}
 
-    DefaultHandlerService(transformer).initialize(CONTEXT)
+    def getitem(key):
+        return properties[key]
+
+    context = MagicMock()
+    context.system_properties.__getitem__.side_effect = getitem
+    DefaultHandlerService(transformer).initialize(context)
 
     assert transformer.validate_and_initialize().called_once()