Merge branch 'aws:master' into master

Author: cj-zhang

src/sagemaker/fw_utils.py

@@ -10,37 +10,37 @@
 # 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.
-"""Utility methods used by framework classes"""
+"""Utility methods used by framework classes."""
 from __future__ import absolute_import
 
 import json
 import logging
 import os
 import re
-import time
 import shutil
 import tempfile
+import time
 from collections import namedtuple
-from typing import List, Optional, Union, Dict
+from typing import Dict, List, Optional, Union
+
 from packaging import version
 
 import sagemaker.image_uris
+import sagemaker.utils
+from sagemaker.deprecations import deprecation_warn_base, renamed_kwargs, renamed_warning
 from sagemaker.instance_group import InstanceGroup
 from sagemaker.s3_utils import s3_path_join
 from sagemaker.session_settings import SessionSettings
-import sagemaker.utils
 from sagemaker.workflow import is_pipeline_variable
-
-from sagemaker.deprecations import renamed_warning, renamed_kwargs
 from sagemaker.workflow.entities import PipelineVariable
-from sagemaker.deprecations import deprecation_warn_base
 
 logger = logging.getLogger(__name__)
 
 _TAR_SOURCE_FILENAME = "source.tar.gz"
 
 UploadedCode = namedtuple("UploadedCode", ["s3_prefix", "script_name"])
 """sagemaker.fw_utils.UploadedCode: An object containing the S3 prefix and script name.
+
 This is for the source code used for the entry point with an ``Estimator``. It can be
 instantiated with positional or keyword arguments.
 """
@@ -211,7 +211,7 @@ def validate_source_code_input_against_pipeline_variables(
     git_config: Optional[Dict[str, str]] = None,
     enable_network_isolation: Union[bool, PipelineVariable] = False,
 ):
-    """Validate source code input against pipeline variables
+    """Validate source code input against pipeline variables.
 
     Args:
         entry_point (str or PipelineVariable): The path to the local Python source file that
@@ -481,7 +481,7 @@ def tar_and_upload_dir(
 
 
 def _list_files_to_compress(script, directory):
-    """Placeholder docstring"""
+    """Placeholder docstring."""
     if directory is None:
         return [script]
 
@@ -585,7 +585,6 @@ def model_code_key_prefix(code_location_key_prefix, model_name, image):
     The location returned is a potential concatenation of 2 parts
         1. code_location_key_prefix if it exists
         2. model_name or a name derived from the image
-
     Args:
         code_location_key_prefix (str): the s3 key prefix from code_location
         model_name (str): the name of the model
@@ -620,8 +619,6 @@ def warn_if_parameter_server_with_multi_gpu(training_instance_type, distribution
                         "enabled": True
                     }
                 }
