fix(mlops): support StepCollection in the pipeline steps compiler

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)}")
 

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"]