change: improve jumpstart retrieve uri unit tests, fix logic for image uris

Author: evakravi

src/sagemaker/estimator.py

@@ -2425,51 +2425,32 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
 
         return init_params
 
-    def training_image_uri(self):
+    def training_image_uri(self, region=None):
         """Return the Docker image to use for training.
 
         The :meth:`~sagemaker.estimator.EstimatorBase.fit` method, which does
         the model training, calls this method to find the image to use for model
         training.
 
+        Args:
+            region: Region to use for image uri.
+                Default: Region associated with SageMaker session.
+
         Returns:
             str: The URI of the Docker image.
         """
-        if self.image_uri:
-            return self.image_uri
-        if hasattr(self, "distribution"):
-            distribution = self.distribution  # pylint: disable=no-member
-        else:
-            distribution = None
-        compiler_config = getattr(self, "compiler_config", None)
-
-        if hasattr(self, "tensorflow_version") or hasattr(self, "pytorch_version"):
-            processor = image_uris._processor(self.instance_type, ["cpu", "gpu"])
-            is_native_huggingface_gpu = processor == "gpu" and not compiler_config
-            container_version = "cu110-ubuntu18.04" if is_native_huggingface_gpu else None
-            if self.tensorflow_version is not None:  # pylint: disable=no-member
-                base_framework_version = (
-                    f"tensorflow{self.tensorflow_version}"  # pylint: disable=no-member
-                )
-            else:
-                base_framework_version = (
-                    f"pytorch{self.pytorch_version}"  # pylint: disable=no-member
-                )
-        else:
-            container_version = None
-            base_framework_version = None
 
-        return image_uris.retrieve(
-            self._framework_name,
-            self.sagemaker_session.boto_region_name,
-            instance_type=self.instance_type,
-            version=self.framework_version,  # pylint: disable=no-member
+        return image_uris.get_training_image_uri(
+            region=region or self.sagemaker_session.boto_region_name,
+            framework=self._framework_name,
+            framework_version=self.framework_version,  # pylint: disable=no-member
             py_version=self.py_version,  # pylint: disable=no-member
-            image_scope="training",
-            distribution=distribution,
-            base_framework_version=base_framework_version,
-            container_version=container_version,
-            training_compiler_config=compiler_config,
+            image_uri=self.image_uri,
+            distribution=getattr(self, "distribution", None),
+            compiler_config=getattr(self, "compiler_config", None),
+            tensorflow_version=getattr(self, "tensorflow_version", None),
+            pytorch_version=getattr(self, "pytorch_version", None),
+            instance_type=self.instance_type,
         )
 
     @classmethod

src/sagemaker/image_uris.py

@@ -17,6 +17,7 @@
 import logging
 import os
 import re
+from typing import Optional
 
 from sagemaker import utils
 from sagemaker.jumpstart.utils import is_jumpstart_model_input
@@ -373,3 +374,68 @@ def _validate_arg(arg, available_options, arg_name):
 def _format_tag(tag_prefix, processor, py_version, container_version):
     """Creates a tag for the image URI."""
     return "-".join(x for x in (tag_prefix, processor, py_version, container_version) if x)
+
+
+def get_training_image_uri(
+    region,
+    framework,
+    framework_version=None,
+    py_version=None,
+    image_uri=None,
+    distribution=None,
+    compiler_config=None,
+    tensorflow_version=None,
+    pytorch_version=None,
+    instance_type=None,
+) -> str:
+    """Retrieve image uri for training.
+
+    Args:
+        region (str): AWS region to use for image URI.
+        framework (str): The framework for which to retrieve an image URI.
+        framework_version (str): The framework version for which to retrieve an
+            image URI (default: None).
+        py_version (str): The python version to use for the image (default: None).
+        image_uri (str): If an image URI is supplied, it will be returned (default: None).
+        distribution (dict): A dictionary with information on how to run distributed
+            training (default: None).
+        compiler_config (:class:`~sagemaker.training_compiler.TrainingCompilerConfig`):
+            A configuration class for the SageMaker Training Compiler
+            (default: None).
+        tensorflow_version (str): Version of tensorflow to use. (default: None)
+        pytorch_version (str): Version of pytorch to use. (default: None)
+        instance_type (str): Instance type fo use. (default: None)
+
+    Returns:
+        str: the image URI string.
+    """
+
+    if image_uri:
+        return image_uri
+
+    base_framework_version: Optional[str] = None
+
+    if tensorflow_version is not None or pytorch_version is not None:
+        processor = _processor(instance_type, ["cpu", "gpu"])
+        is_native_huggingface_gpu = processor == "gpu" and not compiler_config
+        container_version = "cu110-ubuntu18.04" if is_native_huggingface_gpu else None
+        if tensorflow_version is not None:
+            base_framework_version = f"tensorflow{tensorflow_version}"
+        else:
+            base_framework_version = f"pytorch{pytorch_version}"
+    else:
+        container_version = None
+        base_framework_version = None
+
+    return retrieve(
+        framework,
+        region,
+        instance_type=instance_type,
+        version=framework_version,
+        py_version=py_version,
+        image_scope="training",
+        distribution=distribution,
+        base_framework_version=base_framework_version,
+        container_version=container_version,
+        training_compiler_config=compiler_config,
+    )

src/sagemaker/jumpstart/artifacts.py

@@ -19,6 +19,7 @@
     INFERENCE,
     TRAINING,
     SUPPORTED_JUMPSTART_SCOPES,
+    ModelFramework,
 )
 from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
 from sagemaker.jumpstart import accessors as jumpstart_accessors
