[ML][Pipelines] Support assets in parallel_for items. (Azure#28673)

* support yaml cases

* support in SDK

* fix tests

* expose _to_rest_items interface

* remove unnecessary tests

* fix empty items

* fix tests

Author: D-W-

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/pipeline/control_flow_job.py

@@ -4,13 +4,14 @@
 import copy
 import json
 
-from marshmallow import INCLUDE, fields, pre_dump
+from marshmallow import INCLUDE, fields, pre_dump, pre_load
 
 from azure.ai.ml._schema.core.fields import DataBindingStr, NestedField, StringTransformedEnum, UnionField
 from azure.ai.ml._schema.core.schema import PathAwareSchema
 from azure.ai.ml.constants._component import ControlFlowType
 
 from ..job.input_output_entry import OutputSchema
+from ..job.input_output_fields_provider import InputsField
 from ..job.job_limits import DoWhileLimitsSchema
 from .component_job import _resolve_outputs
 from .pipeline_job_io import OutputBindingStr
@@ -73,14 +74,21 @@ def resolve_inputs_outputs(self, data, **kwargs):  # pylint: disable=no-self-use
 
         return result
 
+    @pre_dump
+    def convert_control_flow_body_to_binding_str(self, data, **kwargs):  # pylint: disable=no-self-use, unused-argument
+
+        return super(DoWhileSchema, self).convert_control_flow_body_to_binding_str(data, **kwargs)
+
 
 class ParallelForSchema(BaseLoopSchema):
     type = StringTransformedEnum(allowed_values=[ControlFlowType.PARALLEL_FOR])
     items = UnionField(
         [
+            fields.Dict(keys=fields.Str(), values=InputsField()),
+            fields.List(InputsField()),
+            # put str in last to make sure other type items won't become string when dumps.
+            # TODO: only support binding here
             fields.Str(),
-            fields.Dict(keys=fields.Str(), values=fields.Dict()),
-            fields.List(fields.Dict()),
         ],
         required=True,
     )
@@ -90,19 +98,42 @@ class ParallelForSchema(BaseLoopSchema):
         values=UnionField([OutputBindingStr, NestedField(OutputSchema)], allow_none=True),
     )
 
+    @pre_load
+    def load_items(self, data, **kwargs):  # pylint: disable=no-self-use, unused-argument
+        # load items from json to convert the assets in it to rest
+        try:
+            items = data["items"]
+            if isinstance(items, str):
+                items = json.loads(items)
+            data["items"] = items
+        except Exception: # pylint: disable=broad-except
+            pass
+        return data
+
     @pre_dump
-    def serialize_items(self, data, **kwargs):  # pylint: disable=no-self-use, unused-argument
-        from azure.ai.ml.entities._job.pipeline._io import InputOutputBase
+    def convert_control_flow_body_to_binding_str(self, data, **kwargs):  # pylint: disable=no-self-use, unused-argument
 
-        result = copy.copy(data)
-        if isinstance(result.items, (dict, list)):
-            # use str to serialize input/output builder
-            result._items = json.dumps(result.items, default=lambda x: str(x) if isinstance(x, InputOutputBase) else x)
-        return result
+        return super(ParallelForSchema, self).convert_control_flow_body_to_binding_str(data, **kwargs)
 
     @pre_dump
     def resolve_outputs(self, job, **kwargs):  # pylint: disable=unused-argument
 
         result = copy.copy(job)
         _resolve_outputs(result, job)
         return result
+
+    @pre_dump
+    def serialize_items(self, data, **kwargs):  # pylint: disable=no-self-use, unused-argument
+        # serialize items to json string to avoid being removed by _dump_for_validation
+        from azure.ai.ml.entities._job.pipeline._io import InputOutputBase
+
+        def _binding_handler(obj):
+            if isinstance(obj, InputOutputBase):
+                return str(obj)
+            return repr(obj)
+
+        result = copy.copy(data)
+        if isinstance(result.items, (dict, list)):
+            # use str to serialize input/output builder
+            result._items = json.dumps(result.items, default=_binding_handler)
+        return result

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

