[ML][Pipeline] Support parameter group on pipeline node (Azure#26713)

* Skip check for node input

Signed-off-by: Brynn Yin <[email protected]>

* Support parameter group on pipeline component

Signed-off-by: Brynn Yin <[email protected]>

* Reformat

Signed-off-by: Brynn Yin <[email protected]>

* Fix lint

Signed-off-by: Brynn Yin <[email protected]>

Signed-off-by: Brynn Yin <[email protected]>

Author: brynn-code

sdk/ml/azure-ai-ml/azure/ai/ml/dsl/_pipeline_component_builder.py

@@ -277,7 +277,14 @@ def _map_type(_meta):
         return output_dict
 
     def _get_group_parameter_defaults(self):
-        return {key: copy.deepcopy(val.default) for key, val in self.inputs.items() if isinstance(val, GroupInput)}
+        group_defaults = {}
+        for key, val in self.inputs.items():
+            if not isinstance(val, GroupInput):
+                continue
+            # Copy and insert top-level parameter name into group names for all items
+            group_defaults[key] = copy.deepcopy(val.default)
+            group_defaults[key].insert_group_name_for_items(key)
+        return group_defaults
 
     def _update_nodes_variable_names(self, func_variables: dict):
         """Update nodes list to ordered dict with variable name key and

sdk/ml/azure-ai-ml/azure/ai/ml/dsl/_pipeline_decorator.py

@@ -15,6 +15,8 @@
 from azure.ai.ml.entities import Data, PipelineJob, PipelineJobSettings
 from azure.ai.ml.entities._builders.pipeline import Pipeline
 from azure.ai.ml.entities._inputs_outputs import Input, is_parameter_group
+from azure.ai.ml.entities._job.pipeline._io import NodeOutput, PipelineInput, _GroupAttrDict
+from azure.ai.ml.entities._job.pipeline._pipeline_expression import PipelineExpression
 from azure.ai.ml.exceptions import (
     MissingPositionalArgsError,
     MultipleValueError,
@@ -23,8 +25,6 @@
     UnsupportedParameterKindError,
     UserErrorException,
 )
-from azure.ai.ml.entities._job.pipeline._io import NodeOutput, PipelineInput
-from azure.ai.ml.entities._job.pipeline._pipeline_expression import PipelineExpression
 
 from ._pipeline_component_builder import PipelineComponentBuilder, _is_inside_dsl_pipeline_func
 from ._settings import _dsl_settings_stack
@@ -42,6 +42,8 @@
     bool,
     int,
     float,
+    PipelineExpression,
+    _GroupAttrDict,
 )
 module_logger = logging.getLogger(__name__)
 
@@ -242,7 +244,7 @@ def _validate_args(func, args, kwargs):
 
     def _is_supported_data_type(_data):
         return (
-            isinstance(_data, SUPPORTED_INPUT_TYPES + (PipelineExpression,))
+            isinstance(_data, SUPPORTED_INPUT_TYPES)
             or is_parameter_group(_data)
         )
 

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

@@ -23,7 +23,7 @@
 from azure.ai.ml.entities._job.sweep.search_space import SweepDistribution
 from azure.ai.ml.entities._mixins import YamlTranslatableMixin
 from azure.ai.ml.entities._util import convert_ordered_dict_to_dict, resolve_pipeline_parameters
-from azure.ai.ml.entities._validation import SchemaValidatableMixin, MutableValidationResult
+from azure.ai.ml.entities._validation import MutableValidationResult, SchemaValidatableMixin
 from azure.ai.ml.exceptions import ErrorTarget, ValidationErrorType, ValidationException
 
 module_logger = logging.getLogger(__name__)
@@ -231,7 +231,10 @@ def _parse_io(cls, io_dict: dict, parse_cls):
                 value = value._deepcopy()  # Decoupled input and output
                 io_dict[key] = value
                 value.mode = None
-            elif isinstance(value, dict):
+            elif type(value) == dict: # pylint: disable=unidiomatic-typecheck
+                # Use type comparison instead of is_instance to skip _GroupAttrDict
+                # when loading from yaml io will be a dict,
+                # like {'job_data_path': '${{parent.inputs.pipeline_job_data_path}}'}
                 # parse dict to allowed type
                 io_dict[key] = parse_cls(**value)
 

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_component/pipeline_component.py

@@ -232,7 +232,13 @@ def _validate_binding_inputs(self, node: BaseNode) -> MutableValidationResult:
         not set. Raise error if pipeline input is optional but link to
         required inputs.
         """
-        component_definition_inputs = node.component.inputs
+        component_definition_inputs = {}
+        # Add flattened group input into definition inputs.
+        # e.g. Add {'group_name.item': PipelineInput} for {'group_name': GroupInput}
+        for name, val in node.component.inputs.items():
+            if isinstance(val, GroupInput):
+                component_definition_inputs.update(val.flatten(group_parameter_name=name))
+            component_definition_inputs[name] = val
         # Collect binding relation dict {'pipeline_input': ['node_input']}
         validation_result = self._create_empty_validation_result()
         binding_dict, optional_binding_in_expression_dict = self._get_input_binding_dict(node)

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_inputs_outputs/group_input.py

@@ -1,7 +1,7 @@
 # ---------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
-
+import copy
 from enum import Enum as PyEnum
 
 from azure.ai.ml.constants._component import IOConstants
@@ -38,9 +38,15 @@ def _create_default(self):
         from .._job.pipeline._io import PipelineInput
 
         default_dict = {}
+        # Note: no top-level group names at this time.
         for k, v in self.values.items():
-            # Assign directly if is subgroup, else create PipelineInput object
-            default_dict[k] = v.default if isinstance(v, GroupInput) else PipelineInput(name=k, data=v.default, meta=v)
+            # Create PipelineInput object if not subgroup
+            if not isinstance(v, GroupInput):
+                default_dict[k] = PipelineInput(name=k, data=v.default, meta=v)
+                continue
+            # Copy and insert k into group names for subgroup
+            default_dict[k] = copy.deepcopy(v.default)
+            default_dict[k].insert_group_name_for_items(k)
         return self._create_group_attr_dict(default_dict)
 
     @classmethod

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_job/_input_output_helpers.py

@@ -12,8 +12,7 @@
 from azure.ai.ml._restclient.v2022_10_01_preview.models import JobInput as RestJobInput
 from azure.ai.ml._restclient.v2022_10_01_preview.models import JobInputType
 from azure.ai.ml._restclient.v2022_10_01_preview.models import JobOutput as RestJobOutput
-from azure.ai.ml._restclient.v2022_10_01_preview.models import JobOutputType
-from azure.ai.ml._restclient.v2022_10_01_preview.models import LiteralJobInput
+from azure.ai.ml._restclient.v2022_10_01_preview.models import JobOutputType, LiteralJobInput
 from azure.ai.ml._restclient.v2022_10_01_preview.models import MLFlowModelJobInput as RestMLFlowModelJobInput
 from azure.ai.ml._restclient.v2022_10_01_preview.models import MLFlowModelJobOutput as RestMLFlowModelJobOutput
 from azure.ai.ml._restclient.v2022_10_01_preview.models import MLTableJobInput as RestMLTableJobInput
@@ -191,7 +190,8 @@ def to_rest_dataset_literal_inputs(
     for input_name, input_value in inputs.items():
         if job_type == JobType.PIPELINE:
             validate_pipeline_input_key_contains_allowed_characters(input_name)
-        else:
+        elif job_type:
+            # We pass job_type=None for pipeline node, and want skip this check for nodes.
             validate_key_contains_allowed_characters(input_name)
         if isinstance(input_value, Input):
             if input_value.path and isinstance(input_value.path, str) and is_data_binding_expression(input_value.path):

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

@@ -438,7 +438,6 @@ def __init__(self, name: str, meta: Input, group_names: List[str] = None, **kwar
         """
         super(PipelineInput, self).__init__(name=name, meta=meta, **kwargs)
         self._group_names = group_names if group_names else []
-        self._full_name = "%s.%s" % (".".join(self._group_names), self._name) if self._group_names else self._name
 
     def __str__(self) -> str:
         return self._data_binding()
@@ -451,7 +450,7 @@ def _build_data(self, data, key=None):  # pylint: disable=unused-argument
             msg = "Can not bind input to another component's input."
             raise ValidationException(message=msg, no_personal_data_message=msg, target=ErrorTarget.PIPELINE)
         if isinstance(data, (PipelineInput, NodeOutput)):
-            # If value is input or output, it's a data binding, we require it have a owner so we can convert it to
+            # If value is input or output, it's a data binding, owner is required to convert it to
             # a data binding, eg: ${{parent.inputs.xxx}}
             if isinstance(data, NodeOutput) and data._owner is None:
                 msg = "Setting input binding {} to output without owner is not allowed."
@@ -468,7 +467,8 @@ def _build_data(self, data, key=None):  # pylint: disable=unused-argument
         return data
 
     def _data_binding(self):
-        return f"${{{{parent.inputs.{self._full_name}}}}}"
+        full_name = "%s.%s" % (".".join(self._group_names), self._name) if self._group_names else self._name
+        return f"${{{{parent.inputs.{full_name}}}}}"
 
     def _to_input(self) -> Input:
         """Convert pipeline input to component input for pipeline component."""
@@ -628,6 +628,15 @@ def flatten(self, group_parameter_name):
                 )
         return flattened_parameters
 
+    def insert_group_name_for_items(self, group_name):
+        # Insert one group name for all items.
+        for v in self.values():
+            if isinstance(v, _GroupAttrDict):
+                v.insert_group_name_for_items(group_name)
+            elif isinstance(v, PipelineInput):
+                # Insert group names for pipeline input
+                v._group_names = [group_name] + v._group_names
+
 
 class OutputsAttrDict(dict):
     def __init__(self, outputs: dict, **kwargs):

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

@@ -7,9 +7,6 @@
 
 import pydash
 import pytest
-from pipeline_job.e2etests.test_pipeline_job import assert_job_input_output_types
-from test_utilities.utils import _PYTEST_TIMEOUT_METHOD, omit_with_wildcard
-
 from azure.ai.ml import (
     Input,
     MLClient,
@@ -25,18 +22,21 @@
 from azure.ai.ml.constants._common import AssetTypes, InputOutputModes
 from azure.ai.ml.constants._job.pipeline import PipelineConstants
 from azure.ai.ml.dsl._load_import import to_component
+from azure.ai.ml.dsl._parameter_group_decorator import parameter_group
 from azure.ai.ml.entities import CommandComponent, CommandJob
 from azure.ai.ml.entities import Component
 from azure.ai.ml.entities import Component as ComponentEntity
 from azure.ai.ml.entities import Data, PipelineJob
 from azure.ai.ml.exceptions import ValidationException
 from azure.ai.ml.parallel import ParallelJob, RunFunction, parallel_run_function
+from azure.core.exceptions import HttpResponseError
 from azure.core.polling import LROPoller
+from devtools_testutils import AzureRecordedTestCase
+from pipeline_job.e2etests.test_pipeline_job import assert_job_input_output_types
+from test_utilities.utils import _PYTEST_TIMEOUT_METHOD, omit_with_wildcard
 
 from .._util import _DSL_TIMEOUT_SECOND
 
-from devtools_testutils import AzureRecordedTestCase
-
 tests_root_dir = Path(__file__).parent.parent.parent
 components_dir = tests_root_dir / "test_configs/components/"
 job_input = Input(
@@ -57,6 +57,17 @@
 ]
 
 
+def assert_job_cancel(pipeline, client: MLClient):
+    job = client.jobs.create_or_update(pipeline)
+    try:
+        cancel_poller = client.jobs.begin_cancel(job.name)
+        assert isinstance(cancel_poller, LROPoller)
+        assert cancel_poller.result() is None
+    except HttpResponseError:
+        pass
+    return job
+
+
 @pytest.mark.usefixtures(
     "enable_environment_id_arm_expansion",
     "enable_pipeline_private_preview_features",
@@ -1058,7 +1069,6 @@ def valid_pipeline_func(
             required_input: Input,
             required_param: str,
             node_compute: str = "cpu-cluster",
-            # node_compute: str = 'azureml:cpu-cluster', # both will be supported
         ):
             default_optional_func(
                 required_input=required_input,
@@ -1090,6 +1100,64 @@ def valid_pipeline_func(
             in caplog.messages
         )
 
+    def test_create_pipeline_with_parameter_group(self, client: MLClient) -> None:
+        default_optional_func = load_component(source=str(components_dir / "default_optional_component.yml"))
+
+        @parameter_group
+        class SubGroup:
+            required_param: str
+
+        @parameter_group
+        class Group:
+            sub: SubGroup
+            node_compute: str = "cpu-cluster"
+
+        @dsl.pipeline()
+        def sub_pipeline_func(
+            required_input: Input,
+            group: Group,
+            sub_group: SubGroup,
+        ):
+            default_optional_func(
+                required_input=required_input,
+                required_param=group.sub.required_param,
+            )
+            node2 = default_optional_func(
+                required_input=required_input,
+                required_param=sub_group.required_param,
+            )
+            node2.compute = group.node_compute
+
+        @dsl.pipeline(default_compute="cpu-cluster")
+        def root_pipeline_with_group(
+            r_required_input: Input,
+            r_group: Group,
+        ):
+            sub_pipeline_func(required_input=r_required_input, group=r_group, sub_group=r_group.sub)
+
+        job = root_pipeline_with_group(
+            r_required_input=Input(type="uri_file", path="https://dprepdata.blob.core.windows.net/demo/Titanic.csv"),
+            r_group=Group(sub=SubGroup(required_param="hello")),
+        )
+        rest_job = assert_job_cancel(job, client)
+        assert len(rest_job.inputs) == 2
+        rest_job_dict = rest_job._to_dict()
+        assert rest_job_dict["inputs"] == {
+            "r_required_input": {
+                "mode": "ro_mount",
+                "type": "uri_file",
+                "path": "azureml:https://dprepdata.blob.core.windows.net/demo/Titanic.csv",
+            },
+            "r_group.sub.required_param": "hello",
+            "r_group.node_compute": "cpu-cluster",
+        }
+        assert rest_job_dict["jobs"]["sub_pipeline_func"]["inputs"] == {
+            "required_input": {"path": "${{parent.inputs.r_required_input}}"},
+            "group.sub.required_param": {"path": "${{parent.inputs.r_group.sub.required_param}}"},
+            "group.node_compute": {"path": "${{parent.inputs.r_group.node_compute}}"},
+            "sub_group.required_param": {"path": "${{parent.inputs.r_group.sub.required_param}}"},
+        }
+
     def test_pipeline_with_none_parameter_has_default_optional_true(self, client: MLClient) -> None:
         default_optional_func = load_component(source=str(components_dir / "default_optional_component.yml"))
 
@@ -1525,9 +1593,7 @@ def parallel_in_pipeline(job_data_path, score_model):
         assert_job_input_output_types(pipeline_job)
         assert pipeline_job.settings.default_compute == "cpu-cluster"
 
-    @pytest.mark.skip(
-        "https://dev.azure.com/msdata/Vienna/_workitems/edit/2009659"
-    )
+    @pytest.mark.skip("https://dev.azure.com/msdata/Vienna/_workitems/edit/2009659")
     def test_parallel_components_with_file_input(self, client: MLClient) -> None:
         components_dir = tests_root_dir / "test_configs/dsl_pipeline/parallel_component_with_file_input"
 
@@ -1958,11 +2024,11 @@ def pipeline(job_in_number, job_in_other_number, job_in_path):
             client.jobs.get(child.name)
             client.jobs.get(child.name)._repr_html_()
 
-    @pytest.mark.skip(
-        "https://dev.azure.com/msdata/Vienna/_workitems/edit/2009659"
-    )
+    @pytest.mark.skip("https://dev.azure.com/msdata/Vienna/_workitems/edit/2009659")
     def test_dsl_pipeline_without_setting_binding_node(self, client: MLClient) -> None:
-        from test_configs.dsl_pipeline.pipeline_with_set_binding_output_input.pipeline import pipeline_without_setting_binding_node
+        from test_configs.dsl_pipeline.pipeline_with_set_binding_output_input.pipeline import (
+            pipeline_without_setting_binding_node,
+        )
 
         pipeline = pipeline_without_setting_binding_node()
         pipeline_job = client.jobs.create_or_update(pipeline)
@@ -2011,9 +2077,7 @@ def test_dsl_pipeline_without_setting_binding_node(self, client: MLClient) -> No
         }
         assert expected_job == actual_job
 
-    @pytest.mark.skip(
-        "https://dev.azure.com/msdata/Vienna/_workitems/edit/2009659"
-    )
+    @pytest.mark.skip("https://dev.azure.com/msdata/Vienna/_workitems/edit/2009659")
     def test_dsl_pipeline_with_only_setting_pipeline_level(self, client: MLClient) -> None:
         from test_configs.dsl_pipeline.pipeline_with_set_binding_output_input.pipeline import (
             pipeline_with_only_setting_pipeline_level,
@@ -2066,12 +2130,12 @@ def test_dsl_pipeline_with_only_setting_pipeline_level(self, client: MLClient) -
         }
         assert expected_job == actual_job
 
-    @pytest.mark.skip(
-        "https://dev.azure.com/msdata/Vienna/_workitems/edit/2009659"
-    )
+    @pytest.mark.skip("https://dev.azure.com/msdata/Vienna/_workitems/edit/2009659")
     def test_dsl_pipeline_with_only_setting_binding_node(self, client: MLClient) -> None:
         # Todo: checkout run priority when backend is ready
-        from test_configs.dsl_pipeline.pipeline_with_set_binding_output_input.pipeline import pipeline_with_only_setting_binding_node
+        from test_configs.dsl_pipeline.pipeline_with_set_binding_output_input.pipeline import (
+            pipeline_with_only_setting_binding_node,
+        )
 
         pipeline = pipeline_with_only_setting_binding_node()
         pipeline_job = client.jobs.create_or_update(pipeline)
@@ -2130,9 +2194,7 @@ def test_dsl_pipeline_with_only_setting_binding_node(self, client: MLClient) ->
         }
         assert expected_job == actual_job
 
-    @pytest.mark.skip(
-        "https://dev.azure.com/msdata/Vienna/_workitems/edit/2009659"
-    )
+    @pytest.mark.skip("https://dev.azure.com/msdata/Vienna/_workitems/edit/2009659")
     def test_dsl_pipeline_with_setting_binding_node_and_pipeline_level(self, client: MLClient) -> None:
         from test_configs.dsl_pipeline.pipeline_with_set_binding_output_input.pipeline import (
             pipeline_with_setting_binding_node_and_pipeline_level,