-
-
     """
     if training_instance_type == "local" or distribution is None:
         return
@@ -646,7 +643,7 @@ def warn_if_parameter_server_with_multi_gpu(training_instance_type, distribution
 def profiler_config_deprecation_warning(
     profiler_config, image_uri, framework_name, framework_version
 ):
-    """Put out a deprecation message for if framework profiling is specified TF >= 2.12 and PT >= 2.0"""
+    """Deprecation message if framework profiling is specified TF >= 2.12 and PT >= 2.0."""
     if profiler_config is None or profiler_config.framework_profile_params is None:
         return
 
@@ -692,6 +689,7 @@ def validate_smdistributed(
         framework_name (str): A string representing the name of framework selected.
         framework_version (str): A string representing the framework version selected.
         py_version (str): A string representing the python version selected.
+        Ex: `py38, py39, py310, py311`
         distribution (dict): A dictionary with information to enable distributed training.
             (Defaults to None if distributed training is not enabled.) For example:
 
@@ -763,7 +761,8 @@ def _validate_smdataparallel_args(
         instance_type (str): A string representing the type of training instance selected. Ex: `ml.p3.16xlarge`
         framework_name (str): A string representing the name of framework selected. Ex: `tensorflow`
         framework_version (str): A string representing the framework version selected. Ex: `2.3.1`
-        py_version (str): A string representing the python version selected. Ex: `py3`
+        py_version (str): A string representing the python version selected.
+        Ex: `py38, py39, py310, py311`
         distribution (dict): A dictionary with information to enable distributed training.
             (Defaults to None if distributed training is not enabled.) Ex:
 
@@ -847,6 +846,7 @@ def validate_distribution(
         framework_name (str): A string representing the name of framework selected.
         framework_version (str): A string representing the framework version selected.
         py_version (str): A string representing the python version selected.
+        Ex: `py38, py39, py310, py311`
         image_uri (str): A string representing a Docker image URI.
         kwargs(dict): Additional kwargs passed to this function
 
@@ -953,7 +953,7 @@ def validate_distribution(
 
 
 def validate_distribution_for_instance_type(instance_type, distribution):
-    """Check if the provided distribution strategy is supported for the instance_type
+    """Check if the provided distribution strategy is supported for the instance_type.
 
     Args:
         instance_type (str): A string representing the type of training instance selected.
@@ -1010,6 +1010,7 @@ def validate_torch_distributed_distribution(
                 }
         framework_version (str): A string representing the framework version selected.
         py_version (str): A string representing the python version selected.
+        Ex: `py38, py39, py310, py311`
         image_uri (str): A string representing a Docker image URI.
         entry_point (str or PipelineVariable): The absolute or relative path to the local Python
             source file that should be executed as the entry point to
@@ -1072,7 +1073,7 @@ def validate_torch_distributed_distribution(
 
 
 def _is_gpu_instance(instance_type):
-    """Returns bool indicating whether instance_type supports GPU
+    """Returns bool indicating whether instance_type supports GPU.
 
     Args:
         instance_type (str): Name of the instance_type to check against.
@@ -1091,7 +1092,7 @@ def _is_gpu_instance(instance_type):
 
 
 def _is_trainium_instance(instance_type):
-    """Returns bool indicating whether instance_type is a Trainium instance
+    """Returns bool indicating whether instance_type is a Trainium instance.
 
     Args:
         instance_type (str): Name of the instance_type to check against.
@@ -1107,7 +1108,7 @@ def _is_trainium_instance(instance_type):
 
 
 def python_deprecation_warning(framework, latest_supported_version):
-    """Placeholder docstring"""
+    """Placeholder docstring."""
     return PYTHON_2_DEPRECATION_WARNING.format(
         framework=framework, latest_supported_version=latest_supported_version
     )
@@ -1121,7 +1122,6 @@ def _region_supports_debugger(region_name):
 
     Returns:
         bool: Whether or not the region supports Amazon SageMaker Debugger.
-
     """
     return region_name.lower() not in DEBUGGER_UNSUPPORTED_REGIONS
 
@@ -1134,7 +1134,6 @@ def _region_supports_profiler(region_name):
 
     Returns:
         bool: Whether or not the region supports Amazon SageMaker Debugger profiling feature.
