Merge branch 'fix/make-instance-type-optional' of https://github.com/rahven14/sagemaker-python-sdk into fix/make-instance-type-optional

Author: rahven14

CHANGELOG.md

@@ -1,5 +1,20 @@
 # Changelog
 
+## v2.92.0 (2022-05-26)
+
+### Features
+
+ * add 'Domain' property to RegisterModel step
+
+### Bug Fixes and Other Changes
+
+ * support estimator output path parameterization
+ * Add back Prevent passing PipelineVariable object into image_uris.retrieve
+ * jumpstart amt tracking
+ * fix missing register method params for framework models
+ * fix docstring for decorated functions
+ * Documents: add sagemaker model building pipeline readthedocs
+
 ## v2.91.1 (2022-05-19)
 
 ### Bug Fixes and Other Changes

VERSION

@@ -1 +1 @@
-2.91.2.dev0
+2.92.1.dev0

setup.py

@@ -48,13 +48,13 @@ def read_requirements(filename):
 # Declare minimal set for installation
 required_packages = [
     "attrs==20.3.0",
-    "boto3>=1.20.21",
+    "boto3>=1.20.21,<2.0",
     "google-pasta",
-    "numpy>=1.9.0",
-    "protobuf>=3.1",
-    "protobuf3-to-dict>=0.1.5",
+    "numpy>=1.9.0,<2.0",
+    "protobuf>=3.1,<4.0",
+    "protobuf3-to-dict>=0.1.5,<1.0",
     "smdebug_rulesconfig==1.0.1",
-    "importlib-metadata>=1.4.0",
+    "importlib-metadata>=1.4.0,<2.0",
     "packaging>=20.0",
     "pandas",
     "pathos",

src/sagemaker/estimator.py

@@ -695,26 +695,45 @@ def _stage_user_code_in_s3(self) -> str:
 
         Returns: S3 URI
         """
-        local_mode = self.output_path.startswith("file://")
-
-        if self.code_location is None and local_mode:
-            code_bucket = self.sagemaker_session.default_bucket()
-            code_s3_prefix = "{}/{}".format(self._current_job_name, "source")
-            kms_key = None
-        elif self.code_location is None:
-            code_bucket, _ = parse_s3_url(self.output_path)
-            code_s3_prefix = "{}/{}".format(self._current_job_name, "source")
-            kms_key = self.output_kms_key
-        elif local_mode:
-            code_bucket, key_prefix = parse_s3_url(self.code_location)
-            code_s3_prefix = "/".join(filter(None, [key_prefix, self._current_job_name, "source"]))
-            kms_key = None
+        if is_pipeline_variable(self.output_path):
+            if self.code_location is None:
+                code_bucket = self.sagemaker_session.default_bucket()
+                code_s3_prefix = "{}/{}".format(self._current_job_name, "source")
+                kms_key = None
+            else:
+                code_bucket, key_prefix = parse_s3_url(self.code_location)
+                code_s3_prefix = "/".join(
+                    filter(None, [key_prefix, self._current_job_name, "source"])
+                )
+
+                output_bucket = self.sagemaker_session.default_bucket()
+                kms_key = self.output_kms_key if code_bucket == output_bucket else None
         else:
-            code_bucket, key_prefix = parse_s3_url(self.code_location)
-            code_s3_prefix = "/".join(filter(None, [key_prefix, self._current_job_name, "source"]))
+            local_mode = self.output_path.startswith("file://")
+            if local_mode:
+                if self.code_location is None:
+                    code_bucket = self.sagemaker_session.default_bucket()
+                    code_s3_prefix = "{}/{}".format(self._current_job_name, "source")
+                    kms_key = None
+                else:
+                    code_bucket, key_prefix = parse_s3_url(self.code_location)
+                    code_s3_prefix = "/".join(
+                        filter(None, [key_prefix, self._current_job_name, "source"])
+                    )
+                    kms_key = None
+            else:
+                if self.code_location is None:
+                    code_bucket, _ = parse_s3_url(self.output_path)
+                    code_s3_prefix = "{}/{}".format(self._current_job_name, "source")
+                    kms_key = self.output_kms_key
+                else:
+                    code_bucket, key_prefix = parse_s3_url(self.code_location)
+                    code_s3_prefix = "/".join(
+                        filter(None, [key_prefix, self._current_job_name, "source"])
+                    )
 
-            output_bucket, _ = parse_s3_url(self.output_path)
-            kms_key = self.output_kms_key if code_bucket == output_bucket else None
+                    output_bucket, _ = parse_s3_url(self.output_path)
+                    kms_key = self.output_kms_key if code_bucket == output_bucket else None
 
         return tar_and_upload_dir(
             session=self.sagemaker_session.boto_session,

src/sagemaker/image_uris.py

@@ -23,6 +23,7 @@
 from sagemaker.jumpstart.utils import is_jumpstart_model_input
 from sagemaker.spark import defaults
 from sagemaker.jumpstart import artifacts
+from sagemaker.workflow import is_pipeline_variable
 
 logger = logging.getLogger(__name__)
 
@@ -104,11 +105,17 @@ def retrieve(
 
     Raises:
         NotImplementedError: If the scope is not supported.
-        ValueError: If the combination of arguments specified is not supported.
+        ValueError: If the combination of arguments specified is not supported or
+            any PipelineVariable object is passed in.
         VulnerableJumpStartModelError: If any of the dependencies required by the script have
             known security vulnerabilities.
         DeprecatedJumpStartModelError: If the version of the model is deprecated.
     """
+    args = dict(locals())
+    for name, val in args.items():
+        if is_pipeline_variable(val):
+            raise ValueError("%s should not be a pipeline variable (%s)" % (name, type(val)))
+
     if is_jumpstart_model_input(model_id, model_version):
         return artifacts._retrieve_image_uri(
             model_id,

src/sagemaker/tuner.py

@@ -32,6 +32,7 @@
 from sagemaker.estimator import Framework
 from sagemaker.inputs import TrainingInput
 from sagemaker.job import _Job
+from sagemaker.jumpstart.utils import add_jumpstart_tags, get_jumpstart_base_name_if_jumpstart_model
 from sagemaker.parameter import (
     CategoricalParameter,
     ContinuousParameter,
@@ -319,6 +320,42 @@ def _prepare_for_tuning(self, job_name=None, include_cls_metadata=False):
         """Prepare the tuner instance for tuning (fit)."""
         self._prepare_job_name_for_tuning(job_name=job_name)
         self._prepare_static_hyperparameters_for_tuning(include_cls_metadata=include_cls_metadata)
+        self._prepare_tags_for_tuning()
+
+    def _get_model_uri(
+        self,
+        estimator,
+    ):
+        """Return the model artifact URI used by the Estimator instance.
+
+        This attribute can live in multiple places, and accessing the attribute can
+        raise a TypeError, which needs to be handled.
+        """
+        try:
+            return getattr(estimator, "model_data", None)
+        except TypeError:
+            return getattr(estimator, "model_uri", None)
+
+    def _prepare_tags_for_tuning(self):
+        """Add tags to tuning job (from Estimator and JumpStart tags)."""
+
+        # Add tags from Estimator class
+        estimator = self.estimator or self.estimator_dict[sorted(self.estimator_dict.keys())[0]]
+
+        estimator_tags = getattr(estimator, "tags", []) or []
+
+        if self.tags is None and len(estimator_tags) > 0:
+            self.tags = []
+
+        for tag in estimator_tags:
+            if tag not in self.tags:
+                self.tags.append(tag)
+
+        self.tags = add_jumpstart_tags(
+            tags=self.tags,
+            training_script_uri=getattr(estimator, "source_dir", None),
+            training_model_uri=self._get_model_uri(estimator),
+        )
 
     def _prepare_job_name_for_tuning(self, job_name=None):
         """Set current job name before starting tuning."""
@@ -331,6 +368,12 @@ def _prepare_job_name_for_tuning(self, job_name=None):
                     self.estimator or self.estimator_dict[sorted(self.estimator_dict.keys())[0]]
                 )
                 base_name = base_name_from_image(estimator.training_image_uri())
+
+                jumpstart_base_name = get_jumpstart_base_name_if_jumpstart_model(
+                    getattr(estimator, "source_dir", None),
+                    self._get_model_uri(estimator),
+                )
+                base_name = jumpstart_base_name or base_name
             self._current_job_name = name_from_base(
                 base_name, max_length=self.TUNING_JOB_NAME_MAX_LENGTH, short=True
             )

src/sagemaker/workflow/entities.py

@@ -78,7 +78,11 @@ def __add__(self, other: Union[Expression, PrimitiveType]):
 
     def __str__(self):
         """Override built-in String function for PipelineVariable"""
-        raise TypeError("Pipeline variables do not support __str__ operation.")
+        raise TypeError(
+            "Pipeline variables do not support __str__ operation. "
+            "Please use `.to_string()` to convert it to string type in execution time"
+            "or use `.expr` to translate it to Json for display purpose in Python SDK."
+        )
 
     def __int__(self):
         """Override built-in Integer function for PipelineVariable"""

tests/integ/sagemaker/workflow/test_model_create_and_registration.py

@@ -90,10 +90,12 @@ def test_conditional_pytorch_training_model_registration(
     inputs = TrainingInput(s3_data=input_path)
 
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_type = "ml.m5.xlarge"
     good_enough_input = ParameterInteger(name="GoodEnoughInput", default_value=1)
     in_condition_input = ParameterString(name="Foo", default_value="Foo")
 
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     pytorch_estimator = PyTorch(
         entry_point=entry_point,
         role=role,
@@ -153,7 +155,6 @@ def test_conditional_pytorch_training_model_registration(
             in_condition_input,
             good_enough_input,
             instance_count,
-            instance_type,
         ],
         steps=[step_train, step_cond],
         sagemaker_session=sagemaker_session,
@@ -259,8 +260,10 @@ def test_sklearn_xgboost_sip_model_registration(
     prefix = "sip"
     bucket_name = sagemaker_session.default_bucket()
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_type = "ml.m5.xlarge"
 
+    # The instance_type should not be a pipeline variable
+    # since it is used to retrieve image_uri in compile time (PySDK)
     sklearn_processor = SKLearnProcessor(
         role=role,
         instance_type=instance_type,
@@ -331,6 +334,8 @@ def test_sklearn_xgboost_sip_model_registration(
     source_dir = base_dir
     code_location = "s3://{0}/{1}/code".format(bucket_name, prefix)
 
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     estimator = XGBoost(
         entry_point=entry_point,
         source_dir=source_dir,
@@ -416,7 +421,6 @@ def test_sklearn_xgboost_sip_model_registration(
             train_data_path_param,
             val_data_path_param,
             model_path_param,
-            instance_type,
             instance_count,
             output_path_param,
         ],
@@ -462,7 +466,7 @@ def test_model_registration_with_drift_check_baselines(
     pipeline_name,
 ):
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_type = "ml.m5.xlarge"
 
     # upload model data to s3
     model_local_path = os.path.join(DATA_DIR, "mxnet_mnist/model.tar.gz")
@@ -551,6 +555,9 @@ def test_model_registration_with_drift_check_baselines(
     )
     customer_metadata_properties = {"key1": "value1"}
     domain = "COMPUTER_VISION"
+
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     estimator = XGBoost(
         entry_point="training.py",
         source_dir=os.path.join(DATA_DIR, "sip"),
@@ -581,7 +588,6 @@ def test_model_registration_with_drift_check_baselines(
         parameters=[
             model_uri_param,
             metrics_uri_param,
-            instance_type,
             instance_count,
         ],
         steps=[step_register],
@@ -670,9 +676,11 @@ def test_model_registration_with_model_repack(
     inputs = TrainingInput(s3_data=input_path)
 
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_type = "ml.m5.xlarge"
     good_enough_input = ParameterInteger(name="GoodEnoughInput", default_value=1)
 
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     pytorch_estimator = PyTorch(
         entry_point=entry_point,
         role=role,
@@ -727,7 +735,7 @@ def test_model_registration_with_model_repack(
 
     pipeline = Pipeline(
         name=pipeline_name,
-        parameters=[good_enough_input, instance_count, instance_type],
+        parameters=[good_enough_input, instance_count],
         steps=[step_cond],
         sagemaker_session=sagemaker_session,
     )
@@ -770,8 +778,10 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model(
     inputs = TrainingInput(s3_data=input_path)
 
     instance_count = ParameterInteger(name="InstanceCount", default_value=1)
-    instance_type = ParameterString(name="InstanceType", default_value="ml.m5.xlarge")
+    instance_type = "ml.m5.xlarge"
 
+    # If image_uri is not provided, the instance_type should not be a pipeline variable
+    # since instance_type is used to retrieve image_uri in compile time (PySDK)
     tensorflow_estimator = TensorFlow(
         entry_point=entry_point,
         role=role,
@@ -812,10 +822,7 @@ def test_model_registration_with_tensorflow_model_with_pipeline_model(
 
     pipeline = Pipeline(
         name=pipeline_name,
-        parameters=[
-            instance_count,
-            instance_type,
-        ],
+        parameters=[instance_count],
         steps=[step_train, step_register_model],
         sagemaker_session=sagemaker_session,
     )