Merge branch 'master' into change/remove-setuptools-deprecation

Author: jjerphan

requirements/extras/test_requirements.txt

@@ -24,7 +24,7 @@ pandas>=1.3.5,<1.5
 scikit-learn==1.3.0
 cloudpickle==2.2.1
 scipy==1.10.1
-urllib3>=1.26.8,<1.26.15
+urllib3>=1.26.8,<3.0.0
 docker>=5.0.2,<7.0.0
 PyYAML==6.0
 pyspark==3.3.1

src/sagemaker/chainer/estimator.py

@@ -108,7 +108,7 @@ def __init__(
             framework_version (str): Chainer version you want to use for
                 executing your model training code. Defaults to ``None``. Required unless
                 ``image_uri`` is provided. List of supported versions:
-                https://github.com/aws/sagemaker-python-sdk#chainer-sagemaker-estimators.
+                https://sagemaker.readthedocs.io/en/stable/frameworks/chainer/using_chainer.html#using-chainer-with-the-sagemaker-python-sdk.
             image_uri (str): If specified, the estimator will use this image
                 for training and hosting, instead of selecting the appropriate
                 SageMaker official image based on framework_version and

src/sagemaker/estimator.py

@@ -101,6 +101,7 @@
 )
 from sagemaker.workflow import is_pipeline_variable
 from sagemaker.workflow.entities import PipelineVariable
+from sagemaker.workflow.parameters import ParameterString
 from sagemaker.workflow.pipeline_context import PipelineSession, runnable_by_pipeline
 
 logger = logging.getLogger(__name__)
