feat: jumpstart instance type variants

Author: evakravi

src/sagemaker/image_uris.py

@@ -262,20 +262,6 @@ def retrieve(
     return ECR_URI_TEMPLATE.format(registry=registry, hostname=hostname, repository=repo)
 
 
-def _get_instance_type_family(instance_type):
-    """Return the family of the instance type.
-
-    Regex matches either "ml.<family>.<size>" or "ml_<family>. If input is None
-    or there is no match, return an empty string.
-    """
-    instance_type_family = ""
-    if isinstance(instance_type, str):
-        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
-        if match is not None:
-            instance_type_family = match[1]
-    return instance_type_family
-
-
 def _get_image_tag(
     container_version,
     distribution,
@@ -289,7 +275,7 @@ def _get_image_tag(
     version,
 ):
     """Return image tag based on framework, container, and compute configuration(s)."""
-    instance_type_family = _get_instance_type_family(instance_type)
+    instance_type_family = utils.get_instance_type_family(instance_type)
     if framework in (XGBOOST_FRAMEWORK, SKLEARN_FRAMEWORK):
         if instance_type_family and final_image_scope == INFERENCE_GRAVITON:
             _validate_arg(
@@ -377,7 +363,7 @@ def _config_for_framework_and_scope(framework, image_scope, accelerator_type=Non
 
 def _validate_instance_deprecation(framework, instance_type, version):
     """Check if instance type is deprecated for a certain framework with a certain version"""
-    if _get_instance_type_family(instance_type) == "p2":
+    if utils.get_instance_type_family(instance_type) == "p2":
         if (framework == "pytorch" and Version(version) >= Version("1.13")) or (
             framework == "tensorflow" and Version(version) >= Version("2.12")
         ):
@@ -401,7 +387,7 @@ def _validate_for_suppported_frameworks_and_instance_type(framework, instance_ty
     # Validate for Graviton allowed frameowrks
     if (
         instance_type is not None
-        and _get_instance_type_family(instance_type) in GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY
+        and utils.get_instance_type_family(instance_type) in GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY
         and framework not in GRAVITON_ALLOWED_FRAMEWORKS
     ):
         _validate_framework(framework, GRAVITON_ALLOWED_FRAMEWORKS, "framework", "Graviton")
@@ -418,7 +404,7 @@ def _get_final_image_scope(framework, instance_type, image_scope):
     """Return final image scope based on provided framework and instance type."""
     if (
         framework in GRAVITON_ALLOWED_FRAMEWORKS
-        and _get_instance_type_family(instance_type) in GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY
+        and utils.get_instance_type_family(instance_type) in GRAVITON_ALLOWED_TARGET_INSTANCE_FAMILY
     ):
         return INFERENCE_GRAVITON
     if image_scope is None and framework in (XGBOOST_FRAMEWORK, SKLEARN_FRAMEWORK):
@@ -433,7 +419,7 @@ def _get_final_image_scope(framework, instance_type, image_scope):
 def _get_inference_tool(inference_tool, instance_type):
     """Extract the inference tool name from instance type."""
     if not inference_tool:
-        instance_type_family = _get_instance_type_family(instance_type)
+        instance_type_family = utils.get_instance_type_family(instance_type)
         if instance_type_family.startswith("inf") or instance_type_family.startswith("trn"):
             return "neuron"
     return inference_tool
@@ -517,7 +503,7 @@ def _processor(instance_type, available_processors, serverless_inference_config=
         processor = "neuron"
     else:
         # looks for either "ml.<family>.<size>" or "ml_<family>"
-        family = _get_instance_type_family(instance_type)
+        family = utils.get_instance_type_family(instance_type)
         if family:
             # For some frameworks, we have optimized images for specific families, e.g c5 or p3.
             # In those cases, we use the family name in the image tag. In other cases, we use
@@ -547,7 +533,7 @@ def _should_auto_select_container_version(instance_type, distribution):
     p4d = False
     if instance_type:
         # looks for either "ml.<family>.<size>" or "ml_<family>"
-        family = _get_instance_type_family(instance_type)
+        family = utils.get_instance_type_family(instance_type)
         if family:
             p4d = family == "p4d"
 

src/sagemaker/jumpstart/artifacts/image_uris.py

@@ -111,8 +111,30 @@ def _retrieve_image_uri(
     )
 
     if image_scope == JumpStartScriptScope.INFERENCE:
+        hosting_instance_type_variants = model_specs.hosting_instance_type_variants
+        if hosting_instance_type_variants:
+            image_uri = hosting_instance_type_variants.get_image_uri(
+                instance_type=instance_type, region=region
+            )
+            if image_uri is None:
+                raise ValueError(
+                    f"Inference image uri is unavailable for model id '{model_id}' "
+                    f"with '{instance_type}' instance type in '{region}' region."
+                )
+            return image_uri
         ecr_specs = model_specs.hosting_ecr_specs
     elif image_scope == JumpStartScriptScope.TRAINING:
+        training_instance_type_variants = model_specs.training_instance_type_variants
+        if training_instance_type_variants:
+            image_uri = training_instance_type_variants.get_image_uri(
+                instance_type=instance_type, region=region
+            )
+            if image_uri is None:
+                raise ValueError(
+                    f"Training image uri is unavailable for model id '{model_id}' "
+                    f"with '{instance_type}' instance type in '{region}' region."
+                )
+            return image_uri
         ecr_specs = model_specs.training_ecr_specs
 
     if framework is not None and framework != ecr_specs.framework:

src/sagemaker/jumpstart/types.py

@@ -15,6 +15,7 @@
 from copy import deepcopy
 from enum import Enum
 from typing import Any, Dict, List, Optional, Set, Union
+from sagemaker.utils import get_instance_type_family
 
 
 class JumpStartDataHolderType:
@@ -309,6 +310,69 @@ def to_json(self) -> Dict[str, Any]:
         return json_obj
 
 
+class JumpStartInstanceTypeVariants(JumpStartDataHolderType):
+    """Data class for JumpStart instance type variants."""
+
+    __slots__ = [
+        "aliases",
+        "variants",
+    ]
+
+    def __init__(self, spec: Optional[Dict[str, Any]]):
+        """Initializes a JumpStartInstanceTypeVariants object from its json representation.
+
+        Args:
+            spec (Dict[str, Any]): Dictionary representation of instance type variants.
+        """
+        self.from_json(spec)
+
+    def from_json(self, json_obj: Optional[Dict[str, Any]]) -> None:
+        """Sets fields in object based on json.
+
+        Args:
+            json_obj (Dict[str, Any]): Dictionary representation of instance type variants.
+        """
+
+        if json_obj is None:
+            return
+
+        self.aliases: dict = json_obj["aliases"]
+        self.variant: dict = json_obj["variants"]
+
+    def to_json(self) -> Dict[str, Any]:
+        """Returns json representation of JumpStartPredictorSpecs object."""
+        json_obj = {att: getattr(self, att) for att in self.__slots__ if hasattr(self, att)}
+        return json_obj
+
+    def get_image_uri(self, instance_type: str, region: str) -> Optional[str]:
+        """Returns image uri from instance type and region.
+
+        Returns None if no instance type is available or found.
+        """
+
+        image_uri_alias: Optional[str] = None
+        if instance_type in self.variants:
+            image_uri_alias = self.variants[instance_type]["properties"].get("image_uri")
+        else:
+            instance_type_family = get_instance_type_family(instance_type)
+            image_uri_alias = (
+                self.variants[instance_type_family]["properties"].get("image_uri")
+                if instance_type_family in self.variants
+                else None
+            )
+
+        if image_uri_alias is None:
+            return image_uri_alias
+
+        if not image_uri_alias.startswith("$"):
+            raise TypeError("All image uris should map to an alias an start with '$'.")
+
+        if region not in self.aliases:
+            return None
+        alias_value = self.aliases[region][image_uri_alias[1:]]
+        return alias_value
+
+
 class JumpStartModelSpecs(JumpStartDataHolderType):
     """Data class JumpStart model specs."""
 
@@ -357,6 +421,8 @@ class JumpStartModelSpecs(JumpStartDataHolderType):
         "hosting_model_package_arns",
         "training_model_package_artifact_uris",
         "hosting_use_script_uri",
+        "hosting_instance_type_variants",
+        "training_instance_type_variants",
     ]
 
     def __init__(self, spec: Dict[str, Any]):
@@ -432,6 +498,12 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.hosting_model_package_arns: Optional[Dict] = json_obj.get("hosting_model_package_arns")
         self.hosting_use_script_uri: bool = json_obj.get("hosting_use_script_uri", True)
 
+        self.hosting_instance_type_variants: Optional[JumpStartInstanceTypeVariants] = (
+            JumpStartInstanceTypeVariants(json_obj["hosting_instance_type_variants"])
+            if json_obj.get("hosting_instance_type_variants")
+            else None
+        )
+
         if self.training_supported:
             self.training_ecr_specs: JumpStartECRSpecs = JumpStartECRSpecs(
                 json_obj["training_ecr_specs"]
@@ -453,6 +525,11 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
             self.training_model_package_artifact_uris: Optional[Dict] = json_obj.get(
                 "training_model_package_artifact_uris"
             )
+            self.training_instance_type_variants: Optional[JumpStartInstanceTypeVariants] = (
+                JumpStartInstanceTypeVariants(json_obj["training_instance_type_variants"])
+                if json_obj.get("training_instance_type_variants")
+                else None
+            )
 
     def to_json(self) -> Dict[str, Any]:
         """Returns json representation of JumpStartModelSpecs object."""

src/sagemaker/utils.py

@@ -1435,3 +1435,17 @@ def instance_supports_kms(instance_type: str) -> bool:
         ValueError: If the instance type is improperly formatted.
     """
     return volume_size_supported(instance_type)
+
+
+def get_instance_type_family(instance_type):
+    """Return the family of the instance type.
+
+    Regex matches either "ml.<family>.<size>" or "ml_<family>. If input is None
+    or there is no match, return an empty string.
+    """
+    instance_type_family = ""
+    if isinstance(instance_type, str):
+        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        if match is not None:
+            instance_type_family = match[1]
+    return instance_type_family

tests/unit/sagemaker/image_uris/jumpstart/test_variants.py

@@ -0,0 +1,85 @@
+# 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
+import pytest
+
+from sagemaker import image_uris
+from sagemaker.jumpstart.utils import verify_model_region_and_return_specs
+
+from tests.unit.sagemaker.jumpstart.utils import get_special_model_spec
+
+
+@patch("sagemaker.jumpstart.artifacts.image_uris.verify_model_region_and_return_specs")
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_variants_image_uri(
+    patched_get_model_specs, patched_verify_model_region_and_return_specs
+):
+
+    patched_verify_model_region_and_return_specs.side_effect = verify_model_region_and_return_specs
+    patched_get_model_specs.side_effect = get_special_model_spec
+
+    assert (
+        "763104351884.dkr.ecr.us-west-2.amazonaws.com/"
+        "huggingface-pytorch-inference:1.13.1-transformers4.26.0-gpu-py39-cu117-ubuntu20.04"
+        == image_uris.retrieve(
+            framework=None,
+            region="us-west-2",
+            image_scope="inference",
+            model_id="variant-model",
+            model_version="*",
+            instance_type="ml.p2.xlarge",
+        )
+    )
+
+    assert "867930986793.dkr.us-west-2.amazonaws.com/cpu-blah" == image_uris.retrieve(
+        framework=None,
+        region="us-west-2",
+        image_scope="inference",
+        model_id="variant-model",
+        model_version="*",
+        instance_type="ml.c2.xlarge",
+    )
+
+    with pytest.raises(ValueError):
+        image_uris.retrieve(
+            framework=None,
+            region="us-west-29",
+            image_scope="inference",
+            model_id="variant-model",
+            model_version="*",
+            instance_type="ml.c2.xlarge",
+        )
+
+    with pytest.raises(ValueError):
+        image_uris.retrieve(
+            framework=None,
+            region="us-west-2",
+            image_scope="inference",
+            model_id="variant-model",
+            model_version="*",
+            instance_type="ml.c200000.xlarge",
+        )
+
+    assert (
+        "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.5.0-gpu-py3"
+        == image_uris.retrieve(
+            framework=None,
+            region="us-west-2",
+            image_scope="training",
+            model_id="variant-model",
+            model_version="*",
+            instance_type="ml.g4dn.2xlarge",
+        )
+    )