Skip to content

Commit ce3ac25

Browse files
nayan3107nayan3107
andauthored
feat(mlops): Add run conditions to step dependencies (#6475)
* feat(mlops): Add run conditions to step dependencies Add StepDependency and TriggerCondition so a step can run when an upstream step fails, stops, or reaches any terminal state. Completed dependencies are still emitted as plain step names, so existing pipeline definitions are unchanged. Local mode rejects non-Completed triggers because it stops on the first failed step. Integ test passed against a production region. --- X-AI-Prompt: Add backward compatible Run Conditions support to the PySDK and raise a PR against master X-AI-Tool: kiro-cli * fix(mlops): Treat plain depends_on entries as Completed A plain step or step name listed next to a StepDependency on the same step with a non-Completed trigger was ignored by the conflict check, so the Completed dependency was silently dropped. Plain entries now count as Completed and such combinations raise. --- X-AI-Prompt: Address jkasiraj review on PR 6475 about mixed depends_on conflict cases X-AI-Tool: kiro-cli --------- Co-authored-by: nayan3107 <nayancho@amazon.com>
1 parent 8e39e3b commit ce3ac25

7 files changed

Lines changed: 459 additions & 3 deletions

File tree

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,14 @@ def __init__(self, execution, sagemaker_session):
6161
self.sagemaker_session = sagemaker_session
6262
self.execution = execution
6363
self.pipeline_dag = PipelineGraph.from_pipeline(self.execution.pipeline)
64+
for step in self.pipeline_dag.step_map.values():
65+
# compiled run-condition entries are dicts; this executor stops on the first failure
66+
dependencies = step.depends_on if isinstance(step.depends_on, list) else []
67+
if any(isinstance(dependency, dict) for dependency in dependencies):
68+
raise ValueError(
69+
f"Step {step.name} uses a run condition (StepDependency with a trigger other "
70+
"than Completed), which is not supported in local mode."
71+
)
6472
self.local_sagemaker_client = self.sagemaker_session.sagemaker_client
6573
self._blocked_steps = set()
6674
self._step_executor_factory = _StepExecutorFactory(self)

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,9 @@
6565
# Step collections
6666
from sagemaker.mlops.workflow.step_collections import StepCollection
6767

68+
# Run conditions
69+
from sagemaker.mlops.workflow.step_dependency import StepDependency, TriggerCondition
70+
6871
# Retry policies
6972
from sagemaker.mlops.workflow.retry import (
7073
RetryPolicy,
@@ -124,6 +127,9 @@
124127
"QualityCheckConfig",
125128
# Step collections
126129
"StepCollection",
130+
# Run conditions
131+
"StepDependency",
132+
"TriggerCondition",
127133
# Retry policies
128134
"RetryPolicy",
129135
"StepRetryPolicy",

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

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
from sagemaker.mlops.workflow.function_step import _FunctionStep
2424
from sagemaker.mlops.workflow.steps import Step, StepTypeEnum, PropertyFile
2525
from sagemaker.mlops.workflow.condition_step import ConditionStep
26+
from sagemaker.mlops.workflow.step_dependency import StepDependency
2627

2728
from sagemaker.core.workflow.step_outputs import get_step
2829
from sagemaker.core.workflow.pipeline_definition_config import PipelineDefinitionConfig
@@ -223,6 +224,8 @@ def _simplify_step_list(
223224

224225
resolved_step_set = _StepsSet()
225226
for item in step_list:
227+
if isinstance(item, StepDependency):
228+
item = item.step
226229
if isinstance(item, str):
227230
step_name = item
228231
step = self._input_steps_map.get(item)
@@ -293,6 +296,26 @@ def _get_upstream_steps(self, step: Step, arguments: dict):
293296

294297
return list(depends_on), upstream_step_in_args
295298

299+
def _to_depends_on_request(self, step: Step, depends_on: List[Step]) -> List[Union[str, dict]]:
300+
"""Build the `DependsOn` entries, keeping plain names for Completed dependencies."""
301+
dependencies = {}
302+
for item in step.depends_on or []:
303+
# a plain str/Step/StepOutput entry is an implicit Completed dependency
304+
if not isinstance(item, StepDependency):
305+
item = StepDependency(item)
306+
name = self._simplify_step_list([item])[0].name
307+
if name in dependencies and dependencies[name].trigger != item.trigger:
308+
raise ValueError(
309+
f"Step {step.name} depends on step {name} with conflicting triggers: "
310+
f"{dependencies[name].trigger.value} and {item.trigger.value}."
311+
)
312+
dependencies[name] = item
313+
314+
return [
315+
dependencies[s.name].to_request(s.name) if s.name in dependencies else s.name
316+
for s in depends_on
317+
]
318+
296319
def _set_serialize_output_to_json_flag(self, compiled_steps: List[CompiledStep]):
297320
"""Set the serialize_output_to_json to true in container arguments
298321
@@ -332,7 +355,7 @@ def _build_step(self, step: Step) -> CompiledStep:
332355
step, request_dict["Arguments"]
333356
)
334357
if depends_on:
335-
request_dict["DependsOn"] = [s.name for s in depends_on]
358+
request_dict["DependsOn"] = self._to_depends_on_request(step, depends_on)
336359

337360
self._push_to_build_queue(depends_on)
338361
self._push_to_build_queue(upstream_steps_in_args)
@@ -353,7 +376,7 @@ def _build_condition_step(self, condition_step: ConditionStep) -> ConditionStep:
353376
conditions=list(condition_step.conditions),
354377
if_steps=self._build_steps(condition_step.if_steps),
355378
else_steps=self._build_steps(condition_step.else_steps),
356-
depends_on=[step.name for step in depends_on],
379+
depends_on=self._to_depends_on_request(condition_step, depends_on),
357380
)
358381

359382
return compiled_condition_step
Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
"""Run conditions for step dependencies in SageMaker Pipelines."""
14+
15+
from __future__ import absolute_import
16+
17+
from enum import Enum
18+
from typing import Union, TYPE_CHECKING
19+
20+
from sagemaker.core.workflow.step_outputs import StepOutput, get_step
21+
22+
if TYPE_CHECKING:
23+
from sagemaker.mlops.workflow.steps import Step
24+
25+
26+
class TriggerCondition(Enum):
27+
"""Upstream step outcome that allows a dependent step to run."""
28+
29+
COMPLETED = "Completed"
30+
FAILED = "Failed"
31+
STOPPED = "Stopped"
32+
ANY = "Any"
33+
34+
35+
class StepDependency:
36+
"""A `depends_on` entry that runs the dependent step on a chosen upstream outcome.
37+
38+
Plain entries in `depends_on` (a step or step name) keep their existing meaning: the
39+
dependent step runs only after the upstream step completes. Wrap an entry in
40+
`StepDependency` to run the dependent step when the upstream step fails, stops, or
41+
reaches any terminal state, for example to send a notification when training fails:
42+
43+
.. code-block:: python
44+
45+
notify = LambdaStep(
46+
name="NotifyOnFailure",
47+
lambda_func=notify_lambda,
48+
depends_on=[StepDependency(train, TriggerCondition.FAILED)],
49+
)
50+
51+
A pipeline whose steps run on a failed upstream step ends in the
52+
`CompletedWithFailures` status instead of `Failed`. Run conditions are not supported
53+
in local mode.
54+
"""
55+
56+
def __init__(
57+
self,
58+
step: Union[str, "Step", StepOutput],
59+
trigger: TriggerCondition = TriggerCondition.COMPLETED,
60+
):
61+
"""Create a step dependency with a run condition.
62+
63+
Args:
64+
step (Union[str, Step, StepOutput]): The upstream `Step`, its name, or a
65+
`StepOutput` produced by it.
66+
trigger (TriggerCondition): The upstream outcome that lets the dependent step
67+
run (default: `TriggerCondition.COMPLETED`).
68+
"""
69+
if not isinstance(trigger, TriggerCondition):
70+
raise TypeError(f"trigger must be a TriggerCondition, got {type(trigger).__name__}.")
71+
self.step = get_step(step) if isinstance(step, StepOutput) else step
72+
self.trigger = trigger
73+
74+
@property
75+
def step_name(self) -> str:
76+
"""The name of the upstream step."""
77+
return self.step if isinstance(self.step, str) else self.step.name
78+
79+
def to_request(self, step_name: str = None) -> Union[str, dict]:
80+
"""Get the `DependsOn` entry for this dependency.
81+
82+
A `COMPLETED` dependency is emitted as a plain step name, so definitions that do not
83+
use run conditions are unchanged.
84+
"""
85+
name = step_name or self.step_name
86+
if self.trigger == TriggerCondition.COMPLETED:
87+
return name
88+
return {"Step": name, "Trigger": self.trigger.value}

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535

3636
# Orchestration imports (now in mlops)
3737
from sagemaker.mlops.workflow.retry import RetryPolicy
38+
from sagemaker.mlops.workflow.step_dependency import StepDependency
3839
from sagemaker.core.workflow.step_outputs import StepOutput
3940
from sagemaker.core.workflow.utilities import trim_request_dict
4041
from sagemaker.core.processing import Processor
@@ -90,7 +91,8 @@ def __init__(
9091
step_type (StepTypeEnum): The type of the `Step`.
9192
depends_on (List[Union[str, Step, StepCollection]]): The list of `Step`/`StepCollection`
9293
names or `Step` or `StepCollection`, `StepOutput` instances that the current `Step`
93-
depends on.
94+
depends on. Wrap an entry in `StepDependency` to run the current `Step` when
95+
the upstream step fails, stops, or reaches any terminal state.
9496
"""
9597
self.name = name
9698
self.display_name = display_name
@@ -189,6 +191,11 @@ def _find_dependencies_in_depends_on_list(
189191

190192
dependencies = set()
191193
for step in self.depends_on:
194+
if isinstance(step, StepDependency):
195+
step = step.step
196+
elif isinstance(step, dict):
197+
# compiled run-condition entry: {"Step": <name>, "Trigger": <trigger>}
198+
step = step["Step"]
192199
if isinstance(step, Step):
193200
dependencies.add(step.name)
194201
elif isinstance(step, StepCollection):
Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"). You
4+
# may not use this file except in compliance with the License. A copy of
5+
# the License is located at
6+
#
7+
# http://aws.amazon.com/apache2.0/
8+
#
9+
# or in the "license" file accompanying this file. This file is
10+
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
11+
# ANY KIND, either express or implied. See the License for the specific
12+
# language governing permissions and limitations under the License.
13+
"""Integration test for run conditions (StepDependency triggers).
14+
15+
A Processing step fails on purpose. A step that depends on it with the Failed trigger
16+
runs, and a step with a plain (Completed) dependency does not.
17+
"""
18+
19+
from __future__ import absolute_import
20+
21+
import logging
22+
import time
23+
import uuid
24+
25+
from sagemaker.core import image_uris
26+
from sagemaker.core.processing import Processor
27+
from sagemaker.mlops.workflow import StepDependency, TriggerCondition
28+
from sagemaker.mlops.workflow.pipeline import Pipeline
29+
from sagemaker.mlops.workflow.steps import ProcessingStep
30+
31+
32+
def _shell_step(name, command, image_uri, role, pipeline_session, depends_on=None):
33+
processor = Processor(
34+
role=role,
35+
image_uri=image_uri,
36+
instance_count=1,
37+
instance_type="ml.m5.large",
38+
entrypoint=["/bin/sh", "-c", command],
39+
sagemaker_session=pipeline_session,
40+
)
41+
return ProcessingStep(name=name, step_args=processor.run(), depends_on=depends_on)
42+
43+
44+
def test_failed_trigger_runs_after_upstream_failure(pipeline_session, role, sklearn_latest_version):
45+
stamp = uuid.uuid4().hex[:8]
46+
image_uri = image_uris.retrieve(
47+
framework="sklearn",
48+
region=pipeline_session.boto_region_name,
49+
version=sklearn_latest_version,
50+
py_version="py3",
51+
instance_type="ml.m5.large",
52+
)
53+
54+
work = _shell_step("Work", "exit 1", image_uri, role, pipeline_session)
55+
on_failure = _shell_step(
56+
"OnFailure",
57+
"echo recovered",
58+
image_uri,
59+
role,
60+
pipeline_session,
61+
depends_on=[StepDependency(work, TriggerCondition.FAILED)],
62+
)
63+
on_success = _shell_step(
64+
"OnSuccess", "echo done", image_uri, role, pipeline_session, depends_on=[work]
65+
)
66+
pipeline = Pipeline(
67+
name=f"integ-run-conditions-{stamp}",
68+
steps=[work, on_failure, on_success],
69+
sagemaker_session=pipeline_session,
70+
)
71+
72+
try:
73+
pipeline.upsert(role_arn=role)
74+
execution = pipeline.start()
75+
76+
status = None
77+
deadline = time.time() + 1800
78+
while time.time() < deadline:
79+
status = execution.describe()["PipelineExecutionStatus"]
80+
if status in ("Succeeded", "Failed", "Stopped", "CompletedWithFailures"):
81+
break
82+
time.sleep(30)
83+
84+
step_status = {s["StepName"]: s["StepStatus"] for s in execution.list_steps()}
85+
assert status == "CompletedWithFailures", f"status={status}, steps={step_status}"
86+
assert step_status["Work"] == "Failed"
87+
assert step_status["OnFailure"] == "Succeeded"
88+
assert "OnSuccess" not in step_status
89+
finally:
90+
try:
91+
pipeline.delete()
92+
except Exception as e:
93+
logging.warning("Failed to delete pipeline %s: %s", pipeline.name, e)

0 commit comments

Comments
 (0)