[ML][Pipelines] feat: support v2 style 1p only components (Azure#30443)

* refactor: use pytest.param to shorten test name

* feat: support v2 style 1p only component types

Author: elliotzh

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

@@ -10,7 +10,7 @@ repos:
   hooks:
   - id: interrogate
     types_or: [python]
-    exclude: ^(sdk/ml/azure-ai-ml/tests/|sdk/ml/azure-ai-ml/azure/ai/ml/_restclient|sdk/ml/azure-ai-ml/setup.py)
+    exclude: ^(sdk/ml/azure-ai-ml/tests/|sdk/ml/azure-ai-ml/scripts/|sdk/ml/azure-ai-ml/azure/ai/ml/_restclient|sdk/ml/azure-ai-ml/setup.py)
     # exclude defined here because exclude in pyproject.toml is not being respected
 - repo: https://github.com/streetsidesoftware/cspell-cli
   rev: v6.31.0

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

@@ -30,6 +30,12 @@ class NodeType:
     DATA_TRANSFER = "DataTransferComponent"
     DISTRIBUTED = "DistributedComponent"
     HDI = "HDInsightComponent"
+    SCOPE_V2 = "scope"
+    HDI_V2 = "hdinsight"
+    HEMERA_V2 = "hemera"
+    STARLITE_V2 = "starlite"
+    AE365EXEPOOL_V2 = "ae365exepool"
+    AETHER_BRIDGE_V2 = "aetherbridge"
     PARALLEL = "ParallelComponent"
     SCOPE = "ScopeComponent"
     STARLITE = "StarliteComponent"

sdk/ml/azure-ai-ml/azure/ai/ml/_internal/_setup.py

@@ -70,11 +70,18 @@ def enable_internal_components_in_pipeline(*, force=False):
 
     # redo the registration for those with specific runsettings
     _register_node(NodeType.DATA_TRANSFER, DataTransfer, InternalBaseNodeSchema)
-    _register_node(NodeType.HEMERA, Hemera, InternalBaseNodeSchema)
-    _register_node(NodeType.STARLITE, Starlite, InternalBaseNodeSchema)
     _register_node(NodeType.COMMAND, Command, CommandSchema)
     _register_node(NodeType.DISTRIBUTED, Distributed, DistributedSchema)
-    _register_node(NodeType.SCOPE, Scope, ScopeSchema)
     _register_node(NodeType.PARALLEL, Parallel, ParallelSchema)
+    _register_node(NodeType.HEMERA, Hemera, InternalBaseNodeSchema)
+    _register_node(NodeType.STARLITE, Starlite, InternalBaseNodeSchema)
+    _register_node(NodeType.SCOPE, Scope, ScopeSchema)
     _register_node(NodeType.HDI, HDInsight, HDInsightSchema)
+
+    # register v2 style 1p only components
+    _register_node(NodeType.HEMERA_V2, Hemera, InternalBaseNodeSchema)
+    _register_node(NodeType.STARLITE_V2, Starlite, InternalBaseNodeSchema)
+    _register_node(NodeType.SCOPE_V2, Scope, ScopeSchema)
+    _register_node(NodeType.HDI_V2, HDInsight, HDInsightSchema)
+    # Ae365exepool and AetherBridge have been registered to InternalBaseNode
     _set_registered(True)

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

@@ -5,6 +5,7 @@
 
 import copy
 import logging
+import re
 from enum import Enum
 from os import PathLike, path
 from pathlib import Path
@@ -466,10 +467,7 @@ def _customized_validate(self):
                 message=SPARK_ENVIRONMENT_WARNING_MESSAGE,
             )
         result.merge_with(self._validate_entry_exist(raise_error=False))
-        try:
-            self._validate_fields()
-        except ValidationException as e:
-            result.append_error(yaml_path="*", message=str(e))
+        result.merge_with(self._validate_fields())
         return result
 
     def _validate_entry_exist(self, raise_error=False) -> MutableValidationResult:
