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

Author: jjerphan

CHANGELOG.md

@@ -1,5 +1,42 @@
 # Changelog
 
+## v2.196.0 (2023-10-27)
+
+### Features
+
+ * inference instance type conditioned on training instance type
+
+### Bug Fixes and Other Changes
+
+ * improved jumpstart tagging
+
+## v2.195.1 (2023-10-26)
+
+### Bug Fixes and Other Changes
+
+ * Allow either instance_type or instance_group to be defined in…
+ * enhance image_uris unit tests
+
+## v2.195.0 (2023-10-25)
+
+### Features
+
+ * jumpstart gated model artifacts
+ * jumpstart extract generated text from response
+ * jumpstart contruct payload utility
+
+### Bug Fixes and Other Changes
+
+ * relax upper bound on urllib in local mode requirements
+ * bump urllib3 version
+ * allow smdistributed to be enabled with torch_distributed.
+ * fix URL links
+
+### Documentation Changes
+
+ * remove python 2 reference
+ * update framework version links
+
 ## v2.194.0 (2023-10-19)
 
 ### Features

VERSION

@@ -1 +1 @@
-2.194.1.dev0
+2.196.1.dev0

doc/overview.rst

@@ -32,7 +32,7 @@ After you train a model, you can save it, and then serve the model as an endpoin
 Prepare a Training script
 =========================
 
-Your training script must be a Python 2.7 or 3.6 compatible source file.
+Your training script must be a 3.6 compatible source file.
 
 The training script is very similar to a training script you might run outside of SageMaker, but you can access useful properties about the training environment through various environment variables, including the following:
 

requirements/extras/local_requirements.txt

@@ -1,3 +1,3 @@
-urllib3>=1.26.8,<1.26.15
+urllib3>=1.26.8,<3.0.0
 docker>=5.0.2,<7.0.0
 PyYAML>=5.4.1,<7

src/sagemaker/estimator.py

@@ -71,7 +71,7 @@
 from sagemaker.utils import instance_supports_kms
 from sagemaker.job import _Job
 from sagemaker.jumpstart.utils import (
-    add_jumpstart_tags,
+    add_jumpstart_uri_tags,
     get_jumpstart_base_name_if_jumpstart_model,
     update_inference_tags_with_jumpstart_training_tags,
 )
