remove _adapt_to_ts_format fn

Author: mseth10

src/sagemaker_pytorch_serving_container/torchserve.py

@@ -41,13 +41,11 @@
 DEFAULT_TS_LOG_FILE = pkg_resources.resource_filename(
     sagemaker_pytorch_serving_container.__name__, "/etc/log4j2.xml"
 )
-DEFAULT_TS_MODEL_DIRECTORY = os.path.join(os.getcwd(), ".sagemaker", "ts", "models")
 DEFAULT_TS_MODEL_NAME = "model"
-DEFAULT_TS_CODE_DIR = "code"
 DEFAULT_HANDLER_SERVICE = "sagemaker_pytorch_serving_container.handler_service"
 
 ENABLE_MULTI_MODEL = os.getenv("SAGEMAKER_MULTI_MODEL", "false") == "true"
-MODEL_STORE = "/" if ENABLE_MULTI_MODEL else DEFAULT_TS_MODEL_DIRECTORY
+MODEL_STORE = "/" if ENABLE_MULTI_MODEL else os.path.join(os.getcwd(), ".sagemaker", "ts", "models")
 
 PYTHON_PATH_ENV = "PYTHONPATH"
 REQUIREMENTS_PATH = os.path.join(code_dir, "requirements.txt")
@@ -70,12 +68,10 @@ def start_torchserve(handler_service=DEFAULT_HANDLER_SERVICE):
 
     """
 
-    if ENABLE_MULTI_MODEL:
-        if "SAGEMAKER_HANDLER" not in os.environ:
-            os.environ["SAGEMAKER_HANDLER"] = handler_service
-        _set_python_path()
-    else:
-        _adapt_to_ts_format(handler_service)
+    if ENABLE_MULTI_MODEL and "SAGEMAKER_HANDLER" not in os.environ:
+        os.environ["SAGEMAKER_HANDLER"] = handler_service
+
+    _set_python_path()
 
     _create_torchserve_config_file()
 
@@ -92,7 +88,7 @@ def start_torchserve(handler_service=DEFAULT_HANDLER_SERVICE):
         "--log-config",
         DEFAULT_TS_LOG_FILE,
         "--models",
-        "model.mar"
+        DEFAULT_TS_MODEL_NAME + "=" + environment.model_dir
     ]
 
     print(ts_torchserve_cmd)
@@ -107,30 +103,6 @@ def start_torchserve(handler_service=DEFAULT_HANDLER_SERVICE):
     ts_process.wait()
 
 
-def _adapt_to_ts_format(handler_service):
-    if not os.path.exists(DEFAULT_TS_MODEL_DIRECTORY):
-        os.makedirs(DEFAULT_TS_MODEL_DIRECTORY)
-
-    model_archiver_cmd = [
-        "torch-model-archiver",
-        "--model-name",
-        DEFAULT_TS_MODEL_NAME,
-        "--handler",
-        handler_service,
-        "--export-path",
-        DEFAULT_TS_MODEL_DIRECTORY,
-        "--version",
-        "1",
-        "--extra-files",
-        os.path.join(environment.model_dir)
-    ]
-
-    logger.info(model_archiver_cmd)
-    subprocess.check_call(model_archiver_cmd)
-
-    _set_python_path()
-
-
 def _set_python_path():
     # Torchserve handles code execution by appending the export path, provided
     # to the model archiver, to the PYTHONPATH env var.

test/unit/test_model_server.py

@@ -34,9 +34,7 @@
 @patch("sagemaker_pytorch_serving_container.torchserve._install_requirements")
 @patch("os.path.exists", return_value=True)
 @patch("sagemaker_pytorch_serving_container.torchserve._create_torchserve_config_file")
-@patch("sagemaker_pytorch_serving_container.torchserve._adapt_to_ts_format")
 def test_start_torchserve_default_service_handler(
-    adapt,
     create_config,
     exists,
     install_requirements,
@@ -47,7 +45,6 @@ def test_start_torchserve_default_service_handler(
 ):
     torchserve.start_torchserve()
 
-    adapt.assert_called_once_with(torchserve.DEFAULT_HANDLER_SERVICE)
     create_config.assert_called_once_with()
     exists.assert_called_once_with(REQUIREMENTS_PATH)
     install_requirements.assert_called_once_with()
@@ -62,7 +59,7 @@ def test_start_torchserve_default_service_handler(
         "--log-config",
         torchserve.DEFAULT_TS_LOG_FILE,
         "--models",
-        "model.mar"
+        "model=/opt/ml/model"
     ]
 
     subprocess_popen.assert_called_once_with(ts_model_server_cmd)
@@ -76,9 +73,7 @@ def test_start_torchserve_default_service_handler(
 @patch("sagemaker_pytorch_serving_container.torchserve._install_requirements")
 @patch("os.path.exists", return_value=True)
 @patch("sagemaker_pytorch_serving_container.torchserve._create_torchserve_config_file")
-@patch("sagemaker_pytorch_serving_container.torchserve._adapt_to_ts_format")
 def test_start_torchserve_default_service_handler_multi_model(
-    adapt,
     create_config,
     exists,
     install_requirements,
@@ -104,7 +99,7 @@ def test_start_torchserve_default_service_handler_multi_model(
         "--log-config",
         torchserve.DEFAULT_TS_LOG_FILE,
         "--models",
-        "model.mar"
+        "model=/opt/ml/model"
     ]
 
     subprocess_popen.assert_called_once_with(ts_model_server_cmd)
@@ -116,61 +111,13 @@ def test_start_torchserve_default_service_handler_multi_model(
 @patch("sagemaker_pytorch_serving_container.torchserve._retrieve_ts_server_process")
 @patch("sagemaker_pytorch_serving_container.torchserve._add_sigterm_handler")
 @patch("sagemaker_pytorch_serving_container.torchserve._create_torchserve_config_file")
-@patch("sagemaker_pytorch_serving_container.torchserve._adapt_to_ts_format")
 def test_start_torchserve_custom_handler_service(
-    adapt, create_config, sigterm, retrieve, subprocess_popen, subprocess_call
+    create_config, sigterm, retrieve, subprocess_popen, subprocess_call
 ):
     handler_service = Mock()
 
     torchserve.start_torchserve(handler_service)
 
-    adapt.assert_called_once_with(handler_service)
-
-
-@patch("sagemaker_pytorch_serving_container.torchserve._set_python_path")
-@patch("subprocess.check_call")
-@patch("os.makedirs")
-@patch("os.path.exists", return_value=False)
-def test_adapt_to_ts_format(path_exists, make_dir, subprocess_check_call, set_python_path):
-    handler_service = Mock()
-
-    torchserve._adapt_to_ts_format(handler_service)
-
-    path_exists.assert_called_once_with(torchserve.DEFAULT_TS_MODEL_DIRECTORY)
-    make_dir.assert_called_once_with(torchserve.DEFAULT_TS_MODEL_DIRECTORY)
-
-    model_archiver_cmd = [
-        "torch-model-archiver",
-        "--model-name",
-        torchserve.DEFAULT_TS_MODEL_NAME,
-        "--handler",
-        handler_service,
-        "--export-path",
-        torchserve.DEFAULT_TS_MODEL_DIRECTORY,
-        "--version",
-        "1",
-        "--extra-files",
-        environment.model_dir
-    ]
-
-    subprocess_check_call.assert_called_once_with(model_archiver_cmd)
-    set_python_path.assert_called_once_with()
-
-
-@patch("sagemaker_pytorch_serving_container.torchserve._set_python_path")
-@patch("subprocess.check_call")
-@patch("os.makedirs")
-@patch("os.path.exists", return_value=True)
-def test_adapt_to_ts_format_existing_path(
-    path_exists, make_dir, subprocess_check_call, set_python_path
-):
-    handler_service = Mock()
-
-    torchserve._adapt_to_ts_format(handler_service)
-
-    path_exists.assert_called_once_with(torchserve.DEFAULT_TS_MODEL_DIRECTORY)
-    make_dir.assert_not_called()
-
 
 @patch.dict(os.environ, {torchserve.PYTHON_PATH_ENV: PYTHON_PATH}, clear=True)
 def test_set_existing_python_path():