@@ -115,20 +116,39 @@ def _retrieve_image_uri(
             f"Bad value for container python version for JumpStart model: '{py_version}'."
         )
 
-    if framework == "huggingface":
-        base_framework_version = ecr_specs.framework_version
+    base_framework_version_override = None
+    version_override = None
+    if ecr_specs.framework == ModelFramework.HUGGINGFACE.value:
+        base_framework_version_override = ecr_specs.framework_version
+        version_override = ecr_specs.huggingface_transformers_version
+
+    if image_scope == TRAINING:
+        return image_uris.get_training_image_uri(
+            region=region,
+            framework=ecr_specs.framework,
+            framework_version=version_override or ecr_specs.framework_version,
+            py_version=ecr_specs.py_version,
+            image_uri=None,
+            distribution=None,
+            compiler_config=None,
+            tensorflow_version=None,
+            pytorch_version=base_framework_version_override or base_framework_version,
+            instance_type=instance_type,
+        )
+    if base_framework_version_override is not None:
+        base_framework_version_override = f"pytorch{base_framework_version_override}"
 
     return image_uris.retrieve(
         framework=ecr_specs.framework,
         region=region,
-        version=ecr_specs.framework_version,
+        version=version_override or ecr_specs.framework_version,
         py_version=ecr_specs.py_version,
         instance_type=instance_type,
         accelerator_type=accelerator_type,
         image_scope=image_scope,
         container_version=container_version,
         distribution=distribution,
-        base_framework_version=base_framework_version,
+        base_framework_version=base_framework_version_override or base_framework_version,
         training_compiler_config=training_compiler_config,
     )
 

src/sagemaker/jumpstart/constants.py

@@ -13,6 +13,7 @@
 """This module stores constants related to SageMaker JumpStart."""
 from __future__ import absolute_import
 from typing import Set
+from enum import Enum
 import boto3
 from sagemaker.jumpstart.types import JumpStartLaunchedRegionInfo
 
@@ -118,3 +119,20 @@
 INFERENCE = "inference"
 TRAINING = "training"
 SUPPORTED_JUMPSTART_SCOPES = set([INFERENCE, TRAINING])