@@ -3198,6 +3199,7 @@ class Framework(EstimatorBase):
     """
 
     _framework_name = None
+    UNSUPPORTED_DLC_IMAGE_FOR_SM_PARALLELISM = ("2.0.1-gpu-py310-cu121", "2.0-gpu-py310-cu121")
 
     def __init__(
         self,
@@ -3843,16 +3845,43 @@ def _distribution_configuration(self, distribution):
                 "custom_mpi_options", ""
             )
 
-            if get_mp_parameters(distribution):
-                distribution_config["mp_parameters"] = get_mp_parameters(distribution)
-
-        elif "modelparallel" in distribution.get("smdistributed", {}):
-            raise ValueError("Cannot use Model Parallelism without MPI enabled!")
-
         if "smdistributed" in distribution:
             # smdistributed strategy selected
+            if get_mp_parameters(distribution):
+                distribution_config["mp_parameters"] = get_mp_parameters(distribution)
+            # first make sure torch_distributed is enabled if instance type is p5
+            torch_distributed_enabled = False
+            if "torch_distributed" in distribution:
+                torch_distributed_enabled = distribution.get("torch_distributed").get(
+                    "enabled", False
+                )
             smdistributed = distribution["smdistributed"]
             smdataparallel_enabled = smdistributed.get("dataparallel", {}).get("enabled", False)
+            if isinstance(self.instance_type, ParameterString):
+                p5_enabled = "p5.48xlarge" in self.instance_type.default_value
+            elif isinstance(self.instance_type, str):
+                p5_enabled = "p5.48xlarge" in self.instance_type
+            else:
+                raise ValueError(
+                    "Invalid object type for instance_type argument. Expected "
+                    f"{type(str)} or {type(ParameterString)} but got {type(self.instance_type)}."
+                )
+            img_uri = "" if self.image_uri is None else self.image_uri
+            for unsupported_image in Framework.UNSUPPORTED_DLC_IMAGE_FOR_SM_PARALLELISM:
+                if (
+                    unsupported_image in img_uri and not torch_distributed_enabled
+                ):  # disabling DLC images with CUDA12
+                    raise ValueError(
+                        f"SMDistributed is currently incompatible with DLC image: {img_uri}. "
+                        "(Could be due to CUDA version being greater than 11.)"
+                    )
+            if (
+                not torch_distributed_enabled and p5_enabled
+            ):  # disabling p5 when torch distributed is disabled
+                raise ValueError(
+                    "SMModelParallel and SMDataParallel currently do not support p5 instances."
+                )
+            # smdistributed strategy selected with supported instance type
             distribution_config[self.LAUNCH_SM_DDP_ENV_NAME] = smdataparallel_enabled
             distribution_config[self.INSTANCE_TYPE] = self.instance_type
             if smdataparallel_enabled:

src/sagemaker/huggingface/model.py

@@ -138,12 +138,10 @@ def __init__(
                 unless ``image_uri`` is provided.
             tensorflow_version (str): TensorFlow version you want to use for
                 executing your inference code. Defaults to ``None``. Required unless
-                ``pytorch_version`` is provided. List of supported versions:
-                https://github.com/aws/sagemaker-python-sdk#huggingface-sagemaker-estimators.
+                ``pytorch_version`` is provided. The current supported version is ``2.4.1``.
             pytorch_version (str): PyTorch version you want to use for
                 executing your inference code. Defaults to ``None``. Required unless
-                ``tensorflow_version`` is provided. List of supported versions:
-                https://github.com/aws/sagemaker-python-sdk#huggingface-sagemaker-estimators.
+                ``tensorflow_version`` is provided. The current supported versions are ``1.7.1`` and ``1.6.0``.
             py_version (str): Python version you want to use for executing your
                 model training code. Defaults to ``None``. Required unless
                 ``image_uri`` is provided.

src/sagemaker/jumpstart/accessors.py

@@ -127,6 +127,7 @@ class JumpStartModelsAccessor(object):
     _curr_region = JUMPSTART_DEFAULT_REGION_NAME
 
     _content_bucket: Optional[str] = None
+    _gated_content_bucket: Optional[str] = None
 
     _cache_kwargs: Dict[str, Any] = {}
 
@@ -140,6 +141,16 @@ def get_jumpstart_content_bucket() -> Optional[str]:
         """Returns JumpStart content bucket."""
         return JumpStartModelsAccessor._content_bucket
 
+    @staticmethod
+    def set_jumpstart_gated_content_bucket(gated_content_bucket: str) -> None:
+        """Sets JumpStart gated content bucket."""
+        JumpStartModelsAccessor._gated_content_bucket = gated_content_bucket
+
+    @staticmethod
+    def get_jumpstart_gated_content_bucket() -> Optional[str]:
+        """Returns JumpStart gated content bucket."""
+        return JumpStartModelsAccessor._gated_content_bucket
+
     @staticmethod
     def _validate_and_mutate_region_cache_kwargs(
         cache_kwargs: Optional[Dict[str, Any]] = None, region: Optional[str] = None

src/sagemaker/jumpstart/artifacts/model_uris.py

@@ -25,6 +25,7 @@
 )
 from sagemaker.jumpstart.utils import (
     get_jumpstart_content_bucket,
+    get_jumpstart_gated_content_bucket,
     verify_model_region_and_return_specs,
 )
 from sagemaker.session import Session
@@ -157,9 +158,16 @@ def _retrieve_model_uri(
 
         model_artifact_key = _retrieve_training_artifact_key(model_specs, instance_type)
 
-    bucket = os.environ.get(
-        ENV_VARIABLE_JUMPSTART_MODEL_ARTIFACT_BUCKET_OVERRIDE
-    ) or get_jumpstart_content_bucket(region)
+    default_jumpstart_bucket: str = (
+        get_jumpstart_gated_content_bucket(region)
+        if model_specs.gated_bucket
+        else get_jumpstart_content_bucket(region)
+    )
+
+    bucket = (
+        os.environ.get(ENV_VARIABLE_JUMPSTART_MODEL_ARTIFACT_BUCKET_OVERRIDE)
+        or default_jumpstart_bucket
+    )
 
     model_s3_uri = f"s3://{bucket}/{model_artifact_key}"
 

src/sagemaker/jumpstart/constants.py

@@ -38,82 +38,102 @@
         JumpStartLaunchedRegionInfo(
             region_name="us-west-2",
             content_bucket="jumpstart-cache-prod-us-west-2",
+            gated_content_bucket="jumpstart-private-cache-prod-us-west-2",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="us-east-1",
             content_bucket="jumpstart-cache-prod-us-east-1",
+            gated_content_bucket="jumpstart-private-cache-prod-us-east-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="us-east-2",
             content_bucket="jumpstart-cache-prod-us-east-2",
+            gated_content_bucket="jumpstart-private-cache-prod-us-east-2",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="eu-west-1",
             content_bucket="jumpstart-cache-prod-eu-west-1",
+            gated_content_bucket="jumpstart-private-cache-prod-eu-west-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="eu-central-1",
             content_bucket="jumpstart-cache-prod-eu-central-1",
+            gated_content_bucket="jumpstart-private-cache-prod-eu-central-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="eu-north-1",
             content_bucket="jumpstart-cache-prod-eu-north-1",
+            gated_content_bucket="jumpstart-private-cache-prod-eu-north-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="me-south-1",
             content_bucket="jumpstart-cache-prod-me-south-1",
+            gated_content_bucket="jumpstart-private-cache-prod-me-south-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="ap-south-1",
             content_bucket="jumpstart-cache-prod-ap-south-1",
+            gated_content_bucket="jumpstart-private-cache-prod-ap-south-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="eu-west-3",
             content_bucket="jumpstart-cache-prod-eu-west-3",
+            gated_content_bucket="jumpstart-private-cache-prod-eu-west-3",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="af-south-1",
             content_bucket="jumpstart-cache-prod-af-south-1",
+            gated_content_bucket="jumpstart-private-cache-prod-af-south-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="sa-east-1",
             content_bucket="jumpstart-cache-prod-sa-east-1",
+            gated_content_bucket="jumpstart-private-cache-prod-sa-east-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="ap-east-1",
             content_bucket="jumpstart-cache-prod-ap-east-1",
+            gated_content_bucket="jumpstart-private-cache-prod-ap-east-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="ap-northeast-2",
             content_bucket="jumpstart-cache-prod-ap-northeast-2",
+            gated_content_bucket="jumpstart-private-cache-prod-ap-northeast-2",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="eu-west-2",
             content_bucket="jumpstart-cache-prod-eu-west-2",
+            gated_content_bucket="jumpstart-private-cache-prod-eu-west-2",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="eu-south-1",
             content_bucket="jumpstart-cache-prod-eu-south-1",
+            gated_content_bucket="jumpstart-private-cache-prod-eu-south-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="ap-northeast-1",
             content_bucket="jumpstart-cache-prod-ap-northeast-1",
+            gated_content_bucket="jumpstart-private-cache-prod-ap-northeast-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="us-west-1",
             content_bucket="jumpstart-cache-prod-us-west-1",
+            gated_content_bucket="jumpstart-private-cache-prod-us-west-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="ap-southeast-1",
             content_bucket="jumpstart-cache-prod-ap-southeast-1",
+            gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="ap-southeast-2",
             content_bucket="jumpstart-cache-prod-ap-southeast-2",
+            gated_content_bucket="jumpstart-private-cache-prod-ap-southeast-2",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="ca-central-1",
             content_bucket="jumpstart-cache-prod-ca-central-1",
+            gated_content_bucket="jumpstart-private-cache-prod-ca-central-1",
         ),
         JumpStartLaunchedRegionInfo(
             region_name="cn-north-1",
@@ -128,6 +148,11 @@
 JUMPSTART_REGION_NAME_SET = {region.region_name for region in JUMPSTART_LAUNCHED_REGIONS}
 
 JUMPSTART_BUCKET_NAME_SET = {region.content_bucket for region in JUMPSTART_LAUNCHED_REGIONS}
+JUMPSTART_GATED_BUCKET_NAME_SET = {
+    region.gated_content_bucket
+    for region in JUMPSTART_LAUNCHED_REGIONS
+    if region.gated_content_bucket is not None
+}
 
 JUMPSTART_DEFAULT_REGION_NAME = boto3.session.Session().region_name or "us-west-2"
 

src/sagemaker/jumpstart/types.py

@@ -107,16 +107,21 @@ class JumpStartS3FileType(str, Enum):
 class JumpStartLaunchedRegionInfo(JumpStartDataHolderType):
     """Data class for launched region info."""
 
-    __slots__ = ["content_bucket", "region_name"]
+    __slots__ = ["content_bucket", "region_name", "gated_content_bucket"]
 
-    def __init__(self, content_bucket: str, region_name: str):
+    def __init__(
+        self, content_bucket: str, region_name: str, gated_content_bucket: Optional[str] = None
+    ):
         """Instantiates JumpStartLaunchedRegionInfo object.
 
         Args:
             content_bucket (str): Name of JumpStart s3 content bucket associated with region.
             region_name (str): Name of JumpStart launched region.