@@ -508,14 +506,42 @@ def _validate_entry_exist(self, raise_error=False) -> MutableValidationResult:
                     )
         return validation_result.try_raise(error_target=self._get_validation_error_target(), raise_error=raise_error)
 
-    def _validate_fields(self) -> None:
-        _validate_compute_or_resources(self.compute, self.resources)
-        _validate_input_output_mode(self.inputs, self.outputs)
-        _validate_spark_configurations(self)
-        self._validate_entry()
+    def _validate_fields(self) -> MutableValidationResult:
+        validation_result = self._create_empty_validation_result()
+        try:
+            _validate_compute_or_resources(self.compute, self.resources)
+        except ValidationException as e:
+            validation_result.append_error(message=str(e), yaml_path="resources")
+            validation_result.append_error(message=str(e), yaml_path="compute")
+
+        try:
+            _validate_input_output_mode(self.inputs, self.outputs)
+        except ValidationException as e:
+            msg = str(e)
+            m = re.match(r"(Input|Output) '(\w+)'", msg)
+            if m:
+                io_type, io_name = m.groups()
+                if io_type == "Input":
+                    validation_result.append_error(message=msg, yaml_path=f"inputs.{io_name}")
+                else:
+                    validation_result.append_error(message=msg, yaml_path=f"outputs.{io_name}")
+
+        try:
+            _validate_spark_configurations(self)
+        except ValidationException as e:
+            validation_result.append_error(message=str(e), yaml_path="conf")
+
+        try:
+            self._validate_entry()
+        except ValidationException as e:
+            validation_result.append_error(message=str(e), yaml_path="entry")
 
         if self.args:
-            validate_inputs_for_args(self.args, self.inputs)
+            try:
+                validate_inputs_for_args(self.args, self.inputs)
+            except ValidationException as e:
+                validation_result.append_error(message=str(e), yaml_path="args")
+        return validation_result
 
     def __call__(self, *args, **kwargs) -> "Spark":
         """Call Spark as a function will return a new instance each time."""

sdk/ml/azure-ai-ml/pyproject.toml

@@ -8,7 +8,7 @@ verifytypes = false
 fail-under = 35
 verbose = 2
 ignore-module = true
-exclude = ["setup.py", "tests", "azure/ai/ml/_restclient"]
+exclude = ["setup.py", "tests", "azure/ai/ml/_restclient", "scripts"]
 
 [tool.isort]
 profile = "black"

sdk/ml/azure-ai-ml/scripts/run_tests.py

@@ -3,6 +3,7 @@
 # ---------------------------------------------------------
 import argparse
 import contextlib
+import glob
 import json
 import os
 import re
@@ -88,24 +89,34 @@ def run_simple(
     extra_params,
     *,
     is_live_and_recording,
-    log_file_path=None,
+    log_file_path,
+    log_suffix=None,
     log_in_json=False,
 ):
     print(f"Running {len(tests_to_run)} tests under {working_dir}: ")
     for test_name in tests_to_run:
         print(test_name)
 
-    if log_in_json or log_file_path is None:
+    if log_file_path and log_suffix:
+        log_file_path = log_file_path.with_suffix(log_file_path.suffix + log_suffix)
+
+    if log_in_json:
+        if log_file_path is None:
+            raise ValueError("log_file_path must be specified when log_in_json is True")
         stdout = None
+        json_log_file_path = log_file_path.with_suffix(log_file_path.suffix + ".log")
     else:
-        stdout = open(log_file_path, "wb")
+        stdout = open(log_file_path.with_suffix(log_file_path.suffix + ".txt"), "wb")
+        json_log_file_path = None
+
     with update_dot_env_file(
         {"AZURE_TEST_RUN_LIVE": is_live_and_recording, "AZURE_SKIP_LIVE_RECORDING": not is_live_and_recording},
     ):
         for test_class, keyword_param in reorganize_tests(tests_to_run):
             tmp_extra_params = extra_params + keyword_param
             if log_in_json:
-                temp_log_file_path = log_file_path.with_stem("temp")
+                # use a temp json file to avoid overwriting the final log file
+                temp_log_file_path = json_log_file_path.with_stem("temp")
                 tmp_extra_params += ["--report-log", temp_log_file_path.as_posix()]
 
             subprocess.run(
@@ -120,10 +131,14 @@ def run_simple(
                 stdout=stdout,
             )
             if log_in_json:
-                with open(log_file_path, "a", encoding="utf-8") as f:
+                # append temp json file to the final log file
+                with open(json_log_file_path, "a", encoding="utf-8") as f:
                     f.write(temp_log_file_path.read_text())
     if stdout is not None:
         stdout.close()
+        print(log_file_path.with_suffix(log_file_path.suffix + ".txt").read_text())
+
+    return json_log_file_path
 
 
 def reorganize_tests(tests_to_run):
@@ -168,59 +183,80 @@ def reorganize_tests(tests_to_run):
         yield test_class, keyword_param
 
 
-def run_tests(tests_to_run, extras, *, skip_first_run=False, record_mismatch=False, is_live_and_recording=False):
+def get_base_log_path(working_dir, *, create_new=True):
+    log_dir = working_dir / "scripts" / "tmp"
+    if not create_new:
+        logs = sorted(glob.glob(str(log_dir / "pytest.*.first.log")))
+        if len(logs) == 0:
+            raise RuntimeError("No previous run log file found")
+        return Path(logs[-1][: -len(".first.log")])
+    else:
+        log_file_path = log_dir / "pytest.{}".format(datetime.now().strftime("%Y%m%d%H%M%S"))
+        log_file_path.parent.mkdir(parents=True, exist_ok=True)
+        return log_file_path
+
+
+def get_failed_tests(log_file_path):
+    tests_failed_with_recording_mismatch = []
+    failed_tests = []
+    with open(log_file_path, "r") as f:
+        for line in f:
+            node = json.loads(line)
+            if "outcome" not in node:
+                continue
+            if node["outcome"] != "failed":
+                continue
+            test_name = location_to_test_name(node["location"])
+            failed_tests.append(test_name)
+            msg = node["longrepr"]["reprcrash"]["message"]
+            if "NotFound" in msg:
+                tests_failed_with_recording_mismatch.append(test_name)
+    return failed_tests, tests_failed_with_recording_mismatch
+
+
+def run_tests(tests_to_run, extras, *, skip_first_run=False, record_mismatch=False):
     working_dir = Path(__file__).parent.parent
-    log_file_path = working_dir / "scripts" / "tmp" / "pytest_first_run.log"
-    log_file_path.parent.mkdir(parents=True, exist_ok=True)
-    if record_mismatch and not skip_first_run:
-        # reset the log file
-        log_file_path.unlink(missing_ok=True)
-
-    if not (record_mismatch and skip_first_run):
-        # first run
-        run_simple(
+    log_file_path = get_base_log_path(working_dir, create_new=not skip_first_run)
+
+    if skip_first_run:
+        json_log_file_path = log_file_path.with_suffix(log_file_path.suffix + ".first.log")
+    else:
+        json_log_file_path = run_simple(
             tests_to_run,
             working_dir,
             extras + ["--disable-warnings", "--disable-pytest-warnings"],
-            # first run in record-mismatch mode is always in playback mode
-            is_live_and_recording=is_live_and_recording and not record_mismatch,
-            log_file_path=log_file_path if record_mismatch else None,
-            log_in_json=record_mismatch,
+            # first run is always in playback mode
+            is_live_and_recording=False,
+            log_file_path=log_file_path,
+            log_in_json=True,
+            log_suffix=".first",
         )
 
     if record_mismatch:
-        tests_failed_with_recording_mismatch = []
-        with open(log_file_path, "r") as f:
-            for line in f:
-                node = json.loads(line)
-                if "outcome" not in node:
-                    continue
-                if node["outcome"] != "failed":
-                    continue
-                msg = node["longrepr"]["reprcrash"]["message"]
-                if "ResourceNotFoundError" in msg:
-                    tests_failed_with_recording_mismatch.append(location_to_test_name(node["location"]))
-
+        failed_tests, tests_failed_with_recording_mismatch = get_failed_tests(json_log_file_path)
         if tests_failed_with_recording_mismatch:
             print("Re-do live mode recording for tests: \n", json.dumps(tests_failed_with_recording_mismatch, indent=2))
             run_simple(
                 tests_failed_with_recording_mismatch,
                 working_dir,
                 extra_params=["--tb=line"],
                 is_live_and_recording=True,
+                log_suffix=".record",
+                log_file_path=log_file_path,
             )
 
-            # re-run the original tests to check if they are still failures and output the log
+            print(
+                "Rerun playback mode for failed tests: \n", json.dumps(tests_failed_with_recording_mismatch, indent=2)
+            )
             run_simple(
-                tests_to_run,
+                failed_tests,
                 working_dir,
-                extras + ["--disable-warnings", "--disable-pytest-warnings"],
+                extra_params=extras + ["--disable-warnings", "--disable-pytest-warnings"],
                 is_live_and_recording=False,
-                log_file_path=working_dir
-                / "scripts"
-                / "tmp"
-                / "pytest.{}.log".format(datetime.now().strftime("%Y%m%d%H%M%S")),
+                log_file_path=log_file_path,
+                log_suffix=".final",
             )
+            print(log_file_path.with_suffix(log_file_path.suffix + ".final.log").read_text())
 
 
 if __name__ == "__main__":
@@ -263,6 +299,9 @@ def run_tests(tests_to_run, extras, *, skip_first_run=False, record_mismatch=Fal
         _tests = load_tests_from_file(_args.file)
     elif _args.name:
         _tests = [_args.name]
+    elif _args.skip_first_run and _args.record_mismatch:
+        # load failed tests from last run log
+        _tests = []
     else:
         raise ValueError("Must specify either --file or --name")
     run_tests(

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

@@ -380,7 +380,8 @@ def test_spark_job_with_builder_in_pipeline_with_dynamic_allocation_disabled(
         validation_result = client.jobs.validate(pipeline)
         assert validation_result.passed is False
         assert validation_result.error_messages == {
-            "jobs.add_greeting_column": "Should not specify min or max executors when dynamic allocation is disabled.",
+            "jobs.add_greeting_column.conf": "Should not specify min or max executors "
+            "when dynamic allocation is disabled.",
         }
 
     @pytest.mark.e2etest

sdk/ml/azure-ai-ml/tests/dsl/unittests/test_command_builder.py

@@ -780,7 +780,7 @@ def test_spark_job_with_dynamic_allocation_disabled(self):
         )
         result = node._validate()
         message = "Should not specify min or max executors when dynamic allocation is disabled."
-        assert "*" in result.error_messages and message == result.error_messages["*"]
+        assert "conf" in result.error_messages and message == result.error_messages["conf"]
 
     def test_executor_instances_is_mandatory_when_dynamic_allocation_disabled(self):
         node = spark(
@@ -797,7 +797,7 @@ def test_executor_instances_is_mandatory_when_dynamic_allocation_disabled(self):
             "spark.driver.cores, spark.driver.memory, spark.executor.cores, spark.executor.memory and "
             "spark.executor.instances are mandatory fields."
         )
-        assert "*" in result.error_messages and message == result.error_messages["*"]
+        assert "conf" in result.error_messages and message == result.error_messages["conf"]
 
     def test_executor_instances_is_specified_as_min_executor_if_unset(self):
         node = spark(
@@ -834,7 +834,7 @@ def test_excutor_instances_throw_error_when_out_of_range(self):
             "Executor instances must be a valid non-negative integer and must be between "
             "spark.dynamicAllocation.minExecutors and spark.dynamicAllocation.maxExecutors"
         )
-        assert "*" in result.error_messages and message == result.error_messages["*"]
+        assert "conf" in result.error_messages and message == result.error_messages["conf"]
 
     def test_spark_job_with_additional_conf(self):
         node = spark(