reuse sagemaker-inference's requirements.txt installation logic (#150)

Closes #149

Author: humanzz

src/sagemaker_pytorch_serving_container/torchserve.py

@@ -16,7 +16,6 @@
 import os
 import signal
 import subprocess
-import sys
 
 import pkg_resources
 import psutil
@@ -25,8 +24,7 @@
 
 import sagemaker_pytorch_serving_container
 from sagemaker_pytorch_serving_container import ts_environment
-from sagemaker_inference import environment, utils
-from sagemaker_inference.environment import code_dir
+from sagemaker_inference import environment, utils, model_server
 
 logger = logging.getLogger()
 
@@ -47,7 +45,6 @@
 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")
 TS_NAMESPACE = "org.pytorch.serve.ModelServer"
 
 
@@ -78,8 +75,8 @@ def start_torchserve(handler_service=DEFAULT_HANDLER_SERVICE):
 
     _create_torchserve_config_file(handler_service)
 
-    if os.path.exists(REQUIREMENTS_PATH):
-        _install_requirements()
+    if os.path.exists(model_server.REQUIREMENTS_PATH):
+        model_server._install_requirements()
 
     ts_torchserve_cmd = [
         "torchserve",
@@ -181,17 +178,6 @@ def _terminate(signo, frame):  # pylint: disable=unused-argument
     signal.signal(signal.SIGTERM, _terminate)
 
 
-def _install_requirements():
-    logger.info("installing packages from requirements.txt...")
-    pip_install_cmd = [sys.executable, "-m", "pip", "install", "-r", REQUIREMENTS_PATH]
-
-    try:
-        subprocess.check_call(pip_install_cmd)
-    except subprocess.CalledProcessError:
-        logger.exception("failed to install required packages, exiting")
-        raise ValueError("failed to install required packages")
-
-
 # retry for 10 seconds
 @retry(stop_max_delay=10 * 1000)
 def _retrieve_ts_server_process():

test/unit/test_model_server.py

@@ -19,9 +19,9 @@
 from mock import Mock, patch
 import pytest
 
-from sagemaker_inference import environment
+from sagemaker_inference import environment, model_server
 from sagemaker_pytorch_serving_container import torchserve
-from sagemaker_pytorch_serving_container.torchserve import TS_NAMESPACE, REQUIREMENTS_PATH
+from sagemaker_pytorch_serving_container.torchserve import TS_NAMESPACE
 
 PYTHON_PATH = "python_path"
 DEFAULT_CONFIGURATION = "default_configuration"
@@ -31,7 +31,7 @@
 @patch("subprocess.Popen")
 @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._install_requirements")
+@patch("sagemaker_inference.model_server._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._set_python_path")
@@ -72,7 +72,7 @@ def test_start_torchserve_default_service_handler(
 @patch("subprocess.Popen")
 @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._install_requirements")
+@patch("sagemaker_inference.model_server._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._set_python_path")
@@ -92,7 +92,7 @@ def test_start_torchserve_default_service_handler_multi_model(
 
     set_python_path.assert_called_once_with()
     create_config.assert_called_once_with(torchserve.DEFAULT_HANDLER_SERVICE)
-    exists.assert_called_once_with(REQUIREMENTS_PATH)
+    exists.assert_called_once_with(model_server.REQUIREMENTS_PATH)
     install_requirements.assert_called_once_with()
 
     ts_model_server_cmd = [
@@ -210,15 +210,15 @@ def test_add_sigterm_handler(signal_call):
 
 @patch("subprocess.check_call")
 def test_install_requirements(check_call):
-    torchserve._install_requirements()
+    model_server._install_requirements()
     for i in ['pip', 'install', '-r', '/opt/ml/model/code/requirements.txt']:
         assert i in check_call.call_args.args[0]
 
 
 @patch("subprocess.check_call", side_effect=subprocess.CalledProcessError(0, "cmd"))
 def test_install_requirements_installation_failed(check_call):
     with pytest.raises(ValueError) as e:
-        torchserve._install_requirements()
+        model_server._install_requirements()
     assert "failed to install required packages" in str(e.value)