-
     """
     return region_name.lower() not in PROFILER_UNSUPPORTED_REGIONS
 
@@ -1162,7 +1161,8 @@ def validate_version_or_image_args(framework_version, py_version, image_uri):
 
     Args:
         framework_version (str): The version of the framework.
-        py_version (str): The version of Python.
+        py_version (str): A string representing the python version selected.
+        Ex: `py38, py39, py310, py311`
         image_uri (str): The URI of the image.
 
     Raises:
@@ -1194,9 +1194,8 @@ def create_image_uri(
         instance_type (str): SageMaker instance type. Used to determine device
             type (cpu/gpu/family-specific optimized).
         framework_version (str): The version of the framework.
-        py_version (str): Optional. Python version. If specified, should be one
-            of 'py2' or 'py3'. If not specified, image uri will not include a
-            python component.
+        py_version (str): Optional. Python version Ex: `py38, py39, py310, py311`.
+            If not specified, image uri will not include a python component.
         account (str): AWS account that contains the image. (default:
             '520713654638')
         accelerator_type (str): SageMaker Elastic Inference accelerator type.

src/sagemaker/huggingface/estimator.py

@@ -15,17 +15,13 @@
 
 import logging
 import re
-from typing import Optional, Union, Dict
+from typing import Dict, Optional, Union
 
-from sagemaker.estimator import Framework, EstimatorBase
-from sagemaker.fw_utils import (
-    framework_name_from_image,
-    validate_distribution,
-)
+from sagemaker.estimator import EstimatorBase, Framework
+from sagemaker.fw_utils import framework_name_from_image, validate_distribution
 from sagemaker.huggingface.model import HuggingFaceModel
-from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
-
 from sagemaker.huggingface.training_compiler.config import TrainingCompilerConfig
+from sagemaker.vpc_utils import VPC_CONFIG_DEFAULT
 from sagemaker.workflow.entities import PipelineVariable
 
 logger = logging.getLogger("sagemaker")
@@ -66,7 +62,7 @@ def __init__(
         Args:
             py_version (str): Python version you want to use for executing your model training
                 code. Defaults to ``None``. Required unless ``image_uri`` is provided.  If
-                using PyTorch, the current supported version is ``py36``. If using TensorFlow,
+                using PyTorch, the current supported version is ``py39``. If using TensorFlow,
                 the current supported version is ``py37``.
             entry_point (str or PipelineVariable): Path (absolute or relative) to the Python source
                 file which should be executed as the entry point to training.

src/sagemaker/image_uri_config/pytorch.json

@@ -1479,6 +1479,7 @@
                     "ap-southeast-3": "907027046896",
                     "ap-southeast-4": "457447274322",
                     "ap-southeast-5": "550225433462",
+                    "ap-southeast-7": "590183813437",
                     "ca-central-1": "763104351884",
                     "ca-west-1": "204538143572",
                     "cn-north-1": "727897471807",
@@ -1494,6 +1495,7 @@
                     "il-central-1": "780543022126",
                     "me-central-1": "914824155844",
                     "me-south-1": "217643126080",
+                    "mx-central-1": "637423239942",
                     "sa-east-1": "763104351884",
                     "us-east-1": "763104351884",
                     "us-east-2": "763104351884",

src/sagemaker/jumpstart/accessors.py

@@ -288,6 +288,7 @@ def get_model_specs(
         )
         JumpStartModelsAccessor._set_cache_and_region(region, cache_kwargs)
 
+        # Users only input model id, not contentType, so first try to describe with ModelReference, then with Model
         if hub_arn:
             try:
                 hub_model_arn = construct_hub_model_reference_arn_from_inputs(
@@ -308,11 +309,22 @@ def get_model_specs(
                 hub_model_arn = construct_hub_model_arn_from_inputs(
                     hub_arn=hub_arn, model_name=model_id, version=version
                 )
-                model_specs = JumpStartModelsAccessor._cache.get_hub_model(
-                    hub_model_arn=hub_model_arn
-                )
-                model_specs.set_hub_content_type(HubContentType.MODEL)
-                return model_specs
+
+                # Failed to describe ModelReference, try with Model
+                try:
+                    model_specs = JumpStartModelsAccessor._cache.get_hub_model(
+                        hub_model_arn=hub_model_arn
+                    )
+                    model_specs.set_hub_content_type(HubContentType.MODEL)
+
+                    return model_specs
+                except Exception as ex:
+                    # Failed with both, throw a custom error message
+                    raise RuntimeError(
+                        f"Cannot get details for {model_id} in Hub {hub_arn}. \
+                            {model_id} does not exist as a Model or ModelReference: \n"
+                        + str(ex)
+                    )
 
         return JumpStartModelsAccessor._cache.get_specs(  # type: ignore
             model_id=model_id, version_str=version, model_type=model_type

src/sagemaker/jumpstart/hub/hub.py

@@ -272,18 +272,21 @@ def delete_model_reference(self, model_name: str) -> None:
     def describe_model(
         self, model_name: str, hub_name: Optional[str] = None, model_version: Optional[str] = None
     ) -> DescribeHubContentResponse:
-        """Describe model in the SageMaker Hub."""
+        """Describe Model or ModelReference in a Hub."""
+        hub_name = hub_name or self.hub_name
+
+        # Users only input model id, not contentType, so first try to describe with ModelReference, then with Model
         try:
             model_version = get_hub_model_version(
                 hub_model_name=model_name,
                 hub_model_type=HubContentType.MODEL_REFERENCE.value,
-                hub_name=self.hub_name if not hub_name else hub_name,
+                hub_name=hub_name,
                 sagemaker_session=self._sagemaker_session,
                 hub_model_version=model_version,
             )
 
             hub_content_description: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
-                hub_name=self.hub_name if not hub_name else hub_name,
+                hub_name=hub_name,
                 hub_content_name=model_name,
                 hub_content_version=model_version,
                 hub_content_type=HubContentType.MODEL_REFERENCE.value,
@@ -294,19 +297,32 @@ def describe_model(
                 "Received exeption while calling APIs for ContentType ModelReference, retrying with ContentType Model: "
                 + str(ex)
             )
-            model_version = get_hub_model_version(
-                hub_model_name=model_name,
-                hub_model_type=HubContentType.MODEL.value,
-                hub_name=self.hub_name if not hub_name else hub_name,
-                sagemaker_session=self._sagemaker_session,
-                hub_model_version=model_version,
-            )
 
-            hub_content_description: Dict[str, Any] = self._sagemaker_session.describe_hub_content(
-                hub_name=self.hub_name if not hub_name else hub_name,
-                hub_content_name=model_name,
-                hub_content_version=model_version,
-                hub_content_type=HubContentType.MODEL.value,
-            )
+            # Failed to describe ModelReference, try with Model
+            try:
+                model_version = get_hub_model_version(
+                    hub_model_name=model_name,
+                    hub_model_type=HubContentType.MODEL.value,
+                    hub_name=hub_name,
+                    sagemaker_session=self._sagemaker_session,
+                    hub_model_version=model_version,
+                )
+
+                hub_content_description: Dict[str, Any] = (
+                    self._sagemaker_session.describe_hub_content(
+                        hub_name=hub_name,
+                        hub_content_name=model_name,
+                        hub_content_version=model_version,
+                        hub_content_type=HubContentType.MODEL.value,
+                    )
+                )
+
+            except Exception as ex:
+                # Failed with both, throw a custom error message
+                raise RuntimeError(
+                    f"Cannot get details for {model_name} in Hub {hub_name}. \
+                        {model_name} does not exist as a Model or ModelReference in {hub_name}: \n"
+                    + str(ex)
+                )
 
         return DescribeHubContentResponse(hub_content_description)

src/sagemaker/jumpstart/types.py

@@ -1363,9 +1363,10 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.deploy_kwargs = deepcopy(json_obj.get("deploy_kwargs", {}))
         self.predictor_specs: Optional[JumpStartPredictorSpecs] = (
             JumpStartPredictorSpecs(
-                json_obj["predictor_specs"], is_hub_content=self._is_hub_content
+                json_obj.get("predictor_specs"),
+                is_hub_content=self._is_hub_content,
             )
-            if "predictor_specs" in json_obj
+            if json_obj.get("predictor_specs")
             else None
         )
         self.default_payloads: Optional[Dict[str, JumpStartSerializablePayload]] = (
@@ -1501,6 +1502,9 @@ class JumpStartConfigComponent(JumpStartMetadataBaseFields):
         "incremental_training_supported",
     ]
 
+    # Map of HubContent fields that map to custom names in MetadataBaseFields
+    CUSTOM_FIELD_MAP = {"sage_maker_sdk_predictor_specifications": "predictor_specs"}
+
     __slots__ = slots + JumpStartMetadataBaseFields.__slots__
 
     def __init__(
@@ -1532,6 +1536,11 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
             if field in self.__slots__:
                 setattr(self, field, json_obj[field])
 
+        # Handle custom fields
+        for custom_field, field in self.CUSTOM_FIELD_MAP.items():
+            if custom_field in json_obj:
+                setattr(self, field, json_obj.get(custom_field))
+
 
 class JumpStartMetadataConfig(JumpStartDataHolderType):
     """Data class of JumpStart metadata config."""