+
+
+class ModelFramework(str, Enum):
+    """Enum class for JumpStart model framework.
+
+    The ML framework as referenced in the prefix of the model ID.
+    This value does not necessarily correspond to the container name.
+    """
+
+    PYTORCH = "pytorch"
+    TENSORFLOW = "tensorflow"
+    MXNET = "mxnet"
+    HUGGINGFACE = "huggingface"
+    LIGHTGBM = "lightgbm"
+    CATBOOST = "catboost"
+    XGBOOST = "xgboost"
+    SKLEARN = "sklearn"

src/sagemaker/jumpstart/types.py

@@ -41,8 +41,13 @@ def __eq__(self, other: Any) -> bool:
         if self.__slots__ != other.__slots__:
             return False
         for attribute in self.__slots__:
-            if getattr(self, attribute) != getattr(other, attribute):
+            if (hasattr(self, attribute) and not hasattr(other, attribute)) or (
+                hasattr(other, attribute) and not hasattr(self, attribute)
+            ):
                 return False
+            if hasattr(self, attribute) and hasattr(other, attribute):
+                if getattr(self, attribute) != getattr(other, attribute):
+                    return False
         return True
 
     def __hash__(self) -> int:
@@ -112,7 +117,7 @@ def __init__(self, header: Dict[str, str]):
 
     def to_json(self) -> Dict[str, str]:
         """Returns json representation of JumpStartModelHeader object."""
-        json_obj = {att: getattr(self, att) for att in self.__slots__}
+        json_obj = {att: getattr(self, att) for att in self.__slots__ if hasattr(self, att)}
         return json_obj
 
     def from_json(self, json_obj: Dict[str, str]) -> None:
@@ -134,6 +139,7 @@ class JumpStartECRSpecs(JumpStartDataHolderType):
         "framework",
         "framework_version",
         "py_version",
+        "huggingface_transformers_version",
     }
 
     def __init__(self, spec: Dict[str, Any]):
@@ -154,10 +160,13 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.framework = json_obj["framework"]
         self.framework_version = json_obj["framework_version"]
         self.py_version = json_obj["py_version"]
+        huggingface_transformers_version = json_obj.get("huggingface_transformers_version")
+        if huggingface_transformers_version is not None:
+            self.huggingface_transformers_version = huggingface_transformers_version
 
     def to_json(self) -> Dict[str, Any]:
         """Returns json representation of JumpStartECRSpecs object."""
-        json_obj = {att: getattr(self, att) for att in self.__slots__}
+        json_obj = {att: getattr(self, att) for att in self.__slots__ if hasattr(self, att)}
         return json_obj
 
 
@@ -202,26 +211,23 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.hosting_script_key: str = json_obj["hosting_script_key"]
         self.training_supported: bool = bool(json_obj["training_supported"])
         if self.training_supported:
