[ML][Pipelines] Test: migrate do-while tests (Azure#29593)

* update enable_private_preview_schema_features for pipeline job schema/entity reload

* fix: support input list in do-while mapping

* test: migrate some skip tests in mldesigner

* test: add primitive test and recordings

* test: fix automl in pipeline serverless test

* test: fix wrong registry component

* pylint: remove unused import

* test: re-org tests & add missing recordings

* reformat with black

* fix: consider internal flag for reload schema

Author: zhengfeiwang

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

@@ -68,19 +68,24 @@ def _infer_and_update_body_input_from_mapping():
             # if loop body output type is not specified, skip as we have no place to infer
             if body.outputs[output_name].type is None:
                 continue
-            # if input type is specified, no need to infer and skip
-            if body_input.type is not None:
-                continue
-            inferred_type = body.outputs[output_name].type
-            # update node input
-            body_input._meta._is_inferred_optional = True
-            body_input.type = inferred_type
-            # update node corresponding component input
-            input_name = body_input._meta.name
-            body.component.inputs[input_name]._is_inferred_optional = True
-            body.component.inputs[input_name].type = inferred_type
+            # body input can be a list of inputs, normalize as a list to process
+            if not isinstance(body_input, list):
+                body_input = [body_input]
+            for single_input in body_input:
+                # if input type is specified, no need to infer and skip
+                if single_input.type is not None:
+                    continue
+                inferred_type = body.outputs[output_name].type
+                # update node input
+                single_input._meta._is_inferred_optional = True
+                single_input.type = inferred_type
+                # update node corresponding component input
+                input_name = single_input._meta.name
+                body.component.inputs[input_name]._is_inferred_optional = True
+                body.component.inputs[input_name].type = inferred_type
 
-    # infer and update for dynamic input
-    _infer_and_update_body_input_from_mapping()
+    # when mapping is a dictionary, infer and update for dynamic input
+    if isinstance(mapping, dict):
+        _infer_and_update_body_input_from_mapping()
 
     return do_while_node

sdk/ml/azure-ai-ml/tests/conftest.py

@@ -780,23 +780,33 @@ def enable_pipeline_private_preview_features(mocker: MockFixture):
 @pytest.fixture()
 def enable_private_preview_schema_features():
     """Schemas will be imported at the very beginning, so need to reload related classes."""
+    from azure.ai.ml._internal._setup import _registered, enable_internal_components_in_pipeline
     from azure.ai.ml._schema.component import command_component as command_component_schema
     from azure.ai.ml._schema.component import component as component_schema
     from azure.ai.ml._schema.component import input_output
     from azure.ai.ml._schema.pipeline import pipeline_component as pipeline_component_schema
+    from azure.ai.ml._schema.pipeline import pipeline_job as pipeline_job_schema
     from azure.ai.ml.entities._component import command_component as command_component_entity
     from azure.ai.ml.entities._component import pipeline_component as pipeline_component_entity
+    from azure.ai.ml.entities._job.pipeline import pipeline_job as pipeline_job_entity
 
     def _reload_related_classes():
         reload(component_schema)
         reload(input_output)
         reload(command_component_schema)
         reload(pipeline_component_schema)
+        reload(pipeline_job_schema)
 
         command_component_entity.CommandComponentSchema = command_component_schema.CommandComponentSchema
         pipeline_component_entity.PipelineComponentSchema = pipeline_component_schema.PipelineComponentSchema
+        pipeline_job_entity.PipelineJobSchema = pipeline_job_schema.PipelineJobSchema
+
+        # check internal flag after reload, force register if it is set as True
+        if _registered:
+            enable_internal_components_in_pipeline(force=True)
 
     with patch.dict(os.environ, {AZUREML_PRIVATE_FEATURES_ENV_VAR: "True"}):
+
         _reload_related_classes()
         yield
     _reload_related_classes()

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

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 import pytest
 from azure.ai.ml.dsl._group_decorator import group
 from devtools_testutils import AzureRecordedTestCase, is_live
@@ -395,6 +397,149 @@ def condition_pipeline():
         }
 
 
