feat: jumpstart model artifact instance type variants

Author: evakravi

src/sagemaker/jumpstart/artifacts/model_uris.py

@@ -34,6 +34,7 @@ def _retrieve_model_uri(
     model_id: str,
     model_version: str,
     model_scope: Optional[str] = None,
+    instance_type: Optional[str] = None,
     region: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
@@ -50,6 +51,8 @@ def _retrieve_model_uri(
             artifact S3 URI.
         model_scope (str): The model type, i.e. what it is used for.
             Valid values: "training" and "inference".
+        instance_type (str): An instance type to optionally supply in order to get
+            model artifacts specific for the instance type. (Default: None).
         region (str): Region for which to retrieve model S3 URI. (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
@@ -84,14 +87,55 @@ def _retrieve_model_uri(
         sagemaker_session=sagemaker_session,
     )
 
+    model_artifact_key: str
+
     if model_scope == JumpStartScriptScope.INFERENCE:
+        instance_specific_prepacked_hosting_artifact_key: Optional[str] = (
+            model_specs.hosting_instance_type_variants.get_instance_specific_prepacked_artifact_key(
+                instance_type=instance_type
+            )
+            if instance_type
+            and getattr(model_specs, "hosting_instance_type_variants", None) is not None
+            else None
+        )
+
+        instance_specific_hosting_artifact_key: Optional[str] = (
+            model_specs.hosting_instance_type_variants.get_instance_specific_artifact_key(
+                instance_type=instance_type
+            )
+            if instance_type
+            and getattr(model_specs, "hosting_instance_type_variants", None) is not None
+            else None
+        )
+
+        default_prepacked_hosting_artifact_key: Optional[str] = getattr(
+            model_specs, "hosting_prepacked_artifact_key"
+        )
+
+        default_hosting_artifact_key: str = model_specs.hosting_artifact_key
+
         model_artifact_key = (
-            getattr(model_specs, "hosting_prepacked_artifact_key", None)
-            or model_specs.hosting_artifact_key
+            instance_specific_prepacked_hosting_artifact_key
+            or instance_specific_hosting_artifact_key
+            or default_prepacked_hosting_artifact_key
+            or default_hosting_artifact_key
         )
 
     elif model_scope == JumpStartScriptScope.TRAINING:
-        model_artifact_key = model_specs.training_artifact_key
+        instance_specific_training_artifact_key: Optional[str] = (
+            model_specs.training_instance_type_variants.get_instance_specific_artifact_key(
+                instance_type=instance_type
+            )
+            if instance_type
+            and getattr(model_specs, "training_instance_type_variants", None) is not None
+            else None
+        )
+
+        default_training_artifact_key: str = model_specs.training_artifact_key
+
+        model_artifact_key = (
+            instance_specific_training_artifact_key or default_training_artifact_key
+        )
 
     bucket = os.environ.get(
         ENV_VARIABLE_JUMPSTART_MODEL_ARTIFACT_BUCKET_OVERRIDE

src/sagemaker/jumpstart/factory/estimator.py

@@ -473,6 +473,7 @@ def _add_model_uri_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartE
             tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
             tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
             sagemaker_session=kwargs.sagemaker_session,
+            instance_type=kwargs.instance_type,
         )
 
         if (

src/sagemaker/jumpstart/factory/model.py

@@ -215,6 +215,7 @@ def _add_model_data_to_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpStartMode
         tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
         tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
         sagemaker_session=kwargs.sagemaker_session,
+        instance_type=kwargs.instance_type,
     )
 
     if isinstance(model_data, str) and model_data.startswith("s3://") and model_data.endswith("/"):

src/sagemaker/jumpstart/types.py

@@ -346,6 +346,62 @@ def to_json(self) -> Dict[str, Any]:
         json_obj = {att: getattr(self, att) for att in self.__slots__ if hasattr(self, att)}
         return json_obj
 
+    def get_instance_specific_prepacked_artifact_key(self, instance_type: str) -> Optional[str]:
+        """Returns instance specific model artifact key.
+
+        Returns None if a model, instance type tuple does not have specific
+        artifact key.
+        """
+
+        return self._get_instance_specific_property(
+            instance_type=instance_type, property_name="prepacked_artifact_key"
+        )
+
+    def get_instance_specific_artifact_key(self, instance_type: str) -> Optional[str]:
+        """Returns instance specific model artifact key.
+
+        Returns None if a model, instance type tuple does not have specific
+        artifact key.
+        """
+
+        return self._get_instance_specific_property(
+            instance_type=instance_type, property_name="artifact_key"
+        )
+
+    def _get_instance_specific_property(
+        self, instance_type: str, property_name: str
+    ) -> Optional[str]:
+        """Returns instance specific property.
+
+        If a value exists for both the instance family and instance type,
+        the instance type value is chosen.
+
+        Returns None if a (model, instance type, property name) tuple does not have
+        specific prepacked artifact key.
+        """
+
+        if self.variants is None:
+            return None
+
+        instance_specific_property: Optional[str] = (
+            self.variants.get(instance_type, {}).get("properties", {}).get(property_name, None)
+        )
+
+        if instance_specific_property:
+            return instance_specific_property
+
+        instance_type_family = get_instance_type_family(instance_type)
+
+        instance_family_property: Optional[str] = (
+            self.variants.get(instance_type_family, {})
+            .get("properties", {})
+            .get(property_name, None)
+            if instance_type_family not in {"", None}
+            else None
+        )
+
+        return instance_family_property
+
     def get_instance_specific_environment_variables(self, instance_type: str) -> Dict[str, str]:
         """Returns instance specific environment variables.
 

src/sagemaker/model_uris.py

@@ -30,6 +30,7 @@ def retrieve(
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
     model_scope: Optional[str] = None,
+    instance_type: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
@@ -44,6 +45,8 @@ def retrieve(
             the model artifact S3 URI.
         model_scope (str): The model type.
             Valid values: "training" and "inference".
+        instance_type (str): An instance type to optionally supply in order to get
+            model artifacts specific for the instance type. (Default: None).
         tolerate_vulnerable_model (bool): ``True`` if vulnerable versions of model
             specifications should be tolerated without raising an exception. If ``False``, raises an
             exception if the script used by this version of the model has dependencies with known
@@ -71,11 +74,12 @@ def retrieve(
         )
 
     return artifacts._retrieve_model_uri(
-        model_id,
-        model_version,  # type: ignore
-        model_scope,
-        region,
-        tolerate_vulnerable_model,
-        tolerate_deprecated_model,
+        model_id=model_id,
+        model_version=model_version,  # type: ignore
+        model_scope=model_scope,
+        instance_type=instance_type,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
     )