Accept step object in dependson list

doc/workflows/pipelines/sagemaker.workflow.pipelines.rst

@@ -5,7 +5,6 @@ ConditionStep
 -------------
 
 .. autoclass:: sagemaker.workflow.condition_step.ConditionStep
-
 .. autoclass:: sagemaker.workflow.condition_step.JsonGet
 
 Conditions

src/sagemaker/workflow/_utils.py

@@ -17,8 +17,7 @@
 import shutil
 import tarfile
 import tempfile
-from typing import List
-
+from typing import List, Union
 from sagemaker import image_uris
 from sagemaker.inputs import TrainingInput
 from sagemaker.s3 import (
@@ -61,7 +60,7 @@ def __init__(
         entry_point: str,
         source_dir: str = None,
         dependencies: List = None,
-        depends_on: List[str] = None,
+        depends_on: Union[List[str], List[Step]] = None,
         **kwargs,
     ):
         """Constructs a TrainingStep, given an `EstimatorBase` instance.
@@ -230,7 +229,7 @@ def __init__(
         image_uri=None,
         compile_model_family=None,
         description=None,
-        depends_on: List[str] = None,
+        depends_on: Union[List[str], List[Step]] = None,
         tags=None,
         container_def_list=None,
         **kwargs,
@@ -239,30 +238,35 @@ def __init__(
 
         Args:
             name (str): The name of the training step.
-            step_type (StepTypeEnum): The type of the step with value `StepTypeEnum.Training`.
+            step_type (StepTypeEnum): The type of the step with value
+                `StepTypeEnum.Training`.
             estimator (EstimatorBase): A `sagemaker.estimator.EstimatorBase` instance.
             model_data: the S3 URI to the model data from training.
-            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).
+            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).
             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_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).
-            approval_status (str): Model Approval Status, values can be "Approved", "Rejected",
-                or "PendingManualApproval" (default: "PendingManualApproval").
+            metadata_properties (MetadataProperties): MetadataProperties object
+                (default: None).
+            approval_status (str): Model Approval Status, values can be "Approved",
+                "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).
+            compile_model_family (str): Instance family for compiled model,
+                if specified, a compiled model will be used (default: None).
             description (str): Model Package description (default: None).
-            depends_on (List[str]): A list of step names this `sagemaker.workflow.steps.TrainingStep`
-                depends on
-            container_def_list (list): A list of containers.
+            depends_on (List[str] or List[Step]): A list of step names or instances
+                this step depends on
             **kwargs: additional arguments to `create_model`.
         """
         super(_RegisterModelStep, self).__init__(name, StepTypeEnum.REGISTER_MODEL, depends_on)

src/sagemaker/workflow/callback_step.py

@@ -13,7 +13,7 @@
 """The step definitions for workflow."""
 from __future__ import absolute_import
 
-from typing import List, Dict
+from typing import List, Dict, Union
 from enum import Enum
 
 import attr
@@ -84,7 +84,7 @@ def __init__(
         inputs: dict,
         outputs: List[CallbackOutput],
         cache_config: CacheConfig = None,
-        depends_on: List[str] = None,
+        depends_on: Union[List[str], List[Step]] = None,
     ):
         """Constructs a CallbackStep.
 
@@ -95,8 +95,8 @@ def __init__(
                 in the SQS message body of callback messages.
             outputs (List[CallbackOutput]): Outputs that can be provided when completing a callback.
             cache_config (CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance.
-            depends_on (List[str]): A list of step names this `sagemaker.workflow.steps.TransformStep`
-                depends on
+            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)
         self.sqs_queue_url = sqs_queue_url

src/sagemaker/workflow/condition_step.py

@@ -18,6 +18,12 @@
 import attr
 
 from sagemaker.workflow.conditions import Condition
+from sagemaker.workflow.steps import (
+    Step,
+    StepTypeEnum,
+)
+from sagemaker.workflow.step_collections import StepCollection
+from sagemaker.workflow.utilities import list_to_request
 from sagemaker.workflow.entities import (
     Expression,
     RequestType,
@@ -26,12 +32,6 @@
     Properties,
     PropertyFile,
 )
-from sagemaker.workflow.steps import (
-    Step,
-    StepTypeEnum,
-)
-from sagemaker.workflow.step_collections import StepCollection
-from sagemaker.workflow.utilities import list_to_request
 
 
 class ConditionStep(Step):
@@ -40,7 +40,7 @@ class ConditionStep(Step):
     def __init__(
         self,
         name: str,
-        depends_on: List[str] = None,
+        depends_on: Union[List[str], List[Step]] = None,
         conditions: List[Condition] = None,
         if_steps: List[Union[Step, StepCollection]] = None,
         else_steps: List[Union[Step, StepCollection]] = None,

src/sagemaker/workflow/step_collections.py

@@ -13,7 +13,7 @@
 """The step definitions for workflow."""
 from __future__ import absolute_import
 
-from typing import List
+from typing import List, Union
 
 import attr
 
@@ -61,7 +61,7 @@ def __init__(
         transform_instances,
         estimator: EstimatorBase = None,
         model_data=None,
-        depends_on: List[str] = None,
+        depends_on: Union[List[str], List[Step]] = None,
         model_package_group_name=None,
         model_metrics=None,
         approval_status=None,
@@ -84,8 +84,8 @@ def __init__(
                 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).
-            depends_on (List[str]): The list of step names the first step in the collection
-                depends on
+            depends_on (List[str] or List[Step]): The list of step names or step instances
+                the first step in the collection depends on
             model_package_group_name (str): The Model Package Group name, exclusive to
                 `model_package_name`, using `model_package_group_name` makes the Model Package
                 versioned (default: None).
@@ -226,7 +226,7 @@ def __init__(
         max_payload=None,
         tags=None,
         volume_kms_key=None,
-        depends_on: List[str] = None,
+        depends_on: Union[List[str], List[Step]] = None,
         **kwargs,
     ):
         """Construct steps required for a Transformer step collection:
@@ -263,8 +263,8 @@ def __init__(
                 it will be the format of the batch transform output.
             env (dict): The Environment variables to be set for use during the
                 transform job (default: None).
-            depends_on (List[str]): The list of step names the first step in
-                the collection depends on
+            depends_on (List[str] or List[Step]): The list of step names or step instances
+                the first step in the collection depends on
         """
         steps = []
         if "entry_point" in kwargs:

src/sagemaker/workflow/steps.py

@@ -63,12 +63,13 @@ class Step(Entity):
     Attributes:
         name (str): The name of the step.
         step_type (StepTypeEnum): The type of the step.
-        depends_on (List[str]): The list of step names the current step depends on
+        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)
     step_type: StepTypeEnum = attr.ib(factory=StepTypeEnum.factory)
-    depends_on: List[str] = attr.ib(default=None)
+    depends_on: Union[List[str], List["Step"]] = attr.ib(default=None)
 
     @property
     @abc.abstractmethod
@@ -88,11 +89,13 @@ def to_request(self) -> RequestType:
             "Arguments": self.arguments,
         }
         if self.depends_on:
-            request_dict["DependsOn"] = self.depends_on
+            request_dict["DependsOn"] = self._resolve_depends_on(self.depends_on)
+
         return request_dict
 
-    def add_depends_on(self, step_names: List[str]):
-        """Add step names to the current step depends on list"""
+    def add_depends_on(self, step_names: Union[List[str], List["Step"]]):
+        """Add step names or step instances to the current step depends on list"""
+
         if not step_names:
             return
 
@@ -105,6 +108,19 @@ def ref(self) -> Dict[str, str]:
         """Gets a reference dict for steps"""
         return {"Name": self.name}
 
+    @staticmethod
+    def _resolve_depends_on(depends_on_list: Union[List[str], List["Step"]]):
+        """Resolver the step depends on list"""
+        depends_on = []
+        for step in depends_on_list:
+            if isinstance(step, Step):
+                depends_on.append(step.name)
+            elif isinstance(step, str):
+                depends_on.append(step)
+            else:
+                raise ValueError(f"Invalid input step name: {step}")
+        return depends_on
+
 
 @attr.s
 class CacheConfig:
@@ -147,7 +163,7 @@ def __init__(
         estimator: EstimatorBase,
         inputs: Union[TrainingInput, dict, str, FileSystemInput] = None,
         cache_config: CacheConfig = None,
-        depends_on: List[str] = None,
+        depends_on: Union[List[str], List[Step]] = None,
     ):
         """Construct a TrainingStep, given an `EstimatorBase` instance.
 
@@ -175,8 +191,8 @@ def __init__(
                     the path to the training dataset.
 
             cache_config (CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance.
-            depends_on (List[str]): A list of step names this `sagemaker.workflow.steps.TrainingStep`
-                depends on
+            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)
         self.estimator = estimator
@@ -221,7 +237,11 @@ class CreateModelStep(Step):
     """CreateModel step for workflow."""
 
     def __init__(
-        self, name: str, model: Model, inputs: CreateModelInput, depends_on: List[str] = None
+        self,
+        name: str,
+        model: Model,
+        inputs: CreateModelInput,
+        depends_on: Union[List[str], List[Step]] = None,
     ):
         """Construct a CreateModelStep, given an `sagemaker.model.Model` instance.
 
@@ -233,8 +253,8 @@ def __init__(
             model (Model): A `sagemaker.model.Model` instance.
             inputs (CreateModelInput): A `sagemaker.inputs.CreateModelInput` instance.
                 Defaults to `None`.
-            depends_on (List[str]): A list of step names this `sagemaker.workflow.steps.CreateModelStep`
-                depends on
+            depends_on (List[str] or List[Step]): A list of step names or step instances
+                this `sagemaker.workflow.steps.CreateModelStep` depends on
         """
         super(CreateModelStep, self).__init__(name, StepTypeEnum.CREATE_MODEL, depends_on)
         self.model = model
@@ -279,7 +299,7 @@ def __init__(
         transformer: Transformer,
         inputs: TransformInput,
         cache_config: CacheConfig = None,
-        depends_on: List[str] = None,
+        depends_on: Union[List[str], List[Step]] = None,
     ):
         """Constructs a TransformStep, given an `Transformer` instance.
 
@@ -291,8 +311,8 @@ 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]): A list of step names this `sagemaker.workflow.steps.TransformStep`
-                depends on
+            depends_on (List[str] or List[Step]): A list of step names or step instances
+                this `sagemaker.workflow.steps.TransformStep` depends on
         """
         super(TransformStep, self).__init__(name, StepTypeEnum.TRANSFORM, depends_on)
         self.transformer = transformer
@@ -355,7 +375,7 @@ def __init__(
         code: str = None,
         property_files: List[PropertyFile] = None,
         cache_config: CacheConfig = None,
-        depends_on: List[str] = None,
+        depends_on: Union[List[str], List[Step]] = None,
     ):
         """Construct a ProcessingStep, given a `Processor` instance.
 
@@ -376,8 +396,8 @@ def __init__(
             property_files (List[PropertyFile]): A list of property files that workflow looks
                 for and resolves from the configured processing output list.
             cache_config (CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance.
-            depends_on (List[str]): A list of step names this `sagemaker.workflow.steps.ProcessingStep`
-                depends on
+            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)
         self.processor = processor
@@ -445,7 +465,7 @@ def __init__(
         inputs=None,
         job_arguments: List[str] = None,
         cache_config: CacheConfig = None,
-        depends_on: List[str] = None,
+        depends_on: Union[List[str], List[Step]] = None,
     ):
         """Construct a TuningStep, given a `HyperparameterTuner` instance.
 
@@ -485,8 +505,8 @@ def __init__(
             job_arguments (List[str]): A list of strings to be passed into the processing job.
                 Defaults to `None`.
             cache_config (CacheConfig):  A `sagemaker.workflow.steps.CacheConfig` instance.
-            depends_on (List[str]): A list of step names this `sagemaker.workflow.steps.ProcessingStep`
-                depends on
+            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)
         self.tuner = tuner
@@ -539,7 +559,7 @@ def to_request(self) -> RequestType:
 
         return request_dict
 
-    def get_top_model_s3_uri(self, top_k: int, s3_bucket: str, prefix: str = ""):
+    def get_top_model_s3_uri(self, top_k: int, s3_bucket: str, prefix: str = "") -> Join:
         """Get the model artifact s3 uri from the top performing training jobs.
 
         Args:

tests/integ/test_workflow.py

@@ -1501,7 +1501,7 @@ def test_two_processing_job_depends_on(
 
     step_pyspark_2 = ProcessingStep(
         name="pyspark-process-2",
-        depends_on=[step_pyspark_1.name],
+        depends_on=[step_pyspark_1],
         processor=pyspark_processor,
         inputs=spark_run_args.inputs,
         outputs=spark_run_args.outputs,