-            self.training_ecr_specs: Optional[JumpStartECRSpecs] = JumpStartECRSpecs(
+            self.training_ecr_specs: JumpStartECRSpecs = JumpStartECRSpecs(
                 json_obj["training_ecr_specs"]
             )
-            self.training_artifact_key: Optional[str] = json_obj["training_artifact_key"]
-            self.training_script_key: Optional[str] = json_obj["training_script_key"]
-            self.hyperparameters: Optional[Dict[str, Any]] = json_obj.get("hyperparameters")
-        else:
-            self.training_ecr_specs = (
-                self.training_artifact_key
-            ) = self.training_script_key = self.hyperparameters = None
+            self.training_artifact_key: str = json_obj["training_artifact_key"]
+            self.training_script_key: str = json_obj["training_script_key"]
+            self.hyperparameters: Dict[str, Any] = json_obj.get("hyperparameters", {})
 
     def to_json(self) -> Dict[str, Any]:
         """Returns json representation of JumpStartModelSpecs object."""
         json_obj = {}
         for att in self.__slots__:
-            cur_val = getattr(self, att)
-            if isinstance(cur_val, JumpStartECRSpecs):
-                json_obj[att] = cur_val.to_json()
-            else:
-                json_obj[att] = cur_val
+            if hasattr(self, att):
+                cur_val = getattr(self, att)
+                if isinstance(cur_val, JumpStartECRSpecs):
+                    json_obj[att] = cur_val.to_json()
+                else:
+                    json_obj[att] = cur_val
         return json_obj
 
 

src/sagemaker/script_uris.py

@@ -43,9 +43,7 @@ def retrieve(
         ValueError: If the combination of arguments specified is not supported.
     """
     if not jumpstart_utils.is_jumpstart_model_input(model_id, model_version):
-        raise ValueError(
-            "Must specify `model_id` and `model_version` when retrieving script URIs."
-        )
+        raise ValueError("Must specify `model_id` and `model_version` when retrieving script URIs.")
 
     assert model_id is not None
     assert model_version is not None

src/sagemaker/sklearn/estimator.py

@@ -43,6 +43,7 @@ def __init__(
         source_dir=None,
         hyperparameters=None,
         image_uri=None,
+        image_uri_region=None,
         **kwargs
     ):
         """Creates a SKLearn Estimator for Scikit-learn environment.
@@ -99,6 +100,9 @@ def __init__(
                 If ``framework_version`` or ``py_version`` are ``None``, then
                 ``image_uri`` is required. If also ``None``, then a ``ValueError``
                 will be raised.
+            image_uri_region (str): If ``image_uri` argument is None, the image uri
+                associated with this object will be in this region.
+                    Default: region associated with SageMaker session.
             **kwargs: Additional kwargs passed to the
                 :class:`~sagemaker.estimator.Framework` constructor.
 
@@ -144,7 +148,7 @@ def __init__(
         if image_uri is None:
             self.image_uri = image_uris.retrieve(
                 SKLearn._framework_name,
-                self.sagemaker_session.boto_region_name,
+                image_uri_region or self.sagemaker_session.boto_region_name,
                 version=self.framework_version,
                 py_version=self.py_version,
                 instance_type=instance_type,

src/sagemaker/tensorflow/model.py

@@ -352,14 +352,14 @@ def _get_container_env(self):
         env[self.LOG_LEVEL_PARAM_NAME] = self.LOG_LEVEL_MAP[self._container_log_level]
         return env
 
-    def _get_image_uri(self, instance_type, accelerator_type=None):
+    def _get_image_uri(self, instance_type, accelerator_type=None, region_name=None):
         """Placeholder docstring."""
         if self.image_uri:
             return self.image_uri
 
         return image_uris.retrieve(
             self._framework_name,
-            self.sagemaker_session.boto_region_name,
+            region_name or self.sagemaker_session.boto_region_name,
             version=self.framework_version,
             instance_type=instance_type,
             accelerator_type=accelerator_type,
@@ -383,4 +383,6 @@ def serving_image_uri(
             str: The appropriate image URI based on the given parameters.
 
         """
-        return self._get_image_uri(instance_type=instance_type, accelerator_type=accelerator_type)
+        return self._get_image_uri(
+            instance_type=instance_type, accelerator_type=accelerator_type, region_name=region_name
+        )

src/sagemaker/xgboost/estimator.py

@@ -48,6 +48,7 @@ def __init__(
         hyperparameters=None,
         py_version="py3",
         image_uri=None,
+        image_uri_region=None,
         **kwargs
     ):
         """An estimator that executes an XGBoost-based SageMaker Training Job.
@@ -89,6 +90,9 @@ def __init__(
                 Examples:
                     123.dkr.ecr.us-west-2.amazonaws.com/my-custom-image:1.0
                     custom-image:latest.
+            image_uri_region (str): If ``image_uri` argument is None, the image uri
+                associated with this object will be in this region.
+                    Default: region associated with SageMaker session.
             **kwargs: Additional kwargs passed to the
                 :class:`~sagemaker.estimator.Framework` constructor.
 
@@ -114,7 +118,7 @@ def __init__(
         if image_uri is None:
             self.image_uri = image_uris.retrieve(
                 self._framework_name,
-                self.sagemaker_session.boto_region_name,
+                image_uri_region or self.sagemaker_session.boto_region_name,
                 version=framework_version,
                 py_version=self.py_version,
                 instance_type=instance_type,