Skip to content

Commit 6d00cf2

Browse files
author
nayan3107
committed
fix(mlops): Keep run-condition triggers on ForEachStep dependencies
ForEach compiled DependsOn as plain names, so a StepDependency trigger on the ForEach step was silently dropped. Use the shared run-condition serializer like regular and Condition steps. --- X-AI-Prompt: Rebase the ForEach PR onto the merged Run Conditions PR X-AI-Tool: kiro-cli
1 parent 6497fb6 commit 6d00cf2

2 files changed

Lines changed: 34 additions & 1 deletion

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -405,7 +405,7 @@ def _build_for_each_step(self, for_each_step: ForEachStep) -> ForEachStep:
405405
max_concurrency=for_each_step.max_concurrency,
406406
selector=for_each_step.selector,
407407
for_each_body=self._build_steps(for_each_step.for_each_body),
408-
depends_on=[step.name for step in depends_on],
408+
depends_on=self._to_depends_on_request(for_each_step, depends_on),
409409
)
410410

411411
return compiled_for_each_step

‎sagemaker-mlops/tests/unit/workflow/test_for_each_step.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
LambdaStep,
3737
)
3838
from sagemaker.mlops.workflow.pipeline import Pipeline, PipelineGraph
39+
from sagemaker.mlops.workflow.step_dependency import StepDependency, TriggerCondition
3940
from sagemaker.mlops.workflow.steps import Step, StepTypeEnum
4041

4142

@@ -602,3 +603,35 @@ def test_pipeline_definition_literal_list_items_and_aggregate_reference(mock_ses
602603
"status": {"Get": "Steps.ForEachItems.Status"},
603604
"total": {"Get": "Steps.ForEachItems.TotalStepCount"},
604605
}
606+
607+
608+
def test_pipeline_definition_run_conditions(mock_session):
609+
producer = CustomStep(name="ProduceItem")
610+
first = CustomStep(name="FirstBodyStep", input_data=CurrentItem())
611+
second = CustomStep(
612+
name="SecondBodyStep",
613+
depends_on=[StepDependency(first, TriggerCondition.ANY)],
614+
)
615+
for_each = ForEachStep(
616+
name="ForEachItems",
617+
iterable_items=ParameterString(name="ItemList"),
618+
max_concurrency=2,
619+
for_each_body=[first, second],
620+
depends_on=[StepDependency(producer, TriggerCondition.FAILED)],
621+
)
622+
notify = CustomStep(
623+
name="NotifyOnFailure",
624+
depends_on=[StepDependency(for_each, TriggerCondition.FAILED)],
625+
)
626+
pipeline = Pipeline(
627+
name="test-pipeline",
628+
parameters=[ParameterString(name="ItemList")],
629+
steps=[producer, for_each, notify],
630+
sagemaker_session=mock_session,
631+
)
632+
steps = {s["Name"]: s for s in _definition_steps(pipeline)}
633+
assert steps["ForEachItems"]["DependsOn"] == [{"Step": "ProduceItem", "Trigger": "Failed"}]
634+
assert steps["NotifyOnFailure"]["DependsOn"] == [{"Step": "ForEachItems", "Trigger": "Failed"}]
635+
body = {s["Name"]: s for s in steps["ForEachItems"]["Arguments"]["ForEachBody"]}
636+
assert "DependsOn" not in body["FirstBodyStep"]
637+
assert body["SecondBodyStep"]["DependsOn"] == [{"Step": "FirstBodyStep", "Trigger": "Any"}]

0 commit comments

Comments
 (0)