Feat/foreach step

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

@@ -51,6 +51,13 @@
 from sagemaker.mlops.workflow.endpoint_step import EndpointConfigStep, EndpointStep
 from sagemaker.mlops.workflow.fail_step import FailStep
 from sagemaker.mlops.workflow.inference_component_step import InferenceComponentStep
+from sagemaker.mlops.workflow.for_each_step import (
+    ForEachStep,
+    ForEachSelector,
+    ForEachSelectorPickType,
+    CurrentItem,
+    IterationOutput,
+)
 from sagemaker.mlops.workflow.lambda_step import LambdaStep, LambdaOutput
 from sagemaker.mlops.workflow.lineage_step import (
     LineageAssociation,
@@ -115,6 +122,11 @@
     "EndpointStep",
     "FailStep",
     "InferenceComponentStep",
+    "ForEachStep",
+    "ForEachSelector",
+    "ForEachSelectorPickType",
+    "CurrentItem",
+    "IterationOutput",
     "LambdaStep",
     "LambdaOutput",
     "LineageAssociation",

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.for_each_step import ForEachStep
 
 from sagemaker.core.workflow.step_outputs import get_step
 from sagemaker.core.workflow.pipeline_definition_config import PipelineDefinitionConfig
@@ -207,11 +208,13 @@ def _generate_step_map(steps: Sequence[Union[Step, StepOutput]], step_map: dict)
             if step.name in step_map:
                 raise ValueError(
                     "Pipeline steps cannot have duplicate names. In addition, steps added in "
-                    "the ConditionStep cannot be added in the Pipeline steps list."
+                    "the ConditionStep or ForEachStep cannot be added in the Pipeline steps list."
                 )
             step_map[step.name] = step
             if isinstance(step, ConditionStep):
                 StepsCompiler._generate_step_map(step.if_steps + step.else_steps, step_map)
+            if isinstance(step, ForEachStep):
+                StepsCompiler._generate_step_map(step.for_each_body, step_map)
 
     def _simplify_step_list(
         self,
@@ -247,14 +250,17 @@ def _simplify_step_list(
 
         return list(resolved_step_set)
 
-    def _flatten_condition_step(self, step: ConditionStep) -> List[Step]:
-        """Flatten a ConditionStep into a list of steps."""
+    def _flatten_compound_step(self, step: Union[ConditionStep, ForEachStep]) -> List[Step]:
+        """Flatten a compound step (ConditionStep or ForEachStep) into a list of steps."""
         flattened_steps = _StepsSet()
         flattened_steps.add(step)
-        sub_steps = self._simplify_step_list(step.if_steps + step.else_steps)
+        if isinstance(step, ConditionStep):
+            sub_steps = self._simplify_step_list(step.if_steps + step.else_steps)
+        else:
+            sub_steps = self._simplify_step_list(step.for_each_body)
         for sub_step in sub_steps:
-            if isinstance(sub_step, ConditionStep):
-                flattened_steps.add_list(self._flatten_condition_step(sub_step))
+            if isinstance(sub_step, (ConditionStep, ForEachStep)):
+                flattened_steps.add_list(self._flatten_compound_step(sub_step))
             else:
                 flattened_steps.add(sub_step)
         return list(flattened_steps)
@@ -265,8 +271,8 @@ def _push_to_build_queue(self, steps: List[Step]):
             if step not in self._all_known_steps:
                 self._build_queue.push([step])
 
-            if isinstance(step, ConditionStep):
-                self._all_known_steps.add_list(self._flatten_condition_step(step))
+            if isinstance(step, (ConditionStep, ForEachStep)):
+                self._all_known_steps.add_list(self._flatten_compound_step(step))
             else:
                 self._all_known_steps.add(step)
 
@@ -325,6 +331,8 @@ def _set_serialize_output_to_json_flag(self, compiled_steps: List[CompiledStep])
             if isinstance(step, ConditionStep):
                 self._set_serialize_output_to_json_flag(step.if_steps)
                 self._set_serialize_output_to_json_flag(step.else_steps)
+            elif isinstance(step, ForEachStep):
+                self._set_serialize_output_to_json_flag(step.for_each_body)
             elif step.name in self._steps_need_json_serialization:
                 step_container_args = step._request_dict["Arguments"]["AlgorithmSpecification"][
                     "ContainerArguments"
@@ -381,6 +389,27 @@ def _build_condition_step(self, condition_step: ConditionStep) -> ConditionStep:
 
         return compiled_condition_step
 
+    def _build_for_each_step(self, for_each_step: ForEachStep) -> ForEachStep:
+        """Build a ForEach step."""
+        depends_on, upstream_steps_in_args = self._get_upstream_steps(
+            for_each_step, for_each_step.step_only_arguments
+        )
+        self._push_to_build_queue(depends_on)
+        self._push_to_build_queue(upstream_steps_in_args)
+
+        compiled_for_each_step = ForEachStep(
+            name=for_each_step.name,
+            display_name=for_each_step.display_name,
+            description=for_each_step.description,
+            iterable_items=for_each_step.iterable_items,
+            max_concurrency=for_each_step.max_concurrency,
+            selector=for_each_step.selector,
+            for_each_body=self._build_steps(for_each_step.for_each_body),
+            depends_on=self._to_depends_on_request(for_each_step, depends_on),
+        )
+
+        return compiled_for_each_step
+
     def _build_steps(self, steps: List[Union[Step, StepOutput]]):
         """Build a list of steps."""
         simple_steps = self._simplify_step_list(steps)
@@ -389,6 +418,8 @@ def _build_steps(self, steps: List[Union[Step, StepOutput]]):
         for step in simple_steps:
             if isinstance(step, ConditionStep):
                 compiled_steps.append(self._build_condition_step(step))
+            elif isinstance(step, ForEachStep):
+                compiled_steps.append(self._build_for_each_step(step))
             else:
                 compiled_steps.append(self._build_step(step))
         return compiled_steps
@@ -404,6 +435,8 @@ def _initialize_queue_and_build(self, steps: List[Union[Step, StepOutput]]):
             step = self._build_queue.pop()
             if isinstance(step, ConditionStep):
                 compiled_steps.append(self._build_condition_step(step))
+            elif isinstance(step, ForEachStep):
+                compiled_steps.append(self._build_for_each_step(step))
             else:
                 compiled_steps.append(self._build_step(step))