feat(mlops): Add run conditions to step dependencies

sagemaker-mlops/src/sagemaker/mlops/local/pipeline.py

@@ -61,6 +61,14 @@ def __init__(self, execution, sagemaker_session):
         self.sagemaker_session = sagemaker_session
         self.execution = execution
         self.pipeline_dag = PipelineGraph.from_pipeline(self.execution.pipeline)
+        for step in self.pipeline_dag.step_map.values():
+            # compiled run-condition entries are dicts; this executor stops on the first failure
+            dependencies = step.depends_on if isinstance(step.depends_on, list) else []
+            if any(isinstance(dependency, dict) for dependency in dependencies):
+                raise ValueError(
+                    f"Step {step.name} uses a run condition (StepDependency with a trigger other "
+                    "than Completed), which is not supported in local mode."
+                )
         self.local_sagemaker_client = self.sagemaker_session.sagemaker_client
         self._blocked_steps = set()
         self._step_executor_factory = _StepExecutorFactory(self)

sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py

@@ -65,6 +65,9 @@
 # Step collections
 from sagemaker.mlops.workflow.step_collections import StepCollection
 
+# Run conditions
+from sagemaker.mlops.workflow.step_dependency import StepDependency, TriggerCondition
+
 # Retry policies
 from sagemaker.mlops.workflow.retry import (
     RetryPolicy,
@@ -124,6 +127,9 @@
     "QualityCheckConfig",
     # Step collections
     "StepCollection",
+    # Run conditions
+    "StepDependency",
+    "TriggerCondition",
     # Retry policies
     "RetryPolicy",
     "StepRetryPolicy",

sagemaker-mlops/src/sagemaker/mlops/workflow/_steps_compiler.py

@@ -23,6 +23,7 @@
 from sagemaker.mlops.workflow.function_step import _FunctionStep
 from sagemaker.mlops.workflow.steps import Step, StepTypeEnum, PropertyFile
 from sagemaker.mlops.workflow.condition_step import ConditionStep
+from sagemaker.mlops.workflow.step_dependency import StepDependency
 
 from sagemaker.core.workflow.step_outputs import get_step
 from sagemaker.core.workflow.pipeline_definition_config import PipelineDefinitionConfig
@@ -223,6 +224,8 @@ def _simplify_step_list(
 
         resolved_step_set = _StepsSet()
         for item in step_list:
+            if isinstance(item, StepDependency):
+                item = item.step
             if isinstance(item, str):
                 step_name = item
                 step = self._input_steps_map.get(item)
@@ -293,6 +296,26 @@ def _get_upstream_steps(self, step: Step, arguments: dict):
 
         return list(depends_on), upstream_step_in_args
 
+    def _to_depends_on_request(self, step: Step, depends_on: List[Step]) -> List[Union[str, dict]]:
+        """Build the `DependsOn` entries, keeping plain names for Completed dependencies."""
+        dependencies = {}
+        for item in step.depends_on or []:
+            # a plain str/Step/StepOutput entry is an implicit Completed dependency
+            if not isinstance(item, StepDependency):
+                item = StepDependency(item)
+            name = self._simplify_step_list([item])[0].name
+            if name in dependencies and dependencies[name].trigger != item.trigger:
+                raise ValueError(
+                    f"Step {step.name} depends on step {name} with conflicting triggers: "
+                    f"{dependencies[name].trigger.value} and {item.trigger.value}."
+                )
+            dependencies[name] = item
+
+        return [
+            dependencies[s.name].to_request(s.name) if s.name in dependencies else s.name
+            for s in depends_on
+        ]
+
     def _set_serialize_output_to_json_flag(self, compiled_steps: List[CompiledStep]):
         """Set the serialize_output_to_json to true in container arguments
 
@@ -332,7 +355,7 @@ def _build_step(self, step: Step) -> CompiledStep:
                 step, request_dict["Arguments"]
             )
             if depends_on:
-                request_dict["DependsOn"] = [s.name for s in depends_on]
+                request_dict["DependsOn"] = self._to_depends_on_request(step, depends_on)
 
             self._push_to_build_queue(depends_on)
             self._push_to_build_queue(upstream_steps_in_args)
@@ -353,7 +376,7 @@ def _build_condition_step(self, condition_step: ConditionStep) -> ConditionStep:
             conditions=list(condition_step.conditions),
             if_steps=self._build_steps(condition_step.if_steps),
             else_steps=self._build_steps(condition_step.else_steps),
-            depends_on=[step.name for step in depends_on],
+            depends_on=self._to_depends_on_request(condition_step, depends_on),
         )
 
         return compiled_condition_step

sagemaker-mlops/src/sagemaker/mlops/workflow/step_dependency.py

@@ -0,0 +1,88 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Run conditions for step dependencies in SageMaker Pipelines."""
+
+from __future__ import absolute_import
+
+from enum import Enum
+from typing import Union, TYPE_CHECKING
+
+from sagemaker.core.workflow.step_outputs import StepOutput, get_step
+
+if TYPE_CHECKING:
+    from sagemaker.mlops.workflow.steps import Step
+
+
+class TriggerCondition(Enum):
+    """Upstream step outcome that allows a dependent step to run."""
+
+    COMPLETED = "Completed"
+    FAILED = "Failed"
+    STOPPED = "Stopped"
+    ANY = "Any"
+
+
+class StepDependency:
+    """A `depends_on` entry that runs the dependent step on a chosen upstream outcome.
+
+    Plain entries in `depends_on` (a step or step name) keep their existing meaning: the
+    dependent step runs only after the upstream step completes. Wrap an entry in
+    `StepDependency` to run the dependent step when the upstream step fails, stops, or
+    reaches any terminal state, for example to send a notification when training fails:
+
+    .. code-block:: python
+
+        notify = LambdaStep(
+            name="NotifyOnFailure",
+            lambda_func=notify_lambda,
+            depends_on=[StepDependency(train, TriggerCondition.FAILED)],
+        )
+
+    A pipeline whose steps run on a failed upstream step ends in the
+    `CompletedWithFailures` status instead of `Failed`. Run conditions are not supported
+    in local mode.
+    """
+
+    def __init__(
+        self,
+        step: Union[str, "Step", StepOutput],
+        trigger: TriggerCondition = TriggerCondition.COMPLETED,
+    ):
+        """Create a step dependency with a run condition.
+
+        Args:
+            step (Union[str, Step, StepOutput]): The upstream `Step`, its name, or a
+                `StepOutput` produced by it.
+            trigger (TriggerCondition): The upstream outcome that lets the dependent step
+                run (default: `TriggerCondition.COMPLETED`).
+        """
+        if not isinstance(trigger, TriggerCondition):
+            raise TypeError(f"trigger must be a TriggerCondition, got {type(trigger).__name__}.")
+        self.step = get_step(step) if isinstance(step, StepOutput) else step
+        self.trigger = trigger
+
+    @property
+    def step_name(self) -> str:
+        """The name of the upstream step."""
+        return self.step if isinstance(self.step, str) else self.step.name
+
+    def to_request(self, step_name: str = None) -> Union[str, dict]:
+        """Get the `DependsOn` entry for this dependency.
+
+        A `COMPLETED` dependency is emitted as a plain step name, so definitions that do not
+        use run conditions are unchanged.
+        """
+        name = step_name or self.step_name
+        if self.trigger == TriggerCondition.COMPLETED:
+            return name
+        return {"Step": name, "Trigger": self.trigger.value}

sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py

@@ -35,6 +35,7 @@
 
 # Orchestration imports (now in mlops)
 from sagemaker.mlops.workflow.retry import RetryPolicy
+from sagemaker.mlops.workflow.step_dependency import StepDependency
 from sagemaker.core.workflow.step_outputs import StepOutput
 from sagemaker.core.workflow.utilities import trim_request_dict
 from sagemaker.core.processing import Processor
@@ -90,7 +91,8 @@ def __init__(
             step_type (StepTypeEnum): The type of the `Step`.
             depends_on (List[Union[str, Step, StepCollection]]): The list of `Step`/`StepCollection`
                 names or `Step` or `StepCollection`, `StepOutput` instances that the current `Step`
-                depends on.
+                depends on. Wrap an entry in `StepDependency` to run the current `Step` when
+                the upstream step fails, stops, or reaches any terminal state.
         """
         self.name = name
         self.display_name = display_name
@@ -189,6 +191,11 @@ def _find_dependencies_in_depends_on_list(
 
         dependencies = set()
         for step in self.depends_on:
+            if isinstance(step, StepDependency):
+                step = step.step
+            elif isinstance(step, dict):
+                # compiled run-condition entry: {"Step": <name>, "Trigger": <trigger>}
+                step = step["Step"]
             if isinstance(step, Step):
                 dependencies.add(step.name)
             elif isinstance(step, StepCollection):

sagemaker-mlops/tests/integ/workflow/test_run_conditions.py

@@ -0,0 +1,93 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Integration test for run conditions (StepDependency triggers).
+
+A Processing step fails on purpose. A step that depends on it with the Failed trigger
+runs, and a step with a plain (Completed) dependency does not.
+"""
+
+from __future__ import absolute_import
+
+import logging
+import time
+import uuid
+
+from sagemaker.core import image_uris
+from sagemaker.core.processing import Processor
+from sagemaker.mlops.workflow import StepDependency, TriggerCondition
+from sagemaker.mlops.workflow.pipeline import Pipeline
+from sagemaker.mlops.workflow.steps import ProcessingStep
+
+
+def _shell_step(name, command, image_uri, role, pipeline_session, depends_on=None):
+    processor = Processor(
+        role=role,
+        image_uri=image_uri,
+        instance_count=1,
+        instance_type="ml.m5.large",
+        entrypoint=["/bin/sh", "-c", command],
+        sagemaker_session=pipeline_session,
+    )
+    return ProcessingStep(name=name, step_args=processor.run(), depends_on=depends_on)
+
+
+def test_failed_trigger_runs_after_upstream_failure(pipeline_session, role, sklearn_latest_version):
+    stamp = uuid.uuid4().hex[:8]
+    image_uri = image_uris.retrieve(
+        framework="sklearn",
+        region=pipeline_session.boto_region_name,
+        version=sklearn_latest_version,
+        py_version="py3",
+        instance_type="ml.m5.large",
+    )
+
+    work = _shell_step("Work", "exit 1", image_uri, role, pipeline_session)
+    on_failure = _shell_step(
+        "OnFailure",
+        "echo recovered",
+        image_uri,
+        role,
+        pipeline_session,
+        depends_on=[StepDependency(work, TriggerCondition.FAILED)],
+    )
+    on_success = _shell_step(
+        "OnSuccess", "echo done", image_uri, role, pipeline_session, depends_on=[work]
+    )
+    pipeline = Pipeline(
+        name=f"integ-run-conditions-{stamp}",
+        steps=[work, on_failure, on_success],
+        sagemaker_session=pipeline_session,
+    )
+
+    try:
+        pipeline.upsert(role_arn=role)
+        execution = pipeline.start()
+
+        status = None
+        deadline = time.time() + 1800
+        while time.time() < deadline:
+            status = execution.describe()["PipelineExecutionStatus"]
+            if status in ("Succeeded", "Failed", "Stopped", "CompletedWithFailures"):
+                break
+            time.sleep(30)
+
+        step_status = {s["StepName"]: s["StepStatus"] for s in execution.list_steps()}
+        assert status == "CompletedWithFailures", f"status={status}, steps={step_status}"
+        assert step_status["Work"] == "Failed"
+        assert step_status["OnFailure"] == "Succeeded"
+        assert "OnSuccess" not in step_status
+    finally:
+        try:
+            pipeline.delete()
+        except Exception as e:
+            logging.warning("Failed to delete pipeline %s: %s", pipeline.name, e)