@@ -2,9 +2,10 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 import json
+import os
 from typing import Dict, Union
 
-from azure.ai.ml import Output
+from azure.ai.ml import Output, Input
 from azure.ai.ml._schema import PathAwareSchema
 from azure.ai.ml._schema.pipeline.control_flow_job import ParallelForSchema
 from azure.ai.ml._utils.utils import is_data_binding_expression
@@ -28,7 +29,7 @@ class ParallelFor(LoopNode, NodeIOMixin):
     :param body: Pipeline job for the parallel for loop body.
     :type body: Pipeline
     :param items: The loop body's input which will bind to the loop node.
-    :type items: Union[list, dict, str, PipelineInput, NodeOutput]
+    :type items: typing.Union[list, dict, str, NodeOutput, PipelineInput]
     :param max_concurrency: Maximum number of concurrent iterations to run. All loop body nodes will be executed
         in parallel if not specified.
     :type max_concurrency: int
@@ -105,17 +106,92 @@ def _attr_type_map(cls) -> dict:
             "items": (dict, list, str, PipelineInput, NodeOutput),
         }
 
+    @classmethod
+    def _to_rest_item(cls, item: dict) -> dict:
+        """Convert item to rest object."""
+        primitive_inputs, asset_inputs = {}, {}
+        # validate item
+        for key, val in item.items():
+            if isinstance(val, Input):
+                asset_inputs[key] = val
+            elif isinstance(val, (PipelineInput, NodeOutput)):
+                # convert binding object to string
+                primitive_inputs[key] = str(val)
+            else:
+                primitive_inputs[key] = val
+        return {
+            # asset type inputs will be converted to JobInput dict:
+            # {"asset_param": {"uri": "xxx", "job_input_type": "uri_file"}}
+            **cls._input_entity_to_rest_inputs(input_entity=asset_inputs),
+            # primitive inputs has primitive type value like this
+            # {"int_param": 1}
+            **primitive_inputs
+        }
+
+    @classmethod
+    def _to_rest_items(cls, items: Union[list, dict, str, NodeOutput, PipelineInput]) -> str:
+        """Convert items to rest object."""
+        # validate items.
+        cls._validate_items(items=items, raise_error=True, body_component=None)
+        # convert items to rest object
+        if isinstance(items, list):
+            rest_items = [cls._to_rest_item(item=i) for i in items]
+            rest_items = json.dumps(rest_items)
+        elif isinstance(items, dict):
+            rest_items = {k: cls._to_rest_item(item=v) for k, v in items.items()}
+            rest_items = json.dumps(rest_items)
+        elif isinstance(items, (NodeOutput, PipelineInput)):
+            rest_items = str(items)
+        elif isinstance(items, str):
+            rest_items = items
+        else:
+            raise UserErrorException("Unsupported items type: {}".format(type(items)))
+        return rest_items
+
     def _to_rest_object(self, **kwargs) -> dict:  # pylint: disable=unused-argument
         """Convert self to a rest object for remote call."""
         rest_node = super(ParallelFor, self)._to_rest_object(**kwargs)
-        rest_node.update(dict(outputs=self._to_rest_outputs()))
+        # convert items to rest object
+        rest_items = self._to_rest_items(items=self.items)
+        rest_node.update(dict(
+            items=rest_items,
+            outputs=self._to_rest_outputs()
+        ))
         return convert_ordered_dict_to_dict(rest_node)
 
