support step DependsOn

Author: jerrypeng7773

src/sagemaker/workflow/_utils.py

@@ -59,6 +59,7 @@ def __init__(
         entry_point: str,
         source_dir: str = None,
         dependencies: List = None,
+        depends_on: List[str] = None,
     ):
         """Constructs a TrainingStep, given an `EstimatorBase` instance.
 
@@ -102,7 +103,9 @@ def __init__(
         inputs = TrainingInput(self._model_prefix)
 
         # super!
-        super(_RepackModelStep, self).__init__(name=name, estimator=repacker, inputs=inputs)
+        super(_RepackModelStep, self).__init__(
+            name=name, depends_on=depends_on, estimator=repacker, inputs=inputs
+        )
 
     def _prepare_for_repacking(self):
         """Prepares the source for the estimator."""
@@ -221,6 +224,7 @@ def __init__(
         image_uri=None,
         compile_model_family=None,
         description=None,
+        depends_on: List[str] = None,
         **kwargs,
     ):
         """Constructor of a register model step.
@@ -248,9 +252,10 @@ def __init__(
             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
             **kwargs: additional arguments to `create_model`.
         """
-        super(_RegisterModelStep, self).__init__(name, StepTypeEnum.REGISTER_MODEL)
+        super(_RegisterModelStep, self).__init__(name, StepTypeEnum.REGISTER_MODEL, depends_on)
         self.estimator = estimator
         self.model_data = model_data
         self.content_types = content_types

src/sagemaker/workflow/condition_step.py

@@ -40,6 +40,7 @@ class ConditionStep(Step):
     def __init__(
         self,
         name: str,
+        depends_on: List[str] = None,
         conditions: List[Condition] = None,
         if_steps: List[Union[Step, StepCollection]] = None,
         else_steps: List[Union[Step, StepCollection]] = None,
@@ -60,7 +61,7 @@ def __init__(
                 and `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)
+        super(ConditionStep, self).__init__(name, StepTypeEnum.CONDITION, depends_on)
         self.conditions = conditions or []
         self.if_steps = if_steps or []
         self.else_steps = else_steps or []

src/sagemaker/workflow/step_collections.py

@@ -60,6 +60,7 @@ def __init__(
         response_types,
         inference_instances,
         transform_instances,
+        depends_on: List[str] = None,
         model_package_group_name=None,
         model_metrics=None,
         approval_status=None,
@@ -80,6 +81,7 @@ 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
             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).
@@ -94,12 +96,15 @@ def __init__(
             **kwargs: additional arguments to `create_model`.
         """
         steps: List[Step] = []
+        repack_model = False
         if "entry_point" in kwargs:
+            repack_model = True
             entry_point = kwargs["entry_point"]
             source_dir = kwargs.get("source_dir")
             dependencies = kwargs.get("dependencies")
             repack_model_step = _RepackModelStep(
                 name=f"{name}RepackModel",
+                depends_on=depends_on,
                 estimator=estimator,
                 model_data=model_data,
                 entry_point=entry_point,
@@ -130,6 +135,9 @@ def __init__(
             description=description,
             **kwargs,
         )
+        if not repack_model:
+            register_model_step.add_depends_on(depends_on)
+
         steps.append(register_model_step)
         self.steps = steps
 
@@ -160,6 +168,7 @@ def __init__(
         max_payload=None,
         tags=None,
         volume_kms_key=None,
+        depends_on: List[str] = None,
         **kwargs,
     ):
         """Construct steps required for a Transformer step collection:
@@ -196,6 +205,7 @@ 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
         """
         steps = []
         if "entry_point" in kwargs:
@@ -204,6 +214,7 @@ def __init__(
             dependencies = kwargs.get("dependencies")
             repack_model_step = _RepackModelStep(
                 name=f"{name}RepackModel",
+                depends_on=depends_on,
                 estimator=estimator,
                 model_data=model_data,
                 entry_point=entry_point,
@@ -232,6 +243,9 @@ def predict_wrapper(endpoint, session):
             model=model,
             inputs=model_inputs,
         )
+        if "entry_point" not in kwargs and depends_on:
+            # if the CreateModelStep is the first step in the collection
+            model_step.add_depends_on(depends_on)
         steps.append(model_step)
 
         transformer = Transformer(

src/sagemaker/workflow/steps.py

@@ -63,10 +63,12 @@ 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
     """
 
     name: str = attr.ib(factory=str)
     step_type: StepTypeEnum = attr.ib(factory=StepTypeEnum.factory)
+    depends_on: List[str] = attr.ib(default=None)
 
     @property
     @abc.abstractmethod
@@ -80,11 +82,21 @@ def properties(self):
 
     def to_request(self) -> RequestType:
         """Gets the request structure for workflow service calls."""
-        return {
+        request_dict = {
             "Name": self.name,
             "Type": self.step_type.value,
             "Arguments": self.arguments,
         }
+        if self.depends_on:
+            request_dict["DependsOn"] = self.depends_on
+        return request_dict
+
+    def add_depends_on(self, step_names: List[str]):
+        if not step_names:
+            return
+        if not self.depends_on:
+            self.depends_on = []
+        self.depends_on.extend(step_names)
 
     @property
     def ref(self) -> Dict[str, str]:
@@ -133,6 +145,7 @@ def __init__(
         estimator: EstimatorBase,
         inputs: TrainingInput = None,
         cache_config: CacheConfig = None,
+        depends_on: List[str] = None,
     ):
         """Construct a TrainingStep, given an `EstimatorBase` instance.
 
@@ -144,8 +157,9 @@ def __init__(
             estimator (EstimatorBase): A `sagemaker.estimator.EstimatorBase` instance.
             inputs (TrainingInput): A `sagemaker.inputs.TrainingInput` instance. 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.TrainingStep` depends on
         """
-        super(TrainingStep, self).__init__(name, StepTypeEnum.TRAINING)
+        super(TrainingStep, self).__init__(name, StepTypeEnum.TRAINING, depends_on)
         self.estimator = estimator
         self.inputs = inputs
         self._properties = Properties(
@@ -188,10 +202,7 @@ class CreateModelStep(Step):
     """CreateModel step for workflow."""
 
     def __init__(
-        self,
-        name: str,
-        model: Model,
-        inputs: CreateModelInput,
+        self, name: str, model: Model, inputs: CreateModelInput, depends_on: List[str] = None
     ):
         """Construct a CreateModelStep, given an `sagemaker.model.Model` instance.
 
@@ -203,8 +214,9 @@ 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
         """
-        super(CreateModelStep, self).__init__(name, StepTypeEnum.CREATE_MODEL)
+        super(CreateModelStep, self).__init__(name, StepTypeEnum.CREATE_MODEL, depends_on)
         self.model = model
         self.inputs = inputs or CreateModelInput()
 
@@ -247,6 +259,7 @@ def __init__(
         transformer: Transformer,
         inputs: TransformInput,
         cache_config: CacheConfig = None,
+        depends_on: List[str] = None,
     ):
         """Constructs a TransformStep, given an `Transformer` instance.
 
@@ -258,8 +271,9 @@ 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
         """
-        super(TransformStep, self).__init__(name, StepTypeEnum.TRANSFORM)
+        super(TransformStep, self).__init__(name, StepTypeEnum.TRANSFORM, depends_on)
         self.transformer = transformer
         self.inputs = inputs
         self.cache_config = cache_config
@@ -320,6 +334,7 @@ def __init__(
         code: str = None,
         property_files: List[PropertyFile] = None,
         cache_config: CacheConfig = None,
+        depends_on: List[str] = None,
     ):
         """Construct a ProcessingStep, given a `Processor` instance.
 
@@ -340,8 +355,9 @@ 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
         """
-        super(ProcessingStep, self).__init__(name, StepTypeEnum.PROCESSING)
+        super(ProcessingStep, self).__init__(name, StepTypeEnum.PROCESSING, depends_on)
         self.processor = processor
         self.inputs = inputs
         self.outputs = outputs

tests/integ/test_workflow.py

@@ -28,7 +28,7 @@
     rule_configs,
 )
 from datetime import datetime
-from sagemaker.inputs import CreateModelInput, TrainingInput
+from sagemaker.inputs import CreateModelInput, TrainingInput, TransformInput
 from sagemaker.model import Model
 from sagemaker.processing import ProcessingInput, ProcessingOutput
 from sagemaker.pytorch.estimator import PyTorch
@@ -969,3 +969,110 @@ def test_training_job_with_debugger_and_profiler(
             pipeline.delete()
         except Exception:
             pass
+
+
+def test_two_processing_job_depends_on(
+    sagemaker_session,
+    role,
+    pipeline_name,
+    region_name,
+    cpu_instance_type,
+):
+    instance_count = ParameterInteger(name="InstanceCount", default_value=2)
+    script_path = os.path.join(DATA_DIR, "dummy_script.py")
+
+    pyspark_processor = PySparkProcessor(
+        base_job_name="sm-spark",
+        framework_version="2.4",
+        role=role,
+        instance_count=instance_count,
+        instance_type=cpu_instance_type,
+        max_runtime_in_seconds=1200,
+        sagemaker_session=sagemaker_session,
+    )
+
+    spark_run_args = pyspark_processor.get_run_args(
+        submit_app=script_path,
+        arguments=[
+            "--s3_input_bucket",
+            sagemaker_session.default_bucket(),
+            "--s3_input_key_prefix",
+            "spark-input",
+            "--s3_output_bucket",
+            sagemaker_session.default_bucket(),
+            "--s3_output_key_prefix",
+            "spark-output",
+        ],
+    )
+
+    step_pyspark_1 = ProcessingStep(
+        name="pyspark-process-1",
+        processor=pyspark_processor,
+        inputs=spark_run_args.inputs,
+        outputs=spark_run_args.outputs,
+        job_arguments=spark_run_args.arguments,
+        code=spark_run_args.code,
+    )
+
+    step_pyspark_2 = ProcessingStep(
+        name="pyspark-process-2",
+        depends_on=[step_pyspark_1.name],
+        processor=pyspark_processor,
+        inputs=spark_run_args.inputs,
+        outputs=spark_run_args.outputs,
+        job_arguments=spark_run_args.arguments,
+        code=spark_run_args.code,
+    )
+
+    pipeline = Pipeline(
+        name=pipeline_name,
+        parameters=[instance_count],
+        steps=[step_pyspark_1, step_pyspark_2],
+        sagemaker_session=sagemaker_session,
+    )
+
+    try:
+        response = pipeline.create(role)
+        create_arn = response["PipelineArn"]
+        assert re.match(
+            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
+            create_arn,
+        )
+
+        pipeline.parameters = [ParameterInteger(name="InstanceCount", default_value=1)]
+        response = pipeline.update(role)
+        update_arn = response["PipelineArn"]
+        assert re.match(
+            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}",
+            update_arn,
+        )
+
+        execution = pipeline.start(parameters={})
+        assert re.match(
+            fr"arn:aws:sagemaker:{region_name}:\d{{12}}:pipeline/{pipeline_name}/execution/",
+            execution.arn,
+        )
+
+        response = execution.describe()
+        assert response["PipelineArn"] == create_arn
+
+        try:
+            execution.wait(delay=60)
+        except WaiterError:
+            pass
+
+        execution_steps = execution.list_steps()
+        assert len(execution_steps) == 2
+        time_stamp = {}
+        for execution_step in execution_steps:
+            name = execution_step["StepName"]
+            if name == "pyspark-process-1":
+                time_stamp[name] = execution_step["EndTime"]
+            else:
+                time_stamp[name] = execution_step["StartTime"]
+        assert time_stamp["pyspark-process-1"] < time_stamp["pyspark-process-2"]
+    finally:
+        try:
+            pipeline.delete()
+        except Exception:
+            pass

tests/unit/sagemaker/workflow/test_condition_step.py

@@ -43,11 +43,17 @@ def test_condition_step():
     step1 = CustomStep("MyStep1")
     step2 = CustomStep("MyStep2")
     cond_step = ConditionStep(
-        name="MyConditionStep", conditions=[cond], if_steps=[step1], else_steps=[step2]
+        name="MyConditionStep",
+        depends_on=["TestStep"],
+        conditions=[cond],
+        if_steps=[step1],
+        else_steps=[step2],
     )
+    cond_step.add_depends_on(["SecondTestStep"])
     assert cond_step.to_request() == {
         "Name": "MyConditionStep",
         "Type": "Condition",
+        "DependsOn": ["TestStep", "SecondTestStep"],
         "Arguments": {
             "Conditions": [
                 {