From 5428911bbae426f5d39dce696347956519b9cab3 Mon Sep 17 00:00:00 2001 From: Daniel Castro Date: Wed, 22 Jul 2026 09:12:23 -0700 Subject: [PATCH 1/6] feat(mlops/workflow): Add ForEachStep for iterating over collections Adds a new ForEachStep to sagemaker.mlops.workflow that iterates a body of steps over a collection at pipeline execution time. Supports CurrentItem / ChildOutput references, a Selector for picking values from prior child outputs, and a configurable max_concurrency in the range [1, 50]. - New for_each_step.py exposing ForEachStep, ForEachSelector, ForEachSelectorPickType, CurrentItem, and ChildOutput. - Add FOR_EACH to StepTypeEnum. - Extend StepsCompiler to compile ForEachStep bodies and share duplicate-name detection with ConditionStep by generalizing _flatten_condition_step into _flatten_compound_step. - Track for_each_body dependencies in PipelineGraph and enforce unique step names across ForEachStep bodies. - Re-export the new symbols from sagemaker.mlops.workflow. - Unit tests covering the new step type, compiler paths, and pipeline wire format. --- .../src/sagemaker/mlops/workflow/__init__.py | 12 + .../mlops/workflow/_steps_compiler.py | 49 +- .../sagemaker/mlops/workflow/for_each_step.py | 452 +++++++++++++++ .../src/sagemaker/mlops/workflow/pipeline.py | 10 +- .../src/sagemaker/mlops/workflow/steps.py | 1 + .../tests/unit/workflow/test_for_each_step.py | 548 ++++++++++++++++++ .../unit/workflow/test_steps_compiler.py | 27 +- 7 files changed, 1088 insertions(+), 11 deletions(-) create mode 100644 sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py create mode 100644 sagemaker-mlops/tests/unit/workflow/test_for_each_step.py diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py index 4f50c1be74..d9e49ff6c1 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, + ChildOutput, +) 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", + "ChildOutput", "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..4ddc2db5b5 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=[step.name for step in 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..6e4f7f8652 --- /dev/null +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py @@ -0,0 +1,452 @@ +# 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 + +_CHILD_OUTPUT_NAMESPACE = "ChildOutput" +_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 _TerminalPropertyReference(_ForEachPropertyReference): + """A property reference that rejects further drilling. + + Used for paths that are at their maximum allowed depth (e.g. a single-level + CurrentItem field like ``ForEach.CurrentItem.name``). + """ + + def __getattr__(self, name: str) -> "_ForEachPropertyReference": + """Reject further attribute access.""" + if name.startswith("_"): + raise AttributeError(name) + raise AttributeError( + f"Cannot drill deeper into '{self._path}'. The service only " + f"supports one level of field access on CurrentItem " + f"(e.g. CurrentItem().name, not CurrentItem().a.b)." + ) + + def __getitem__(self, item: Union[int, str]) -> "_ForEachPropertyReference": + """Reject indexing on a terminal reference.""" + raise TypeError( + f"Cannot drill deeper into '{self._path}'. The service only " + f"supports one level of field access on CurrentItem." + ) + + +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 one level into a top-level field:: + + item = CurrentItem() + item.expr # {"Get": "ForEach.CurrentItem"} + item.name.expr # {"Get": "ForEach.CurrentItem.name"} + + Only one level of field access is supported. Deeper paths + (``item.a.b``) and bracket indexing (``item["score"]``) are rejected. + """ + + def __init__(self): + """Initialize a CurrentItem reference.""" + super().__init__(_CURRENT_ITEM_PATH) + + def __getattr__(self, name: str) -> "_ForEachPropertyReference": + """Allow one level of field access, returning a terminal reference.""" + if name.startswith("_"): + raise AttributeError(name) + return _TerminalPropertyReference(f"{self._path}.{name}", self._refs) + + def __getitem__(self, item: Union[int, str]) -> "_ForEachPropertyReference": + """Reject indexing; the CurrentItem grammar only supports dot paths.""" + raise TypeError( + f"CurrentItem does not support indexing ('[{item!r}]'). Use attribute " + f"access instead, e.g. CurrentItem().{item} for field access." + ) + + +class ChildOutput(_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:: + + ChildOutput(process_step).OutputParameters["InstanceCount"].expr + # {"Get": "ChildOutput.ProcessItem.OutputParameters['InstanceCount']"} + + Args: + body_step (Union[str, Step]): The body step (or its name) whose output + the selector compares across children. + """ + + def __init__(self, body_step: Union[str, Step]): + """Initialize a ChildOutput reference rooted at the given body step.""" + step_name = body_step.name if isinstance(body_step, Step) else body_step + super().__init__(f"{_CHILD_OUTPUT_NAMESPACE}.{step_name}") + + +class _ChildOutputsAccessor: + """Indexed accessor for a ForEach step's per-child outputs. + + ``ChildOutputs`` MUST be indexed (``properties.ChildOutputs[0].``); + a bare ``ChildOutputs.`` reference is invalid, so this accessor is + intentionally not a ``PipelineVariable`` itself. + """ + + def __init__(self, step: Step): + """Initialize a _ChildOutputsAccessor for the given ForEach step.""" + self._step = step + + def __getitem__(self, index: int) -> _ForEachPropertyReference: + """Get a drillable reference to the output of the child at ``index``.""" + if not isinstance(index, int): + raise TypeError( + f"ChildOutputs of ForEach step '{self._step.name}' must be indexed " + f"with an integer, e.g. properties.ChildOutputs[0]." + ) + return _ForEachPropertyReference( + f"Steps.{self._step.name}.ChildOutputs[{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 child's output is exposed to downstream steps via + ``Steps..Selected.``. + + Args: + pick (Union[str, ForEachSelectorPickType]): The pick strategy, + ``"MaxBy"`` or ``"MinBy"``. + field (ChildOutput): A ``ChildOutput`` 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' (a ChildOutput reference).") + if not isinstance(field, PipelineVariable): + raise ValueError( + "ForEachSelector 'field' must be a ChildOutput reference, e.g. " + 'ChildOutput(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 once per item as child steps, enforces a concurrency cap, and exposes + per-child outputs plus aggregate results to downstream steps. + + 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: + + - ``properties.ChildOutputs[i].`` — per-child output (delegates to + the body step's property schema; must be indexed) + - ``properties.Selected.`` — output of the selector-picked child + - ``properties.Status`` (String: Succeeded | PartiallySucceeded | Failed) + - ``properties.InputItemCount``, ``properties.TotalStepCount``, + ``properties.SucceededStepCount``, ``properties.FailedStepCount``, + ``properties.StoppedStepCount`` (Integers) + """ + + _AGGREGATE_PROPERTY_MEMBERS = ( + "Status", + "InputItemCount", + "TotalStepCount", + "SucceededStepCount", + "FailedStepCount", + "StoppedStepCount", + ) + + def __init__( + self, + name: str, + iterable_items: PipelineVariable = 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 (PipelineVariable): REQUIRED. The items to iterate + over. A property reference (e.g. a `ParameterString` or a step + output) whose value must resolve to a JSON list string at + runtime, such as ``'["item-a", "item-b"]'``. The service + currently only accepts the property-reference form; literal + values (including literal lists) are rejected ("a literal value + cannot resolve to a list"). 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. v1 supports exactly ONE body step; the list signature is + future-proofing for multi-step bodies. 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 + child by comparing a ``ChildOutput`` field across children. + 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 = 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) + root_prop.__dict__["ChildOutputs"] = _ChildOutputsAccessor(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." + ) + 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; " + f"a 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." + ) + if len(for_each_body) != 1: + raise ValueError( + f"ForEach step '{name}': for_each_body currently supports exactly " + f"one step; got {len(for_each_body)}. Multi-step bodies are not " + f"yet supported by the service." + ) + 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"{_CHILD_OUTPUT_NAMESPACE}." + if not isinstance(field_path, str) or not field_path.startswith(prefix): + raise ValueError( + f"ForEach step '{name}': Selector.Field must be a ChildOutput " + f"reference (e.g. ChildOutput(body_step).); other " + f"namespaces are constant across children. 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 aggregate results. + + Members: ``Status``, ``InputItemCount``, ``TotalStepCount``, + ``SucceededStepCount``, ``FailedStepCount``, ``StoppedStepCount``, + ``ChildOutputs[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..9cfef502f8 --- /dev/null +++ b/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py @@ -0,0 +1,548 @@ +# 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 ( + ChildOutput, + CurrentItem, + ForEachSelector, + ForEachSelectorPickType, + ForEachStep, +) +from sagemaker.mlops.workflow.lambda_step import ( + LambdaOutput, + LambdaOutputTypeEnum, + LambdaStep, +) +from sagemaker.mlops.workflow.pipeline import Pipeline, PipelineGraph +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 / ChildOutput 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"} + + +def test_current_item_multi_level_drill_rejected(): + # the service only supports one level of field access + ref = CurrentItem().a + with pytest.raises(AttributeError, match="Cannot drill deeper"): + ref.b + + +def test_current_item_indexing_rejected(): + # the CurrentItem grammar only supports dot paths, not bracket indexing + with pytest.raises(TypeError, match="does not support indexing"): + CurrentItem()["score"] + + +def test_current_item_has_no_referenced_steps(): + assert CurrentItem().name._referenced_steps == [] + + +def test_child_output_expr_from_step_and_str(): + body = make_lambda_body_step() + ref = ChildOutput(body).OutputParameters["ProcessedScore"] + assert ref.expr == {"Get": "ChildOutput.ProcessItem.OutputParameters['ProcessedScore']"} + ref_from_str = ChildOutput("ProcessItem").Loss + assert ref_from_str.expr == {"Get": "ChildOutput.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_array_items_rejected(): + # the service currently only accepts the property-reference form + literal_item = '{"name": "lit-c", "score": 5}' + ref_item = ParameterString(name="ItemOne") + with pytest.raises(ValueError, match="literal value cannot resolve to a list"): + make_for_each_step(iterable_items=[ref_item, literal_item]) + + +def test_for_each_step_to_request_with_selector(): + body = make_lambda_body_step() + selector = ForEachSelector( + pick="MaxBy", + field=ChildOutput(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": "ChildOutput.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_aggregate_properties(): + step = make_for_each_step() + props = step.properties + assert props.Status.expr == {"Get": "Steps.ForEachItems.Status"} + assert props.InputItemCount.expr == {"Get": "Steps.ForEachItems.InputItemCount"} + assert props.TotalStepCount.expr == {"Get": "Steps.ForEachItems.TotalStepCount"} + assert props.SucceededStepCount.expr == {"Get": "Steps.ForEachItems.SucceededStepCount"} + assert props.FailedStepCount.expr == {"Get": "Steps.ForEachItems.FailedStepCount"} + assert props.StoppedStepCount.expr == {"Get": "Steps.ForEachItems.StoppedStepCount"} + + +def test_for_each_step_child_outputs_property(): + step = make_for_each_step() + ref = step.properties.ChildOutputs[3].OutputParameters["ResultPath"] + assert ref.expr == {"Get": "Steps.ForEachItems.ChildOutputs[3].OutputParameters['ResultPath']"} + assert ref._referenced_steps == [step] + + +def test_for_each_step_child_outputs_must_be_indexed(): + step = make_for_each_step() + with pytest.raises(TypeError, match="must be indexed"): + step.properties.ChildOutputs["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_rejected(): + with pytest.raises(ValueError, match="exactly one step"): + ForEachStep( + name="ForEachItems", + iterable_items=ParameterString(name="ItemList"), + max_concurrency=2, + for_each_body=[CustomStep(name="A"), CustomStep(name="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=ChildOutput("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=ChildOutput("ProcessItem").Loss) + + +def test_selector_pick_enum(): + selector = ForEachSelector( + pick=ForEachSelectorPickType.MIN_BY, field=ChildOutput("ProcessItem").Loss + ) + assert selector.to_request()["Pick"] == "MinBy" + + +def test_selector_field_must_be_child_output(): + body = make_lambda_body_step() + with pytest.raises(ValueError, match="must be a ChildOutput 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=ChildOutput("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=ChildOutput(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=ChildOutput(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": "ChildOutput.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=ChildOutput(body).OutputParameters["ProcessedScore"] + ), + ) + downstream = CustomStep( + name="Consumer", + input_data={ + "status": for_each.properties.Status, + "input_item_count": for_each.properties.InputItemCount, + "total": for_each.properties.TotalStepCount, + "succeeded": for_each.properties.SucceededStepCount, + "failed": for_each.properties.FailedStepCount, + "stopped": for_each.properties.StoppedStepCount, + "child_0": for_each.properties.ChildOutputs[0].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"] == { + "status": {"Get": "Steps.ForEachItems.Status"}, + "input_item_count": {"Get": "Steps.ForEachItems.InputItemCount"}, + "total": {"Get": "Steps.ForEachItems.TotalStepCount"}, + "succeeded": {"Get": "Steps.ForEachItems.SucceededStepCount"}, + "failed": {"Get": "Steps.ForEachItems.FailedStepCount"}, + "stopped": {"Get": "Steps.ForEachItems.StoppedStepCount"}, + "child_0": {"Get": "Steps.ForEachItems.ChildOutputs[0].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"] 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( From 959f8ab9c1dc092facf0b1aebbada504ea6d513a Mon Sep 17 00:00:00 2001 From: Daniel Castro Date: Tue, 28 Jul 2026 10:04:17 -0700 Subject: [PATCH 2/6] feat(workflow): Remove aggregate properties from ForEachStep Remove Status, InputItemCount, TotalStepCount, SucceededStepCount, FailedStepCount, and StoppedStepCount from ForEachStep.properties. Only ChildOutputs[i] and Selected remain as downstream-referenceable properties. This aligns with the backend policy change: any child failure now causes the entire ForEach step to fail (no partial success). The PartiallySucceeded status is removed, making aggregate counters no longer meaningful to downstream steps. --- .../sagemaker/mlops/workflow/for_each_step.py | 21 ++--------------- .../tests/unit/workflow/test_for_each_step.py | 23 ------------------- 2 files changed, 2 insertions(+), 42 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py index 6e4f7f8652..e0b8b8ca57 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py @@ -256,21 +256,8 @@ class ForEachStep(Step): - ``properties.ChildOutputs[i].`` — per-child output (delegates to the body step's property schema; must be indexed) - ``properties.Selected.`` — output of the selector-picked child - - ``properties.Status`` (String: Succeeded | PartiallySucceeded | Failed) - - ``properties.InputItemCount``, ``properties.TotalStepCount``, - ``properties.SucceededStepCount``, ``properties.FailedStepCount``, - ``properties.StoppedStepCount`` (Integers) """ - _AGGREGATE_PROPERTY_MEMBERS = ( - "Status", - "InputItemCount", - "TotalStepCount", - "SucceededStepCount", - "FailedStepCount", - "StoppedStepCount", - ) - def __init__( self, name: str, @@ -324,8 +311,6 @@ def __init__( 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) root_prop.__dict__["ChildOutputs"] = _ChildOutputsAccessor(self) root_prop.__dict__["Selected"] = _ForEachPropertyReference(f"Steps.{name}.Selected", [self]) self._properties = root_prop @@ -443,10 +428,8 @@ def step_only_arguments(self) -> RequestType: @property def properties(self): - """A Properties object exposing the ForEach step's aggregate results. + """A Properties object exposing the ForEach step's results. - Members: ``Status``, ``InputItemCount``, ``TotalStepCount``, - ``SucceededStepCount``, ``FailedStepCount``, ``StoppedStepCount``, - ``ChildOutputs[i]`` and ``Selected``. + Members: ``ChildOutputs[i]`` and ``Selected``. """ return self._properties diff --git a/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py b/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py index 9cfef502f8..e9dedcc009 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py +++ b/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py @@ -209,17 +209,6 @@ def test_step_only_arguments_excludes_body(): # --------------------------------------------------------------------------- -def test_for_each_step_aggregate_properties(): - step = make_for_each_step() - props = step.properties - assert props.Status.expr == {"Get": "Steps.ForEachItems.Status"} - assert props.InputItemCount.expr == {"Get": "Steps.ForEachItems.InputItemCount"} - assert props.TotalStepCount.expr == {"Get": "Steps.ForEachItems.TotalStepCount"} - assert props.SucceededStepCount.expr == {"Get": "Steps.ForEachItems.SucceededStepCount"} - assert props.FailedStepCount.expr == {"Get": "Steps.ForEachItems.FailedStepCount"} - assert props.StoppedStepCount.expr == {"Get": "Steps.ForEachItems.StoppedStepCount"} - - def test_for_each_step_child_outputs_property(): step = make_for_each_step() ref = step.properties.ChildOutputs[3].OutputParameters["ResultPath"] @@ -434,12 +423,6 @@ def test_pipeline_definition_downstream_references(mock_session): downstream = CustomStep( name="Consumer", input_data={ - "status": for_each.properties.Status, - "input_item_count": for_each.properties.InputItemCount, - "total": for_each.properties.TotalStepCount, - "succeeded": for_each.properties.SucceededStepCount, - "failed": for_each.properties.FailedStepCount, - "stopped": for_each.properties.StoppedStepCount, "child_0": for_each.properties.ChildOutputs[0].OutputParameters["ResultPath"], "selected": for_each.properties.Selected.OutputParameters["ResultPath"], }, @@ -453,12 +436,6 @@ def test_pipeline_definition_downstream_references(mock_session): steps = _definition_steps(pipeline) consumer = next(s for s in steps if s["Name"] == "Consumer") assert consumer["Arguments"]["input_data"] == { - "status": {"Get": "Steps.ForEachItems.Status"}, - "input_item_count": {"Get": "Steps.ForEachItems.InputItemCount"}, - "total": {"Get": "Steps.ForEachItems.TotalStepCount"}, - "succeeded": {"Get": "Steps.ForEachItems.SucceededStepCount"}, - "failed": {"Get": "Steps.ForEachItems.FailedStepCount"}, - "stopped": {"Get": "Steps.ForEachItems.StoppedStepCount"}, "child_0": {"Get": "Steps.ForEachItems.ChildOutputs[0].OutputParameters['ResultPath']"}, "selected": {"Get": "Steps.ForEachItems.Selected.OutputParameters['ResultPath']"}, } From 73771200167eb58b9317a8ae6167c47d8818ff04 Mon Sep 17 00:00:00 2001 From: Daniel Castro Date: Tue, 28 Jul 2026 11:20:53 -0700 Subject: [PATCH 3/6] feat(workflow): Support deep paths and bracket indexing for CurrentItem Remove _TerminalPropertyReference class and CurrentItem overrides. CurrentItem now inherits full drilling from _ForEachPropertyReference: - Deep dot paths: CurrentItem().a.b.c - Bracket indexing: CurrentItem()['score'], CurrentItem()[0] All forms produce valid {"Get": "ForEach.CurrentItem..."} expressions. The service resolves these at expansion time for each child lane. --- .../sagemaker/mlops/workflow/for_each_step.py | 50 +++---------------- .../tests/unit/workflow/test_for_each_step.py | 18 +++---- 2 files changed, 13 insertions(+), 55 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py index e0b8b8ca57..f589bc10ba 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py @@ -76,63 +76,25 @@ def _referenced_steps(self) -> List[Any]: return self._refs -class _TerminalPropertyReference(_ForEachPropertyReference): - """A property reference that rejects further drilling. - - Used for paths that are at their maximum allowed depth (e.g. a single-level - CurrentItem field like ``ForEach.CurrentItem.name``). - """ - - def __getattr__(self, name: str) -> "_ForEachPropertyReference": - """Reject further attribute access.""" - if name.startswith("_"): - raise AttributeError(name) - raise AttributeError( - f"Cannot drill deeper into '{self._path}'. The service only " - f"supports one level of field access on CurrentItem " - f"(e.g. CurrentItem().name, not CurrentItem().a.b)." - ) - - def __getitem__(self, item: Union[int, str]) -> "_ForEachPropertyReference": - """Reject indexing on a terminal reference.""" - raise TypeError( - f"Cannot drill deeper into '{self._path}'. The service only " - f"supports one level of field access on CurrentItem." - ) - - 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 one level into a top-level field:: + 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"} - - Only one level of field access is supported. Deeper paths - (``item.a.b``) and bracket indexing (``item["score"]``) are rejected. + 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) - def __getattr__(self, name: str) -> "_ForEachPropertyReference": - """Allow one level of field access, returning a terminal reference.""" - if name.startswith("_"): - raise AttributeError(name) - return _TerminalPropertyReference(f"{self._path}.{name}", self._refs) - - def __getitem__(self, item: Union[int, str]) -> "_ForEachPropertyReference": - """Reject indexing; the CurrentItem grammar only supports dot paths.""" - raise TypeError( - f"CurrentItem does not support indexing ('[{item!r}]'). Use attribute " - f"access instead, e.g. CurrentItem().{item} for field access." - ) - class ChildOutput(_ForEachPropertyReference): """Reference to a ForEach body step's output, for use in ``ForEachSelector.field``. diff --git a/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py b/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py index e9dedcc009..21efe8d598 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py +++ b/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py @@ -115,19 +115,15 @@ def test_current_item_bare_expr(): 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_multi_level_drill_rejected(): - # the service only supports one level of field access - ref = CurrentItem().a - with pytest.raises(AttributeError, match="Cannot drill deeper"): - ref.b - - -def test_current_item_indexing_rejected(): - # the CurrentItem grammar only supports dot paths, not bracket indexing - with pytest.raises(TypeError, match="does not support indexing"): - CurrentItem()["score"] +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(): From 872d45f5b1d83d4f69af131e5d521a0641fff533 Mon Sep 17 00:00:00 2001 From: Daniel Castro Date: Mon, 3 Aug 2026 14:11:50 -0700 Subject: [PATCH 4/6] feat(workflow): Support multi-step ForEach bodies MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove the "exactly one body step" restriction so ForEach bodies can contain any number of steps. Each iteration spawns one child instance of every body step. Rename ChildOutput / ChildOutputs to IterationOutput / IterationOutputs across the public API. Property reference changes: - properties.IterationOutputs[i].. now requires a body step name qualifier, since an iteration may span multiple body steps. - properties.Selected. resolves directly to the selector- picked iteration's output — no body step name qualifier, because the selector's field already identifies the target body step. --- .../src/sagemaker/mlops/workflow/__init__.py | 4 +- .../sagemaker/mlops/workflow/for_each_step.py | 114 +++++++++--------- .../tests/unit/workflow/test_for_each_step.py | 92 ++++++++------ 3 files changed, 118 insertions(+), 92 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py index d9e49ff6c1..2500f40269 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/__init__.py @@ -56,7 +56,7 @@ ForEachSelector, ForEachSelectorPickType, CurrentItem, - ChildOutput, + IterationOutput, ) from sagemaker.mlops.workflow.lambda_step import LambdaStep, LambdaOutput from sagemaker.mlops.workflow.lineage_step import ( @@ -126,7 +126,7 @@ "ForEachSelector", "ForEachSelectorPickType", "CurrentItem", - "ChildOutput", + "IterationOutput", "LambdaStep", "LambdaOutput", "LineageAssociation", diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py index f589bc10ba..3affd15e59 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py @@ -26,7 +26,7 @@ MAX_CONCURRENCY_MIN = 1 MAX_CONCURRENCY_MAX = 50 -_CHILD_OUTPUT_NAMESPACE = "ChildOutput" +_ITERATION_OUTPUT_NAMESPACE = "IterationOutput" _CURRENT_ITEM_PATH = "ForEach.CurrentItem" # Step types that may not appear in a ForEachBody. @@ -96,47 +96,49 @@ def __init__(self): super().__init__(_CURRENT_ITEM_PATH) -class ChildOutput(_ForEachPropertyReference): +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:: - ChildOutput(process_step).OutputParameters["InstanceCount"].expr - # {"Get": "ChildOutput.ProcessItem.OutputParameters['InstanceCount']"} + 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 children. + the selector compares across iterations. """ def __init__(self, body_step: Union[str, Step]): - """Initialize a ChildOutput reference rooted at the given body 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"{_CHILD_OUTPUT_NAMESPACE}.{step_name}") + super().__init__(f"{_ITERATION_OUTPUT_NAMESPACE}.{step_name}") -class _ChildOutputsAccessor: - """Indexed accessor for a ForEach step's per-child outputs. +class _IterationOutputsAccessor: + """Indexed accessor for a ForEach step's per-iteration outputs. - ``ChildOutputs`` MUST be indexed (``properties.ChildOutputs[0].``); - a bare ``ChildOutputs.`` reference is invalid, so this accessor is - intentionally not a ``PipelineVariable`` itself. + ``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 a _ChildOutputsAccessor for the given ForEach step.""" + """Initialize an _IterationOutputsAccessor for the given ForEach step.""" self._step = step def __getitem__(self, index: int) -> _ForEachPropertyReference: - """Get a drillable reference to the output of the child at ``index``.""" + """Get a drillable reference to the outputs of the iteration at ``index``.""" if not isinstance(index, int): raise TypeError( - f"ChildOutputs of ForEach step '{self._step.name}' must be indexed " - f"with an integer, e.g. properties.ChildOutputs[0]." + 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}.ChildOutputs[{index}]", [self._step] + f"Steps.{self._step.name}.IterationOutputs[{index}]", [self._step] ) @@ -150,14 +152,17 @@ class ForEachSelectorPickType(Enum): class ForEachSelector: """Selects a single child of a ForEach step by comparing a field across children. - The selected child's output is exposed to downstream steps via - ``Steps..Selected.``. + 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 (ChildOutput): A ``ChildOutput`` reference targeting a body step - of the same ForEach step, drilled at least one member deep. + field (IterationOutput): An ``IterationOutput`` reference targeting a + body step of the same ForEach step, drilled at least one member + deep. """ def __init__( @@ -183,11 +188,11 @@ def __init__( f"got: {type(pick)}." ) if field is None: - raise ValueError("ForEachSelector requires 'field' (a ChildOutput reference).") + raise ValueError("ForEachSelector requires 'field' (an IterationOutput reference).") if not isinstance(field, PipelineVariable): raise ValueError( - "ForEachSelector 'field' must be a ChildOutput reference, e.g. " - 'ChildOutput(body_step).OutputParameters["MyOutput"], ' + "ForEachSelector 'field' must be an IterationOutput reference, e.g. " + 'IterationOutput(body_step).OutputParameters["MyOutput"], ' f"got: {type(field)}." ) self.pick = pick @@ -205,19 +210,25 @@ 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 once per item as child steps, enforces a concurrency cap, and exposes - per-child outputs plus aggregate results to downstream steps. - - 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: - - - ``properties.ChildOutputs[i].`` — per-child output (delegates to - the body step's property schema; must be indexed) - - ``properties.Selected.`` — output of the selector-picked child + 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) """ def __init__( @@ -247,12 +258,13 @@ def __init__( 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. v1 supports exactly ONE body step; the list signature is - future-proofing for multi-step bodies. Disallowed body step - types: Condition, Fail, ForEach (no nesting). Body steps may - reference the current item via :class:`CurrentItem`. + 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 - child by comparing a ``ChildOutput`` field across children. + 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 @@ -273,7 +285,7 @@ def __init__( self.selector = selector root_prop = Properties(step_name=name) - root_prop.__dict__["ChildOutputs"] = _ChildOutputsAccessor(self) + root_prop.__dict__["IterationOutputs"] = _IterationOutputsAccessor(self) root_prop.__dict__["Selected"] = _ForEachPropertyReference(f"Steps.{name}.Selected", [self]) self._properties = root_prop @@ -316,12 +328,6 @@ def _validate_for_each_body(name: str, for_each_body: Any): raise ValueError( f"ForEach step '{name}': for_each_body must be a non-empty list " f"of steps." ) - if len(for_each_body) != 1: - raise ValueError( - f"ForEach step '{name}': for_each_body currently supports exactly " - f"one step; got {len(for_each_body)}. Multi-step bodies are not " - f"yet supported by the service." - ) for body_step in for_each_body: if not isinstance(body_step, Step): raise ValueError( @@ -348,12 +354,12 @@ def _validate_selector( f"instance. Got: {type(selector)}." ) field_path = selector.field.expr.get("Get") - prefix = f"{_CHILD_OUTPUT_NAMESPACE}." + 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 a ChildOutput " - f"reference (e.g. ChildOutput(body_step).); other " - f"namespaces are constant across children. Got: {field_path}." + 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) @@ -392,6 +398,6 @@ def step_only_arguments(self) -> RequestType: def properties(self): """A Properties object exposing the ForEach step's results. - Members: ``ChildOutputs[i]`` and ``Selected``. + Members: ``IterationOutputs[i].`` and ``Selected``. """ return self._properties diff --git a/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py b/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py index 21efe8d598..fd2693bc55 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py +++ b/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py @@ -24,11 +24,11 @@ from sagemaker.mlops.workflow.condition_step import ConditionStep from sagemaker.mlops.workflow.fail_step import FailStep from sagemaker.mlops.workflow.for_each_step import ( - ChildOutput, CurrentItem, ForEachSelector, ForEachSelectorPickType, ForEachStep, + IterationOutput, ) from sagemaker.mlops.workflow.lambda_step import ( LambdaOutput, @@ -105,7 +105,7 @@ def make_for_each_step(body_step=None, selector=None, **kwargs): # --------------------------------------------------------------------------- -# CurrentItem / ChildOutput references +# CurrentItem / IterationOutput references # --------------------------------------------------------------------------- @@ -130,12 +130,17 @@ def test_current_item_has_no_referenced_steps(): assert CurrentItem().name._referenced_steps == [] -def test_child_output_expr_from_step_and_str(): +# --------------------------------------------------------------------------- +# IterationOutput references +# --------------------------------------------------------------------------- + + +def test_iteration_output_expr_from_step_and_str(): body = make_lambda_body_step() - ref = ChildOutput(body).OutputParameters["ProcessedScore"] - assert ref.expr == {"Get": "ChildOutput.ProcessItem.OutputParameters['ProcessedScore']"} - ref_from_str = ChildOutput("ProcessItem").Loss - assert ref_from_str.expr == {"Get": "ChildOutput.ProcessItem.Loss"} + 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 == [] @@ -170,13 +175,13 @@ def test_for_each_step_to_request_with_selector(): body = make_lambda_body_step() selector = ForEachSelector( pick="MaxBy", - field=ChildOutput(body).OutputParameters["ProcessedScore"], + 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": "ChildOutput.ProcessItem.OutputParameters['ProcessedScore']" + "Get": "IterationOutput.ProcessItem.OutputParameters['ProcessedScore']" } @@ -205,23 +210,27 @@ def test_step_only_arguments_excludes_body(): # --------------------------------------------------------------------------- -def test_for_each_step_child_outputs_property(): +def test_for_each_step_iteration_outputs_property(): step = make_for_each_step() - ref = step.properties.ChildOutputs[3].OutputParameters["ResultPath"] - assert ref.expr == {"Get": "Steps.ForEachItems.ChildOutputs[3].OutputParameters['ResultPath']"} + ref = step.properties.IterationOutputs[3].ProcessItem.OutputParameters["ResultPath"] + assert ref.expr == { + "Get": "Steps.ForEachItems.IterationOutputs[3].ProcessItem.OutputParameters['ResultPath']" + } assert ref._referenced_steps == [step] -def test_for_each_step_child_outputs_must_be_indexed(): +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.ChildOutputs["ResultPath"] + 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.expr == { + "Get": "Steps.ForEachItems.Selected.OutputParameters['ProcessedScore']" + } assert ref._referenced_steps == [step] @@ -258,14 +267,18 @@ def test_empty_body_rejected(body): make_for_each_step(for_each_body=body) -def test_multi_step_body_rejected(): - with pytest.raises(ValueError, match="exactly one step"): - ForEachStep( - name="ForEachItems", - iterable_items=ParameterString(name="ItemList"), - max_concurrency=2, - for_each_body=[CustomStep(name="A"), CustomStep(name="B")], - ) +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(): @@ -279,26 +292,26 @@ def test_disallowed_body_step_types(): def test_selector_requires_pick_and_field(): with pytest.raises(ValueError, match="requires 'pick'"): - ForEachSelector(pick=None, field=ChildOutput("ProcessItem").Loss) + 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=ChildOutput("ProcessItem").Loss) + ForEachSelector(pick="Median", field=IterationOutput("ProcessItem").Loss) def test_selector_pick_enum(): selector = ForEachSelector( - pick=ForEachSelectorPickType.MIN_BY, field=ChildOutput("ProcessItem").Loss + pick=ForEachSelectorPickType.MIN_BY, field=IterationOutput("ProcessItem").Loss ) assert selector.to_request()["Pick"] == "MinBy" -def test_selector_field_must_be_child_output(): +def test_selector_field_must_be_iteration_output(): body = make_lambda_body_step() - with pytest.raises(ValueError, match="must be a ChildOutput reference"): + 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")), @@ -310,7 +323,7 @@ def test_selector_field_must_target_body_step(): with pytest.raises(ValueError, match="must target a body step"): make_for_each_step( body_step=body, - selector=ForEachSelector(pick="MaxBy", field=ChildOutput("OtherStep").Loss), + selector=ForEachSelector(pick="MaxBy", field=IterationOutput("OtherStep").Loss), ) @@ -319,7 +332,7 @@ def test_selector_field_must_drill_into_member(): with pytest.raises(ValueError, match="drill at least"): make_for_each_step( body_step=body, - selector=ForEachSelector(pick="MaxBy", field=ChildOutput(body)), + selector=ForEachSelector(pick="MaxBy", field=IterationOutput(body)), ) @@ -336,7 +349,7 @@ def test_pipeline_definition_wire_format(mock_session): body = make_lambda_body_step() selector = ForEachSelector( pick="MaxBy", - field=ChildOutput(body).OutputParameters["ProcessedScore"], + field=IterationOutput(body).OutputParameters["ProcessedScore"], ) for_each = ForEachStep( name="ForEachItems", @@ -359,7 +372,7 @@ def test_pipeline_definition_wire_format(mock_session): assert steps[0]["Arguments"]["MaxConcurrency"] == 2 assert steps[0]["Arguments"]["Selector"] == { "Pick": "MaxBy", - "Field": {"Get": "ChildOutput.ProcessItem.OutputParameters['ProcessedScore']"}, + "Field": {"Get": "IterationOutput.ProcessItem.OutputParameters['ProcessedScore']"}, } body_request = steps[0]["Arguments"]["ForEachBody"][0] assert body_request["Name"] == "ProcessItem" @@ -413,13 +426,15 @@ def test_pipeline_definition_downstream_references(mock_session): max_concurrency=2, for_each_body=[body], selector=ForEachSelector( - pick="MinBy", field=ChildOutput(body).OutputParameters["ProcessedScore"] + pick="MinBy", field=IterationOutput(body).OutputParameters["ProcessedScore"] ), ) downstream = CustomStep( name="Consumer", input_data={ - "child_0": for_each.properties.ChildOutputs[0].OutputParameters["ResultPath"], + "iter_0": for_each.properties.IterationOutputs[0].ProcessItem.OutputParameters[ + "ResultPath" + ], "selected": for_each.properties.Selected.OutputParameters["ResultPath"], }, ) @@ -432,8 +447,13 @@ def test_pipeline_definition_downstream_references(mock_session): steps = _definition_steps(pipeline) consumer = next(s for s in steps if s["Name"] == "Consumer") assert consumer["Arguments"]["input_data"] == { - "child_0": {"Get": "Steps.ForEachItems.ChildOutputs[0].OutputParameters['ResultPath']"}, - "selected": {"Get": "Steps.ForEachItems.Selected.OutputParameters['ResultPath']"}, + "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 From 812b351d45c46bd8d986fecfbeb1f0fd0dc840af Mon Sep 17 00:00:00 2001 From: nayan3107 Date: Tue, 13 Oct 2026 15:57:48 -0700 Subject: [PATCH 5/6] feat(mlops/workflow): Align ForEachStep with service The service still resolves the ForEach aggregate members (Status, InputItemCount and the step counts) as downstream references, and accepts a literal IterableItems list of up to 100 scalars or property references. Expose both and add ForEach to the step-type guard test. --- X-AI-Prompt: Fix the ForEachStep gaps against the Tioga backend X-AI-Tool: kiro-cli --- .../sagemaker/mlops/workflow/for_each_step.py | 73 ++++++++++++--- .../tests/unit/workflow/test_for_each_step.py | 91 ++++++++++++++++--- .../tests/unit/workflow/test_job_step.py | 1 + 3 files changed, 136 insertions(+), 29 deletions(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py index 3affd15e59..13b1882c29 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/for_each_step.py @@ -25,6 +25,7 @@ MAX_CONCURRENCY_MIN = 1 MAX_CONCURRENCY_MAX = 50 +MAX_LITERAL_ITERABLE_ITEMS = 100 _ITERATION_OUTPUT_NAMESPACE = "IterationOutput" _CURRENT_ITEM_PATH = "ForEach.CurrentItem" @@ -229,12 +230,29 @@ class ForEachStep(Step): - ``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: PipelineVariable = None, + 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, @@ -246,15 +264,18 @@ def __init__( Args: name (str): The name of the `ForEachStep`. - iterable_items (PipelineVariable): REQUIRED. The items to iterate - over. A property reference (e.g. a `ParameterString` or a step - output) whose value must resolve to a JSON list string at - runtime, such as ``'["item-a", "item-b"]'``. The service - currently only accepts the property-reference form; literal - values (including literal lists) are rejected ("a literal value - cannot resolve to a list"). Items are commonly JSON-object - strings whose fields are addressed via ``CurrentItem().`` - inside the body step. + 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 @@ -279,12 +300,16 @@ def __init__( self._validate_for_each_body(name, for_each_body) self._validate_selector(name, selector, for_each_body) - self.iterable_items = iterable_items + 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 @@ -296,13 +321,29 @@ def _validate_iterable_items(name: str, iterable_items: Any): 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." + 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; " - f"a literal value cannot resolve to a list. " + 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)}." ) @@ -398,6 +439,8 @@ def step_only_arguments(self) -> RequestType: def properties(self): """A Properties object exposing the ForEach step's results. - Members: ``IterationOutputs[i].`` and ``Selected``. + Members: ``Status``, ``InputItemCount``, ``TotalStepCount``, + ``SucceededStepCount``, ``FailedStepCount``, ``StoppedStepCount``, + ``IterationOutputs[i].`` and ``Selected``. """ return self._properties diff --git a/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py b/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py index fd2693bc55..180e01a731 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py +++ b/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py @@ -121,9 +121,7 @@ def test_current_item_drilled_expr(): 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']" - } + assert CurrentItem()["nested"]["key"].expr == {"Get": "ForEach.CurrentItem['nested']['key']"} def test_current_item_has_no_referenced_steps(): @@ -163,12 +161,23 @@ def test_for_each_step_to_request_reference_items(item_list_parameter): assert args["ForEachBody"][0]["Type"] == "Lambda" -def test_for_each_step_literal_array_items_rejected(): - # the service currently only accepts the property-reference form - literal_item = '{"name": "lit-c", "score": 5}' +def test_for_each_step_literal_list_items_accepted(): ref_item = ParameterString(name="ItemOne") - with pytest.raises(ValueError, match="literal value cannot resolve to a list"): - make_for_each_step(iterable_items=[ref_item, literal_item]) + 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(): @@ -219,6 +228,24 @@ def test_for_each_step_iteration_outputs_property(): 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"): @@ -228,9 +255,7 @@ def test_for_each_step_iteration_outputs_must_be_indexed(): 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.expr == {"Get": "Steps.ForEachItems.Selected.OutputParameters['ProcessedScore']"} assert ref._referenced_steps == [step] @@ -451,9 +476,7 @@ def test_pipeline_definition_downstream_references(mock_session): "Get": "Steps.ForEachItems.IterationOutputs[0].ProcessItem" ".OutputParameters['ResultPath']" }, - "selected": { - "Get": "Steps.ForEachItems.Selected.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 @@ -539,3 +562,43 @@ def test_pipeline_graph_includes_body_step_edges(mock_session): 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"}, + } 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", } From 8dd22eca9989c8b2fa9e80a00935ada30f541763 Mon Sep 17 00:00:00 2001 From: nayan3107 Date: Fri, 16 Oct 2026 11:59:54 -0700 Subject: [PATCH 6/6] 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 --- .../mlops/workflow/_steps_compiler.py | 2 +- .../tests/unit/workflow/test_for_each_step.py | 33 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/sagemaker-mlops/src/sagemaker/mlops/workflow/_steps_compiler.py b/sagemaker-mlops/src/sagemaker/mlops/workflow/_steps_compiler.py index 4ddc2db5b5..7ce1b926ae 100644 --- a/sagemaker-mlops/src/sagemaker/mlops/workflow/_steps_compiler.py +++ b/sagemaker-mlops/src/sagemaker/mlops/workflow/_steps_compiler.py @@ -405,7 +405,7 @@ def _build_for_each_step(self, for_each_step: ForEachStep) -> ForEachStep: 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=[step.name for step in depends_on], + depends_on=self._to_depends_on_request(for_each_step, depends_on), ) return compiled_for_each_step diff --git a/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py b/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py index 180e01a731..c83e576ff7 100644 --- a/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py +++ b/sagemaker-mlops/tests/unit/workflow/test_for_each_step.py @@ -36,6 +36,7 @@ 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 @@ -602,3 +603,35 @@ def test_pipeline_definition_literal_list_items_and_aggregate_reference(mock_ses "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"}]