+    @classmethod
+    def _from_rest_item(cls, rest_item):
+        """Convert rest item to item."""
+        primitive_inputs, asset_inputs = {}, {}
+        for key, val in rest_item.items():
+            if isinstance(val, dict) and val.get("job_input_type"):
+                asset_inputs[key] = val
+            else:
+                primitive_inputs[key] = val
+        return {
+            **cls._from_rest_inputs(inputs=asset_inputs),
+            **primitive_inputs
+        }
+
+    @classmethod
+    def _from_rest_items(cls, rest_items: str) -> Union[dict, list, str]:
+        """Convert items from rest object."""
+        try:
+            items = json.loads(rest_items)
+        except json.JSONDecodeError:
+            # return original items when failed to load
+            return rest_items
+        if isinstance(items, list):
+            return [cls._from_rest_item(rest_item=i) for i in items]
+        if isinstance(items, dict):
+            return {k: cls._from_rest_item(rest_item=v) for k, v in items.items()}
+        return rest_items
+
     @classmethod
     def _from_rest_object(cls, obj: dict, pipeline_jobs: dict) -> "ParallelFor":
         # pylint: disable=protected-access
-
         obj = BaseNode._from_rest_object_to_init_params(obj)
+        obj["items"] = cls._from_rest_items(rest_items=obj.get("items", ""))
         return cls._create_instance_from_schema_dict(pipeline_jobs=pipeline_jobs, loaded_data=obj)
 
     @classmethod
@@ -149,11 +225,21 @@ def _convert_output_meta(self, outputs):
             aggregate_outputs[name] = resolved_output
         return aggregate_outputs
 
-    def _validate_items(self, raise_error=True):
-        validation_result = self._create_empty_validation_result()
-        if self.items is not None:
-            items = self.items
+    def _customized_validate(self):
+        """Customized validation for parallel for node."""
+        # pylint: disable=protected-access
+        validation_result = self._validate_body(raise_error=False)
+        validation_result.merge_with(
+            self._validate_items(items=self.items, raise_error=False, body_component=self.body._component)
+        )
+        return validation_result
+
+    @classmethod
+    def _validate_items(cls, items, raise_error=True, body_component=None):
+        validation_result = cls._create_empty_validation_result()
+        if items is not None:
             if isinstance(items, str):
+                # TODO: remove the validation
                 # try to deserialize str if it's a json string
                 try:
                     items = json.loads(items)
@@ -168,7 +254,7 @@ def _validate_items(self, raise_error=True):
                 items = list(items.values())
             if isinstance(items, list):
                 if len(items) > 0:
-                    self._validate_items_list(items, validation_result)
+                    cls._validate_items_list(items, validation_result, body_component=body_component)
                 else:
                     validation_result.append_error(
                         yaml_path="items",
@@ -179,17 +265,12 @@ def _validate_items(self, raise_error=True):
                 message="Items is required for parallel_for node",
             )
         return validation_result.try_raise(
-            self._get_validation_error_target(),
+            cls._get_validation_error_target(),
             raise_error=raise_error,
         )
 
-    def _customized_validate(self):
-        """Customized validation for parallel for node."""
-        validation_result = self._validate_body(raise_error=False)
-        validation_result.merge_with(self._validate_items(raise_error=False))
-        return validation_result
-
-    def _validate_items_list(self, items: list, validation_result):
+    @classmethod
+    def _validate_items_list(cls, items: list, validation_result, body_component=None):
         # pylint: disable=protected-access
         meta = {}
         # all items have to be dict and have matched meta
@@ -213,10 +294,41 @@ def _validate_items_list(self, items: list, validation_result):
                             message=msg
                         )
                 # items' keys should appear in body's inputs
-                body_component = self.body._component
                 if isinstance(body_component, Component) and (not item.keys() <= body_component.inputs.keys()):
                     msg = f"Item {item} got unmatched inputs with loop body component inputs {body_component.inputs}."
                     validation_result.append_error(
                         yaml_path="items",
                         message=msg
                     )
