tests for in process mode

Author: Bryannah Hernandez

src/sagemaker/serve/mode/in_process_mode.py

@@ -13,6 +13,7 @@
 from sagemaker.serve.utils.exceptions import LocalDeepPingException
 from sagemaker.serve.model_server.multi_model_server.server import InProcessMultiModelServer
 from sagemaker.session import Session
+from datetime import datetime, timedelta
 
 logger = logging.getLogger(__name__)
 
@@ -45,19 +46,15 @@ def __init__(
         self.session = session
         self.schema_builder = schema_builder
         self.model_server = model_server
-        self.client = None
-        self.container = None
-        self.secret_key = None
-        self._invoke_serving = None
         self._ping_container = None
 
     def load(self, model_path: str = None):
         """Loads model path, checks that path exists"""
         path = Path(model_path if model_path else self.model_path)
         if not path.exists():
-            raise Exception("model_path does not exist")
+            raise ValueError("model_path does not exist")
         if not path.is_dir():
-            raise Exception("model_path is not a valid directory")
+            raise ValueError("model_path is not a valid directory")
 
         return self.inference_spec.load(str(path))
 
@@ -69,15 +66,18 @@ def create_server(
         predictor: PredictorBase,
     ):
         """Creating the server and checking ping health."""
-
-        # self.destroy_server()
-
         logger.info("Waiting for model server %s to start up...", self.model_server)
 
         if self.model_server == ModelServer.MMS:
             self._ping_container = self._multi_model_server_deep_ping
 
-        while True:
+        time_limit = datetime.now() + timedelta(seconds=5)
+        while self._ping_container is not None:
+            final_pull = datetime.now() > time_limit
+
+            if final_pull:
+                break
+
             time.sleep(10)
 
             healthy, response = self._ping_container(predictor)

tests/unit/sagemaker/serve/mode/test_in_process_mode.py

@@ -0,0 +1,152 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# 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 __future__ import absolute_import
+
+import unittest
+from unittest.mock import MagicMock, patch, Mock, mock_open
+
+from sagemaker.serve.mode.in_process_mode import InProcessMode
+from sagemaker.serve import SchemaBuilder
+from sagemaker.serve.utils.types import ModelServer
+from sagemaker.serve.utils.exceptions import LocalDeepPingException
+
+
+mock_prompt = "Hello, I'm a language model,"
+mock_response = "Hello, I'm a language model, and I'm here to help you with your English."
+mock_sample_input = {"inputs": mock_prompt, "parameters": {}}
+mock_sample_output = [{"generated_text": mock_response}]
+
+
+class TestInProcessMode(unittest.TestCase):
+
+    @patch("sagemaker.serve.mode.in_process_mode.Path")
+    @patch("sagemaker.serve.spec.inference_spec.InferenceSpec")
+    @patch("sagemaker.session.Session")
+    def test_load_happy(self, mock_session, mock_inference_spec, mock_path):
+        mock_path.return_value.exists.side_effect = lambda *args, **kwargs: True
+        mock_path.return_value.is_dir.side_effect = lambda *args, **kwargs: True
+
+        mock_inference_spec.load.side_effect = lambda *args, **kwargs: "Dummy load"
+
+        mock_schema_builder = SchemaBuilder(mock_sample_input, mock_sample_output)
+        in_process_mode = InProcessMode(
+            model_server=ModelServer.MMS,
+            inference_spec=mock_inference_spec,
+            schema_builder=mock_schema_builder,
+            session=mock_session,
+            model_path="model_path",
+            env_vars={"key": "val"},
+        )
+
+        res = in_process_mode.load(model_path="/tmp/model-builder/code/")
+
+        self.assertEqual(res, "Dummy load")
+        self.assertEqual(in_process_mode.inference_spec, mock_inference_spec)
+        self.assertEqual(in_process_mode.schema_builder, mock_schema_builder)
+        self.assertEqual(in_process_mode.model_path, "model_path")
+        self.assertEqual(in_process_mode.env_vars, {"key": "val"})
+
+    @patch("sagemaker.serve.mode.in_process_mode.Path")
+    @patch("sagemaker.serve.spec.inference_spec.InferenceSpec")
+    @patch("sagemaker.session.Session")
+    def test_load_ex(self, mock_session, mock_inference_spec, mock_path):
+        mock_path.return_value.exists.side_effect = lambda *args, **kwargs: False
+        mock_path.return_value.is_dir.side_effect = lambda *args, **kwargs: True
+
+        mock_inference_spec.load.side_effect = lambda *args, **kwargs: "Dummy load"
+
+        mock_schema_builder = SchemaBuilder(mock_sample_input, mock_sample_output)
+        in_process_mode = InProcessMode(
+            model_server=ModelServer.MMS,
+            inference_spec=mock_inference_spec,
+            schema_builder=mock_schema_builder,
+            session=mock_session,
+            model_path="model_path",
+        )
+
+        self.assertRaises(ValueError, in_process_mode.load, "/tmp/model-builder/code/")
+
+        mock_path.return_value.exists.side_effect = lambda *args, **kwargs: True
+        mock_path.return_value.is_dir.side_effect = lambda *args, **kwargs: False
+
+        mock_inference_spec.load.side_effect = lambda *args, **kwargs: "Dummy load"
+        mock_schema_builder = SchemaBuilder(mock_sample_input, mock_sample_output)
+        in_process_mode = InProcessMode(
+            model_server=ModelServer.MMS,
+            inference_spec=mock_inference_spec,
+            schema_builder=mock_schema_builder,
+            session=mock_session,
+            model_path="model_path",
+        )
+
+        self.assertRaises(ValueError, in_process_mode.load, "/tmp/model-builder/code/")
+
+    @patch("sagemaker.serve.mode.in_process_mode.logger")
+    @patch("sagemaker.base_predictor.PredictorBase")
+    @patch("sagemaker.serve.spec.inference_spec.InferenceSpec")
+    @patch("sagemaker.session.Session")
+    def test_create_server_happy(
+        self, mock_session, mock_inference_spec, mock_predictor, mock_logger
+    ):
+        mock_response = "Fake response"
+        mock_multi_model_server_deep_ping = Mock()
+        mock_multi_model_server_deep_ping.side_effect = lambda *args, **kwargs: (
+            True,
+            mock_response,
+        )
+
+        in_process_mode = InProcessMode(
+            model_server=ModelServer.MMS,
+            inference_spec=mock_inference_spec,
+            schema_builder=SchemaBuilder(mock_sample_input, mock_sample_output),
+            session=mock_session,
+            model_path="model_path",
+        )
+
+        in_process_mode._multi_model_server_deep_ping = mock_multi_model_server_deep_ping
+
+        in_process_mode.create_server(predictor=mock_predictor)
+
+        mock_logger.info.assert_called_once_with(
+            "Waiting for model server %s to start up...", ModelServer.MMS
+        )
+        mock_logger.debug.assert_called_once_with(
+            "Ping health check has passed. Returned %s", str(mock_response)
+        )
+
+    @patch("sagemaker.base_predictor.PredictorBase")
+    @patch("sagemaker.serve.spec.inference_spec.InferenceSpec")
+    @patch("sagemaker.session.Session")
+    def test_create_server_ex(
+        self,
+        mock_session,
+        mock_inference_spec,
+        mock_predictor,
+    ):
+        mock_multi_model_server_deep_ping = Mock()
+        mock_multi_model_server_deep_ping.side_effect = lambda *args, **kwargs: (
+            False,
+            None,
+        )
+
+        in_process_mode = InProcessMode(
+            model_server=ModelServer.MMS,
+            inference_spec=mock_inference_spec,
+            schema_builder=SchemaBuilder(mock_sample_input, mock_sample_output),
+            session=mock_session,
+            model_path="model_path",
+        )
+
+        in_process_mode._multi_model_server_deep_ping = mock_multi_model_server_deep_ping
+
+        self.assertRaises(LocalDeepPingException, in_process_mode.create_server, mock_predictor)