[ML][Pipelines] redo feat: flow in pipeline (Azure#31914)

* [ML][Pipelines] redo feat: load component from flow
(Azure#31632) (Azure#31911)
This reverts commit c6be14f.

* fix: remove extra-index-url from requirements.txt

* fix: resolve comments

Author: elliotzh

.vscode/cspell.json

@@ -88,6 +88,7 @@
     "sdk/ml/azure-ai-ml/tests/**",
     "sdk/ml/azure-ai-ml/swagger/**",
     "sdk/ml/azure-ai-ml/NOTICE.txt",
+    "sdk/ml/azure-ai-ml/.pre-commit-config.yaml",
     "sdk/loadtestservice/azure-developer-loadtesting/**",
     "sdk/translation/azure-ai-translation-text/azure/ai/translation/text/_serialization.py",
     "sdk/translation/azure-ai-translation-text/tests/test_break_sentence.py",

sdk/ml/azure-ai-ml/.pre-commit-config.yaml

@@ -26,8 +26,8 @@ repos:
     args: ['--settings-file', 'sdk/ml/azure-ai-ml/pyproject.toml']
 - repo: local
   hooks:
-    - id: pylint
-      name: pylint
+    - id: pylint-dependencies-check
+      name: pylint-dependencies-check
       entry: python
       language: system
       types: [python]
@@ -69,9 +69,9 @@ repos:
             print(f"Please run the following command to install the correct version of {packagename}")
             print(f"\tpython -m pip install {packagename}=={required_version} {' '.join(install_args)}")
             sys.exit(1)
-
-        # Run pylint
-        os.execl(sys.executable, sys.executable, "-m", "pylint", *sys.argv[1:])
-      # cspell:disable-next-line
-      - "--rcfile=pylintrc"
-      - "--output-format=parseable"
+    - id: pylint
+      name: pylint
+      entry: python
+      language: system
+      args: [ -m, pylint, --rcfile=pylintrc, --output-format=parseable ]
+      types: [python]

sdk/ml/azure-ai-ml/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/ml/azure-ai-ml",
-  "Tag": "python/ml/azure-ai-ml_d05969c2b4"
+  "Tag": "python/ml/azure-ai-ml_039ca40c3d"
 }

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/component/__init__.py

@@ -5,21 +5,21 @@
 __path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore
 
 from .command_component import AnonymousCommandComponentSchema, CommandComponentSchema, ComponentFileRefField
-from .component import ComponentSchema
-from .import_component import AnonymousImportComponentSchema, ImportComponentFileRefField, ImportComponentSchema
-from .parallel_component import AnonymousParallelComponentSchema, ParallelComponentFileRefField, ParallelComponentSchema
-from .spark_component import AnonymousSparkComponentSchema, SparkComponentFileRefField, SparkComponentSchema
+from .component import ComponentSchema, ComponentYamlRefField
 from .data_transfer_component import (
     AnonymousDataTransferCopyComponentSchema,
+    AnonymousDataTransferExportComponentSchema,
+    AnonymousDataTransferImportComponentSchema,
     DataTransferCopyComponentFileRefField,
     DataTransferCopyComponentSchema,
-    AnonymousDataTransferImportComponentSchema,
-    DataTransferImportComponentFileRefField,
-    DataTransferImportComponentSchema,
-    AnonymousDataTransferExportComponentSchema,
     DataTransferExportComponentFileRefField,
     DataTransferExportComponentSchema,
+    DataTransferImportComponentFileRefField,
+    DataTransferImportComponentSchema,
 )
+from .import_component import AnonymousImportComponentSchema, ImportComponentFileRefField, ImportComponentSchema
+from .parallel_component import AnonymousParallelComponentSchema, ParallelComponentFileRefField, ParallelComponentSchema
+from .spark_component import AnonymousSparkComponentSchema, SparkComponentFileRefField, SparkComponentSchema
 
 __all__ = [
     "ComponentSchema",
@@ -44,4 +44,5 @@
     "AnonymousDataTransferExportComponentSchema",
     "DataTransferExportComponentFileRefField",
     "DataTransferExportComponentSchema",
+    "ComponentYamlRefField",
 ]

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/component/component.py

@@ -1,8 +1,10 @@
 # ---------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
+from pathlib import Path
 
-from marshmallow import fields, post_dump, pre_dump, pre_load
+from marshmallow import ValidationError, fields, post_dump, pre_dump, pre_load
+from marshmallow.fields import Field
 
 from azure.ai.ml._schema.component.input_output import InputPortSchema, OutputPortSchema, ParameterSchema
 from azure.ai.ml._schema.core.fields import (
@@ -13,8 +15,8 @@
     UnionField,
 )
 from azure.ai.ml._schema.core.intellectual_property import IntellectualPropertySchema
-from azure.ai.ml._utils.utils import is_private_preview_enabled
-from azure.ai.ml.constants._common import AzureMLResourceType
+from azure.ai.ml._utils.utils import is_private_preview_enabled, load_yaml
+from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, AzureMLResourceType
 
 from .._utils.utils import _resolve_group_inputs_for_component
 from ..assets.asset import AssetSchema
@@ -26,6 +28,52 @@ def _get_field_name(self):
         return "Component"
 
 
+class ComponentYamlRefField(Field):
+    """Allows you to nest a :class:`Schema <marshmallow.Schema>`
+    inside a yaml ref field.
+    """
+
+    def _jsonschema_type_mapping(self):
+        schema = {"type": "string"}
+        if self.name is not None:
+            schema["title"] = self.name
+        if self.dump_only:
+            schema["readonly"] = True
+        return schema
+
+    def _deserialize(self, value, attr, data, **kwargs):
+        if not isinstance(value, str):
+            raise ValidationError(f"Nested yaml ref field expected a string but got {type(value)}.")
+
+        base_path = Path(self.context[BASE_PATH_CONTEXT_KEY])
+
+        source_path = Path(value)
+        # raise if the string is not a valid path, like "azureml:xxx"
+        try:
+            source_path.resolve()
+        except OSError as ex:
+            raise ValidationError(f"Nested file ref field expected a local path but got {value}.") from ex
+
+        if not source_path.is_absolute():
+            source_path = base_path / source_path
+
+        if not source_path.is_file():
+            raise ValidationError(
+                f"Nested yaml ref field expected a local path but can't find {value} based on {base_path.as_posix()}."
+            )
+
+        loaded_value = load_yaml(source_path)
+
+        # local import to avoid circular import
+        from azure.ai.ml.entities import Component
+
+        component = Component._load(data=loaded_value, yaml_path=source_path)  # pylint: disable=protected-access
+        return component
+
+    def _serialize(self, value, attr, obj, **kwargs):
+        raise ValidationError("Serialize on RefField is not supported.")
+
+
 class ComponentSchema(AssetSchema):
     schema = fields.Str(data_key="$schema", attribute="_schema")
     name = ComponentNameStr(required=True)

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/component/flow.py

@@ -0,0 +1,89 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# pylint: disable=unused-argument,protected-access,no-member
+
+from marshmallow import fields
+
+from azure.ai.ml._schema import YamlFileSchema
+from azure.ai.ml._schema.component import ComponentSchema
+from azure.ai.ml._schema.component.component import ComponentNameStr
+from azure.ai.ml._schema.core.fields import ArmVersionedStr, LocalPathField, StringTransformedEnum, UnionField
+from azure.ai.ml._schema.core.schema_meta import PatchedSchemaMeta
+from azure.ai.ml.constants._common import AzureMLResourceType
+from azure.ai.ml.constants._component import NodeType
+
+
+class _ComponentMetadataSchema(metaclass=PatchedSchemaMeta):
+    """Schema to recognize metadata of a flow as a component."""
+
+    name = ComponentNameStr()
+    version = fields.Str()
+    display_name = fields.Str()
+    description = fields.Str()
+    tags = fields.Dict(keys=fields.Str(), values=fields.Str())
+    is_deterministic = fields.Bool()
+
+
+class _FlowAttributesSchema(metaclass=PatchedSchemaMeta):
+    """Schema to recognize attributes of a flow."""
+
+    variant = fields.Str()
+    column_mappings = fields.Dict(
+        fields.Str(),
+        fields.Str(),
+    )
+    connections = fields.Dict(
+        keys=fields.Str(),
+        values=fields.Dict(
+            keys=fields.Str(),
+            values=fields.Str(),
+        ),
+    )
+    environment_variables = fields.Dict(
+        fields.Str(),
+        fields.Str(),
+    )
+
+
+class FlowSchema(YamlFileSchema, _ComponentMetadataSchema):
+    """Schema for flow.dag.yaml file."""
+
+    additional_includes = fields.List(LocalPathField())
+
+
+class RunSchema(YamlFileSchema, _ComponentMetadataSchema, _FlowAttributesSchema):
+    """Schema for run.yaml file."""
+
+    flow = LocalPathField(required=True)
+
+
+class FlowComponentSchema(ComponentSchema, _FlowAttributesSchema):
+    """FlowSchema and FlowRunSchema are used to load flow while FlowComponentSchema is used to dump flow."""
+
+    class Meta:
+        """Override this to exclude inputs & outputs as component doesn't have them."""
+
+        exclude = ["inputs", "outputs"]  # component doesn't have inputs & outputs
+
+    # TODO: name should be required?
+    name = ComponentNameStr()
+
+    type = StringTransformedEnum(allowed_values=[NodeType.FLOW_PARALLEL], required=True)
+
+    # name, version, tags, display_name and is_deterministic are inherited from ComponentSchema
+    properties = fields.Dict(
+        fields.Str(),
+        fields.Str(),
+    )
+
+    # this is different from regular CodeField
+    code = UnionField(
+        [
+            LocalPathField(),
+            ArmVersionedStr(azureml_type=AzureMLResourceType.CODE),
+        ],
+        metadata={"description": "A local path or http:, https:, azureml: url pointing to a remote location."},
+    )
+    additional_includes = fields.List(LocalPathField(), load_only=True)

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

@@ -15,6 +15,7 @@
     AnonymousParallelComponentSchema,
     AnonymousSparkComponentSchema,
     ComponentFileRefField,
+    ComponentYamlRefField,
     DataTransferCopyComponentFileRefField,
     ImportComponentFileRefField,
     ParallelComponentFileRefField,
@@ -28,6 +29,7 @@
 from ...exceptions import ValidationException
 from .._sweep.parameterized_sweep import ParameterizedSweepSchema
 from .._utils.data_binding_expression import support_data_binding_expression_for_fields
+from ..component.flow import FlowComponentSchema
 from ..core.fields import (
     ArmVersionedStr,
     ComputeField,
@@ -252,6 +254,10 @@ class ParallelSchema(BaseNodeSchema, ParameterizedParallelSchema):
                 # component file reference
                 ParallelComponentFileRefField(),
             ],
+            NodeType.FLOW_PARALLEL: [
+                NestedField(FlowComponentSchema, unknown=INCLUDE, dump_only=True),
+                ComponentYamlRefField(),
+            ],
         },
         plain_union_fields=[
             # for registry type assets

sdk/ml/azure-ai-ml/azure/ai/ml/constants/_common.py

@@ -381,6 +381,18 @@ class CommonYamlFields:
     """Schema."""
 
 
+class SchemaUrl:
+    """Schema urls.
+
+    Schema urls will be used in VSCode extension to validate yaml files. It will also be used to identify the
+    corresponding entity type of a yaml file, especially for some internal yaml files.
+    """
+
+    PROMPTFLOW_PREFIX = "https://azuremlschemas.azureedge.net/promptflow/"
+    PROMPTFLOW_FLOW = PROMPTFLOW_PREFIX + "latest/Flow.schema.json"
+    PROMPTFLOW_RUN = PROMPTFLOW_PREFIX + "latest/Run.schema.json"
+
+
 class GitProperties:
     """GitProperties is a class that defines the constants used by the SDK/CLI for Git operations.
 

sdk/ml/azure-ai-ml/azure/ai/ml/constants/_component.py

@@ -23,6 +23,7 @@ class NodeType(object):
     IMPORT = "import"
     SPARK = "spark"
     DATA_TRANSFER = "data_transfer"
+    FLOW_PARALLEL = "promptflow_parallel"
     # Note: container is not a real component type,
     # only used to mark component from container data.
     _CONTAINER = "_container"

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

@@ -17,6 +17,7 @@
 from ...constants._common import ARM_ID_PREFIX
 from ...constants._component import NodeType
 from .._component.component import Component
+from .._component.flow import FlowComponent
 from .._component.parallel_component import ParallelComponent
 from .._inputs_outputs import Input, Output
 from .._job.job_resource_configuration import JobResourceConfiguration
@@ -122,15 +123,28 @@ def __init__(
         validate_attribute_type(attrs_to_check=locals(), attr_type_map=self._attr_type_map())
         kwargs.pop("type", None)
 
-        BaseNode.__init__(
-            self,
-            type=NodeType.PARALLEL,
-            component=component,
-            inputs=inputs,
-            outputs=outputs,
-            compute=compute,
-            **kwargs,
-        )
+        if isinstance(component, FlowComponent):
+            # make input definition fit actual inputs for flow component
+            with component._inputs._fit_inputs(inputs):  # pylint: disable=protected-access
+                BaseNode.__init__(
+                    self,
+                    type=NodeType.PARALLEL,
+                    component=component,
+                    inputs=inputs,
+                    outputs=outputs,
+                    compute=compute,
+                    **kwargs,
+                )
+        else:
+            BaseNode.__init__(
+                self,
+                type=NodeType.PARALLEL,
+                component=component,
+                inputs=inputs,
+                outputs=outputs,
+                compute=compute,
+                **kwargs,
+            )
         # init mark for _AttrDict
         self._init = True
 
@@ -309,7 +323,7 @@ def set_resources(
     @classmethod
     def _attr_type_map(cls) -> dict:
         return {
-            "component": (str, ParallelComponent),
+            "component": (str, ParallelComponent, FlowComponent),
             "retry_settings": (dict, RetrySettings),
             "resources": (dict, JobResourceConfiguration),
             "task": (dict, ParallelTask),