feature: Write handler to config file and serve model directly from o… (#115)

* feature: Write handler to config file and serve model directly from original model artifact directory to save disk space.

* Remove Python 2.7 support

Author: davidthomas426

buildspec-release.yml

@@ -22,7 +22,7 @@ phases:
       - tox -e black-check
 
       # run unit tests
-      - tox -e py27,py36,py37 -- test/unit
+      - tox -e py36,py37 -- test/unit
 
       # build dummy container
       - python3 setup.py sdist
@@ -37,7 +37,7 @@ phases:
       - cd ../..
 
       # run local integration tests
-      - IGNORE_COVERAGE=- tox -e py27,py36,py37 -- test/integration/local
+      - IGNORE_COVERAGE=- tox -e py36,py37 -- test/integration/local
 
       # publish the release to github
       - git-release --publish

buildspec.yml

@@ -20,7 +20,7 @@ phases:
       - tox -e twine
 
       # run unit tests
-      - tox -e py27,py36,py37 -- test/unit
+      - tox -e py36,py37 -- test/unit
 
       # build dummy container
       - python setup.py sdist
@@ -35,4 +35,4 @@ phases:
       - cd ../..
 
       # run local integration tests
-      - IGNORE_COVERAGE=- tox -e py27,py36,py37 -- test/integration/local 
+      - IGNORE_COVERAGE=- tox -e py36,py37 -- test/integration/local

setup.py

@@ -56,7 +56,6 @@ def read_version():
         "Natural Language :: English",
         "License :: OSI Approved :: Apache Software License",
         "Programming Language :: Python",
-        "Programming Language :: Python :: 2.7",
         "Programming Language :: Python :: 3.6",
         "Programming Language :: Python :: 3.7",
     ],

src/sagemaker_inference/model_server.py

@@ -41,11 +41,11 @@
 DEFAULT_MMS_LOG_FILE = pkg_resources.resource_filename(
     sagemaker_inference.__name__, "/etc/log4j2.xml"
 )
-DEFAULT_MMS_MODEL_DIRECTORY = os.path.join(os.getcwd(), ".sagemaker/mms/models")
+DEFAULT_MMS_MODEL_EXPORT_DIRECTORY = os.path.join(os.getcwd(), ".sagemaker/mms/models")
 DEFAULT_MMS_MODEL_NAME = "model"
 
 ENABLE_MULTI_MODEL = os.getenv("SAGEMAKER_MULTI_MODEL", "false") == "true"
-MODEL_STORE = "/" if ENABLE_MULTI_MODEL else DEFAULT_MMS_MODEL_DIRECTORY
+MODEL_STORE = "/"
 
 PYTHON_PATH_ENV = "PYTHONPATH"
 REQUIREMENTS_PATH = os.path.join(code_dir, "requirements.txt")
@@ -68,15 +68,16 @@ 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)
+    if ENABLE_MULTI_MODEL and not os.getenv("SAGEMAKER_HANDLER"):
+        os.environ["SAGEMAKER_HANDLER"] = handler_service
+
+    _set_python_path()
 
     env = environment.Environment()
-    _create_model_server_config_file(env)
+
+    # Note: multi-model default config already sets default_service_handler
+    handler_service_for_config = None if ENABLE_MULTI_MODEL else handler_service
+    _create_model_server_config_file(env, handler_service_for_config)
 
     if os.path.exists(REQUIREMENTS_PATH):
         _install_requirements()
@@ -91,6 +92,8 @@ def start_model_server(handler_service=DEFAULT_HANDLER_SERVICE):
         "--log-config",
         DEFAULT_MMS_LOG_FILE,
     ]
+    if not ENABLE_MULTI_MODEL:
+        multi_model_server_cmd += ["--models", DEFAULT_MMS_MODEL_NAME + "=" + environment.model_dir]
 
     logger.info(multi_model_server_cmd)
     subprocess.Popen(multi_model_server_cmd)
@@ -104,9 +107,12 @@ def start_model_server(handler_service=DEFAULT_HANDLER_SERVICE):
     mms_process.wait()
 
 
+# Note: this legacy function is still here for backwards compatibility.
+# It should not normally need to be used, since the model artifact can be used
+# straight from the original model directory
 def _adapt_to_mms_format(handler_service):
-    if not os.path.exists(DEFAULT_MMS_MODEL_DIRECTORY):
-        os.makedirs(DEFAULT_MMS_MODEL_DIRECTORY)
+    if not os.path.exists(DEFAULT_MMS_MODEL_EXPORT_DIRECTORY):
+        os.makedirs(DEFAULT_MMS_MODEL_EXPORT_DIRECTORY)
 
     model_archiver_cmd = [
         "model-archiver",
@@ -117,7 +123,7 @@ def _adapt_to_mms_format(handler_service):
         "--model-path",
         environment.model_dir,
         "--export-path",
-        DEFAULT_MMS_MODEL_DIRECTORY,
+        DEFAULT_MMS_MODEL_EXPORT_DIRECTORY,
         "--archive-format",
         "no-archive",
     ]
