feat: combined model + script artifact

evakravi · evakravi · commit 8bb9cbadf8a1 · 2023-03-13T20:25:33.000Z
diff --git a/src/sagemaker/jumpstart/artifacts.py b/src/sagemaker/jumpstart/artifacts.py
@@ -173,10 +173,11 @@ def _retrieve_image_uri(
 def _retrieve_model_uri(
     model_id: str,
     model_version: str,
-    model_scope: Optional[str],
-    region: Optional[str],
-    tolerate_vulnerable_model: bool,
-    tolerate_deprecated_model: bool,
+    model_scope: Optional[str] = None,
+    region: Optional[str] = None,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+    include_script: bool = False,
 ):
     """Retrieves the model artifact S3 URI for the model matching the given arguments.
 
@@ -197,6 +198,8 @@ def _retrieve_model_uri(
         tolerate_deprecated_model (bool): True if deprecated versions of model
             specifications should be tolerated (exception not raised). If False, raises
             an exception if the version of the model is deprecated.
+        include_script (bool): True if script artifact should be packaged alongside model
+            tarball. (Default: False).
     Returns:
         str: the model artifact S3 URI for the corresponding model.
 
@@ -205,6 +208,8 @@ def _retrieve_model_uri(
         VulnerableJumpStartModelError: If any of the dependencies required by the script have
             known security vulnerabilities.
         DeprecatedJumpStartModelError: If the version of the model is deprecated.
+        NotImplementedError: If the combination of arguments doesn't support combined model
+            and script artifact.
     """
     if region is None:
         region = JUMPSTART_DEFAULT_REGION_NAME
@@ -218,10 +223,24 @@ def _retrieve_model_uri(
         tolerate_deprecated_model=tolerate_deprecated_model,
     )
 
+    error_msg_no_combined_artifact = (
+        "No combined script + model tarball available "
+        f"for {model_id} with version {model_version} for {model_scope}."
+    )
+
     if model_scope == JumpStartScriptScope.INFERENCE:
-        model_artifact_key = model_specs.hosting_artifact_key
+        if not include_script:
+            model_artifact_key = model_specs.hosting_artifact_key
+        else:
+            model_artifact_key = getattr(model_specs, "hosting_prepacked_artifact_key", None)
+            if model_artifact_key is None:
+                raise NotImplementedError(error_msg_no_combined_artifact)
+
     elif model_scope == JumpStartScriptScope.TRAINING:
-        model_artifact_key = model_specs.training_artifact_key
+        if not include_script:
+            model_artifact_key = model_specs.training_artifact_key
+        else:
+            raise NotImplementedError(error_msg_no_combined_artifact)
 
     bucket = os.environ.get(
         ENV_VARIABLE_JUMPSTART_MODEL_ARTIFACT_BUCKET_OVERRIDE
diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
@@ -293,6 +293,7 @@ class JumpStartModelSpecs(JumpStartDataHolderType):
         "training_vulnerabilities",
         "deprecated",
         "metrics",
+        "hosting_prepacked_artifact_key",
     ]
 
     def __init__(self, spec: Dict[str, Any]):
@@ -330,6 +331,9 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.training_vulnerabilities: List[str] = json_obj["training_vulnerabilities"]
         self.deprecated: bool = bool(json_obj["deprecated"])
         self.metrics: Optional[List[Dict[str, str]]] = json_obj.get("metrics", None)
+        self.hosting_prepacked_artifact_key: Optional[str] = json_obj.get(
+            "hosting_prepacked_artifact_key", None
+        )
 
         if self.training_supported:
             self.training_ecr_specs: JumpStartECRSpecs = JumpStartECRSpecs(
diff --git a/src/sagemaker/model_uris.py b/src/sagemaker/model_uris.py
@@ -30,6 +30,7 @@ def retrieve(
     model_scope: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
+    include_script: bool = False,
 ) -> str:
     """Retrieves the model artifact Amazon S3 URI for the model matching the given arguments.
 
@@ -48,6 +49,8 @@ def retrieve(
         tolerate_deprecated_model (bool): ``True`` if deprecated versions of model
             specifications should be tolerated without raising an exception. If ``False``, raises
             an exception if the version of the model is deprecated. (Default: False).
+        include_script (bool): True if script artifact should be packaged alongside model
+            tarball. (Default: False).
     Returns:
         str: The model artifact S3 URI for the corresponding model.
 
@@ -57,6 +60,8 @@ def retrieve(
         VulnerableJumpStartModelError: If any of the dependencies required by the script have
             known security vulnerabilities.
         DeprecatedJumpStartModelError: If the version of the model is deprecated.
+        NotImplementedError: If the combination of arguments doesn't support combined model
+            and script artifact.
     """
     if not jumpstart_utils.is_jumpstart_model_input(model_id, model_version):
         raise ValueError("Must specify `model_id` and `model_version` when retrieving model URIs.")
@@ -68,4 +73,5 @@ def retrieve(
         region,
         tolerate_vulnerable_model,
         tolerate_deprecated_model,
+        include_script,
     )
diff --git a/tests/unit/sagemaker/jumpstart/constants.py b/tests/unit/sagemaker/jumpstart/constants.py
@@ -1070,6 +1070,88 @@
             },
         ],
     },
