diff --git a/sagemaker-mlops/src/sagemaker/mlops/local/pipeline.py b/sagemaker-mlops/src/sagemaker/mlops/local/pipeline.py index 9172659447..438914c9e9 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/local/pipeline.py +++ b/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) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py index d438f4f1db..4f50c1be74 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py +++ b/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", diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/_steps_compiler.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/_steps_compiler.py index 69caf025b4..673cad8c1b 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/_steps_compiler.py +++ b/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 diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/step_dependency.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/step_dependency.py new file mode 100644 index 0000000000..70a297115e --- /dev/null +++ b/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} diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py index 333f32616c..b5b45e77a4 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py +++ b/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": , "Trigger": } + step = step["Step"] if isinstance(step, Step): dependencies.add(step.name) elif isinstance(step, StepCollection): diff --git a/sagemaker-mlops/tests/integ/workflow/test_run_conditions.py b/sagemaker-mlops/tests/integ/workflow/test_run_conditions.py new file mode 100644 index 0000000000..a0f1d1d3a7 --- /dev/null +++ b/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) diff --git a/sagemaker-mlops/tests/unit/workflow/test_step_dependency.py b/sagemaker-mlops/tests/unit/workflow/test_step_dependency.py new file mode 100644 index 0000000000..f5e5e9405c --- /dev/null +++ b/sagemaker-mlops/tests/unit/workflow/test_step_dependency.py @@ -0,0 +1,231 @@ +# 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. +from __future__ import absolute_import + +import json +from unittest.mock import Mock + +import pytest + +from sagemaker.core.workflow.conditions import ConditionEquals +from sagemaker.mlops.local.pipeline import LocalPipelineExecutor +from sagemaker.mlops.workflow import StepDependency, TriggerCondition +from sagemaker.mlops.workflow.condition_step import ConditionStep +from sagemaker.mlops.workflow.pipeline import Pipeline, PipelineGraph +from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + + +class CustomStep(Step): + def __init__(self, name, depends_on=None): + super(CustomStep, self).__init__( + name=name, step_type=StepTypeEnum.TRAINING, depends_on=depends_on + ) + + @property + def arguments(self): + return {} + + @property + def properties(self): + return None + + +@pytest.fixture +def session(): + session = Mock() + session.local_mode = False + session.sagemaker_config = {} + return session + + +def _depends_on_by_step(steps, session): + definition = json.loads(Pipeline(name="P", steps=steps, sagemaker_session=session).definition()) + return {step["Name"]: step.get("DependsOn") for step in definition["Steps"]} + + +def test_to_request_uses_plain_name_for_completed(): + assert StepDependency("A").to_request() == "A" + assert StepDependency("A", TriggerCondition.COMPLETED).to_request() == "A" + + +def test_to_request_uses_object_form_for_other_triggers(): + assert StepDependency("A", TriggerCondition.FAILED).to_request() == { + "Step": "A", + "Trigger": "Failed", + } + assert StepDependency("A", TriggerCondition.STOPPED).to_request()["Trigger"] == "Stopped" + assert StepDependency("A", TriggerCondition.ANY).to_request()["Trigger"] == "Any" + + +def test_step_name_from_step_object(): + assert StepDependency(CustomStep("A"), TriggerCondition.ANY).step_name == "A" + + +def test_rejects_non_enum_trigger(): + with pytest.raises(TypeError): + StepDependency("A", "Failed") + + +def test_completed_dependency_produces_same_definition_as_plain_entry(session): + a = CustomStep("A") + plain = Pipeline( + name="P", steps=[a, CustomStep("B", depends_on=[a, "A"])], sagemaker_session=session + ).definition() + + a = CustomStep("A") + wrapped = Pipeline( + name="P", + steps=[a, CustomStep("B", depends_on=[StepDependency(a), StepDependency("A")])], + sagemaker_session=session, + ).definition() + + assert wrapped == plain + + +def test_definition_mixes_plain_and_object_entries(session): + a = CustomStep("A") + b = CustomStep("B", depends_on=[StepDependency(a, TriggerCondition.FAILED)]) + c = CustomStep("C", depends_on=["A", StepDependency(b, TriggerCondition.ANY)]) + + assert _depends_on_by_step([a, b, c], session) == { + "A": None, + "B": [{"Step": "A", "Trigger": "Failed"}], + "C": ["A", {"Step": "B", "Trigger": "Any"}], + } + + +def test_upstream_step_is_compiled_when_only_referenced_by_step_dependency(session): + a = CustomStep("A") + b = CustomStep("B", depends_on=[StepDependency(a, TriggerCondition.FAILED)]) + + assert set(_depends_on_by_step([b], session)) == {"A", "B"} + + +def test_condition_step_depends_on_object_form(session): + a = CustomStep("A") + condition = ConditionStep( + name="Check", + conditions=[ConditionEquals(left=1, right=1)], + if_steps=[CustomStep("IfStep")], + else_steps=[], + depends_on=[StepDependency(a, TriggerCondition.STOPPED)], + ) + + assert _depends_on_by_step([a, condition], session)["Check"] == [ + {"Step": "A", "Trigger": "Stopped"} + ] + + +def test_same_trigger_listed_twice_is_deduplicated(session): + a = CustomStep("A") + b = CustomStep( + "B", + depends_on=[ + StepDependency(a, TriggerCondition.FAILED), + StepDependency("A", TriggerCondition.FAILED), + ], + ) + + assert _depends_on_by_step([a, b], session)["B"] == [{"Step": "A", "Trigger": "Failed"}] + + +def test_conflicting_triggers_raise(session): + a = CustomStep("A") + b = CustomStep( + "B", + depends_on=[ + StepDependency(a, TriggerCondition.FAILED), + StepDependency(a, TriggerCondition.ANY), + ], + ) + + with pytest.raises(ValueError, match="conflicting triggers"): + Pipeline(name="P", steps=[a, b], sagemaker_session=session).definition() + + +FAILED_A = {"Step": "A", "Trigger": "Failed"} + + +@pytest.mark.parametrize( + "depends_on, expected", + [ + (lambda a: ["A"], ["A"]), + (lambda a: [a], ["A"]), + (lambda a: [a, "A"], ["A"]), + (lambda a: [StepDependency(a)], ["A"]), + (lambda a: [StepDependency(a, TriggerCondition.COMPLETED)], ["A"]), + (lambda a: [StepDependency(a), "A"], ["A"]), + (lambda a: [StepDependency(a, TriggerCondition.FAILED)], [FAILED_A]), + (lambda a: [StepDependency(a, TriggerCondition.ANY)], [{"Step": "A", "Trigger": "Any"}]), + ( + lambda a: [ + StepDependency(a, TriggerCondition.FAILED), + StepDependency("A", TriggerCondition.FAILED), + ], + [FAILED_A], + ), + ], +) +def test_depends_on_combinations(session, depends_on, expected): + a = CustomStep("A") + b = CustomStep("B", depends_on=depends_on(a)) + + assert _depends_on_by_step([a, b], session)["B"] == expected + + +@pytest.mark.parametrize( + "depends_on", + [ + lambda a: [ + StepDependency(a, TriggerCondition.COMPLETED), + StepDependency(a, TriggerCondition.FAILED), + ], + lambda a: ["A", StepDependency(a, TriggerCondition.FAILED)], + lambda a: [a, StepDependency(a, TriggerCondition.FAILED)], + lambda a: [StepDependency(a, TriggerCondition.FAILED), "A"], + lambda a: [StepDependency("A", TriggerCondition.ANY), a], + ], +) +def test_plain_entry_conflicts_with_non_completed_trigger(session, depends_on): + a = CustomStep("A") + b = CustomStep("B", depends_on=depends_on(a)) + + with pytest.raises(ValueError, match="conflicting triggers: (Completed|Failed|Any)"): + Pipeline(name="P", steps=[a, b], sagemaker_session=session).definition() + + +def test_pipeline_graph_keeps_object_form_edges(session): + a = CustomStep("A") + b = CustomStep("B", depends_on=[StepDependency(a, TriggerCondition.ANY)]) + pipeline = Pipeline(name="P", steps=[a, b], sagemaker_session=session) + + assert PipelineGraph.from_pipeline(pipeline).adjacency_list["A"] == ["B"] + + +def test_local_mode_rejects_run_conditions(session): + a = CustomStep("A") + b = CustomStep("B", depends_on=[StepDependency(a, TriggerCondition.FAILED)]) + execution = Mock() + execution.pipeline = Pipeline(name="P", steps=[a, b], sagemaker_session=session) + + with pytest.raises(ValueError, match="not supported in local mode"): + LocalPipelineExecutor(execution, Mock()) + + +def test_local_mode_allows_completed_step_dependency(session): + a = CustomStep("A") + b = CustomStep("B", depends_on=[StepDependency(a)]) + execution = Mock() + execution.pipeline = Pipeline(name="P", steps=[a, b], sagemaker_session=session) + + LocalPipelineExecutor(execution, Mock())