@@ -577,9 +577,7 @@ def __init__(
         self.entry_point = entry_point
         self.dependencies = dependencies or []
         self.uploaded_code: Optional[UploadedCode] = None
-        self.tags = add_jumpstart_tags(
-            tags=tags, training_model_uri=self.model_uri, training_script_uri=self.source_dir
-        )
+
         if self.instance_type in ("local", "local_gpu"):
             if self.instance_type == "local_gpu" and self.instance_count > 1:
                 raise RuntimeError("Distributed Training in Local GPU is not supported")
@@ -592,6 +590,15 @@ def __init__(
         else:
             self.sagemaker_session = sagemaker_session or Session()
 
+        self.tags = (
+            add_jumpstart_uri_tags(
+                tags=tags, training_model_uri=self.model_uri, training_script_uri=self.source_dir
+            )
+            if getattr(self.sagemaker_session, "settings", None) is not None
+            and self.sagemaker_session.settings.include_jumpstart_tags
+            else tags
+        )
+
         self.base_job_name = base_job_name
         self._current_job_name = None
         if (
@@ -3818,6 +3825,7 @@ def _distribution_configuration(self, distribution):
 
         mpi_enabled = False
         smdataparallel_enabled = False
+        p5_enabled = False
         if "instance_groups" in distribution:
             distribution_config["sagemaker_distribution_instance_groups"] = distribution[
                 "instance_groups"
@@ -3862,10 +3870,11 @@ def _distribution_configuration(self, distribution):
             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)}."
-                )
+                for instance in self.instance_groups:
+                    if "p5.48xlarge" in instance._to_request_dict().get("InstanceType", ()):
+                        p5_enabled = True
+                        break
+
             img_uri = "" if self.image_uri is None else self.image_uri
             for unsupported_image in Framework.UNSUPPORTED_DLC_IMAGE_FOR_SM_PARALLELISM:
                 if (

src/sagemaker/instance_types.py

@@ -33,6 +33,7 @@ def retrieve_default(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    training_instance_type: Optional[str] = None,
 ) -> str:
     """Retrieves the default instance type for the model matching the given arguments.
 
@@ -56,6 +57,11 @@ def retrieve_default(
             object, used for SageMaker interactions. If not
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
+        training_instance_type (str): In the case of a model fine-tuned on SageMaker, the training
+            instance type used for the training job that produced the fine-tuned weights.
+            Optionally supply this to get a inference instance type conditioned
+            on the training instance, to ensure compatability of training artifact to inference
+            instance. (Default: None).
     Returns:
         str: The default instance type to use for the model.
 
@@ -78,6 +84,7 @@ def retrieve_default(
         tolerate_vulnerable_model,
         tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
+        training_instance_type=training_instance_type,
     )
 
 
@@ -89,6 +96,7 @@ def retrieve(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    training_instance_type: Optional[str] = None,
 ) -> List[str]:
     """Retrieves the supported training instance types for the model matching the given arguments.
 
@@ -110,6 +118,12 @@ def retrieve(
             object, used for SageMaker interactions. If not
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
+        training_instance_type (str): In the case of a model fine-tuned on SageMaker, the training
+            instance type used for the training job that produced the fine-tuned weights.
+            Optionally supply this to get a inference instance type conditioned
+            on the training instance, to ensure compatability of training artifact to inference
+            instance. (Default: None).
+
     Returns:
         list: The supported instance types to use for the model.
 
@@ -132,4 +146,5 @@ def retrieve(
         tolerate_vulnerable_model,
         tolerate_deprecated_model,
         sagemaker_session=sagemaker_session,
+        training_instance_type=training_instance_type,
     )

src/sagemaker/jumpstart/artifacts/instance_types.py

@@ -37,6 +37,7 @@ def _retrieve_default_instance_type(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    training_instance_type: Optional[str] = None,
 ) -> str:
     """Retrieves the default instance type for the model.
 
@@ -60,6 +61,11 @@ def _retrieve_default_instance_type(
             object, used for SageMaker interactions. If not
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
+        training_instance_type (str): In the case of a model fine-tuned on SageMaker, the training
+            instance type used for the training job that produced the fine-tuned weights.
+            Optionally supply this to get a inference instance type conditioned
+            on the training instance, to ensure compatability of training artifact to inference
+            instance. (Default: None).
     Returns:
         str: the default instance type to use for the model or None.
 
@@ -82,7 +88,21 @@ def _retrieve_default_instance_type(
     )
 
     if scope == JumpStartScriptScope.INFERENCE:
-        default_instance_type = model_specs.default_inference_instance_type
+        instance_specific_default_instance_type = (
+            (
+                model_specs.training_instance_type_variants.get_instance_specific_default_inference_instance_type(  # pylint: disable=C0301 # noqa: E501
+                    training_instance_type
+                )
+            )
+            if training_instance_type is not None
+            and getattr(model_specs, "training_instance_type_variants", None) is not None
+            else None
+        )
+        default_instance_type = (
+            instance_specific_default_instance_type
+            if instance_specific_default_instance_type is not None
+            else model_specs.default_inference_instance_type
+        )
     elif scope == JumpStartScriptScope.TRAINING:
         default_instance_type = model_specs.default_training_instance_type
     else:
@@ -103,6 +123,7 @@ def _retrieve_instance_types(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    training_instance_type: Optional[str] = None,
 ) -> List[str]:
     """Retrieves the supported instance types for the model.
 
@@ -126,6 +147,11 @@ def _retrieve_instance_types(
             object, used for SageMaker interactions. If not
             specified, one is created using the default AWS configuration
             chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
+        training_instance_type (str): In the case of a model fine-tuned on SageMaker, the training
+            instance type used for the training job that produced the fine-tuned weights.
+            Optionally supply this to get a inference instance type conditioned
+            on the training instance, to ensure compatability of training artifact to inference
+            instance. (Default: None).
     Returns:
         list: the supported instance types to use for the model or None.
 
@@ -148,8 +174,24 @@ def _retrieve_instance_types(
     )
 
     if scope == JumpStartScriptScope.INFERENCE:
-        instance_types = model_specs.supported_inference_instance_types
+        default_instance_types = model_specs.supported_inference_instance_types or []
+        instance_specific_instance_types = (
+            model_specs.training_instance_type_variants.get_instance_specific_supported_inference_instance_types(  # pylint: disable=C0301 # noqa: E501
+                training_instance_type
+            )
+            if training_instance_type is not None
+            and getattr(model_specs, "training_instance_type_variants", None) is not None
+            else []
+        )
+        instance_types = (
+            instance_specific_instance_types
+            if len(instance_specific_instance_types) > 0
+            else default_instance_types
+        )
+
     elif scope == JumpStartScriptScope.TRAINING:
+        if training_instance_type is not None:
+            raise ValueError("Cannot use `training_instance_type` argument " "with training scope.")
         instance_types = model_specs.supported_training_instance_types
     else:
         raise NotImplementedError(

src/sagemaker/jumpstart/constants.py

@@ -154,6 +154,10 @@
     if region.gated_content_bucket is not None
 }
 
+JUMPSTART_GATED_AND_PUBLIC_BUCKET_NAME_SET = JUMPSTART_BUCKET_NAME_SET.union(
+    JUMPSTART_GATED_BUCKET_NAME_SET
+)
+
 JUMPSTART_DEFAULT_REGION_NAME = boto3.session.Session().region_name or "us-west-2"
 
 JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"

src/sagemaker/jumpstart/enums.py

@@ -76,6 +76,9 @@ class JumpStartTag(str, Enum):
     TRAINING_MODEL_URI = "aws-jumpstart-training-model-uri"
     TRAINING_SCRIPT_URI = "aws-jumpstart-training-script-uri"
 
+    MODEL_ID = "sagemaker-sdk:jumpstart-model-id"
+    MODEL_VERSION = "sagemaker-sdk:jumpstart-model-version"
+
 
 class SerializerType(str, Enum):
     """Enum class for serializers associated with JumpStart models."""

src/sagemaker/jumpstart/estimator.py

@@ -988,7 +988,6 @@ def deploy(
             use_compiled_model (bool): Flag to select whether to use compiled
                 (optimized) model. (Default: False).
         """
-
         self.orig_predictor_cls = predictor_cls
 
         sagemaker_session = sagemaker_session or self.sagemaker_session
@@ -1039,6 +1038,7 @@ def deploy(
             dependencies=dependencies,
             git_config=git_config,
             use_compiled_model=use_compiled_model,
+            training_instance_type=self.instance_type,
         )
 
         predictor = super(JumpStartEstimator, self).deploy(

src/sagemaker/jumpstart/factory/estimator.py

@@ -60,8 +60,10 @@
     JumpStartModelInitKwargs,
 )
 from sagemaker.jumpstart.utils import (
+    add_jumpstart_model_id_version_tags,
     update_dict_if_key_not_present,
     resolve_estimator_sagemaker_config_field,
+    verify_model_region_and_return_specs,
 )
 
 
@@ -196,6 +198,7 @@ def get_init_kwargs(
     estimator_init_kwargs = _add_estimator_extra_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_role_to_kwargs(estimator_init_kwargs)
     estimator_init_kwargs = _add_env_to_kwargs(estimator_init_kwargs)
+    estimator_init_kwargs = _add_tags_to_kwargs(estimator_init_kwargs)
 
     return estimator_init_kwargs
 
@@ -277,6 +280,7 @@ def get_deploy_kwargs(
     tolerate_vulnerable_model: Optional[bool] = None,
     use_compiled_model: Optional[bool] = None,
     model_name: Optional[str] = None,
+    training_instance_type: Optional[str] = None,
 ) -> JumpStartEstimatorDeployKwargs:
     """Returns kwargs required to call `deploy` on `sagemaker.estimator.Estimator` object."""
 
@@ -310,7 +314,7 @@ def get_deploy_kwargs(
         model_id=model_id,
         model_from_estimator=True,
         model_version=model_version,
-        instance_type=model_deploy_kwargs.instance_type,  # prevent excess logging
+        instance_type=model_deploy_kwargs.instance_type if training_instance_type is None else None,
         region=region,
         image_uri=image_uri,
         source_dir=source_dir,
@@ -330,6 +334,7 @@ def get_deploy_kwargs(
         git_config=git_config,
         tolerate_vulnerable_model=tolerate_vulnerable_model,
         tolerate_deprecated_model=tolerate_deprecated_model,
+        training_instance_type=training_instance_type,
     )
 
     estimator_deploy_kwargs: JumpStartEstimatorDeployKwargs = JumpStartEstimatorDeployKwargs(
@@ -439,6 +444,26 @@ def _add_instance_type_and_count_to_kwargs(
     return kwargs
 
 
+def _add_tags_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartEstimatorInitKwargs:
+    """Sets tags in kwargs based on default or override, returns full kwargs."""
+
+    full_model_version = verify_model_region_and_return_specs(
+        model_id=kwargs.model_id,
+        version=kwargs.model_version,
+        scope=JumpStartScriptScope.TRAINING,
+        region=kwargs.region,
+        tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
+        tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
+        sagemaker_session=kwargs.sagemaker_session,
+    ).version
+
+    if kwargs.sagemaker_session.settings.include_jumpstart_tags:
+        kwargs.tags = add_jumpstart_model_id_version_tags(
+            kwargs.tags, kwargs.model_id, full_model_version
+        )
+    return kwargs
+
+
 def _add_image_uri_to_kwargs(kwargs: JumpStartEstimatorInitKwargs) -> JumpStartEstimatorInitKwargs:
     """Sets image uri in kwargs based on default or override, returns full kwargs."""