+    "huggingface-text2text-flan-t5-xxl-fp16": {
+        "model_id": "huggingface-text2text-flan-t5-xxl-fp16",
+        "url": "https://huggingface.co/google/flan-t5-xxl",
+        "version": "1.0.0",
+        "min_sdk_version": "2.130.0",
+        "training_supported": False,
+        "incremental_training_supported": False,
+        "hosting_ecr_specs": {
+            "framework": "pytorch",
+            "framework_version": "1.12.0",
+            "py_version": "py38",
+            "huggingface_transformers_version": "4.17.0",
+        },
+        "hosting_artifact_key": "huggingface-infer/infer-huggingface-text2text-flan-t5-xxl-fp16.tar.gz",
+        "hosting_script_key": "source-directory-tarballs/huggingface/inference/text2text/v1.0.2/sourcedir.tar.gz",
+        "hosting_prepacked_artifact_key": "huggingface-infer/prepack/v1.0.0/infer-prepack-huggingface-"
+        "text2text-flan-t5-xxl-fp16.tar.gz",
+        "hosting_prepacked_artifact_version": "1.0.0",
+        "inference_vulnerable": False,
+        "inference_dependencies": [
+            "accelerate==0.16.0",
+            "bitsandbytes==0.37.0",
+            "filelock==3.9.0",
+            "huggingface-hub==0.12.0",
+            "regex==2022.7.9",
+            "tokenizers==0.13.2",
+            "transformers==4.26.0",
+        ],
+        "inference_vulnerabilities": [],
+        "training_vulnerable": False,
+        "training_dependencies": [],
+        "training_vulnerabilities": [],
+        "deprecated": False,
+        "inference_environment_variables": [
+            {
+                "name": "SAGEMAKER_PROGRAM",
+                "type": "text",
+                "default": "inference.py",
+                "scope": "container",
+            },
+            {
+                "name": "SAGEMAKER_SUBMIT_DIRECTORY",
+                "type": "text",
+                "default": "/opt/ml/model/code",
+                "scope": "container",
+            },
+            {
+                "name": "SAGEMAKER_CONTAINER_LOG_LEVEL",
+                "type": "text",
+                "default": "20",
+                "scope": "container",
+            },
+            {
+                "name": "MODEL_CACHE_ROOT",
+                "type": "text",
+                "default": "/opt/ml/model",
+                "scope": "container",
+            },
+            {"name": "SAGEMAKER_ENV", "type": "text", "default": "1", "scope": "container"},
+            {
+                "name": "SAGEMAKER_MODEL_SERVER_WORKERS",
+                "type": "text",
+                "default": "1",
+                "scope": "container",
+            },
+            {
+                "name": "SAGEMAKER_MODEL_SERVER_TIMEOUT",
+                "type": "text",
+                "default": "3600",
+                "scope": "container",
+            },
+        ],
+        "metrics": [],
+        "default_inference_instance_type": "ml.g5.12xlarge",
+        "supported_inference_instance_types": [
+            "ml.g5.12xlarge",
+            "ml.g5.24xlarge",
+            "ml.p3.8xlarge",
+            "ml.p3.16xlarge",
+            "ml.g4dn.12xlarge",
+        ],
+    },
 }
 
 BASE_SPEC = {
diff --git a/tests/unit/sagemaker/model_uris/jumpstart/test_combined_artifact.py b/tests/unit/sagemaker/model_uris/jumpstart/test_combined_artifact.py
@@ -0,0 +1,60 @@
+# 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 mock.mock import patch
+
+from sagemaker import model_uris
+import pytest
+
+from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
+
+
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_combined_artifacts(patched_get_model_specs):
+
+    patched_get_model_specs.side_effect = get_prototype_model_spec
+
+    model_id_combined_model_artifact = "huggingface-text2text-flan-t5-xxl-fp16"
+
+    uri = model_uris.retrieve(
+        region="us-west-2",
+        model_scope="inference",
+        model_id=model_id_combined_model_artifact,
+        model_version="*",
+        include_script=True,
+    )
+    assert (
+        uri == "s3://jumpstart-cache-prod-us-west-2/huggingface-infer/"
+        "prepack/v1.0.0/infer-prepack-huggingface-text2text-flan-t5-xxl-fp16.tar.gz"
+    )
+
+    with pytest.raises(NotImplementedError):
+        model_uris.retrieve(
+            region="us-west-2",
+            model_scope="transfer_learning",
+            model_id=model_id_combined_model_artifact,
+            model_version="*",
+            include_script=True,
+        )
+
+    model_id_combined_model_artifact_unsupported = "xgboost-classification-model"
+
+    with pytest.raises(NotImplementedError):
+        model_uris.retrieve(
+            region="us-west-2",
+            model_scope="inference",
+            model_id=model_id_combined_model_artifact_unsupported,
+            model_version="*",
+            include_script=True,
+        )
