change: Clean code and add retry to integ tests

Author: Dewen Qi

src/sagemaker/model.py

@@ -405,6 +405,7 @@ def create(
                 https://boto3.amazonaws.com/v1/documentation
                 /api/latest/reference/services/sagemaker.html#SageMaker.Client.add_tags
         """
+        # TODO: we should replace _create_sagemaker_model() with create()
         self._create_sagemaker_model(
             instance_type=instance_type,
             accelerator_type=accelerator_type,

src/sagemaker/workflow/_utils.py

@@ -76,8 +76,7 @@ def __init__(
                     endpoints use this role to access training data and model
                     artifacts. After the endpoint is created, the inference code
                     might use the IAM role, if it needs to access an AWS resource.
-            model_data (str): The S3 location of a SageMaker model data
-                    ``.tar.gz`` file.
+            model_data (str): The S3 location of a SageMaker model data `.tar.gz` file.
             entry_point (str): Path (absolute or relative) to the local Python
                     source file which should be executed as the entry point to
                     inference. If ``source_dir`` is specified, then ``entry_point``
@@ -94,6 +93,8 @@ def __init__(
                         >>>         |----- test.py
 
                         You can assign entry_point='src/train.py'.
+            display_name (str): The display name of this `_RepackModelStep` step (default: None).
+            description (str): The description of this `_RepackModelStep` (default: None).
             source_dir (str): A relative location to a directory with other training
                 or model hosting source code dependencies aside from the entry point
                 file in the Git repo (default: None). Structure within this
@@ -124,12 +125,14 @@ def __init__(
 
                     This is not supported with "local code" in Local Mode.
             depends_on (List[str] or List[Step]): A list of step names or instances
-                    this step depends on
+                    this step depends on (default: None).
             retry_policies (List[RetryPolicy]): The list of retry policies for the current step
+                (default: None).
             subnets (list[str]): List of subnet ids. If not specified, the re-packing
-                    job will be created without VPC config.
+                    job will be created without VPC config (default: None).
             security_group_ids (list[str]): List of security group ids. If not
-                specified, the re-packing job will be created without VPC config.
+                specified, the re-packing job will be created without VPC config (default: None).
+            **kwargs: additional arguments for the repacking job.
         """
         self._model_data = model_data
         self.sagemaker_session = sagemaker_session
@@ -283,36 +286,38 @@ def __init__(
 
         Args:
             name (str): The name of the training step.
-            step_args (dict): The arguments for the `_RegisterModelStep` definition (default: None).
-            estimator (EstimatorBase): A `sagemaker.estimator.EstimatorBase` instance
+            step_args (dict): The arguments for this `_RegisterModelStep` definition
                 (default: None).
-            model_data: the S3 URI to the model data from training (default: None).
-            content_types (list): The supported MIME types for the
-                input data (default: None).
-            response_types (list): The supported MIME types for
-                the output data (default: None).
+            content_types (list): The supported MIME types for the input data (default: None).
+            response_types (list): The supported MIME types for the output data (default: None).
             inference_instances (list): A list of the instance types that are used to
                 generate inferences in real-time (default: None).
             transform_instances (list): A list of the instance types on which a
                 transformation job can be run or on which an endpoint
                 can be deployed (default: None).
+            estimator (EstimatorBase): A `sagemaker.estimator.EstimatorBase` instance
+                (default: None).
+            model_data: the S3 URI to the model data from training (default: None).
             model_package_group_name (str): Model Package Group name, exclusive to
                 `model_package_name`, using `model_package_group_name`
                 makes the Model Package versioned (default: None).
             model_metrics (ModelMetrics): ModelMetrics object (default: None).
-            metadata_properties (MetadataProperties): MetadataProperties object
-                (default: None).
+            metadata_properties (MetadataProperties): MetadataProperties object (default: None).
             approval_status (str): Model Approval Status, values can be "Approved",
-                "Rejected", or "PendingManualApproval"
-                (default: "PendingManualApproval").
+                "Rejected", or "PendingManualApproval" (default: "PendingManualApproval").
             image_uri (str): The container image uri for Model Package, if not specified,
                 Estimator's training container image will be used (default: None).
             compile_model_family (str): Instance family for compiled model,
                 if specified, a compiled model will be used (default: None).
+            display_name (str): The display name of this `_RegisterModelStep` step (default: None).
             description (str): Model Package description (default: None).
             depends_on (List[str] or List[Step]): A list of step names or instances
                 this step depends on (default: None).
             retry_policies (List[RetryPolicy]): The list of retry policies for the current step
+                (default: None).
+            tags (List[dict[str, str]]): A list of dictionaries containing key-value pairs used to
+                configure the create model package request (default: None).
+            container_def_list (list): A list of container definitions (default: None).
             drift_check_baselines (DriftCheckBaselines): DriftCheckBaselines object (default: None).
             customer_metadata_properties (dict[str, str]): A dictionary of key-value paired
                 metadata properties (default: None).

src/sagemaker/workflow/model_step.py

@@ -55,7 +55,7 @@ def __init__(
                 that this `ModelStep` depends on.
                 If a listed `Step` name does not exist, an error is returned (default: None).
             retry_policies (List[RetryPolicy] or Dict[str, List[RetryPolicy]]): The list of retry
-                policies for the `ModelStep`.
+                policies for the `ModelStep` (default: None).
             display_name (str): The display name of the `ModelStep`.
                 The display name provides better UI readability. (default: None).
             description (str): The description of the `ModelStep` (default: None).
@@ -107,8 +107,9 @@ def __init__(
             self._create_model_retry_policies = retry_policies
             self._register_model_retry_policies = retry_policies
             self._repack_model_retry_policies = retry_policies
-        self._repack()
 
+        if self._need_runtime_repack:
+            self._append_repack_model_step()
         if self._register_model_args:
             self._append_register_model_step()
         else:
@@ -140,12 +141,8 @@ def _append_create_model_step(self):
             create_model_step.add_depends_on(self.depends_on)
         self.steps.append(create_model_step)
 
-    def _repack(self):
-        """Add `_RepackModelStep` for the runtime repack"""
-
-        if not self._need_runtime_repack:
-            return
-
+    def _append_repack_model_step(self):
+        """Create and append a `_RepackModelStep` for the runtime repack"""
         if isinstance(self._model, PipelineModel):
             model_list = self._model.models
         elif isinstance(self._model, Model):

src/sagemaker/workflow/pipeline_context.py

@@ -133,7 +133,8 @@ def runnable_by_pipeline(run_func):
     """
 
     def wrapper(*args, **kwargs):
-        if isinstance(args[0].sagemaker_session, PipelineSession):
+        self_instance = args[0]
+        if isinstance(self_instance.sagemaker_session, PipelineSession):
             run_func_params = inspect.signature(run_func).parameters
             arg_list = list(args)
 
@@ -161,14 +162,14 @@ def wrapper(*args, **kwargs):
                 UserWarning,
             )
             if run_func.__name__ in ["register", "create"]:
-                args[0].sagemaker_session.init_step_arguments(args[0])
+                self_instance.sagemaker_session.init_step_arguments(self_instance)
                 run_func(*args, **kwargs)
-                context = args[0].sagemaker_session.context
-                args[0].sagemaker_session.context = None
+                context = self_instance.sagemaker_session.context
+                self_instance.sagemaker_session.context = None
                 return context
 
             run_func(*args, **kwargs)
-            return args[0].sagemaker_session.context
+            return self_instance.sagemaker_session.context
 
         return run_func(*args, **kwargs)
 

src/sagemaker/workflow/step_collections.py

@@ -13,6 +13,7 @@
 """The step definitions for workflow."""
 from __future__ import absolute_import
 
+import warnings
 from typing import List, Union
 
 import attr
@@ -243,6 +244,16 @@ def __init__(
         steps.append(register_model_step)
         self.steps = steps
 
+        # TODO: add public document link here once ready
+        warnings.warn(
+            (
+                "We are deprecating the use of RegisterModel. "
+                "Instead, please use the ModelStep, which simply takes in the step arguments "
+                "generated by model.register()."
+            ),
+            DeprecationWarning,
+        )
+
 
 RegisterModel = deprecated_class(RegisterModel, "RegisterModel")