address PR comments

Author: Jonathan Makunga

src/sagemaker/image_uri_config/tasks.json

@@ -1,16 +1,12 @@
 {
-	"description": "Sample Task Inputs and Outputs",
 	"fill-mask": {
-		"ref": "https://huggingface.co/tasks/fill-mask",
-        "inputs": {
-			"ref": "https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/tasks/fill-mask/spec/input.json",
+        "sample_inputs": {
 			"properties": {
 				"inputs": "Paris is the <mask> of France.",
 				"parameters": {}
 			}
         },
-        "outputs": {
-			"ref": "https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/tasks/fill-mask/spec/output.json",
+        "sample_outputs": {
 			"properties": [
 				{
 					"sequence": "Paris is the capital of France.",
@@ -20,16 +16,13 @@
 		}
     },
 	"question-answering": {
-		"ref": "https://huggingface.co/tasks/question-answering",
-        "inputs": {
-			"ref": "https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/tasks/question-answering/spec/input.json",
+        "sample_inputs": {
 			"properties": {
 				"context": "I have a German Shepherd dog, named Coco.",
 				"question": "What is my dog's breed?"
 			}
         },
-        "outputs": {
-			"ref": "https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/tasks/question-answering/spec/output.json",
+        "sample_outputs": {
 			"properties": [
 				{
 					"answer": "German Shepherd",
@@ -41,16 +34,13 @@
 		}
     },
 	"text-classification": {
-		"ref": "https://huggingface.co/tasks/text-classification",
-        "inputs": {
-			"ref": "https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/tasks/text-classification/spec/input.json",
+        "sample_inputs": {
 			"properties": {
 				"inputs": "Where is the capital of France?, Paris is the capital of France.",
 				"parameters": {}
 			}
         },
-        "outputs": {
-			"ref": "https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/tasks/text-classification/spec/output.json",
+        "sample_outputs": {
 			"properties": [
 				{
 					"label": "entailment",
@@ -60,16 +50,13 @@
 		}
     },
     "text-generation": {
-		"ref": "https://huggingface.co/tasks/text-generation",
-        "inputs": {
-			"ref": "https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/tasks/text-generation/spec/input.json",
+        "sample_inputs": {
 			"properties": {
 				"inputs": "Hello, I'm a language model",
 				"parameters": {}
 			}
         },
-        "outputs": {
-			"ref": "https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/src/tasks/text-generation/spec/output.json",
+        "sample_outputs": {
 			"properties": [
 				{
 					"generated_text": "Hello, I'm a language modeler. So while writing this, when I went out to meet my wife or come home she told me that my"

src/sagemaker/serve/builder/model_builder.py

@@ -20,7 +20,7 @@
 
 from pathlib import Path
 
-from sagemaker import Session, task
+from sagemaker import Session
 from sagemaker.model import Model
 from sagemaker.base_predictor import PredictorBase
 from sagemaker.serializers import NumpySerializer, TorchTensorSerializer
@@ -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 import task
 from sagemaker.serve.utils.exceptions import TaskNotFoundException
 from sagemaker.serve.utils.predictors import _get_local_mode_predictor
 from sagemaker.serve.detector.image_detector import (
@@ -617,7 +618,7 @@ def build(
                 )
 
                 model_task = hf_model_md.get("pipeline_tag")
-                if self.schema_builder is None:
+                if self.schema_builder is None and model_task:
                     self._schema_builder_init(model_task)
 
                 if model_task == "text-generation":  # pylint: disable=R1705

src/sagemaker/serve/utils/exceptions.py

@@ -63,7 +63,7 @@ def __init__(self, message):
 
 
 class TaskNotFoundException(ModelBuilderException):
-    """Raise when task could not be found"""
+    """Raise when HuggingFace task could not be found"""
 
     fmt = "Error Message: {message}"
 

src/sagemaker/task.py renamed to src/sagemaker/serve/utils/task.py

@@ -31,7 +31,9 @@ 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_io_config_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")
+    c = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
+    task_io_config_path = os.path.join(c, "image_uri_config", "tasks.json")
     try:
         with open(task_io_config_path) as f:
             task_io_config = json.load(f)
@@ -41,8 +43,8 @@ def retrieve_local_schemas(task: str) -> Tuple[Any, Any]:
                 raise ValueError(f"Could not find {task} I/O schema.")
 
             sample_schema = (
-                task_io_schemas["inputs"]["properties"],
-                task_io_schemas["outputs"]["properties"],
+                task_io_schemas["sample_inputs"]["properties"],
+                task_io_schemas["sample_outputs"]["properties"],
             )
         return sample_schema
     except FileNotFoundError:

tests/integ/sagemaker/serve/test_schema_builder.py

@@ -12,8 +12,19 @@
 # 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
+from sagemaker.serve.utils import task
+
+import pytest
+
+from sagemaker.serve.utils.exceptions import TaskNotFoundException
+from tests.integ.sagemaker.serve.constants import (
+    PYTHON_VERSION_IS_NOT_310,
+    SERVE_SAGEMAKER_ENDPOINT_TIMEOUT,
+)
+
+from tests.integ.timeout import timeout
+from tests.integ.utils import cleanup_model_resources
 
 import logging
 
@@ -33,7 +44,13 @@ def test_model_builder_happy_path_with_only_model_id_fill_mask(sagemaker_session
     assert model_builder.schema_builder.sample_output == outputs
 
 
-def test_model_builder_happy_path_with_only_model_id_question_answering(sagemaker_session):
+@pytest.mark.skipif(
+    PYTHON_VERSION_IS_NOT_310,
+    reason="Testing Schema Builder Simplification feature",
+)
+def test_model_builder_happy_path_with_only_model_id_question_answering(
+    sagemaker_session, gpu_instance_type
+):
     model_builder = ModelBuilder(model="bert-large-uncased-whole-word-masking-finetuned-squad")
 
     model = model_builder.build(sagemaker_session=sagemaker_session)
@@ -44,3 +61,41 @@ def test_model_builder_happy_path_with_only_model_id_question_answering(sagemake
     inputs, outputs = task.retrieve_local_schemas("question-answering")
     assert model_builder.schema_builder.sample_input == inputs
     assert model_builder.schema_builder.sample_output == outputs
+
+    with timeout(minutes=SERVE_SAGEMAKER_ENDPOINT_TIMEOUT):
+        caught_ex = None
+        try:
+            iam_client = sagemaker_session.boto_session.client("iam")
+            role_arn = iam_client.get_role(RoleName="JarvisTest")["Role"]["Arn"]
+
+            logger.info("Deploying and predicting in SAGEMAKER_ENDPOINT mode...")
+            predictor = model.deploy(
+                role=role_arn, instance_count=1, instance_type=gpu_instance_type
+            )
+
+            predicted_outputs = predictor.predict(inputs)
+            assert predicted_outputs is not None
+
+        except Exception as e:
+            caught_ex = e
+        finally:
+            cleanup_model_resources(
+                sagemaker_session=model_builder.sagemaker_session,
+                model_name=model.name,
+                endpoint_name=model.endpoint_name,
+            )
+            if caught_ex:
+                logger.exception(caught_ex)
+                assert (
+                    False
+                ), f"{caught_ex} was thrown when running transformers sagemaker endpoint test"
+
+
+def test_model_builder_negative_path(sagemaker_session):
+    model_builder = ModelBuilder(model="CompVis/stable-diffusion-v1-4")
+
+    with pytest.raises(
+        TaskNotFoundException,
+        match="Error Message: Schema builder for text-to-image could not be found.",
+    ):
+        model_builder.build(sagemaker_session=sagemaker_session)

tests/unit/sagemaker/test_task.py renamed to tests/unit/sagemaker/serve/utils/test_task.py

@@ -12,11 +12,15 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+from unittest.mock import patch
+
 import pytest
-from sagemaker import task
+
+from sagemaker.serve.utils import task
 
 EXPECTED_INPUTS = {"inputs": "Paris is the <mask> of France.", "parameters": {}}
 EXPECTED_OUTPUTS = [{"sequence": "Paris is the capital of France.", "score": 0.7}]
+HF_INVALID_TASK = "not-present-task"
 
 
 def test_retrieve_local_schemas_success():
@@ -34,5 +38,12 @@ def test_retrieve_local_schemas_text_generation_success():
 
 
 def test_retrieve_local_schemas_throws():
-    with pytest.raises(ValueError):
-        task.retrieve_local_schemas("not-present-task")
+    with pytest.raises(ValueError, match=f"Could not find {HF_INVALID_TASK} I/O schema."):
+        task.retrieve_local_schemas(HF_INVALID_TASK)
+
+
+@patch("builtins.open")
+def test_retrieve_local_schemas_file_not_found(mock_open):
+    mock_open.side_effect = FileNotFoundError
+    with pytest.raises(ValueError, match="Could not find tasks config file."):
+        task.retrieve_local_schemas(HF_INVALID_TASK)