doc: improvements on the docstring of ModelStep (aws#3137)

Co-authored-by: Zhankui Lu <[email protected]>

Author: navaj0

src/sagemaker/model.py

@@ -341,7 +341,9 @@ def register(
                 "MACHINE_LEARNING" (default: None).
 
         Returns:
-            A `sagemaker.model.ModelPackage` instance.
+            A `sagemaker.model.ModelPackage` instance or pipeline step arguments
+            in case the Model instance is built with
+            :class:`~sagemaker.workflow.pipeline_context.PipelineSession`
         """
         if self.model_data is None:
             raise ValueError("SageMaker Model Package cannot be created without model data.")
@@ -398,15 +400,22 @@ def create(
                 attach to an endpoint for model loading and inference, for
                 example, 'ml.eia1.medium'. If not specified, no Elastic
                 Inference accelerator will be attached to the endpoint (default: None).
-            serverless_inference_config (sagemaker.serverless.ServerlessInferenceConfig):
+            serverless_inference_config (ServerlessInferenceConfig):
                 Specifies configuration related to serverless endpoint. Instance type is
                 not provided in serverless inference. So this is used to find image URIs
                 (default: None).
             tags (List[Dict[str, str]]): The list of tags to add to
-                the model (default: None). Example: >>> tags = [{'Key': 'tagname', 'Value':
-                'tagvalue'}] For more information about tags, see
-                https://boto3.amazonaws.com/v1/documentation
-                /api/latest/reference/services/sagemaker.html#SageMaker.Client.add_tags
+                the model (default: None). Example::
+
+                    tags = [{'Key': 'tagname', 'Value':'tagvalue'}]
+
+                For more information about tags, see
+                `boto3 documentation <https://boto3.amazonaws.com/v1/documentation/\
+api/latest/reference/services/sagemaker.html#SageMaker.Client.add_tags>`_
+
+        Returns:
+            None or pipeline step arguments in case the Model instance is built with
+            :class:`~sagemaker.workflow.pipeline_context.PipelineSession`
         """
         # TODO: we should replace _create_sagemaker_model() with create()
         self._create_sagemaker_model(

src/sagemaker/workflow/README.rst

@@ -50,10 +50,17 @@ def __init__(
         Args:
             name (str): The name of the `ModelStep`. A name is required and must be
                 unique within a pipeline.
-            step_args (_ModelStepArguments): The arguments for the `ModelStep` definition.
-            depends_on (List[Union[str, Step, StepCollection]]): A list of `Step`/`StepCollection`
-                names or `Step` instances or `StepCollection` instances that the first step,
-                in this `ModelStep` collection, depends on.
+            step_args (_ModelStepArguments): The arguments for the `ModelStep` definition,
+                generated by invoking the :func:`~sagemaker.model.Model.register` or
+                :func:`~sagemaker.model.Model.create`
+                under the :class:`~sagemaker.workflow.pipeline_context.PipelineSession`. Example::
+
+                     model = Model(sagemaker_session=PipelineSession())
+                     model_step = ModelStep(step_args=model.register())
+
+            depends_on (List[Union[str, Step, StepCollection]]):
+                A list of `Step` or `StepCollection`
+                names or `Step` instances or `StepCollection` instances that it 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` (default: None).

src/sagemaker/workflow/model_step.py

@@ -18,7 +18,6 @@
 
 import attr
 
-from sagemaker.deprecations import deprecated_class
 from sagemaker.estimator import EstimatorBase
 from sagemaker.model import Model
 from sagemaker import PipelineModel
@@ -262,18 +261,14 @@ def __init__(
         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(). For more, see: "
+                "Please use the ModelStep instead. For more, see: "
                 "https://sagemaker.readthedocs.io/en/stable/"
                 "amazon_sagemaker_model_building_pipeline.html#model-step"
             ),
             DeprecationWarning,
         )
 
 
-RegisterModel = deprecated_class(RegisterModel, "RegisterModel")
-
-
 class EstimatorTransformer(StepCollection):
     """Creates a Transformer step collection for workflow."""