Add Integ tests

Author: Jonathan Makunga

src/sagemaker/serve/builder/model_builder.py

@@ -38,6 +38,7 @@
 from sagemaker.predictor import Predictor
 from sagemaker.serve.save_retrive.version_1_0_0.metadata.metadata import Metadata
 from sagemaker.serve.spec.inference_spec import InferenceSpec
+from sagemaker.serve.utils.exceptions import TaskNotFoundException
 from sagemaker.serve.utils.predictors import _get_local_mode_predictor
 from sagemaker.serve.detector.image_detector import (
     auto_detect_container,
@@ -609,20 +610,17 @@ def build(
             if self._is_jumpstart_model_id():
                 return self._build_for_jumpstart()
             if self._is_djl():  # pylint: disable=R1705
-                if self.schema_builder is None:
-                    self._schema_builder_init("text-generation")
-
                 return self._build_for_djl()
             else:
                 hf_model_md = get_huggingface_model_metadata(
                     self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
                 )
 
-                hf_task = hf_model_md.get("pipeline_tag")
+                model_task = hf_model_md.get("pipeline_tag")
                 if self.schema_builder is None:
-                    self._schema_builder_init(hf_task)
+                    self._schema_builder_init(model_task)
 
-                if hf_task == "text-generation":  # pylint: disable=R1705
+                if model_task == "text-generation":  # pylint: disable=R1705
                     return self._build_for_tgi()
                 else:
                     return self._build_for_transformers()
@@ -682,17 +680,16 @@ def validate(self, model_dir: str) -> Type[bool]:
         return get_metadata(model_dir)
 
     def _schema_builder_init(self, model_task: str):
-        """Initialize the"""
-        sample_inputs, sample_outputs = None, None
+        """Initialize the schema builder
 
+        Args:
+            model_task (str): Required, the task name
+
+        Raises:
+            TaskNotFoundException: If the I/O schema for the given task is not found.
+        """
         try:
             sample_inputs, sample_outputs = task.retrieve_local_schemas(model_task)
-        except ValueError:
-            # TODO: try to retrieve schemas remotely
-            pass
-
-        if sample_inputs and sample_outputs:
             self.schema_builder = SchemaBuilder(sample_inputs, sample_outputs)
-        else:
-            # TODO: Raise ClientError
-            pass
+        except ValueError:
+            raise TaskNotFoundException(f"Schema builder for {model_task} could not be found.")

src/sagemaker/serve/utils/exceptions.py

@@ -60,3 +60,12 @@ class SkipTuningComboException(ModelBuilderException):
 
     def __init__(self, message):
         super().__init__(message=message)
+
+
+class TaskNotFoundException(ModelBuilderException):
+    """Raise when task could not be found"""
+
+    fmt = "Error Message: {message}"
+
+    def __init__(self, message):
+        super().__init__(message=message)

src/sagemaker/task.py

@@ -15,7 +15,6 @@
 
 import json
 import os
-from enum import Enum
 from typing import Any, Tuple
 
 
@@ -32,20 +31,19 @@ def retrieve_local_schemas(task: str) -> Tuple[Any, Any]:
     Raises:
         ValueError: If no tasks config found or the task does not exist in the local config.
     """
-    task_path = os.path.join(os.path.dirname(__file__), "image_uri_config", "tasks.json")
+    task_io_config_path = os.path.join(os.path.dirname(__file__), "image_uri_config", "tasks.json")
     try:
-        with open(task_path) as f:
-            task_config = json.load(f)
-            task_schema = task_config.get(task, None)
+        with open(task_io_config_path) as f:
+            task_io_config = json.load(f)
+            task_io_schemas = task_io_config.get(task, None)
 
-            if task_schema is None:
-                raise ValueError(f"Could not find {task} task schema.")
+            if task_io_schemas is None:
+                raise ValueError(f"Could not find {task} I/O schema.")
 
             sample_schema = (
-                task_schema["inputs"]["properties"],
-                task_schema["outputs"]["properties"],
+                task_io_schemas["inputs"]["properties"],
+                task_io_schemas["outputs"]["properties"],
             )
         return sample_schema
-
     except FileNotFoundError:
         raise ValueError("Could not find tasks config file.")

tests/integ/sagemaker/serve/test_schema_builder.py

@@ -0,0 +1,46 @@
+# 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
+
+from sagemaker import task
+from sagemaker.serve.builder.model_builder import ModelBuilder
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def test_model_builder_happy_path_with_only_model_id_fill_mask(sagemaker_session):
+    model_builder = ModelBuilder(model="bert-base-uncased")
+
+    model = model_builder.build(sagemaker_session=sagemaker_session)
+
+    assert model is not None
+    assert model_builder.schema_builder is not None
+
+    inputs, outputs = task.retrieve_local_schemas("fill-mask")
+    assert model_builder.schema_builder.sample_input == inputs
+    assert model_builder.schema_builder.sample_output == outputs
+
+
+def test_model_builder_happy_path_with_only_model_id_question_answering(sagemaker_session):
+    model_builder = ModelBuilder(model="bert-large-uncased-whole-word-masking-finetuned-squad")
+
+    model = model_builder.build(sagemaker_session=sagemaker_session)
+
+    assert model is not None
+    assert model_builder.schema_builder is not None
+
+    inputs, outputs = task.retrieve_local_schemas("question-answering")
+    assert model_builder.schema_builder.sample_input == inputs
+    assert model_builder.schema_builder.sample_output == outputs

tests/unit/sagemaker/test_task.py

@@ -26,6 +26,13 @@ def test_retrieve_local_schemas_success():
     assert outputs == EXPECTED_OUTPUTS
 
 
+def test_retrieve_local_schemas_text_generation_success():
+    inputs, outputs = task.retrieve_local_schemas("text-generation")
+
+    assert inputs is not None
+    assert outputs is not None
+
+
 def test_retrieve_local_schemas_throws():
     with pytest.raises(ValueError):
-        task.retrieve_local_schemas("invalid-task")
+        task.retrieve_local_schemas("not-present-task")