+            gated_content_bucket (Optional[str[]): Name of JumpStart gated s3 content bucket
+                optionally associated with region.
         """
         self.content_bucket = content_bucket
+        self.gated_content_bucket = gated_content_bucket
         self.region_name = region_name
 
 
@@ -691,6 +696,7 @@ class JumpStartModelSpecs(JumpStartDataHolderType):
         "hosting_instance_type_variants",
         "training_instance_type_variants",
         "default_payloads",
+        "gated_bucket",
     ]
 
     def __init__(self, spec: Dict[str, Any]):
@@ -767,6 +773,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
             if json_obj.get("default_payloads")
             else None
         )
+        self.gated_bucket = json_obj.get("gated_bucket", False)
         self.inference_volume_size: Optional[int] = json_obj.get("inference_volume_size")
         self.inference_enable_network_isolation: bool = json_obj.get(
             "inference_enable_network_isolation", False

src/sagemaker/jumpstart/utils.py

@@ -63,13 +63,64 @@ def get_jumpstart_launched_regions_message() -> str:
     return f"JumpStart is available in {formatted_launched_regions_str} regions."
 
 
+def get_jumpstart_gated_content_bucket(
+    region: str = constants.JUMPSTART_DEFAULT_REGION_NAME,
+) -> str:
+    """Returns regionalized private content bucket name for JumpStart.
+
+    Raises:
+        ValueError: If JumpStart is not launched in ``region`` or private content
+            unavailable in that region.
+    """
+
+    old_gated_content_bucket: Optional[
+        str
+    ] = accessors.JumpStartModelsAccessor.get_jumpstart_gated_content_bucket()
+
+    info_logs: List[str] = []
+
+    gated_bucket_to_return: Optional[str] = None
+    if (
+        constants.ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE in os.environ
+        and len(os.environ[constants.ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE]) > 0
+    ):
+        gated_bucket_to_return = os.environ[
+            constants.ENV_VARIABLE_JUMPSTART_CONTENT_BUCKET_OVERRIDE
+        ]
+        info_logs.append(f"Using JumpStart private bucket override: '{gated_bucket_to_return}'")
+    else:
+        try:
+            gated_bucket_to_return = constants.JUMPSTART_REGION_NAME_TO_LAUNCHED_REGION_DICT[
+                region
+            ].gated_content_bucket
+            if gated_bucket_to_return is None:
+                raise ValueError(
+                    f"No private content bucket for JumpStart exists in {region} region."
+                )
+        except KeyError:
+            formatted_launched_regions_str = get_jumpstart_launched_regions_message()
+            raise ValueError(
+                f"Unable to get private content bucket for JumpStart in {region} region. "
+                f"{formatted_launched_regions_str}"
+            )
+
+    accessors.JumpStartModelsAccessor.set_jumpstart_gated_content_bucket(gated_bucket_to_return)
+
+    if gated_bucket_to_return != old_gated_content_bucket:
+        accessors.JumpStartModelsAccessor.reset_cache()
+        for info_log in info_logs:
+            constants.JUMPSTART_LOGGER.info(info_log)
+
+    return gated_bucket_to_return
+
+
 def get_jumpstart_content_bucket(
     region: str = constants.JUMPSTART_DEFAULT_REGION_NAME,
 ) -> str:
     """Returns regionalized content bucket name for JumpStart.
 
     Raises:
-        RuntimeError: If JumpStart is not launched in ``region``.
+        ValueError: If JumpStart is not launched in ``region``.
     """
 
     old_content_bucket: Optional[

src/sagemaker/mxnet/estimator.py

@@ -77,7 +77,7 @@ def __init__(
             framework_version (str): MXNet version you want to use for executing
                 your model training code. Defaults to `None`. Required unless
                 ``image_uri`` is provided. List of supported versions.
-                https://github.com/aws/sagemaker-python-sdk#mxnet-sagemaker-estimators.
+                https://aws.amazon.com/releasenotes/available-deep-learning-containers-images/.
             py_version (str): Python version you want to use for executing your
                 model training code. One of 'py2' or 'py3'. Defaults to ``None``. Required
                 unless ``image_uri`` is provided.

src/sagemaker/pytorch/estimator.py

@@ -326,6 +326,9 @@ def _pytorch_distribution_configuration(self, distribution):
             if self.instance_type is not None:
                 distribution_config[self.INSTANCE_TYPE_ENV_NAME] = self.instance_type
         elif torch_distributed_enabled:
+            if "smdistributed" in distribution:
+                # Enable torch_distributed for smdistributed.
+                distribution_config = self._distribution_configuration(distribution=distribution)
             distribution_config[self.LAUNCH_TORCH_DISTRIBUTED_ENV_NAME] = torch_distributed_enabled
             if self.instance_type is not None:
                 distribution_config[self.INSTANCE_TYPE_ENV_NAME] = self.instance_type