diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py index 4f50c1be74..2500f40269 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py +++ b/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", diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/_steps_compiler.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/_steps_compiler.py index 673cad8c1b..7ce1b926ae 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.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)) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py new file mode 100644 index 0000000000..13b1882c29 --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py @@ -0,0 +1,446 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""The `ForEach` step definitions for SageMaker Pipelines Workflows.""" + +from __future__ import absolute_import + +import re +from enum import Enum +from typing import Any, List, Optional, Union + +from sagemaker.core.helper.pipeline_variable import PipelineVariable, RequestType +from sagemaker.core.workflow.properties import Properties +from sagemaker.core.workflow.utilities import list_to_request +from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + +MAX_CONCURRENCY_MIN = 1 +MAX_CONCURRENCY_MAX = 50 +MAX_LITERAL_ITERABLE_ITEMS = 100 + +_ITERATION_OUTPUT_NAMESPACE = "IterationOutput" +_CURRENT_ITEM_PATH = "ForEach.CurrentItem" + +# Step types that may not appear in a ForEachBody. +_DISALLOWED_BODY_STEP_TYPES = frozenset( + {StepTypeEnum.CONDITION, StepTypeEnum.FAIL, StepTypeEnum.FOR_EACH} +) + + +class _ForEachPropertyReference(PipelineVariable): + """A drillable ``{"Get": }`` reference used by the ForEach step. + + Supports attribute access (``ref.Member``) and indexing (``ref[0]``, + ``ref["key"]``) to build up a property path, mirroring the drill-down + behavior of :class:`~sagemaker.core.workflow.properties.Properties`. + """ + + def __init__(self, path: str, referenced_steps: Optional[List[Any]] = None): + """Initialize a _ForEachPropertyReference. + + Args: + path (str): The full ``Get`` expression path. + referenced_steps (List): Steps this reference depends on. + """ + self._path = path + self._refs = list(referenced_steps) if referenced_steps else [] + + def __getattr__(self, name: str) -> "_ForEachPropertyReference": + """Drill into a member of the referenced structure.""" + if name.startswith("_"): + raise AttributeError(name) + return _ForEachPropertyReference(f"{self._path}.{name}", self._refs) + + def __getitem__(self, item: Union[int, str]) -> "_ForEachPropertyReference": + """Drill into a list index or map key of the referenced structure.""" + if isinstance(item, int): + return _ForEachPropertyReference(f"{self._path}[{item}]", self._refs) + return _ForEachPropertyReference(f"{self._path}['{item}']", self._refs) + + @property + def expr(self) -> RequestType: + """The 'Get' expression dict for this reference.""" + return {"Get": self._path} + + @property + def _referenced_steps(self) -> List[Any]: + """List of steps that this reference depends on.""" + return self._refs + + +class CurrentItem(_ForEachPropertyReference): + """Reference to the current iteration item inside a ForEach body step. + + Valid ONLY within the steps of a ``ForEachStep``'s ``for_each_body``. + The bare reference resolves to the whole item (for scalar items). For + JSON-object items, drill into fields via attribute access or bracket + indexing to any depth:: + + item = CurrentItem() + item.expr # {"Get": "ForEach.CurrentItem"} + item.name.expr # {"Get": "ForEach.CurrentItem.name"} + item.a.b.c.expr # {"Get": "ForEach.CurrentItem.a.b.c"} + item["score"].expr # {"Get": "ForEach.CurrentItem['score']"} + """ + + def __init__(self): + """Initialize a CurrentItem reference.""" + super().__init__(_CURRENT_ITEM_PATH) + + +class IterationOutput(_ForEachPropertyReference): + """Reference to a ForEach body step's output, for use in ``ForEachSelector.field``. + + Valid ONLY as the ``field`` of a :class:`ForEachSelector`, and must target a + body step of the same ForEach step, drilling at least one member deep:: + + IterationOutput(process_step).OutputParameters["InstanceCount"].expr + # {"Get": "IterationOutput.ProcessItem.OutputParameters['InstanceCount']"} + + Args: + body_step (Union[str, Step]): The body step (or its name) whose output + the selector compares across iterations. + """ + + def __init__(self, body_step: Union[str, Step]): + """Initialize an IterationOutput reference rooted at the given body step.""" + step_name = body_step.name if isinstance(body_step, Step) else body_step + super().__init__(f"{_ITERATION_OUTPUT_NAMESPACE}.{step_name}") + + +class _IterationOutputsAccessor: + """Indexed accessor for a ForEach step's per-iteration outputs. + + ``IterationOutputs`` MUST be indexed (``properties.IterationOutputs[0]..``); + a bare ``IterationOutputs.`` reference is invalid, so this accessor + is intentionally not a ``PipelineVariable`` itself. The indexed reference + must then be drilled by body step name, since a ForEach body may contain + multiple steps. + """ + + def __init__(self, step: Step): + """Initialize an _IterationOutputsAccessor for the given ForEach step.""" + self._step = step + + def __getitem__(self, index: int) -> _ForEachPropertyReference: + """Get a drillable reference to the outputs of the iteration at ``index``.""" + if not isinstance(index, int): + raise TypeError( + f"IterationOutputs of ForEach step '{self._step.name}' must be indexed " + f"with an integer, e.g. properties.IterationOutputs[0].." + ) + return _ForEachPropertyReference( + f"Steps.{self._step.name}.IterationOutputs[{index}]", [self._step] + ) + + +class ForEachSelectorPickType(Enum): + """Enum of pick strategies for a :class:`ForEachSelector`.""" + + MAX_BY = "MaxBy" + MIN_BY = "MinBy" + + +class ForEachSelector: + """Selects a single child of a ForEach step by comparing a field across children. + + The selected iteration's output is exposed to downstream steps via + ``Steps..Selected.`` — no body step name qualifier + is required, because the selector's ``field`` already identifies which + body step's output was compared. + + Args: + pick (Union[str, ForEachSelectorPickType]): The pick strategy, + ``"MaxBy"`` or ``"MinBy"``. + field (IterationOutput): An ``IterationOutput`` reference targeting a + body step of the same ForEach step, drilled at least one member + deep. + """ + + def __init__( + self, + pick: Union[str, ForEachSelectorPickType], + field: PipelineVariable, + ): + """Initialize a ForEachSelector.""" + if pick is None: + raise ValueError("ForEachSelector requires 'pick' (MaxBy or MinBy).") + if isinstance(pick, str): + try: + pick = ForEachSelectorPickType(pick) + except ValueError: + raise ValueError( + f"Invalid ForEachSelector pick type: '{pick}'. " + f"Valid values: " + f"{[member.value for member in ForEachSelectorPickType]}." + ) + if not isinstance(pick, ForEachSelectorPickType): + raise ValueError( + f"ForEachSelector 'pick' must be a str or ForEachSelectorPickType, " + f"got: {type(pick)}." + ) + if field is None: + raise ValueError("ForEachSelector requires 'field' (an IterationOutput reference).") + if not isinstance(field, PipelineVariable): + raise ValueError( + "ForEachSelector 'field' must be an IterationOutput reference, e.g. " + 'IterationOutput(body_step).OutputParameters["MyOutput"], ' + f"got: {type(field)}." + ) + self.pick = pick + self.field = field + + def to_request(self) -> RequestType: + """Get the request structure for workflow service calls.""" + return { + "Pick": self.pick.value, + "Field": self.field, + } + + +class ForEachStep(Step): + """ForEach step for dynamic fan-out in SageMaker Pipelines Workflows. + + A ForEach step takes a list of items (resolved at runtime), runs its body + step(s) once per item as child steps, enforces a concurrency cap, and + exposes per-iteration outputs plus aggregate results to downstream steps. + + A body may contain multiple steps; each item spawns one child instance of + every body step. Child steps get IDs of the form ``-`` + (e.g. ``ProcessItem-0``). Because of this, body step names must be unique + across all ForEach bodies in the pipeline, and no top-level step name may + look like ``-`` — the service rejects such + collisions. + + Downstream steps can reference the parent step's properties, drilling by + body step name to reach a specific output: + + - ``properties.IterationOutputs[i]..`` — per-iteration + output for a specific body step (must be indexed; body step name is + required since the body may contain multiple steps) + - ``properties.Selected.`` — output of the selector-picked + iteration (no body step name qualifier: the selector's ``field`` already + identifies which body step's output was compared) + - ``properties.Status`` (String: Succeeded | Failed | Stopped) + - ``properties.InputItemCount`` (Integer: number of items iterated) and + ``properties.TotalStepCount``, ``properties.SucceededStepCount``, + ``properties.FailedStepCount``, ``properties.StoppedStepCount`` + (Integers: child step counts across all iterations, so up to + ``InputItemCount`` times the number of body steps) + """ + + _AGGREGATE_PROPERTY_MEMBERS = ( + "Status", + "InputItemCount", + "TotalStepCount", + "SucceededStepCount", + "FailedStepCount", + "StoppedStepCount", + ) + + def __init__( + self, + name: str, + iterable_items: Union[ + PipelineVariable, List[Union[PipelineVariable, str, int, float, bool]] + ] = None, + max_concurrency: int = None, + for_each_body: List[Step] = None, + selector: Optional[ForEachSelector] = None, + display_name: Optional[str] = None, + description: Optional[str] = None, + depends_on: Optional[List[Union[str, Step]]] = None, + ): + """Construct a ForEachStep. + + Args: + name (str): The name of the `ForEachStep`. + iterable_items (Union[PipelineVariable, List]): REQUIRED. The items + to iterate over, in one of two forms: + + - A property reference (e.g. a `ParameterString` or a step + output) whose value resolves to a JSON list string at + runtime, such as ``'["item-a", "item-b"]'``. + - A non-empty list of at most 100 elements, each a property + reference or a literal scalar (str, int, float, bool). Each + element is one iteration item. + + Items are commonly JSON-object strings whose fields are + addressed via ``CurrentItem().`` inside the body step. + max_concurrency (int): REQUIRED. The maximum number of child steps + that run concurrently. Range: [1, 50]. + for_each_body (List[Step]): REQUIRED. The step(s) to run once per + item. May contain one or more body steps; each item spawns one + child instance of every body step. Disallowed body step types: + Condition, Fail, ForEach (no nesting). Body steps may reference + the current item via :class:`CurrentItem`. + selector (ForEachSelector): Optional selector that picks a single + iteration by comparing an ``IterationOutput`` field across + iterations. + display_name (str): The display name of the `ForEachStep`. + description (str): The description of the `ForEachStep`. + depends_on (List[Union[str, Step]]): The list of `Step` names or + `Step` instances that the current `Step` depends on. + """ + super(ForEachStep, self).__init__( + name, display_name, description, StepTypeEnum.FOR_EACH, depends_on + ) + + self._validate_iterable_items(name, iterable_items) + self._validate_max_concurrency(name, max_concurrency) + self._validate_for_each_body(name, for_each_body) + self._validate_selector(name, selector, for_each_body) + + self.iterable_items = ( + list(iterable_items) if isinstance(iterable_items, list) else iterable_items + ) + self.max_concurrency = max_concurrency + self.for_each_body = list(for_each_body) + self.selector = selector + + root_prop = Properties(step_name=name) + for member in ForEachStep._AGGREGATE_PROPERTY_MEMBERS: + root_prop.__dict__[member] = Properties(step_name=name, path=member, step=self) + root_prop.__dict__["IterationOutputs"] = _IterationOutputsAccessor(self) + root_prop.__dict__["Selected"] = _ForEachPropertyReference(f"Steps.{name}.Selected", [self]) + self._properties = root_prop + + @staticmethod + def _validate_iterable_items(name: str, iterable_items: Any): + """Validate the iterable_items argument.""" + if iterable_items is None: + raise ValueError( + f"ForEach step '{name}': iterable_items is required. Provide a " + f"property reference that resolves to a JSON list string at " + f"runtime, or a list of items." + ) + if isinstance(iterable_items, list): + if not 1 <= len(iterable_items) <= MAX_LITERAL_ITERABLE_ITEMS: + raise ValueError( + f"ForEach step '{name}': an iterable_items list must have between 1 " + f"and {MAX_LITERAL_ITERABLE_ITEMS} elements. Got: {len(iterable_items)}." + ) + for element in iterable_items: + if element is None or not isinstance( + element, (PipelineVariable, str, int, float, bool) + ): + raise ValueError( + f"ForEach step '{name}': iterable_items list elements must be " + f"property references or literal scalars (str, int, float, bool). " + f"Got: {type(element)}." + ) + return + if not isinstance(iterable_items, PipelineVariable): + raise ValueError( + f"ForEach step '{name}': iterable_items must be a property " + f"reference that resolves to a JSON list string at runtime, or a " + f"list of items; a single literal value cannot resolve to a list. " + f"Got: {type(iterable_items)}." + ) + + @staticmethod + def _validate_max_concurrency(name: str, max_concurrency: Any): + """Validate the max_concurrency argument.""" + if ( + not isinstance(max_concurrency, int) + or isinstance(max_concurrency, bool) + or not MAX_CONCURRENCY_MIN <= max_concurrency <= MAX_CONCURRENCY_MAX + ): + raise ValueError( + f"ForEach step '{name}': max_concurrency is required and must be " + f"an integer in the range " + f"[{MAX_CONCURRENCY_MIN}, {MAX_CONCURRENCY_MAX}]. " + f"Got: {max_concurrency}." + ) + + @staticmethod + def _validate_for_each_body(name: str, for_each_body: Any): + """Validate the for_each_body argument.""" + if not for_each_body: + raise ValueError( + f"ForEach step '{name}': for_each_body must be a non-empty list " f"of steps." + ) + for body_step in for_each_body: + if not isinstance(body_step, Step): + raise ValueError( + f"ForEach step '{name}': for_each_body entries must be Step " + f"instances. Got: {type(body_step)}." + ) + if body_step.step_type in _DISALLOWED_BODY_STEP_TYPES: + raise ValueError( + f"ForEach step '{name}': step type " + f"'{body_step.step_type.value}' is not allowed in " + f"for_each_body. Disallowed types: Condition, Fail, ForEach." + ) + + @staticmethod + def _validate_selector( + name: str, selector: Optional[ForEachSelector], for_each_body: List[Step] + ): + """Validate that the selector field targets a body step of this ForEach.""" + if selector is None: + return + if not isinstance(selector, ForEachSelector): + raise ValueError( + f"ForEach step '{name}': selector must be a ForEachSelector " + f"instance. Got: {type(selector)}." + ) + field_path = selector.field.expr.get("Get") + prefix = f"{_ITERATION_OUTPUT_NAMESPACE}." + if not isinstance(field_path, str) or not field_path.startswith(prefix): + raise ValueError( + f"ForEach step '{name}': Selector.Field must be an IterationOutput " + f"reference (e.g. IterationOutput(body_step).); other " + f"namespaces are constant across iterations. Got: {field_path}." + ) + remainder = field_path[len(prefix) :] + match = re.match(r"^([^.\[\]]+)(.+)$", remainder) + body_step_names = [step.name for step in for_each_body] + if not match or match.group(1) not in body_step_names: + raise ValueError( + f"ForEach step '{name}': Selector.Field must target a body step " + f"of this ForEach step ({body_step_names}) and drill at least " + f"one member deep. Got: {field_path}." + ) + + @property + def arguments(self) -> RequestType: + """The arguments dict that is used to define the ForEach fan-out.""" + args = { + "IterableItems": self.iterable_items, + "MaxConcurrency": self.max_concurrency, + } + if self.selector is not None: + args["Selector"] = self.selector.to_request() + args["ForEachBody"] = list_to_request(self.for_each_body) + return args + + @property + def step_only_arguments(self) -> RequestType: + """Argument dict pertaining to the step only, and not the body steps.""" + args = { + "IterableItems": self.iterable_items, + "MaxConcurrency": self.max_concurrency, + } + if self.selector is not None: + args["Selector"] = self.selector.to_request() + return args + + @property + def properties(self): + """A Properties object exposing the ForEach step's results. + + Members: ``Status``, ``InputItemCount``, ``TotalStepCount``, + ``SucceededStepCount``, ``FailedStepCount``, ``StoppedStepCount``, + ``IterationOutputs[i].`` and ``Selected``. + """ + return self._properties diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py index 05622e9ca0..997c7fda75 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/pipeline.py @@ -77,6 +77,7 @@ from sagemaker.core.workflow.step_outputs import StepOutput from sagemaker.mlops.workflow.steps import Step from sagemaker.mlops.workflow.condition_step import ConditionStep +from sagemaker.mlops.workflow.for_each_step import ForEachStep from sagemaker.mlops.workflow.triggers import ( PipelineSchedule, Trigger, @@ -1005,11 +1006,13 @@ def _generate_step_map(steps: Sequence[Step], 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): _generate_step_map(step.if_steps + step.else_steps, step_map) + if isinstance(step, ForEachStep): + _generate_step_map(step.for_each_body, step_map) @attr.s @@ -1292,6 +1295,11 @@ def _initialize_adjacency_list(self) -> Dict[str, List[str]]: if isinstance(child_step, Step): dependency_list[child_step.name].add(step.name) + if isinstance(step, ForEachStep): + for child_step in step.for_each_body: + if isinstance(child_step, Step): + dependency_list[child_step.name].add(step.name) + adjacency_list = {} for step in dependency_list: for step_dependency in dependency_list[step]: diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py index b5b45e77a4..2bee11ceba 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/steps.py @@ -69,6 +69,7 @@ class StepTypeEnum(Enum): INFERENCE_COMPONENT = "InferenceComponent" LINEAGE = "Lineage" JOB = "Job" + FOR_EACH = "ForEach" class Step(Entity): diff --git a/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py b/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py new file mode 100644 index 0000000000..c83e576ff7 --- /dev/null +++ b/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py @@ -0,0 +1,637 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Unit tests for workflow for_each_step.""" + +from __future__ import absolute_import + +import json + +import pytest +from unittest.mock import Mock + +from sagemaker.core.workflow.parameters import ParameterString +from sagemaker.core.workflow.properties import Properties +from sagemaker.mlops.workflow.condition_step import ConditionStep +from sagemaker.mlops.workflow.fail_step import FailStep +from sagemaker.mlops.workflow.for_each_step import ( + CurrentItem, + ForEachSelector, + ForEachSelectorPickType, + ForEachStep, + IterationOutput, +) +from sagemaker.mlops.workflow.lambda_step import ( + LambdaOutput, + LambdaOutputTypeEnum, + LambdaStep, +) +from sagemaker.mlops.workflow.pipeline import Pipeline, PipelineGraph +from sagemaker.mlops.workflow.step_dependency import StepDependency, TriggerCondition +from sagemaker.mlops.workflow.steps import Step, StepTypeEnum + + +class CustomStep(Step): + """A minimal concrete step for testing.""" + + def __init__(self, name, step_type=StepTypeEnum.TRAINING, input_data=None, depends_on=None): + super(CustomStep, self).__init__(name, None, None, step_type, depends_on) + self._input_data = input_data + self._properties = Properties(step_name=name, step=self) + + @property + def arguments(self): + return {"input_data": self._input_data} if self._input_data is not None else {} + + @property + def properties(self): + return self._properties + + +@pytest.fixture +def mock_session(): + session = Mock() + session.boto_session.client.return_value = Mock() + session.sagemaker_client = Mock() + session.local_mode = False + session.sagemaker_config = {} + return session + + +@pytest.fixture +def item_list_parameter(): + return ParameterString(name="ItemList") + + +def make_lambda_body_step(function_arn="arn:aws:lambda:us-west-2:123456789012:function:body"): + lambda_func = Mock() + lambda_func.function_arn = function_arn + lambda_func.zipped_code_dir = None + lambda_func.script = None + current_item = CurrentItem() + return LambdaStep( + name="ProcessItem", + lambda_func=lambda_func, + inputs={ + "item_name": current_item.name, + "item_score": current_item.score, + }, + outputs=[ + LambdaOutput(output_name="ProcessedScore", output_type=LambdaOutputTypeEnum.Integer), + LambdaOutput(output_name="ResultPath", output_type=LambdaOutputTypeEnum.String), + ], + ) + + +def make_for_each_step(body_step=None, selector=None, **kwargs): + body_step = body_step or make_lambda_body_step() + default_kwargs = dict( + name="ForEachItems", + iterable_items=ParameterString(name="ItemList"), + max_concurrency=2, + for_each_body=[body_step], + selector=selector, + ) + default_kwargs.update(kwargs) + return ForEachStep(**default_kwargs) + + +# --------------------------------------------------------------------------- +# CurrentItem / IterationOutput references +# --------------------------------------------------------------------------- + + +def test_current_item_bare_expr(): + assert CurrentItem().expr == {"Get": "ForEach.CurrentItem"} + + +def test_current_item_drilled_expr(): + assert CurrentItem().name.expr == {"Get": "ForEach.CurrentItem.name"} + assert CurrentItem().a.b.c.expr == {"Get": "ForEach.CurrentItem.a.b.c"} + + +def test_current_item_bracket_indexing(): + assert CurrentItem()["score"].expr == {"Get": "ForEach.CurrentItem['score']"} + assert CurrentItem().items[0].expr == {"Get": "ForEach.CurrentItem.items[0]"} + assert CurrentItem()["nested"]["key"].expr == {"Get": "ForEach.CurrentItem['nested']['key']"} + + +def test_current_item_has_no_referenced_steps(): + assert CurrentItem().name._referenced_steps == [] + + +# --------------------------------------------------------------------------- +# IterationOutput references +# --------------------------------------------------------------------------- + + +def test_iteration_output_expr_from_step_and_str(): + body = make_lambda_body_step() + ref = IterationOutput(body).OutputParameters["ProcessedScore"] + assert ref.expr == {"Get": "IterationOutput.ProcessItem.OutputParameters['ProcessedScore']"} + ref_from_str = IterationOutput("ProcessItem").Loss + assert ref_from_str.expr == {"Get": "IterationOutput.ProcessItem.Loss"} + assert ref._referenced_steps == [] + + +# --------------------------------------------------------------------------- +# to_request serialization +# --------------------------------------------------------------------------- + + +def test_for_each_step_to_request_reference_items(item_list_parameter): + step = make_for_each_step(iterable_items=item_list_parameter) + request = step.to_request() + assert request["Name"] == "ForEachItems" + assert request["Type"] == "ForEach" + args = request["Arguments"] + assert args["IterableItems"] is item_list_parameter + assert args["MaxConcurrency"] == 2 + assert "Selector" not in args + assert len(args["ForEachBody"]) == 1 + assert args["ForEachBody"][0]["Name"] == "ProcessItem" + assert args["ForEachBody"][0]["Type"] == "Lambda" + + +def test_for_each_step_literal_list_items_accepted(): + ref_item = ParameterString(name="ItemOne") + items = [ref_item, '{"name": "lit-c", "score": 5}', 42, 1.5, True] + step = make_for_each_step(iterable_items=items) + assert step.to_request()["Arguments"]["IterableItems"] == items + + +@pytest.mark.parametrize("items", [[], ["x"] * 101]) +def test_for_each_step_literal_list_size_out_of_range(items): + with pytest.raises(ValueError, match="between 1 and 100 elements"): + make_for_each_step(iterable_items=items) + + +@pytest.mark.parametrize("element", [None, {"a": 1}, ["nested"]]) +def test_for_each_step_literal_list_invalid_element(element): + with pytest.raises(ValueError, match="literal scalars"): + make_for_each_step(iterable_items=["ok", element]) + + +def test_for_each_step_to_request_with_selector(): + body = make_lambda_body_step() + selector = ForEachSelector( + pick="MaxBy", + field=IterationOutput(body).OutputParameters["ProcessedScore"], + ) + step = make_for_each_step(body_step=body, selector=selector) + args = step.to_request()["Arguments"] + assert args["Selector"]["Pick"] == "MaxBy" + assert args["Selector"]["Field"].expr == { + "Get": "IterationOutput.ProcessItem.OutputParameters['ProcessedScore']" + } + + +def test_for_each_step_to_request_depends_on_and_metadata(): + upstream = CustomStep(name="UpstreamStep") + step = make_for_each_step( + depends_on=[upstream], + display_name="display", + description="desc", + ) + request = step.to_request() + assert request["DependsOn"] == [upstream] + assert request["DisplayName"] == "display" + assert request["Description"] == "desc" + + +def test_step_only_arguments_excludes_body(): + step = make_for_each_step() + assert "ForEachBody" not in step.step_only_arguments + assert "IterableItems" in step.step_only_arguments + assert "MaxConcurrency" in step.step_only_arguments + + +# --------------------------------------------------------------------------- +# Parent properties +# --------------------------------------------------------------------------- + + +def test_for_each_step_iteration_outputs_property(): + step = make_for_each_step() + ref = step.properties.IterationOutputs[3].ProcessItem.OutputParameters["ResultPath"] + assert ref.expr == { + "Get": "Steps.ForEachItems.IterationOutputs[3].ProcessItem.OutputParameters['ResultPath']" + } + assert ref._referenced_steps == [step] + + +@pytest.mark.parametrize( + "member", + [ + "Status", + "InputItemCount", + "TotalStepCount", + "SucceededStepCount", + "FailedStepCount", + "StoppedStepCount", + ], +) +def test_for_each_step_aggregate_properties(member): + step = make_for_each_step() + ref = getattr(step.properties, member) + assert ref.expr == {"Get": f"Steps.ForEachItems.{member}"} + assert ref._referenced_steps == [step] + + +def test_for_each_step_iteration_outputs_must_be_indexed(): + step = make_for_each_step() + with pytest.raises(TypeError, match="must be indexed"): + step.properties.IterationOutputs["ResultPath"] + + +def test_for_each_step_selected_property(): + step = make_for_each_step() + ref = step.properties.Selected.OutputParameters["ProcessedScore"] + assert ref.expr == {"Get": "Steps.ForEachItems.Selected.OutputParameters['ProcessedScore']"} + assert ref._referenced_steps == [step] + + +# --------------------------------------------------------------------------- +# Validation +# --------------------------------------------------------------------------- + + +def test_iterable_items_required(): + with pytest.raises(ValueError, match="iterable_items is required"): + make_for_each_step(iterable_items=None) + + +def test_iterable_items_literal_scalar_rejected(): + with pytest.raises(ValueError, match="literal value cannot resolve to a list"): + make_for_each_step(iterable_items="just-a-string") + + +@pytest.mark.parametrize("max_concurrency", [None, 0, 51, "2", 2.0, True]) +def test_max_concurrency_invalid(max_concurrency): + with pytest.raises(ValueError, match="max_concurrency is required"): + make_for_each_step(max_concurrency=max_concurrency) + + +@pytest.mark.parametrize("max_concurrency", [1, 25, 50]) +def test_max_concurrency_valid(max_concurrency): + step = make_for_each_step(max_concurrency=max_concurrency) + assert step.max_concurrency == max_concurrency + + +@pytest.mark.parametrize("body", [None, []]) +def test_empty_body_rejected(body): + with pytest.raises(ValueError, match="non-empty"): + make_for_each_step(for_each_body=body) + + +def test_multi_step_body_accepted(): + step_a = CustomStep(name="A") + step_b = CustomStep(name="B") + for_each = ForEachStep( + name="ForEachItems", + iterable_items=ParameterString(name="ItemList"), + max_concurrency=2, + for_each_body=[step_a, step_b], + ) + assert for_each.for_each_body == [step_a, step_b] + request_body = for_each.to_request()["Arguments"]["ForEachBody"] + assert [entry["Name"] for entry in request_body] == ["A", "B"] + + +def test_disallowed_body_step_types(): + condition_step = ConditionStep(name="Cond", conditions=[Mock()], if_steps=[], else_steps=[]) + fail_step = FailStep(name="Failure") + nested_for_each = make_for_each_step(name="Inner") + for bad_body in (condition_step, fail_step, nested_for_each): + with pytest.raises(ValueError, match="not allowed in"): + make_for_each_step(body_step=bad_body) + + +def test_selector_requires_pick_and_field(): + with pytest.raises(ValueError, match="requires 'pick'"): + ForEachSelector(pick=None, field=IterationOutput("ProcessItem").Loss) + with pytest.raises(ValueError, match="requires 'field'"): + ForEachSelector(pick="MinBy", field=None) + + +def test_selector_invalid_pick(): + with pytest.raises(ValueError, match="Invalid ForEachSelector pick type"): + ForEachSelector(pick="Median", field=IterationOutput("ProcessItem").Loss) + + +def test_selector_pick_enum(): + selector = ForEachSelector( + pick=ForEachSelectorPickType.MIN_BY, field=IterationOutput("ProcessItem").Loss + ) + assert selector.to_request()["Pick"] == "MinBy" + + +def test_selector_field_must_be_iteration_output(): + body = make_lambda_body_step() + with pytest.raises(ValueError, match="must be an IterationOutput reference"): + make_for_each_step( + body_step=body, + selector=ForEachSelector(pick="MaxBy", field=ParameterString(name="Other")), + ) + + +def test_selector_field_must_target_body_step(): + body = make_lambda_body_step() + with pytest.raises(ValueError, match="must target a body step"): + make_for_each_step( + body_step=body, + selector=ForEachSelector(pick="MaxBy", field=IterationOutput("OtherStep").Loss), + ) + + +def test_selector_field_must_drill_into_member(): + body = make_lambda_body_step() + with pytest.raises(ValueError, match="drill at least"): + make_for_each_step( + body_step=body, + selector=ForEachSelector(pick="MaxBy", field=IterationOutput(body)), + ) + + +# --------------------------------------------------------------------------- +# Pipeline definition (interpolated wire format) +# --------------------------------------------------------------------------- + + +def _definition_steps(pipeline): + return json.loads(pipeline.definition())["Steps"] + + +def test_pipeline_definition_wire_format(mock_session): + body = make_lambda_body_step() + selector = ForEachSelector( + pick="MaxBy", + field=IterationOutput(body).OutputParameters["ProcessedScore"], + ) + for_each = ForEachStep( + name="ForEachItems", + iterable_items=ParameterString(name="ItemList"), + max_concurrency=2, + for_each_body=[body], + selector=selector, + ) + pipeline = Pipeline( + name="test-pipeline", + parameters=[ParameterString(name="ItemList")], + steps=[for_each], + sagemaker_session=mock_session, + ) + steps = _definition_steps(pipeline) + assert len(steps) == 1 + assert steps[0]["Name"] == "ForEachItems" + assert steps[0]["Type"] == "ForEach" + assert steps[0]["Arguments"]["IterableItems"] == {"Get": "Parameters.ItemList"} + assert steps[0]["Arguments"]["MaxConcurrency"] == 2 + assert steps[0]["Arguments"]["Selector"] == { + "Pick": "MaxBy", + "Field": {"Get": "IterationOutput.ProcessItem.OutputParameters['ProcessedScore']"}, + } + body_request = steps[0]["Arguments"]["ForEachBody"][0] + assert body_request["Name"] == "ProcessItem" + assert body_request["Type"] == "Lambda" + assert body_request["Arguments"] == { + "item_name": {"Get": "ForEach.CurrentItem.name"}, + "item_score": {"Get": "ForEach.CurrentItem.score"}, + } + assert body_request["OutputParameters"] == [ + {"OutputName": "ProcessedScore", "OutputType": "Integer"}, + {"OutputName": "ResultPath", "OutputType": "String"}, + ] + + +def test_pipeline_definition_step_output_reference_items(mock_session): + body = CustomStep(name="ProcessItem", input_data=CurrentItem()) + producer = CustomStep(name="ProduceItem") + producer_output = Properties( + step_name="ProduceItem", path="OutputParameters['ItemList']", step=producer + ) + for_each = ForEachStep( + name="ForEachItems", + iterable_items=producer_output, + max_concurrency=1, + for_each_body=[body], + ) + pipeline = Pipeline( + name="test-pipeline", + steps=[for_each, producer], + sagemaker_session=mock_session, + ) + steps = _definition_steps(pipeline) + for_each_request = next(s for s in steps if s["Name"] == "ForEachItems") + assert for_each_request["Arguments"]["IterableItems"] == { + "Get": "Steps.ProduceItem.OutputParameters['ItemList']" + } + # dependency is inferred by the service from the Get expression; the + # compiler does not add an explicit DependsOn for property references + assert "DependsOn" not in for_each_request + # the referenced producer step must still be compiled as a top-level step + assert {s["Name"] for s in steps} == {"ForEachItems", "ProduceItem"} + body_request = for_each_request["Arguments"]["ForEachBody"][0] + assert body_request["Arguments"] == {"input_data": {"Get": "ForEach.CurrentItem"}} + + +def test_pipeline_definition_downstream_references(mock_session): + body = make_lambda_body_step() + for_each = ForEachStep( + name="ForEachItems", + iterable_items=ParameterString(name="ItemList"), + max_concurrency=2, + for_each_body=[body], + selector=ForEachSelector( + pick="MinBy", field=IterationOutput(body).OutputParameters["ProcessedScore"] + ), + ) + downstream = CustomStep( + name="Consumer", + input_data={ + "iter_0": for_each.properties.IterationOutputs[0].ProcessItem.OutputParameters[ + "ResultPath" + ], + "selected": for_each.properties.Selected.OutputParameters["ResultPath"], + }, + ) + pipeline = Pipeline( + name="test-pipeline", + parameters=[ParameterString(name="ItemList")], + steps=[for_each, downstream], + sagemaker_session=mock_session, + ) + steps = _definition_steps(pipeline) + consumer = next(s for s in steps if s["Name"] == "Consumer") + assert consumer["Arguments"]["input_data"] == { + "iter_0": { + "Get": "Steps.ForEachItems.IterationOutputs[0].ProcessItem" + ".OutputParameters['ResultPath']" + }, + "selected": {"Get": "Steps.ForEachItems.Selected.OutputParameters['ResultPath']"}, + } + # dependency is inferred by the service from the Get expressions; + # no explicit DependsOn is added for property references + assert "DependsOn" not in consumer + + +def test_pipeline_definition_for_each_inside_condition_step(mock_session): + body = CustomStep(name="ProcessItem", input_data=CurrentItem().name) + for_each = ForEachStep( + name="ForEachItems", + iterable_items=ParameterString(name="ItemList"), + max_concurrency=2, + for_each_body=[body], + ) + condition = Mock() + condition.to_request.return_value = {"Type": "Equals", "LeftValue": 1, "RightValue": 1} + condition._referenced_steps = [] + condition_step = ConditionStep( + name="CheckSomething", + conditions=[condition], + if_steps=[for_each], + else_steps=[], + ) + pipeline = Pipeline( + name="test-pipeline", + parameters=[ParameterString(name="ItemList")], + steps=[condition_step], + sagemaker_session=mock_session, + ) + steps = _definition_steps(pipeline) + assert len(steps) == 1 + assert steps[0]["Name"] == "CheckSomething" + if_steps = steps[0]["Arguments"]["IfSteps"] + assert len(if_steps) == 1 + assert if_steps[0]["Type"] == "ForEach" + assert if_steps[0]["Arguments"]["ForEachBody"][0]["Name"] == "ProcessItem" + + +def test_body_steps_are_not_top_level_steps(mock_session): + body = CustomStep(name="ProcessItem", input_data=CurrentItem()) + for_each = ForEachStep( + name="ForEachItems", + iterable_items=ParameterString(name="ItemList"), + max_concurrency=2, + for_each_body=[body], + ) + pipeline = Pipeline( + name="test-pipeline", + parameters=[ParameterString(name="ItemList")], + steps=[for_each], + sagemaker_session=mock_session, + ) + steps = _definition_steps(pipeline) + assert [s["Name"] for s in steps] == ["ForEachItems"] + + +def test_duplicate_body_step_name_rejected(mock_session): + body = CustomStep(name="ProcessItem") + for_each = ForEachStep( + name="ForEachItems", + iterable_items=ParameterString(name="ItemList"), + max_concurrency=2, + for_each_body=[body], + ) + duplicate = CustomStep(name="ProcessItem") + pipeline = Pipeline( + name="test-pipeline", + steps=[for_each, duplicate], + sagemaker_session=mock_session, + ) + with pytest.raises(ValueError, match="duplicate names"): + pipeline.definition() + + +def test_pipeline_graph_includes_body_step_edges(mock_session): + body = CustomStep(name="ProcessItem") + for_each = ForEachStep( + name="ForEachItems", + iterable_items=ParameterString(name="ItemList"), + max_concurrency=2, + for_each_body=[body], + ) + graph = PipelineGraph([for_each]) + assert "ProcessItem" in graph.step_map + assert "ProcessItem" in graph.adjacency_list["ForEachItems"] + + +def test_pipeline_definition_literal_list_items_and_aggregate_reference(mock_session): + body = CustomStep(name="ProcessItem", input_data=CurrentItem()) + producer = CustomStep(name="ProduceItem") + producer_output = Properties( + step_name="ProduceItem", path="OutputParameters['ItemA']", step=producer + ) + for_each = ForEachStep( + name="ForEachItems", + iterable_items=[producer_output, "s3://my-bucket/item-b", 42, True], + max_concurrency=2, + for_each_body=[body], + ) + consumer = CustomStep( + name="Consumer", + input_data={ + "status": for_each.properties.Status, + "total": for_each.properties.TotalStepCount, + }, + ) + pipeline = Pipeline( + name="test-pipeline", + steps=[producer, for_each, consumer], + sagemaker_session=mock_session, + ) + steps = _definition_steps(pipeline) + for_each_request = next(s for s in steps if s["Name"] == "ForEachItems") + # same shape as the service's iterable_literal_list.json fixture + assert for_each_request["Arguments"]["IterableItems"] == [ + {"Get": "Steps.ProduceItem.OutputParameters['ItemA']"}, + "s3://my-bucket/item-b", + 42, + True, + ] + consumer_request = next(s for s in steps if s["Name"] == "Consumer") + assert consumer_request["Arguments"]["input_data"] == { + "status": {"Get": "Steps.ForEachItems.Status"}, + "total": {"Get": "Steps.ForEachItems.TotalStepCount"}, + } + + +def test_pipeline_definition_run_conditions(mock_session): + producer = CustomStep(name="ProduceItem") + first = CustomStep(name="FirstBodyStep", input_data=CurrentItem()) + second = CustomStep( + name="SecondBodyStep", + depends_on=[StepDependency(first, TriggerCondition.ANY)], + ) + for_each = ForEachStep( + name="ForEachItems", + iterable_items=ParameterString(name="ItemList"), + max_concurrency=2, + for_each_body=[first, second], + depends_on=[StepDependency(producer, TriggerCondition.FAILED)], + ) + notify = CustomStep( + name="NotifyOnFailure", + depends_on=[StepDependency(for_each, TriggerCondition.FAILED)], + ) + pipeline = Pipeline( + name="test-pipeline", + parameters=[ParameterString(name="ItemList")], + steps=[producer, for_each, notify], + sagemaker_session=mock_session, + ) + steps = {s["Name"]: s for s in _definition_steps(pipeline)} + assert steps["ForEachItems"]["DependsOn"] == [{"Step": "ProduceItem", "Trigger": "Failed"}] + assert steps["NotifyOnFailure"]["DependsOn"] == [{"Step": "ForEachItems", "Trigger": "Failed"}] + body = {s["Name"]: s for s in steps["ForEachItems"]["Arguments"]["ForEachBody"]} + assert "DependsOn" not in body["FirstBodyStep"] + assert body["SecondBodyStep"]["DependsOn"] == [{"Step": "FirstBodyStep", "Trigger": "Any"}] diff --git a/sagemaker-mlops/tests/unit/workflow/test_job_step.py b/sagemaker-mlops/tests/unit/workflow/test_job_step.py index 4e6e572409..194dc65650 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_job_step.py +++ b/sagemaker-mlops/tests/unit/workflow/test_job_step.py @@ -50,6 +50,7 @@ "INFERENCE_COMPONENT": "InferenceComponent", "LINEAGE": "Lineage", "JOB": "Job", + "FOR_EACH": "ForEach", } diff --git a/sagemaker-mlops/tests/unit/workflow/test_steps_compiler.py b/sagemaker-mlops/tests/unit/workflow/test_steps_compiler.py index 54bd834208..7023496068 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_steps_compiler.py +++ b/sagemaker-mlops/tests/unit/workflow/test_steps_compiler.py @@ -25,6 +25,8 @@ ) from sagemaker.mlops.workflow.steps import Step, StepTypeEnum from sagemaker.mlops.workflow.condition_step import ConditionStep +from sagemaker.mlops.workflow.for_each_step import ForEachStep +from sagemaker.core.workflow.step_outputs import StepOutput class TestCompiledStep: @@ -413,7 +415,7 @@ def test_build_raises_error_on_second_call(self, mock_session, mock_step): assert "more than once" in str(exc_info.value) def test_flatten_condition_step(self, mock_session): - """Test _flatten_condition_step flattens nested condition steps.""" + """Test _flatten_compound_step flattens nested condition steps.""" step1 = Mock(spec=Step) step1.name = "step1" @@ -428,11 +430,32 @@ def test_flatten_condition_step(self, mock_session): steps=[condition_step], ) - result = compiler._flatten_condition_step(condition_step) + result = compiler._flatten_compound_step(condition_step) assert len(result) >= 1 assert condition_step in result + def test_flatten_compound_step_for_each(self, mock_session): + """Test _flatten_compound_step flattens ForEach steps.""" + body_step = Mock(spec=Step) + body_step.name = "body-step" + + for_each_step = Mock(spec=ForEachStep) + for_each_step.name = "for-each" + for_each_step.for_each_body = [body_step] + + compiler = StepsCompiler( + pipeline_name="test-pipeline", + sagemaker_session=mock_session, + steps=[for_each_step], + ) + + result = compiler._flatten_compound_step(for_each_step) + + assert len(result) == 2 + assert for_each_step in result + assert body_step in result + def test_push_to_build_queue(self, mock_session, mock_step): """Test _push_to_build_queue adds steps to queue.""" compiler = StepsCompiler(