Replace the for loop in the Scatter Gather node with a do_while loop. (Azure#29633)

* add temp files

* more stuff

* p-for

* add files before switching branches

* finish up skeleton code

* save progress

* lots of stuff

* remove pdb statements

* temp nodes

* progress pre-main-merge

* make things work

* Changes to aggreagate outputs

* Merge, and lots of other stuff

* remove test nodes and add real output

* mark node as experimental

* nit fixes

* Removing dependency on PyYaml

* nits

* linting nits general pass

* lint pass v2

* disable validation

* better datstore anchoring and testing

* Update sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/fl_scatter_gather.py

Co-authored-by: kdestin <[email protected]>

* PR comments

* testing

* remove pdb

* lint and dev require mldesigner

* spell write

* lint

* more lint

* fix lint-suggested breaking changes

* regenerate test recording

* minor test fix

* run black formatter but for real

* cl and package dep warning

* move test constants to fixtures, make mldesigner import checked in code

* formatting

* rerun black after merge

* implement do_while

* import from dsl module

* add pipeline comp

* update unitest and e2etest

* remove redundant comments

* set telemetry flag for SG node

* formatting + add e2etesting and

* resolve pylint issues

* fix pylint and black issues

---------

Co-authored-by: Miles Holland <[email protected]>
Co-authored-by: MilesHolland <[email protected]>
Co-authored-by: singankit <[email protected]>
Co-authored-by: kdestin <[email protected]>

Author: 4 people

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_builders/fl_scatter_gather.py

@@ -16,13 +16,16 @@
 from azure.ai.ml.entities._assets.federated_learning_silo import FederatedLearningSilo
 from azure.ai.ml.entities._component.component import Component
 from azure.ai.ml.entities._validation import MutableValidationResult
-from .subcomponents import aggregate_output
+from azure.ai.ml.dsl._do_while import do_while
+from azure.ai.ml.dsl import pipeline
+from .subcomponents import create_scatter_output_table
+
 
 # TODO 2293610: add support for more types of outputs besides uri_folder and mltable
 # Likely types that ought to be mergeable: string, int, uri_file
 MERGE_COMPONENT_MAPPING = {
-    "mltable": aggregate_output,
-    "uri_folder": aggregate_output,
+    "mltable": create_scatter_output_table,
+    "uri_folder": create_scatter_output_table,
 }
 
 
@@ -34,7 +37,7 @@
 # big TODO: For some reason, surfacing this file in  __init__.py causes
 #  a circular import exception on the first attempted import
 # In notebooks, the second import succeeds, but then causes a silent failure where the
-# MLDesigner component created by the subcomponents.aggregate_output function
+# MLDesigner component created by the subcomponents.create_scatter_output_table function
 # will produce a ComponentExecutor object instead of the actual component.
 # TODO 2293541: Add telemetry of some sort
 # pylint: disable=too-many-instance-attributes
@@ -101,37 +104,20 @@ def __init__(
         self.silo_to_aggregation_argument_map = silo_to_aggregation_argument_map
         self.aggregation_to_silo_argument_map = aggregation_to_silo_argument_map
         self.max_iterations = max_iterations
-        # subgraph is a list of all iteration's executed steps, stored in a dictionary for each iteration
-        self.subgraph = []
         self._init = True  # Needed by parent class to work properly
 
-        executed_aggregation_step = None
-        # TODO 2293573: replace this for-loop with a do-while node.
-        for _ in range(self.max_iterations):
-            # Create inputs for silo components
-            silo_inputs = {}
-            # Start with static inputs
-            silo_inputs.update(self.shared_silo_kwargs)
-            # merge in inputs passed in from previous iteration's aggregate step if they exist
-            if executed_aggregation_step is not None and self.aggregation_to_silo_argument_map is not None:
-                silo_inputs.update(
-                    FLScatterGather._extract_outputs(
-                        executed_aggregation_step.outputs, self.aggregation_to_silo_argument_map
-                    )
-                )
-            # per-silo inputs added in during scatter_gather
-            # Run scatter-gather iteration.
-            self.subgraph.append(self.scatter_gather(silo_inputs))
-            # re-assign previous agg step to extract outputs for next iter's inputs
-            executed_aggregation_step = self.subgraph[-1]["aggregation"]
+        self.scatter_gather_graph = self.scatter_gather()
+
+        # set SG node flag for telemetry
+        self.scatter_gather_graph.properties["azureml.telemetry.attribute"] = "scatter-gather"
 
         # set output to final aggregation step's output
-        self._outputs = self.subgraph[-1]["aggregation"].outputs
+        self._outputs = self.scatter_gather_graph.outputs
         super(FLScatterGather, self).__init__(
             type=JobType.COMPONENT,  # pylint: disable=redefined-builtin
             component=None,
             inputs=None,
-            outputs=self.subgraph[-1]["aggregation"].outputs,
+            outputs=self.scatter_gather_graph.outputs,
             name=None,
             display_name=None,
             description=None,
@@ -142,14 +128,117 @@ def __init__(
             experiment_name=None,
         )
 
+    def scatter_gather(self):
+        @pipeline(
+            name="Scatter gather",
+            description="It includes all steps that need to be executed in silo and aggregation",
+        )
+        def scatter_gather_iteration_body(**silo_inputs):
+            """
+                Performs a scatter-gather iteration by running copies of the silo step on different
+            computes/datstores according to this node's silo configs. The outputs of these
+            silo components are then merged by an internal helper component. The merged values
+            are then inputted into the user-provided aggregation component. Returns the executed aggregation component.
+            Args:
+                silo_inputs (dict): A dictionary of names and Inputs to be injected into each executed silo step.
+                    This dictionary is merged with silo-specific inputs before each executed.
+            """
+
+            silo_outputs = []
+            # TODO 2293586 replace this for-loop with a parallel-for node
+            for silo_config in self.silo_configs:
+                silo_inputs.update(silo_config.inputs)
+                executed_silo_component = self.silo_component(**silo_inputs)
+                for v, k in executed_silo_component.inputs.items():
+                    if v in silo_config.inputs and k.type == "uri_folder":
+                        k.mode = "ro_mount"
+                FLScatterGather._anchor_step(
+                    pipeline_step=executed_silo_component,
+                    compute=silo_config.compute,
+                    internal_datastore=silo_config.datastore,
+                    orchestrator_datastore=self.aggregation_datastore,
+                )
+                # add to silo outputs list
+                silo_outputs.append(executed_silo_component)
+
+            # produce internal argument-merging components and record them in local subgraph
+            merge_comp_mapping = self._inject_merge_components(silo_outputs)
+
+            # produce aggregate step inputs by merging static kwargs and mapped arguments from
+            # internal merge components
+            agg_inputs = {}
+            agg_inputs.update(self.aggregation_kwargs)
+            internal_merge_outputs = {
+                self._get_aggregator_input_name(k): v.outputs.aggregated_output for k, v in merge_comp_mapping.items()
+            }
+            agg_inputs.update(internal_merge_outputs)
+
+            # run the user aggregation step
+            executed_aggregation_component = self.aggregation_component(**agg_inputs)
+            # Set mode of aggregated mltable inputs as eval mount to allow files referenced within the table
+            # to be accessible by the component
+            for name, agg_input in executed_aggregation_component.inputs.items():
+                if name in self.silo_to_aggregation_argument_map.keys() and agg_input.type == "mltable":
+                    agg_input.mode = "eval_download"
+
+            # Anchor both the internal merge components and the user-supplied aggregation step
+            # to the aggregation compute and datastore
+            if self.aggregation_compute is not None and self.aggregation_datastore is not None:
+                # internal merge component is also siloed to wherever the aggregation component lives.
+                for executed_merge_component in merge_comp_mapping.values():
+                    FLScatterGather._anchor_step(
+                        pipeline_step=executed_merge_component,
+                        compute=self.aggregation_compute,
+                        internal_datastore=self.aggregation_datastore,
+                        orchestrator_datastore=self.aggregation_datastore,
+                    )
+                FLScatterGather._anchor_step(
+                    pipeline_step=executed_aggregation_component,
+                    compute=self.aggregation_compute,
+                    internal_datastore=self.aggregation_datastore,
+                    orchestrator_datastore=self.aggregation_datastore,
+                )
+            return executed_aggregation_component.outputs
+
+        @pipeline(name="Scatter gather graph")
+        def create_scatter_gather_graph():
+            """
+            Creates a scatter-gather graph by executing the scatter_gather_iteration_body
+            function in a do-while loop. The loop terminates when the user-supplied
+            termination condition is met.
+            """
+
+            silo_inputs = {}
+            # Start with static inputs
+            silo_inputs.update(self.shared_silo_kwargs)
+
+            # merge in inputs passed in from previous iteration's aggregate step)
+            if self.aggregation_to_silo_argument_map is not None:
+                silo_inputs.update({v: None for v in self.aggregation_to_silo_argument_map.values()})
+
+            scatter_gather_body = scatter_gather_iteration_body(**silo_inputs)
+
+            # map aggregation outputs to scatter inputs
+            do_while_mapping = {
+                k: getattr(scatter_gather_body.inputs, v) for k, v in self.aggregation_to_silo_argument_map.items()
+            }
+
+            do_while(
+                body=scatter_gather_body,
+                mapping=do_while_mapping,
+                max_iteration_count=self.max_iterations,
+            )
+            return scatter_gather_body.outputs
+
+        return create_scatter_gather_graph()
+
     # TODO potential set default fail_on_missing value to false
     @classmethod
     def _extract_outputs(cls, component_output: Output, argument_map: Dict, fail_on_missing=False):
         """
             Pulls values from a component_output, as specified by the keys of the
         inputted argument_map, and groups in a new dictionary. The keys of the new dictionary
         are specified by the items of the argument_map dictionary.
-
         Example
             component_output = {"one" : 1, "two": 2, "three" : 3}
             argument_map = {"one" : "red", "two" : "two"}
@@ -172,99 +261,13 @@ def _extract_outputs(cls, component_output: Output, argument_map: Dict, fail_on_
                 result[v] = component_output[k]
         return result
 
-    def scatter_gather(self, silo_inputs: Dict):
-        """
-            Performs a scatter-gather iteration by running copies of the silo step on different
-        computes/datstores according to this node's silo configs. The outputs of these
-        silo components are then merged by an internal helper component. The merged values
-        are then inputted into the user-provided aggregation component. Returns the executed aggregation component.
-
-        Args:
-            silo_inputs (dict): A dictionary of names and Inputs to be injected into each executed silo step.
-                This dictionary is merged with silo-specific inputs before each executed.
-        Returns:
-            sg_graph (dict): A dictionary containing the subgraph of this scatter-gather iteration.
-                Contains three indexes which contain a list of executed silo steps, a list of internally
-                created and executed merge components, and the executed aggregation step.
-        """
-
-        sg_graph = {"silo_steps": []}
-        siloed_outputs = {}
-        # TODO 2293586 replace this for-loop with a parallel-for node
-        # pylint: disable=consider-using-enumerate
-        for i in range(len(self.silo_configs)):
-            silo_config = self.silo_configs[i]
-            silo_inputs.update(silo_config.inputs)
-            executed_silo_component = self.silo_component(**silo_inputs)
-            for v, k in executed_silo_component.inputs.items():
-                if v in silo_config.inputs and k.type == "uri_folder":
-                    k.mode = "ro_mount"
-            FLScatterGather._anchor_step(
-                pipeline_step=executed_silo_component,
-                compute=silo_config.compute,
-                internal_datastore=silo_config.datastore,
-                orchestrator_datastore=self.aggregation_datastore,
-            )
-            # include executed silo steps in recorded subgraph
-            sg_graph["silo_steps"].append(executed_silo_component)
-
-            # Extract user-specified outputs from the silo component, rename them as needed,
-            # annotate them with the silo's index, then jam them all into the
-            # variable-length internal component's input list.
-            siloed_outputs.update(
-                {
-                    "{}_{}".format(k, i): v
-                    for k, v in FLScatterGather._extract_outputs(
-                        executed_silo_component.outputs, self.silo_to_aggregation_argument_map
-                    ).items()
-                }
-            )
-
-        # produce internal argument-merging components and record them in local subgraph
-        merge_comp_mapping = self._inject_merge_components(sg_graph["silo_steps"])
-        sg_graph["mergers"] = list(merge_comp_mapping.values())
-
-        # produce aggregate step inputs by merging static kwargs and mapped arguments from
-        # internal merge components
-        agg_inputs = {}
-        agg_inputs.update(self.aggregation_kwargs)
-        internal_merge_outputs = {
-            self._get_aggregator_input_name(k): v.outputs.aggregated_output for k, v in merge_comp_mapping.items()
-        }
-        agg_inputs.update(internal_merge_outputs)
-
-        # run the user aggregation step
-        executed_aggregation_component = self.aggregation_component(**agg_inputs)
-        # Set mode of aggregated mltable inputs as eval mount to allow files referenced within the table
-        # to be accessible by the component
-        for name, agg_input in executed_aggregation_component.inputs.items():
-            if name in self.silo_to_aggregation_argument_map.keys() and agg_input.type == "mltable":
-                agg_input.mode = "eval_download"
-        # record aggregation step in subgraph
-        sg_graph["aggregation"] = executed_aggregation_component
-
-        # Anchor both the internal merge components and the user-supplied aggregation step
-        # to the aggregation compute and datastore
-        if self.aggregation_compute is not None and self.aggregation_datastore is not None:
-            # internal merge component is also siloed to wherever the aggregation component lives.
-            for executed_merge_component in merge_comp_mapping.values():
-                FLScatterGather._anchor_step(
-                    pipeline_step=executed_merge_component,
-                    compute=self.aggregation_compute,
-                    internal_datastore=self.aggregation_datastore,
-                    orchestrator_datastore=self.aggregation_datastore,
-                )
-            FLScatterGather._anchor_step(
-                pipeline_step=executed_aggregation_component,
-                compute=self.aggregation_compute,
-                internal_datastore=self.aggregation_datastore,
-                orchestrator_datastore=self.aggregation_datastore,
-            )
-        return sg_graph
-
     @classmethod
     def _get_fl_datastore_path(
-        cls, datastore_name: str, output_name: str, unique_id: str = "${{name}}", iteration_num: Optional[int] = None
+        cls,
+        datastore_name: str,
+        output_name: str,
+        unique_id: str = "${{name}}",
+        iteration_num: Optional[int] = None,
     ) -> str:
         """Construct a path string using the inputted values. The important aspect is that this produces a
         path with a specified datastore.
@@ -434,7 +437,10 @@ def _anchor_step(
                 if output.type in ANCHORABLE_OUTPUT_TYPES:
                     validation_result.merge_with(
                         cls._check_or_set_datastore(
-                            name=name, output=output, target_datastore=orchestrator_datastore, iteration_num=iteration
+                            name=name,
+                            output=output,
+                            target_datastore=orchestrator_datastore,
+                            iteration_num=iteration,
                         )
                     )
         else:
@@ -463,7 +469,6 @@ def validate_inputs(
         max_iterations: int,
         raise_error=False,
     ) -> MutableValidationResult:
-
         validation_result = cls._create_empty_validation_result()
 
         # saved values for validation later on
@@ -474,7 +479,8 @@ def validate_inputs(
         # validate silo component
         if silo_component is None:
             validation_result.append_error(
-                yaml_path="silo_component", message="silo_component is a required argument for the scatter gather node."
+                yaml_path="silo_component",
+                message="silo_component is a required argument for the scatter gather node.",
             )
         else:
             # ensure that silo component has both inputs and outputs
@@ -522,10 +528,14 @@ def validate_inputs(
         # validate silos configs
         if silo_configs is None:
             validation_result.append_error(
-                yaml_path="silo_configs", message="silo_configs is a required argument for the scatter gather node."
+                yaml_path="silo_configs",
+                message="silo_configs is a required argument for the scatter gather node.",
             )
         elif len(silo_configs) == 0:
-            validation_result.append_error(yaml_path="silo_configs", message="silo_configs cannot be an empty list.")
+            validation_result.append_error(
+                yaml_path="silo_configs",
+                message="silo_configs cannot be an empty list.",
+            )
         else:
             first_silo = silo_configs[0]
             expected_inputs = []
@@ -646,7 +656,8 @@ def validate_inputs(
 
         if max_iterations < 1:
             validation_result.append_error(
-                yaml_path="max_iterations", message=f"max_iterations must be a positive value, not '{max_iterations}'."
+                yaml_path="max_iterations",
+                message=f"max_iterations must be a positive value, not '{max_iterations}'.",
             )
 
         return validation_result.try_raise(
@@ -656,7 +667,11 @@ def validate_inputs(
 
     @classmethod
     def _custom_fl_data_path(
-        cls, datastore_name, output_name, unique_id="${{name}}", iteration_num="${{iteration_num}}"
+        cls,
+        datastore_name,
+        output_name,
+        unique_id="${{name}}",
+        iteration_num="${{iteration_num}}",
     ):
         """Produces a path to store the data during FL training.
         Args:
@@ -682,7 +697,11 @@ def _get_aggregator_input_name(self, silo_output_name):
 
     @classmethod
     def _try_create_default_mappings(
-        cls, silo_comp: Component, agg_comp: Component, silo_agg_map: Dict, agg_silo_map: Dict
+        cls,
+        silo_comp: Component,
+        agg_comp: Component,
+        silo_agg_map: Dict,
+        agg_silo_map: Dict,
     ):
         """
         This function tries to produce dictionaries that link the silo and aggregation
@@ -739,7 +758,10 @@ def _inject_merge_components(self, executed_silo_components):
         executed_component = executed_silo_components[0]
 
         merge_comp_mapping = {}
-        for silo_output_argument_name, _ in self.silo_to_aggregation_argument_map.items():
+        for (
+            silo_output_argument_name,
+            _,
+        ) in self.silo_to_aggregation_argument_map.items():
             merge_comp = self._get_merge_component(executed_component.outputs[silo_output_argument_name].type)
             merge_component_inputs = {
                 silo_output_argument_name