Merge branch 'master' into tgi-optimum-0.0.18

Author: jinyoung-lim

MANIFEST.in

@@ -1,6 +1,7 @@
 recursive-include src/sagemaker *.py
 
 include src/sagemaker/image_uri_config/*.json
+include src/sagemaker/serve/schema/*.json
 include src/sagemaker/serve/requirements.txt
 recursive-include requirements *
 

src/sagemaker/local/image.py

@@ -40,6 +40,7 @@
 import sagemaker.local.data
 import sagemaker.local.utils
 import sagemaker.utils
+from sagemaker.utils import check_tarfile_data_filter_attribute
 
 CONTAINER_PREFIX = "algo"
 STUDIO_HOST_NAME = "sagemaker-local"
@@ -686,7 +687,8 @@ def _prepare_serving_volumes(self, model_location):
         for filename in model_data_source.get_file_list():
             if tarfile.is_tarfile(filename):
                 with tarfile.open(filename) as tar:
-                    tar.extractall(path=model_data_source.get_root_dir())
+                    check_tarfile_data_filter_attribute()
+                    tar.extractall(path=model_data_source.get_root_dir(), filter="data")
 
         volumes.append(_Volume(model_data_source.get_root_dir(), "/opt/ml/model"))
 

src/sagemaker/serve/builder/model_builder.py

@@ -38,6 +38,8 @@
 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 (
     auto_detect_container,
@@ -605,7 +607,7 @@ def build(
 
         self.serve_settings = self._get_serve_setting()
 
-        self._is_custom_image_uri = self.image_uri is None
+        self._is_custom_image_uri = self.image_uri is not None
 
         if isinstance(self.model, str):
             if self._is_jumpstart_model_id():
@@ -616,7 +618,12 @@ def build(
                 hf_model_md = get_huggingface_model_metadata(
                     self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
                 )
-                if hf_model_md.get("pipeline_tag") == "text-generation":  # pylint: disable=R1705
+
+                model_task = hf_model_md.get("pipeline_tag")
+                if self.schema_builder is None and model_task:
+                    self._schema_builder_init(model_task)
+
+                if model_task == "text-generation":  # pylint: disable=R1705
                     return self._build_for_tgi()
                 else:
                     return self._build_for_transformers()
@@ -674,3 +681,18 @@ def validate(self, model_dir: str) -> Type[bool]:
         """
 
         return get_metadata(model_dir)
+
+    def _schema_builder_init(self, model_task: str):
+        """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)
+            self.schema_builder = SchemaBuilder(sample_inputs, sample_outputs)
+        except ValueError:
+            raise TaskNotFoundException(f"Schema builder for {model_task} could not be found.")

src/sagemaker/serve/model_server/djl_serving/prepare.py

@@ -20,7 +20,7 @@
 from typing import List
 from pathlib import Path
 
-from sagemaker.utils import _tmpdir
+from sagemaker.utils import _tmpdir, check_tarfile_data_filter_attribute
 from sagemaker.s3 import S3Downloader
 from sagemaker.djl_inference import DJLModel
 from sagemaker.djl_inference.model import _read_existing_serving_properties
@@ -53,7 +53,8 @@ def _extract_js_resource(js_model_dir: str, js_id: str):
     """Uncompress the jumpstart resource"""
     tmp_sourcedir = Path(js_model_dir).joinpath(f"infer-prepack-{js_id}.tar.gz")
     with tarfile.open(str(tmp_sourcedir)) as resources:
-        resources.extractall(path=js_model_dir)
+        check_tarfile_data_filter_attribute()
+        resources.extractall(path=js_model_dir, filter="data")
 
 
 def _copy_jumpstart_artifacts(model_data: str, js_id: str, code_dir: Path):

src/sagemaker/serve/model_server/tgi/prepare.py

@@ -19,7 +19,7 @@
 from pathlib import Path
 
 from sagemaker.serve.utils.local_hardware import _check_disk_space, _check_docker_disk_usage
-from sagemaker.utils import _tmpdir
+from sagemaker.utils import _tmpdir, check_tarfile_data_filter_attribute
 from sagemaker.s3 import S3Downloader
 
 logger = logging.getLogger(__name__)
@@ -29,7 +29,8 @@ def _extract_js_resource(js_model_dir: str, code_dir: Path, js_id: str):
     """Uncompress the jumpstart resource"""
     tmp_sourcedir = Path(js_model_dir).joinpath(f"infer-prepack-{js_id}.tar.gz")
     with tarfile.open(str(tmp_sourcedir)) as resources:
-        resources.extractall(path=code_dir)
+        check_tarfile_data_filter_attribute()
+        resources.extractall(path=code_dir, filter="data")
 
 
 def _copy_jumpstart_artifacts(model_data: str, js_id: str, code_dir: Path) -> bool:

src/sagemaker/serve/schema/task.json

@@ -0,0 +1,67 @@
+{
+	"fill-mask": {
+        "sample_inputs": {
+			"properties": {
+				"inputs": "Paris is the <mask> of France.",
+				"parameters": {}
+			}
+        },
+        "sample_outputs": {
+			"properties": [
+				{
+					"sequence": "Paris is the capital of France.",
+					"score": 0.7
+				}
+			]
+		}
+    },
+	"question-answering": {
+        "sample_inputs": {
+			"properties": {
+				"context": "I have a German Shepherd dog, named Coco.",
+				"question": "What is my dog's breed?"
+			}
+        },
+        "sample_outputs": {
+			"properties": [
+				{
+					"answer": "German Shepherd",
+					"score": 0.972,
+					"start": 9,
+					"end": 24
+				}
+			]
+		}
+    },
+	"text-classification": {
+        "sample_inputs": {
+			"properties": {
+				"inputs": "Where is the capital of France?, Paris is the capital of France.",
+				"parameters": {}
+			}
+        },
+        "sample_outputs": {
+			"properties": [
+				{
+					"label": "entailment",
+					"score": 0.997
+				}
+			]
+		}
+    },
+    "text-generation": {
+        "sample_inputs": {
+			"properties": {
+				"inputs": "Hello, I'm a language model",
+				"parameters": {}
+			}
+        },
+        "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/utils/exceptions.py

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

src/sagemaker/serve/utils/task.py

@@ -0,0 +1,50 @@
+# 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.
+"""Accessors to retrieve task fallback input/output schema"""
+from __future__ import absolute_import
+
+import json
+import os
+from typing import Any, Tuple
+
+
+def retrieve_local_schemas(task: str) -> Tuple[Any, Any]:
+    """Retrieves task sample inputs and outputs locally.
+
+    Args:
+        task (str): Required, the task name
+
+    Returns:
+        Tuple[Any, Any]: A tuple that contains the sample input,
+        at index 0, and output schema, at index 1.
+
+    Raises:
+        ValueError: If no tasks config found or the task does not exist in the local config.
+    """
+    config_dir = os.path.dirname(os.path.dirname(__file__))
+    task_io_config_path = os.path.join(config_dir, "schema", "task.json")
+    try:
+        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_io_schemas is None:
+                raise ValueError(f"Could not find {task} I/O schema.")
+
+            sample_schema = (
+                task_io_schemas["sample_inputs"]["properties"],
+                task_io_schemas["sample_outputs"]["properties"],
+            )
+        return sample_schema
+    except FileNotFoundError:
+        raise ValueError("Could not find tasks config file.")

src/sagemaker/utils.py

@@ -22,6 +22,7 @@
 import random
 import re
 import shutil
+import sys
 import tarfile
 import tempfile
 import time
@@ -591,7 +592,8 @@ def _create_or_update_code_dir(
         download_file_from_url(source_directory, local_code_path, sagemaker_session)
 
         with tarfile.open(name=local_code_path, mode="r:gz") as t:
-            t.extractall(path=code_dir)
+            check_tarfile_data_filter_attribute()
+            t.extractall(path=code_dir, filter="data")
 
     elif source_directory:
         if os.path.exists(code_dir):
@@ -628,7 +630,8 @@ def _extract_model(model_uri, sagemaker_session, tmp):
     else:
         local_model_path = model_uri.replace("file://", "")
     with tarfile.open(name=local_model_path, mode="r:gz") as t:
-        t.extractall(path=tmp_model_dir)
+        check_tarfile_data_filter_attribute()
+        t.extractall(path=tmp_model_dir, filter="data")
     return tmp_model_dir
 
 
@@ -1489,3 +1492,25 @@ def format_tags(tags: Tags) -> List[TagsDict]:
         return [{"Key": str(k), "Value": str(v)} for k, v in tags.items()]
 
     return tags
+
+
+class PythonVersionError(Exception):
+    """Raise when a secure [/patched] version of Python is not used."""
+
+
+def check_tarfile_data_filter_attribute():
+    """Check if tarfile has data_filter utility.
+
+    Tarfile-data_filter utility has guardrails against untrusted de-serialisation.
+
+    Raises:
+        PythonVersionError: if `tarfile.data_filter` is not available.
+    """
+    # The function and it's usages can be deprecated post support of python >= 3.12
+    if not hasattr(tarfile, "data_filter"):
+        raise PythonVersionError(
+            f"Since tarfile extraction is unsafe the operation is prohibited "
+            f"per PEP-721. Please update your Python [{sys.version}] "
+            f"to latest patch [refer to https://www.python.org/downloads/] "
+            f"to consume the security patch"
+        )

src/sagemaker/workflow/_utils.py

@@ -32,7 +32,12 @@
     Step,
     ConfigurableRetryStep,
 )
-from sagemaker.utils import _save_model, download_file_from_url, format_tags
+from sagemaker.utils import (
+    _save_model,
+    download_file_from_url,
+    format_tags,
+    check_tarfile_data_filter_attribute,
+)
 from sagemaker.workflow.retry import RetryPolicy
 from sagemaker.workflow.utilities import trim_request_dict
 
@@ -257,7 +262,8 @@ def _inject_repack_script_and_launcher(self):
                 download_file_from_url(self._source_dir, old_targz_path, self.sagemaker_session)
 
                 with tarfile.open(name=old_targz_path, mode="r:gz") as t:
-                    t.extractall(path=targz_contents_dir)
+                    check_tarfile_data_filter_attribute()
+                    t.extractall(path=targz_contents_dir, filter="data")
 
                 shutil.copy2(fname, os.path.join(targz_contents_dir, REPACK_SCRIPT))
                 with open(

tests/integ/s3_utils.py

@@ -19,6 +19,8 @@
 import boto3
 from six.moves.urllib.parse import urlparse
 
+from sagemaker.utils import check_tarfile_data_filter_attribute
+
 
 def assert_s3_files_exist(sagemaker_session, s3_url, files):
     parsed_url = urlparse(s3_url)
@@ -55,4 +57,5 @@ def extract_files_from_s3(s3_url, tmpdir, sagemaker_session):
     s3.Bucket(parsed_url.netloc).download_file(parsed_url.path.lstrip("/"), model)
 
     with tarfile.open(model, "r") as tar_file:
-        tar_file.extractall(tmpdir)
+        check_tarfile_data_filter_attribute()
+        tar_file.extractall(tmpdir, filter="data")