+class TestDoWhilePipeline(TestControlFlowPipeline):
+    @property
+    def _basic_component_func(self):
+        return load_component("./tests/test_configs/dsl_pipeline/do_while/basic_component/component.yml")
+
+    def test_do_while_pipeline(self, client: MLClient):
+        @pipeline
+        def do_while_body_pipeline_component(
+            component_in_number: Input(type="integer", optional=True),
+            component_in_path: Input(type="uri_folder"),
+        ):
+            # Call component obj as function: apply given inputs & parameters to create a node in pipeline
+            basic_component = self._basic_component_func(
+                component_in_number=component_in_number,
+                component_in_path=component_in_path,
+            )
+            return basic_component.outputs
+
+        @pipeline
+        def pipeline_with_do_while(
+            component_in_number: Input(type="integer"),
+            component_in_path: Input(type="uri_folder"),
+        ):
+            # Do while node with pipeline component
+            do_while_body_pipeline = do_while_body_pipeline_component(
+                component_in_number=component_in_number,
+                component_in_path=component_in_path,
+            )
+            do_while_with_pipeline = do_while(  # noqa: F841
+                body=do_while_body_pipeline,
+                condition=do_while_body_pipeline.outputs.is_number_larger_than_zero,
+                mapping={
+                    do_while_body_pipeline.outputs.output_in_path: do_while_body_pipeline.inputs.component_in_path,
+                },
+                max_iteration_count=5,
+            )
+
+            command_component = self._basic_component_func(
+                component_in_number=component_in_number,
+                component_in_path=component_in_path,
+            )
+            do_while_with_command_component = do_while(  # noqa: F841
+                body=command_component,
+                condition=command_component.outputs.is_number_larger_than_zero,
+                mapping={
+                    "output_in_number": command_component.inputs.component_in_number,
+                    command_component.outputs.output_in_path: command_component.inputs.component_in_path,
+                },
+                max_iteration_count=5,
+            )
+
+            # Use the outputs of do_while node
+            basic_component = do_while_body_pipeline_component(
+                component_in_number=None,
+                component_in_path=command_component.outputs.output_in_path,
+            )
+            return {"output_in_path": basic_component.outputs.output_in_path}
+
+        pipeline_job = pipeline_with_do_while(
+            component_in_number=2,
+            component_in_path=Input(type="uri_folder", path=str(Path(__file__).parent)),
+        )
+        # set pipeline level compute
+        pipeline_job.settings.default_compute = "cpu-cluster"
+
+        assert_job_cancel(pipeline_job, client)
+
+    def test_do_while_pipeline_with_primitive_inputs(self, client: MLClient):
+        @pipeline
+        def do_while_body_pipeline_component(
+            component_in_number: Input(type="integer"),
+            component_in_number_1: Input(type="integer"),
+            component_in_path: Input(type="uri_folder"),
+        ):
+            """E2E dummy train-score-eval pipeline with components defined via yaml."""
+            # Call component obj as function: apply given inputs & parameters to create a node in pipeline
+            train_with_sample_data = self._basic_component_func(
+                component_in_number=component_in_number,
+                component_in_number_1=component_in_number_1,
+                component_in_path=component_in_path,
+            )
+            return train_with_sample_data.outputs
+
+        @pipeline
+        def pipeline_with_do_while(
+            component_in_number: Input(type="integer"),
+            component_in_path: Input(type="uri_folder"),
+        ):
+            # Do while node with pipeline component
+            do_while_body_pipeline = do_while_body_pipeline_component(
+                component_in_number=component_in_number,
+                component_in_number_1=component_in_number,
+                component_in_path=component_in_path,
+            )
+            do_while_with_pipeline = do_while(  # noqa: F841
+                body=do_while_body_pipeline,
+                condition=do_while_body_pipeline.outputs.is_number_larger_than_zero,
+                mapping={
+                    do_while_body_pipeline.outputs.output_in_number: [
+                        do_while_body_pipeline.inputs.component_in_number,
+                        do_while_body_pipeline.inputs.component_in_number_1,
+                    ],
+                    "output_in_path": do_while_body_pipeline.inputs.component_in_path,
+                },
+                max_iteration_count=5,
+            )
+
+            command_component = self._basic_component_func(
+                component_in_number=component_in_number,
+                component_in_number_1=component_in_number,
+                component_in_path=component_in_path,
+            )
+            do_while_with_command_component = do_while(  # noqa: F841
+                body=command_component,
+                condition=command_component.outputs.is_number_larger_than_zero,
+                mapping={
+                    "output_in_number": [
+                        command_component.inputs.component_in_number,
+                        command_component.inputs.component_in_number_1,
+                    ],
+                    "output_in_path": command_component.inputs.component_in_path,
+                },
+                max_iteration_count=5,
+            )
+
+            # Use the outputs of do_while node
+            basic_component = do_while_body_pipeline_component(
+                component_in_number=do_while_body_pipeline.outputs.output_in_number,
+                component_in_number_1=command_component.outputs.output_in_number,
+                component_in_path=do_while_body_pipeline.outputs.output_in_path,
+            )
+            return {"output_in_path": basic_component.outputs.output_in_path}
+
+        pipeline_job = pipeline_with_do_while(
+            component_in_number=2,
+            component_in_path=Input(type="uri_folder", path=str(Path(__file__).parent)),
+        )
+        # set pipeline level compute
+        pipeline_job.settings.default_compute = "cpu-cluster"
+
+        assert_job_cancel(pipeline_job, client)
+
+
 @pytest.mark.skipif(
     condition=is_live(),
     # TODO: reopen live test when parallel_for deployed to canary

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

@@ -3027,3 +3027,41 @@ def my_pipeline(timeout) -> PipelineJob:
         pipeline_job = assert_job_cancel(pipeline, client)
         assert pipeline_job.jobs["node_0"].limits.timeout == "${{parent.inputs.timeout}}"
         assert pipeline_job.jobs["node_1"].limits.timeout == 1
+
+    def test_pipeline_component_primitive_type_consumption(self, client: MLClient):
+        baisc_component_func = load_component(
+            "./tests/test_configs/dsl_pipeline/primitive_type_components/basic_component.yml"
+        )
+        boolean_func = load_component("./tests/test_configs/dsl_pipeline/primitive_type_components/boolean.yml")
+        integer_func = load_component("./tests/test_configs/dsl_pipeline/primitive_type_components/integer.yml")
+        number_func = load_component("./tests/test_configs/dsl_pipeline/primitive_type_components/number.yml")
+        string_func = load_component("./tests/test_configs/dsl_pipeline/primitive_type_components/string.yml")
+
+        @dsl.pipeline
+        def pipeline_component_func(bool_param: bool, int_param: int, float_param: float, str_param: str):
+            baisc_component_func(
+                bool_param=bool_param,
+                int_param=int_param,
+                float_param=float_param,
+                str_param=str_param,
+            )
+
+        @dsl.pipeline
+        def pipeline_func():
+            # components return primitive type outputs
+            bool_node = boolean_func()
+            int_node = integer_func()
+            float_node = number_func()
+            str_node = string_func()
+            # pipeline component consume above primitive type outputs
+            pipeline_node = pipeline_component_func(  # noqa: F841
+                bool_param=bool_node.outputs.output,
+                int_param=int_node.outputs.output,
+                float_param=float_node.outputs.output,
+                str_param=str_node.outputs.output,
+            )
+
+        pipeline_job = pipeline_func()
+        pipeline_job.settings.default_compute = "cpu-cluster"
+
+        assert_job_cancel(pipeline_job, client)