feature: support displayName and description for pipeline steps (#2580)

Co-authored-by: Ahsan Khan <[email protected]>
Co-authored-by: Shreya Pandit <[email protected]>

Author: 3 people

src/sagemaker/workflow/_utils.py

@@ -55,6 +55,8 @@ def __init__(
         role,
         model_data: str,
         entry_point: str,
+        display_name: str = None,
+        description: str = None,
         source_dir: str = None,
         dependencies: List = None,
         depends_on: Union[List[str], List[Step]] = None,
@@ -165,7 +167,12 @@ def __init__(
 
         # super!
         super(_RepackModelStep, self).__init__(
-            name=name, depends_on=depends_on, estimator=repacker, inputs=inputs
+            name=name,
+            display_name=display_name,
+            description=description,
+            depends_on=depends_on,
+            estimator=repacker,
+            inputs=inputs,
         )
 
     def _prepare_for_repacking(self):
@@ -285,6 +292,7 @@ def __init__(
         approval_status="PendingManualApproval",
         image_uri=None,
         compile_model_family=None,
+        display_name: str = None,
         description=None,
         depends_on: Union[List[str], List[Step]] = None,
         tags=None,
@@ -326,7 +334,9 @@ def __init__(
                 this step depends on
             **kwargs: additional arguments to `create_model`.
         """
-        super(_RegisterModelStep, self).__init__(name, StepTypeEnum.REGISTER_MODEL, depends_on)
+        super(_RegisterModelStep, self).__init__(
+            name, display_name, description, StepTypeEnum.REGISTER_MODEL, depends_on
+        )
         self.estimator = estimator
         self.model_data = model_data
         self.content_types = content_types

src/sagemaker/workflow/callback_step.py

@@ -83,6 +83,8 @@ def __init__(
         sqs_queue_url: str,
         inputs: dict,
         outputs: List[CallbackOutput],
+        display_name: str = None,
+        description: str = None,
         cache_config: CacheConfig = None,
         depends_on: Union[List[str], List[Step]] = None,
     ):
@@ -94,11 +96,15 @@ def __init__(
             inputs (dict): Input arguments that will be provided
                 in the SQS message body of callback messages.
             outputs (List[CallbackOutput]): Outputs that can be provided when completing a callback.
+            display_name (str): The display name of the callback step.
+            description (str): The description of the callback step.
             cache_config (CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance.
             depends_on (List[str] or List[Step]): A list of step names or step instances
                 this `sagemaker.workflow.steps.CallbackStep` depends on
         """
-        super(CallbackStep, self).__init__(name, StepTypeEnum.CALLBACK, depends_on)
+        super(CallbackStep, self).__init__(
+            name, display_name, description, StepTypeEnum.CALLBACK, depends_on
+        )
         self.sqs_queue_url = sqs_queue_url
         self.outputs = outputs
         self.cache_config = cache_config

src/sagemaker/workflow/condition_step.py

@@ -42,6 +42,8 @@ def __init__(
         self,
         name: str,
         depends_on: Union[List[str], List[Step]] = None,
+        display_name: str = None,
+        description: str = None,
         conditions: List[Condition] = None,
         if_steps: List[Union[Step, StepCollection]] = None,
         else_steps: List[Union[Step, StepCollection]] = None,
@@ -53,6 +55,9 @@ def __init__(
         execution.
 
         Args:
+            name (str): The name of the condition step.
+            display_name (str): The display name of the condition step.
+            description (str): The description of the condition step.
             conditions (List[Condition]): A list of `sagemaker.workflow.conditions.Condition`
                 instances.
             if_steps (List[Union[Step, StepCollection]]): A list of `sagemaker.workflow.steps.Step`
@@ -62,7 +67,9 @@ def __init__(
                 or `sagemaker.workflow.step_collections.StepCollection` instances that are
                 marked as ready for execution if the list of conditions evaluates to False.
         """
-        super(ConditionStep, self).__init__(name, StepTypeEnum.CONDITION, depends_on)
+        super(ConditionStep, self).__init__(
+            name, display_name, description, StepTypeEnum.CONDITION, depends_on
+        )
         self.conditions = conditions or []
         self.if_steps = if_steps or []
         self.else_steps = else_steps or []

src/sagemaker/workflow/lambda_step.py

@@ -82,6 +82,8 @@ def __init__(
         self,
         name: str,
         lambda_func: Lambda,
+        display_name: str = None,
+        description: str = None,
         inputs: dict = None,
         outputs: List[LambdaOutput] = None,
         cache_config: CacheConfig = None,
@@ -91,6 +93,8 @@ def __init__(
 
         Args:
             name (str): The name of the lambda step.
+            display_name (str): The display name of the Lambda step.
+            description (str): The description of the Lambda step.
             lambda_func (str): An instance of sagemaker.lambda_helper.Lambda.
                 If lambda arn is specified in the instance, LambdaStep just invokes the function,
                 else lambda function will be created while creating the pipeline.
@@ -101,7 +105,9 @@ def __init__(
             depends_on (List[str]): A list of step names this `sagemaker.workflow.steps.LambdaStep`
                 depends on
         """
-        super(LambdaStep, self).__init__(name, StepTypeEnum.LAMBDA, depends_on)
+        super(LambdaStep, self).__init__(
+            name, display_name, description, StepTypeEnum.LAMBDA, depends_on
+        )
         self.lambda_func = lambda_func
         self.outputs = outputs if outputs is not None else []
         self.cache_config = cache_config

src/sagemaker/workflow/step_collections.py

@@ -67,6 +67,7 @@ def __init__(
         approval_status=None,
         image_uri=None,
         compile_model_family=None,
+        display_name=None,
         description=None,
         tags=None,
         model: Union[Model, PipelineModel] = None,
@@ -138,6 +139,8 @@ def __init__(
                 tags=tags,
                 subnets=subnets,
                 security_group_ids=security_group_ids,
+                description=description,
+                display_name=display_name,
                 **kwargs,
             )
             steps.append(repack_model_step)
@@ -179,6 +182,8 @@ def __init__(
                         tags=tags,
                         subnets=subnets,
                         security_group_ids=security_group_ids,
+                        description=description,
+                        display_name=display_name,
                         **kwargs,
                     )
                     steps.append(repack_model_step)
@@ -208,6 +213,7 @@ def __init__(
             image_uri=image_uri,
             compile_model_family=compile_model_family,
             description=description,
+            display_name=display_name,
             tags=tags,
             container_def_list=self.container_def_list,
             **kwargs,
@@ -231,6 +237,8 @@ def __init__(
         instance_count,
         instance_type,
         transform_inputs,
+        description: str = None,
+        display_name: str = None,
         # model arguments
         image_uri=None,
         predictor_cls=None,
@@ -302,6 +310,8 @@ def __init__(
                 tags=tags,
                 subnets=estimator.subnets,
                 security_group_ids=estimator.security_group_ids,
+                description=description,
+                display_name=display_name,
             )
             steps.append(repack_model_step)
             model_data = repack_model_step.properties.ModelArtifacts.S3ModelArtifacts
@@ -324,6 +334,8 @@ def predict_wrapper(endpoint, session):
             name=f"{name}CreateModelStep",
             model=model,
             inputs=model_inputs,
+            description=description,
+            display_name=display_name,
         )
         if "entry_point" not in kwargs and depends_on:
             # if the CreateModelStep is the first step in the collection
@@ -351,6 +363,8 @@ def predict_wrapper(endpoint, session):
             name=f"{name}TransformStep",
             transformer=transformer,
             inputs=transform_inputs,
+            description=description,
+            display_name=display_name,
         )
         steps.append(transform_step)
 

src/sagemaker/workflow/steps.py

@@ -63,12 +63,16 @@ class Step(Entity):
 
     Attributes:
         name (str): The name of the step.
+        display_name (str): The display name of the step.
+        description (str): The description of the step.
         step_type (StepTypeEnum): The type of the step.
         depends_on (List[str] or List[Step]): The list of step names or step
             instances the current step depends on
     """
 
     name: str = attr.ib(factory=str)
+    display_name: str = attr.ib(default=None)
+    description: str = attr.ib(default=None)
     step_type: StepTypeEnum = attr.ib(factory=StepTypeEnum.factory)
     depends_on: Union[List[str], List["Step"]] = attr.ib(default=None)
 
@@ -91,7 +95,10 @@ def to_request(self) -> RequestType:
         }
         if self.depends_on:
             request_dict["DependsOn"] = self._resolve_depends_on(self.depends_on)
-
+        if self.display_name:
+            request_dict["DisplayName"] = self.display_name
+        if self.description:
+            request_dict["Description"] = self.description
         return request_dict
 
     def add_depends_on(self, step_names: Union[List[str], List["Step"]]):
@@ -168,6 +175,8 @@ def __init__(
         self,
         name: str,
         estimator: EstimatorBase,
+        display_name: str = None,
+        description: str = None,
         inputs: Union[TrainingInput, dict, str, FileSystemInput] = None,
         cache_config: CacheConfig = None,
         depends_on: Union[List[str], List[Step]] = None,
@@ -180,6 +189,8 @@ def __init__(
         Args:
             name (str): The name of the training step.
             estimator (EstimatorBase): A `sagemaker.estimator.EstimatorBase` instance.
+            display_name (str): The display name of the training step.
+            description (str): The description of the training step.
             inputs (Union[str, dict, TrainingInput, FileSystemInput]): Information
                 about the training data. This can be one of three types:
 
@@ -200,7 +211,9 @@ def __init__(
             depends_on (List[str] or List[Step]): A list of step names or step instances
                 this `sagemaker.workflow.steps.TrainingStep` depends on
         """
-        super(TrainingStep, self).__init__(name, StepTypeEnum.TRAINING, depends_on)
+        super(TrainingStep, self).__init__(
+            name, display_name, description, StepTypeEnum.TRAINING, depends_on
+        )
         self.estimator = estimator
         self.inputs = inputs
         self._properties = Properties(
@@ -248,6 +261,8 @@ def __init__(
         model: Model,
         inputs: CreateModelInput,
         depends_on: Union[List[str], List[Step]] = None,
+        display_name: str = None,
+        description: str = None,
     ):
         """Construct a CreateModelStep, given an `sagemaker.model.Model` instance.
 
@@ -261,8 +276,12 @@ def __init__(
                 Defaults to `None`.
             depends_on (List[str] or List[Step]): A list of step names or step instances
                 this `sagemaker.workflow.steps.CreateModelStep` depends on
+            display_name (str): The display name of the CreateModel step.
+            description (str): The description of the CreateModel step.
         """
-        super(CreateModelStep, self).__init__(name, StepTypeEnum.CREATE_MODEL, depends_on)
+        super(CreateModelStep, self).__init__(
+            name, display_name, description, StepTypeEnum.CREATE_MODEL, depends_on
+        )
         self.model = model
         self.inputs = inputs or CreateModelInput()
 
@@ -304,6 +323,8 @@ def __init__(
         name: str,
         transformer: Transformer,
         inputs: TransformInput,
+        display_name: str = None,
+        description: str = None,
         cache_config: CacheConfig = None,
         depends_on: Union[List[str], List[Step]] = None,
     ):
@@ -317,10 +338,14 @@ def __init__(
             transformer (Transformer): A `sagemaker.transformer.Transformer` instance.
             inputs (TransformInput): A `sagemaker.inputs.TransformInput` instance.
             cache_config (CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance.
-            depends_on (List[str] or List[Step]): A list of step names or step instances
-                this `sagemaker.workflow.steps.TransformStep` depends on
+            display_name (str): The display name of the transform step.
+            description (str): The description of the transform step.
+            depends_on (List[str]): A list of step names this `sagemaker.workflow.steps.TransformStep`
+                depends on
         """
-        super(TransformStep, self).__init__(name, StepTypeEnum.TRANSFORM, depends_on)
+        super(TransformStep, self).__init__(
+            name, display_name, description, StepTypeEnum.TRANSFORM, depends_on
+        )
         self.transformer = transformer
         self.inputs = inputs
         self.cache_config = cache_config
@@ -375,6 +400,8 @@ def __init__(
         self,
         name: str,
         processor: Processor,
+        display_name: str = None,
+        description: str = None,
         inputs: List[ProcessingInput] = None,
         outputs: List[ProcessingOutput] = None,
         job_arguments: List[str] = None,
@@ -391,6 +418,8 @@ def __init__(
         Args:
             name (str): The name of the processing step.
             processor (Processor): A `sagemaker.processing.Processor` instance.
+            display_name (str): The display name of the processing step.
+            description (str): The description of the processing step.
             inputs (List[ProcessingInput]): A list of `sagemaker.processing.ProcessorInput`
                 instances. Defaults to `None`.
             outputs (List[ProcessingOutput]): A list of `sagemaker.processing.ProcessorOutput`
@@ -405,7 +434,9 @@ def __init__(
             depends_on (List[str] or List[Step]): A list of step names or step instance
                 this `sagemaker.workflow.steps.ProcessingStep` depends on
         """
-        super(ProcessingStep, self).__init__(name, StepTypeEnum.PROCESSING, depends_on)
+        super(ProcessingStep, self).__init__(
+            name, display_name, description, StepTypeEnum.PROCESSING, depends_on
+        )
         self.processor = processor
         self.inputs = inputs
         self.outputs = outputs
@@ -468,6 +499,8 @@ def __init__(
         self,
         name: str,
         tuner: HyperparameterTuner,
+        display_name: str = None,
+        description: str = None,
         inputs=None,
         job_arguments: List[str] = None,
         cache_config: CacheConfig = None,
@@ -481,6 +514,8 @@ def __init__(
         Args:
             name (str): The name of the tuning step.
             tuner (HyperparameterTuner): A `sagemaker.tuner.HyperparameterTuner` instance.
+            display_name (str): The display name of the tuning step.
+            description (str): The description of the tuning step.
             inputs: Information about the training data. Please refer to the
                 ``fit()`` method of the associated estimator, as this can take
                 any of the following forms:
@@ -514,7 +549,9 @@ def __init__(
             depends_on (List[str] or List[Step]): A list of step names or step instance
                 this `sagemaker.workflow.steps.ProcessingStep` depends on
         """
-        super(TuningStep, self).__init__(name, StepTypeEnum.TUNING, depends_on)
+        super(TuningStep, self).__init__(
+            name, display_name, description, StepTypeEnum.TUNING, depends_on
+        )
         self.tuner = tuner
         self.inputs = inputs
         self.job_arguments = job_arguments