+                # validate item value type
+                cls._validate_item_value_type(item=item, validation_result=validation_result)
+
+    @classmethod
+    def _validate_item_value_type(cls, item: dict, validation_result):
+        # pylint: disable=protected-access
+        supported_types = (Input, str, bool, int, float, PipelineInput)
+        for _, val in item.items():
+            if not isinstance(val, supported_types):
+                validation_result.append_error(
+                    yaml_path="items",
+                    message="Unsupported type {} in parallel_for items. Supported types are: {}".format(
+                        type(val), supported_types
+                    )
+                )
+            if isinstance(val, Input):
+                cls._validate_input_item_value(entry=val, validation_result=validation_result)
+
+    @classmethod
+    def _validate_input_item_value(cls, entry: Input, validation_result):
+        if not isinstance(entry, Input):
+            return
+        if not entry.path:
+            validation_result.append_error(
+                yaml_path="items",
+                message=f"Input path not provided for {entry}.",
+            )
+        if isinstance(entry.path, str) and os.path.exists(entry.path):
+            validation_result.append_error(
+                yaml_path="items",
+                message=f"Local file input {entry} is not supported, please create it as a dataset.",
+            )

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/pipeline/_io/mixin.py

@@ -154,10 +154,14 @@ def _to_rest_inputs(self) -> Dict[str, Dict]:
         }
         """
         built_inputs = self._build_inputs()
+        return self._input_entity_to_rest_inputs(input_entity=built_inputs)
+
+    @classmethod
+    def _input_entity_to_rest_inputs(cls, input_entity: Dict[str, Input]) -> Dict[str, Dict]:
 
         # Convert io entity to rest io objects
         input_bindings, dataset_literal_inputs = process_sdk_component_job_io(
-            built_inputs, [ComponentJobConstants.INPUT_PATTERN]
+            input_entity, [ComponentJobConstants.INPUT_PATTERN]
         )
 
         # parse input_bindings to InputLiteral(value=str(binding))

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_util.py

@@ -205,7 +205,6 @@ def is_empty_target(obj):
 
 def convert_ordered_dict_to_dict(target_object: Union[Dict, List], remove_empty=True) -> Union[Dict, List]:
     """Convert ordered dict to dict. Remove keys with None value.
-
     This is a workaround for rest request must be in dict instead of
     ordered dict.
     """

sdk/ml/azure-ai-ml/tests/dsl/e2etests/test_controlflow_pipeline.py

@@ -645,3 +645,42 @@ def parallel_for_pipeline():
                                                                  'value': '${{parent.outputs.component_out_path}}'}},
                               'type': 'parallel_for'}
         }
+
+    def test_parallel_for_pipeline_with_asset_items(self, client: MLClient):
+        hello_world_component = load_component(
+            source="./tests/test_configs/components/helloworld_component.yml"
+        )
+
+        @pipeline
+        def parallel_for_pipeline():
+            parallel_body = hello_world_component()
+            parallel_node = parallel_for(
+                body=parallel_body,
+                items=[
+                    {"component_in_number": 1, "component_in_path": test_input},
+                    {"component_in_number": 2, "component_in_path": test_input},
+                ]
+            )
+            after_node = hello_world_component(
+                component_in_path=parallel_node.outputs.component_out_path,
+            )
+
+        pipeline_job = parallel_for_pipeline()
+        pipeline_job.settings.default_compute = "cpu-cluster"
+
+        with include_private_preview_nodes_in_pipeline():
+            pipeline_job = assert_job_cancel(pipeline_job, client)
+
+        dsl_pipeline_job_dict = omit_with_wildcard(pipeline_job._to_rest_object().as_dict(), *omit_fields)
+        assert dsl_pipeline_job_dict["properties"]["jobs"]["parallel_node"] == {
+            'body': '${{parent.jobs.parallel_body}}',
+            'items': '[{"component_in_path": {"uri": '
+                     '"https://dprepdata.blob.core.windows.net/demo/Titanic.csv", '
+                     '"job_input_type": "uri_file"}, '
+                     '"component_in_number": 1}, {"component_in_path": '
+                     '{"uri": '
+                     '"https://dprepdata.blob.core.windows.net/demo/Titanic.csv", '
+                     '"job_input_type": "uri_file"}, '
+                     '"component_in_number": 2}]',
+            'type': 'parallel_for'
+        }