@@ -141,20 +147,23 @@ def _set_python_path():
         os.environ[PYTHON_PATH_ENV] = code_dir_path
 
 
-def _create_model_server_config_file(env):
-    configuration_properties = _generate_mms_config_properties(env)
+def _create_model_server_config_file(env, handler_service=None):
+    configuration_properties = _generate_mms_config_properties(env, handler_service)
 
     utils.write_file(MMS_CONFIG_FILE, configuration_properties)
 
 
-def _generate_mms_config_properties(env):
+def _generate_mms_config_properties(env, handler_service=None):
     user_defined_configuration = {
         "default_response_timeout": env.model_server_timeout,
         "default_workers_per_model": env.model_server_workers,
         "inference_address": "http://0.0.0.0:{}".format(env.inference_http_port),
         "management_address": "http://0.0.0.0:{}".format(env.management_http_port),
         "vmargs": "-XX:-UseContainerSupport",
     }
+    # If provided, add handler service to user config
+    if handler_service:
+        user_defined_configuration["default_service_handler"] = handler_service
 
     custom_configuration = str()
 

test/unit/test_model_server.py

@@ -49,20 +49,23 @@ def test_start_model_server_default_service_handler(
 
     model_server.start_model_server()
 
-    adapt.assert_called_once_with(model_server.DEFAULT_HANDLER_SERVICE)
-    create_config.assert_called_once_with(env.return_value)
+    adapt.assert_not_called()
+
+    create_config.assert_called_once_with(env.return_value, model_server.DEFAULT_HANDLER_SERVICE)
     exists.assert_called_once_with(REQUIREMENTS_PATH)
     install_requirements.assert_called_once_with()
 
     multi_model_server_cmd = [
         "multi-model-server",
         "--start",
         "--model-store",
-        model_server.DEFAULT_MMS_MODEL_DIRECTORY,
+        model_server.MODEL_STORE,
         "--mms-config",
         model_server.MMS_CONFIG_FILE,
         "--log-config",
         model_server.DEFAULT_MMS_LOG_FILE,
+        "--models",
+        "{}={}".format(model_server.DEFAULT_MMS_MODEL_NAME, environment.model_dir),
     ]
 
     subprocess_popen.assert_called_once_with(multi_model_server_cmd)
@@ -75,14 +78,16 @@ def test_start_model_server_default_service_handler(
 @patch("sagemaker_inference.model_server._add_sigterm_handler")
 @patch("sagemaker_inference.model_server._create_model_server_config_file")
 @patch("sagemaker_inference.model_server._adapt_to_mms_format")
+@patch("sagemaker_inference.environment.Environment")
 def test_start_model_server_custom_handler_service(
-    adapt, create_config, sigterm, retrieve, subprocess_popen, subprocess_call
+    env, adapt, create_config, sigterm, retrieve, subprocess_popen, subprocess_call
 ):
     handler_service = Mock()
 
     model_server.start_model_server(handler_service)
 
-    adapt.assert_called_once_with(handler_service)
+    adapt.assert_not_called()
+    create_config.assert_called_once_with(env.return_value, handler_service)
 
 
 @patch("sagemaker_inference.model_server._set_python_path")
@@ -94,8 +99,8 @@ def test_adapt_to_mms_format(path_exists, make_dir, subprocess_check_call, set_p
 
     model_server._adapt_to_mms_format(handler_service)
 
-    path_exists.assert_called_once_with(model_server.DEFAULT_MMS_MODEL_DIRECTORY)
-    make_dir.assert_called_once_with(model_server.DEFAULT_MMS_MODEL_DIRECTORY)
+    path_exists.assert_called_once_with(model_server.DEFAULT_MMS_MODEL_EXPORT_DIRECTORY)
+    make_dir.assert_called_once_with(model_server.DEFAULT_MMS_MODEL_EXPORT_DIRECTORY)
 
     model_archiver_cmd = [
         "model-archiver",
@@ -106,7 +111,7 @@ def test_adapt_to_mms_format(path_exists, make_dir, subprocess_check_call, set_p
         "--model-path",
         environment.model_dir,
         "--export-path",
-        model_server.DEFAULT_MMS_MODEL_DIRECTORY,
+        model_server.DEFAULT_MMS_MODEL_EXPORT_DIRECTORY,
         "--archive-format",
         "no-archive",
     ]
@@ -126,7 +131,7 @@ def test_adapt_to_mms_format_existing_path(
 
     model_server._adapt_to_mms_format(handler_service)
 
-    path_exists.assert_called_once_with(model_server.DEFAULT_MMS_MODEL_DIRECTORY)
+    path_exists.assert_called_once_with(model_server.DEFAULT_MMS_MODEL_EXPORT_DIRECTORY)
     make_dir.assert_not_called()
 
 

tox.ini

@@ -4,7 +4,7 @@
 # and then run "tox" from this directory.
 
 [tox]
-envlist = black-format,flake8,pylint,twine,py27,py36,py37
+envlist = black-format,flake8,pylint,twine,py36,py37
 
 skip_missing_interpreters = False