Repository navigation
Feat/foreach step - #6473
Feat/foreach step#6473
Conversation
GitHub Actions
Error
GitHub Actions
Error
GitHub Actions
Error
839b04b to
8dd22ec
Compare
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
🤖 AI code review
This PR adds a new ForEachStep workflow step (plus ForEachSelector, CurrentItem, IterationOutput, and supporting references) to sagemaker.mlops.workflow, and wires it into StepsCompiler and PipelineGraph following the established ConditionStep compound-step pattern. The implementation is consistent with existing conventions: it overrides step_only_arguments so body references are not treated as top-level dependencies, mirrors _flatten_compound_step/_build_*/_generate_step_map handling, and interpolates correctly as a PipelineVariable. Validation is comprehensive and the test suite is thorough (wire format, downstream references, run conditions, nesting in ConditionStep, duplicate-name enforcement). I found no correctness or security issues; only one low-severity edge case worth a glance.
Reviewed commit 8dd22ec. Automated review; verify before acting.
| f"IterationOutputs of ForEach step '{self._step.name}' must be indexed " | ||
| f"with an integer, e.g. properties.IterationOutputs[0].<bodyStepName>." | ||
| ) | ||
| return _ForEachPropertyReference( |
There was a problem hiding this comment.
🟡 Low: __getitem__ uses isinstance(item, int) to decide between list-index and map-key formatting. Because bool is a subclass of int, IterationOutputs[True] (or ref[False]) is silently treated as index 1/0 rather than rejected. The _IterationOutputsAccessor.__getitem__ already explicitly guards against non-int keys; you may want the same care here (e.g. isinstance(item, int) and not isinstance(item, bool)) so booleans don't accidentally produce an integer index. Minor, since passing a bool here is unusual.
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.
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.
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.
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].<bodyStepName>.<member> now requires a body step name qualifier, since an iteration may span multiple body steps. - properties.Selected.<member> 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.
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
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
8dd22ec to
18a890d
Compare
There was a problem hiding this comment.
🤖 AI code review
This PR adds ForEachStep (plus ForEachSelector, ForEachSelectorPickType, CurrentItem, IterationOutput) to sagemaker.mlops.workflow, wires it into the public __init__ exports, the StepsCompiler, and PipelineGraph, and adds a new ForEach StepTypeEnum. The implementation closely mirrors the existing ConditionStep compound-step patterns (flattening, step-map recursion, JSON-serialization flag propagation, adjacency edges) and is backed by a thorough unit test suite covering validation, wire format, downstream references, and run-condition triggers. The code is well-structured and I found no correctness or compatibility bugs; only minor points below.
Reviewed commit 18a890d. Automated review; verify before acting.
| root_prop.__dict__["Selected"] = _ForEachPropertyReference(f"Steps.{name}.Selected", [self]) | ||
| self._properties = root_prop | ||
|
|
||
| @staticmethod |
There was a problem hiding this comment.
🟡 Low: _validate_selector indexes selector.field.expr.get("Get"), but field is only checked to be a PipelineVariable in ForEachSelector.__init__, not an IterationOutput. A PipelineVariable subclass whose expr is not a dict (e.g. returns a list or string) would raise AttributeError here instead of the intended ValueError. Consider guarding with isinstance(selector.field.expr, dict) before .get, or validating the field type up front. Low severity since the public IterationOutput always returns a dict expr.
| @staticmethod | ||
| def _validate_selector( | ||
| name: str, selector: Optional[ForEachSelector], for_each_body: List[Step] | ||
| ): |
There was a problem hiding this comment.
🟡 Low: arguments and step_only_arguments duplicate the IterableItems/MaxConcurrency/Selector construction. step_only_arguments could be built once and arguments extend it with ForEachBody, which keeps the two in sync (compare ConditionStep.step_only_arguments reusing Conditions). Minor maintenance nit.
This branch was successfully deployed
3 participants
What: Add ForEachStep to sagemaker.mlops.workflow. A ForEach step runs its body step(s) once per item of a list resolved at execution time, with a concurrency cap of 1-50.
iterable_items: a property reference that resolves to a JSON list, or a literal list of up to 100 scalars/property references.
for_each_body: one or more steps (Condition, Fail and nested ForEach are not allowed). Body steps read the current item via CurrentItem(), including deep paths and bracket indexing (CurrentItem().a.b, CurrentItem()["score"]).
ForEachSelector(pick="MaxBy"|"MinBy", field=IterationOutput(body_step).) picks one iteration.
Downstream references: properties.IterationOutputs[i].., properties.Selected., and the aggregates Status, InputItemCount, TotalStepCount, SucceededStepCount, FailedStepCount, StoppedStepCount.
Run conditions: depends_on on the ForEach step and on its body steps accepts StepDependency(step, TriggerCondition.<...>), serialized the same way as for other steps.
StepsCompiler and PipelineGraph compile ForEach bodies and enforce unique step names across bodies.
Why: SageMaker Pipelines adds a ForEach step type; this exposes it in the PySDK.
Test: Unit tests for the new step, compiler paths and pipeline wire format, including full Pipeline.definition() checks and run-condition triggers on the ForEach step, its body and downstream steps; all 441 tests in sagemaker-mlops/tests/unit/workflow pass.
flake8 and pylint still haven't been run locally, so CI will be the first lint check.