diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/_steps_compiler.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/_steps_compiler.py index 673cad8c1b..97776a681a 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/_steps_compiler.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/_steps_compiler.py @@ -24,6 +24,7 @@ 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.mlops.workflow.step_collections import StepCollection from sagemaker.core.workflow.step_outputs import get_step from sagemaker.core.workflow.pipeline_definition_config import PipelineDefinitionConfig @@ -212,6 +213,8 @@ def _generate_step_map(steps: Sequence[Union[Step, StepOutput]], step_map: dict) step_map[step.name] = step if isinstance(step, ConditionStep): StepsCompiler._generate_step_map(step.if_steps + step.else_steps, step_map) + if isinstance(step, StepCollection): + StepsCompiler._generate_step_map(step.steps, step_map) def _simplify_step_list( self, @@ -241,7 +244,10 @@ def _simplify_step_list( if isinstance(step, Step): if step not in resolved_step_set: resolved_step_set.add(step) - + elif isinstance(step, StepCollection): + for sub_step in step.steps: + if sub_step not in resolved_step_set: + resolved_step_set.add(sub_step) else: raise ValueError(f"Invalid input step type: {type(step)}") diff --git a/sagemaker-mlops/tests/unit/workflow/test_steps_compiler.py b/sagemaker-mlops/tests/unit/workflow/test_steps_compiler.py index 54bd834208..a506188421 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_steps_compiler.py +++ b/sagemaker-mlops/tests/unit/workflow/test_steps_compiler.py @@ -459,3 +459,42 @@ def test_push_to_build_queue_duplicate(self, mock_session, mock_step): # Should only be added once to queue assert len(compiler._build_queue) == 1 + + +class TestStepCollectionSupport: + """StepCollections (e.g. MonitorBatchTransformStep) are flattened into their sub-steps.""" + + @staticmethod + def _definition_step_names(steps): + import json + from sagemaker.mlops.workflow.pipeline import Pipeline + + session = Mock(boto_region_name="us-west-2", local_mode=False) + definition = json.loads( + Pipeline(name="p", steps=steps, sagemaker_session=session).definition() + ) + return {step["Name"]: step for step in definition["Steps"]} + + def test_pipeline_with_step_collection(self): + from sagemaker.mlops.workflow.fail_step import FailStep + from sagemaker.mlops.workflow.step_collections import StepCollection + + first = FailStep("first") + second = FailStep("second", depends_on=[first]) + + steps = self._definition_step_names([StepCollection(name="coll", steps=[first, second])]) + + assert set(steps) == {"first", "second"} + assert steps["second"]["DependsOn"] == ["first"] + + @pytest.mark.parametrize("by_name", [False, True]) + def test_depends_on_step_collection(self, by_name): + from sagemaker.mlops.workflow.fail_step import FailStep + from sagemaker.mlops.workflow.step_collections import StepCollection + + collection = StepCollection(name="coll", steps=[FailStep("first"), FailStep("second")]) + after = FailStep("after", depends_on=["coll" if by_name else collection]) + + steps = self._definition_step_names([collection, after]) + + assert sorted(steps["after"]["DependsOn"]) == ["first", "second"]