[ML][Pipelines] fix: pipeline job schema validation fail for data-binding expression (Azure#28517)

* fix: pipeline job schema validation fail for data-binding expression

* fix: test public json schema

Author: elliotzh

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/core/fields.py

@@ -455,6 +455,7 @@ def __init__(self, union_fields: List[fields.Field], is_strict=False, **kwargs):
                 resolve_field_instance(cls_or_instance)
                 for cls_or_instance in union_fields
             ]
+            # TODO: make serialization/de-serialization work in the same way as json schema when is_strict is True
             self.is_strict = is_strict  # S\When True, combine fields with oneOf instead of anyOf at schema generation
         except FieldInstanceResolutionError as error:
             raise ValueError(

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/job/input_output_fields_provider.py

@@ -4,6 +4,7 @@
 
 from marshmallow import fields
 
+from azure.ai.ml._schema._utils.data_binding_expression import support_data_binding_expression_for_fields
 from azure.ai.ml._schema.core.fields import NestedField, PrimitiveValueField, UnionField
 from azure.ai.ml._schema.job.input_output_entry import (
     DataInputSchema,
@@ -14,20 +15,30 @@
 )
 
 
-def InputsField(**kwargs):
+def InputsField(*, support_databinding: bool = False, **kwargs):
+    value_fields = [
+        NestedField(DataInputSchema),
+        NestedField(ModelInputSchema),
+        NestedField(MLTableInputSchema),
+        NestedField(InputLiteralValueSchema),
+        PrimitiveValueField(is_strict=False),
+        # This ordering of types for the values keyword is intentional. The ordering of types
+        # determines what order schema values are matched and cast in. Changing the current ordering can
+        # result in values being mis-cast such as 1.0 translating into True.
+    ]
+
+    # As is_strict is set to True, 1 and only 1 value field must be matched.
+    # root level data-binding expression has already been covered by PrimitiveValueField;
+    # If support_databinding is True, we should only add data-binding expression support for nested fields.
+    if support_databinding:
+        for field_obj in value_fields:
+            if isinstance(field_obj, NestedField):
+                support_data_binding_expression_for_fields(field_obj.schema)
+
     return fields.Dict(
         keys=fields.Str(),
         values=UnionField(
-            [
-                NestedField(DataInputSchema),
-                NestedField(ModelInputSchema),
-                NestedField(MLTableInputSchema),
-                NestedField(InputLiteralValueSchema),
-                PrimitiveValueField(is_strict=False),
-                # This ordering of types for the values keyword is intentional. The ordering of types
-                # determines what order schema values are matched and cast in. Changing the current ordering can
-                # result in values being mis-cast such as 1.0 translating into True.
-            ],
+            value_fields,
             metadata={"description": "Inputs to a job."},
             is_strict=True,
             **kwargs

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

@@ -48,7 +48,7 @@
 class BaseNodeSchema(PathAwareSchema):
     unknown = INCLUDE
 
-    inputs = InputsField()
+    inputs = InputsField(support_databinding=True)
     outputs = fields.Dict(
         keys=fields.Str(),
         values=UnionField([OutputBindingStr, NestedField(OutputSchema)], allow_none=True),
@@ -61,7 +61,7 @@ def __init__(self, *args, **kwargs):
         # data binding expression is not supported inside component field, while validation error
         # message will be very long when component is an object as error message will include
         # str(component), so just add component to skip list. The same to trial in Sweep.
-        support_data_binding_expression_for_fields(self, ["type", "component", "trial"])
+        support_data_binding_expression_for_fields(self, ["type", "component", "trial", "inputs"])
 
     @post_dump(pass_original=True)
     def add_user_setting_attr_dict(self, data, original_data, **kwargs):  # pylint: disable=unused-argument

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

@@ -8,7 +8,6 @@
 
 from marshmallow import INCLUDE, ValidationError, post_load, pre_dump, pre_load
 
-from azure.ai.ml._schema._utils.data_binding_expression import _add_data_binding_to_field
 from azure.ai.ml._schema.core.fields import (
     ArmVersionedStr,
     ComputeField,
@@ -36,7 +35,8 @@ class PipelineJobSchema(BaseJobSchema):
     type = StringTransformedEnum(allowed_values=[JobType.PIPELINE])
     compute = ComputeField()
     settings = NestedField(PipelineJobSettingsSchema, unknown=INCLUDE)
-    inputs = InputsField()
+    # Support databinding in inputs as we support macro like ${{name}}
+    inputs = InputsField(support_databinding=True)
     outputs = OutputsField()
     jobs = PipelineJobsField()
     component = UnionField(
@@ -50,11 +50,6 @@ class PipelineJobSchema(BaseJobSchema):
         ],
     )
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        # Support databinding in inputs as we support macro like ${{name}}
-        _add_data_binding_to_field(self.load_fields["inputs"], [], [])
-
     @pre_dump()
     def backup_jobs_and_remove_component(self, job, **kwargs):
         # pylint: disable=protected-access

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

@@ -48,6 +48,8 @@
 from azure.core.pipeline.transport import HttpTransport
 from azure.identity import AzureCliCredential, ClientSecretCredential, DefaultAzureCredential
 
+from test_utilities.utils import reload_schema_for_nodes_in_pipeline_job
+
 E2E_TEST_LOGGING_ENABLED = "E2E_TEST_LOGGING_ENABLED"
 test_folder = Path(os.path.abspath(__file__)).parent.absolute()
 
@@ -840,14 +842,28 @@ def pytest_configure(config):
 
 @pytest.fixture()
 def enable_private_preview_pipeline_node_types():
-    # Update the node types in pipeline jobs to include the private preview node types
-    from azure.ai.ml._schema.pipeline import pipeline_job
+    with reload_schema_for_nodes_in_pipeline_job():
+        yield
 
-    schema = pipeline_job.PipelineJobSchema
-    original_jobs = schema._declared_fields["jobs"]
-    schema._declared_fields["jobs"] = pipeline_job.PipelineJobsField()
 
-    try:
+@pytest.fixture()
+def disable_internal_components():
+    """Some global changes are made in enable_internal_components, so we need to explicitly disable it.
+    It's not recommended to use this fixture along with other related fixtures like enable_internal_components
+    and enable_private_preview_features, as the execution order of fixtures is not guaranteed.
+    """
+    from azure.ai.ml._internal._schema.component import NodeType
+    from azure.ai.ml._internal._util import _set_registered
+    from azure.ai.ml.entities._component.component_factory import component_factory
+    from azure.ai.ml.entities._job.pipeline._load_component import pipeline_node_factory
+
+    for _type in NodeType.all_values():
+        pipeline_node_factory._create_instance_funcs.pop(_type, None)  # pylint: disable=protected-access
+        pipeline_node_factory._load_from_rest_object_funcs.pop(_type, None)  # pylint: disable=protected-access
+        component_factory._create_instance_funcs.pop(_type, None)  # pylint: disable=protected-access
+        component_factory._create_schema_funcs.pop(_type, None)  # pylint: disable=protected-access
+
+    _set_registered(False)
+
+    with reload_schema_for_nodes_in_pipeline_job(revert_after_yield=False):
         yield
-    finally:
-        schema._declared_fields["jobs"] = original_jobs

sdk/ml/azure-ai-ml/tests/internal/_utils.py

@@ -282,18 +282,3 @@ def extract_non_primitive(obj):
     if isinstance(obj, (float, int, str)):
         return None
     return obj
-
-
-def unregister_internal_components():
-    from azure.ai.ml._internal._schema.component import NodeType
-    from azure.ai.ml._internal._util import _set_registered
-    from azure.ai.ml.entities._component.component_factory import component_factory
-    from azure.ai.ml.entities._job.pipeline._load_component import pipeline_node_factory
-
-    for _type in NodeType.all_values():
-        pipeline_node_factory._create_instance_funcs.pop(_type, None)  # pylint: disable=protected-access
-        pipeline_node_factory._load_from_rest_object_funcs.pop(_type, None)  # pylint: disable=protected-access
-        component_factory._create_instance_funcs.pop(_type, None)  # pylint: disable=protected-access
-        component_factory._create_schema_funcs.pop(_type, None)  # pylint: disable=protected-access
-
-    _set_registered(False)

sdk/ml/azure-ai-ml/tests/internal/unittests/test_internal_disabled.py

@@ -0,0 +1,24 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+import pytest
+from azure.ai.ml.constants._common import AZUREML_INTERNAL_COMPONENTS_ENV_VAR
+from azure.ai.ml.dsl._utils import environment_variable_overwrite
+from azure.ai.ml.exceptions import ValidationException
+
+
+@pytest.mark.usefixtures("disable_internal_components")
+@pytest.mark.unittest
+@pytest.mark.pipeline_test
+class TestInternalDisabled:
+    def test_load_pipeline_job_with_internal_nodes_from_rest(self):
+        # this is a simplified test case which avoid constructing a complete pipeline job rest object
+        from azure.ai.ml.entities._job.pipeline._load_component import pipeline_node_factory
+
+        internal_node_type = "CommandComponent"
+        with environment_variable_overwrite(AZUREML_INTERNAL_COMPONENTS_ENV_VAR, "False"):
+            with pytest.raises(ValidationException, match=f"Unsupported component type: {internal_node_type}."):
+                pipeline_node_factory.get_load_from_rest_object_func(internal_node_type)
+
+        with environment_variable_overwrite(AZUREML_INTERNAL_COMPONENTS_ENV_VAR, "True"):
+            pipeline_node_factory.get_load_from_rest_object_func(internal_node_type)

sdk/ml/azure-ai-ml/tests/internal/unittests/test_pipeline_job.py

@@ -32,12 +32,10 @@
     TargetSelector,
 )
 from azure.ai.ml._internal.entities import InternalBaseNode, InternalComponent, Scope
-from azure.ai.ml.constants._common import AZUREML_INTERNAL_COMPONENTS_ENV_VAR, AssetTypes
+from azure.ai.ml.constants._common import AssetTypes
 from azure.ai.ml.constants._job.job import JobComputePropertyFields
 from azure.ai.ml.dsl import pipeline
-from azure.ai.ml.dsl._utils import environment_variable_overwrite
 from azure.ai.ml.entities import CommandComponent, Data, PipelineJob
-from azure.ai.ml.exceptions import ValidationException
 
 from .._utils import (
     DATA_VERSION,
@@ -46,7 +44,6 @@
     extract_non_primitive,
     get_expected_runsettings_items,
     set_run_settings,
-    unregister_internal_components,
 )
 
 
@@ -602,15 +599,3 @@ def test_job_properties(self):
             assert len(node_dict["properties"]) == 1
             assert "AZURE_ML_PathOnCompute_" in list(node_dict["properties"].keys())[0]
             assert node_dict["properties"] == rest_node_dict["properties"]
-
-    def test_load_pipeline_job_with_internal_nodes_from_rest(self):
-        # this is a simplified test case which avoid constructing a complete pipeline job rest object
-        from azure.ai.ml.entities._job.pipeline._load_component import pipeline_node_factory
-
-        unregister_internal_components()
-        internal_node_type = "CommandComponent"
-        with environment_variable_overwrite(AZUREML_INTERNAL_COMPONENTS_ENV_VAR, "False"):
-            with pytest.raises(ValidationException, match=f"Unsupported component type: {internal_node_type}."):
-                pipeline_node_factory.get_load_from_rest_object_func(internal_node_type)
-
-        pipeline_node_factory.get_load_from_rest_object_func(internal_node_type)

sdk/ml/azure-ai-ml/tests/pipeline_job/unittests/test_private_preview_disabled.py

@@ -0,0 +1,26 @@
+from pathlib import Path
+import pytest
+import yaml
+from jsonschema.validators import validate
+from azure.ai.ml.entities import PipelineJob
+from test_utilities.json_schema import PatchedJSONSchema
+
+from .._util import _PIPELINE_JOB_TIMEOUT_SECOND
+
+
+# schema of nodes will be reloaded with private preview features disabled in unregister_internal_components
+@pytest.mark.usefixtures("disable_internal_components")
+@pytest.mark.timeout(_PIPELINE_JOB_TIMEOUT_SECOND)
+@pytest.mark.unittest
+@pytest.mark.pipeline_test
+class TestPrivatePreviewDisabled:
+    def test_public_json_schema(self):
+        # public json schema is the json schema to be saved in
+        # https://azuremlschemas.azureedge.net/latest/pipelineJob.schema.json
+        base_dir = Path("./tests/test_configs/pipeline_jobs/json_schema_validation")
+        target_schema = PatchedJSONSchema().dump(PipelineJob._create_schema_for_validation(context={"base_path": "./"}))
+
+        with open(base_dir.joinpath("component_spec.yaml"), "r") as f:
+            yaml_data = yaml.safe_load(f.read())
+
+        validate(yaml_data, target_schema)

sdk/ml/azure-ai-ml/tests/test_configs/pipeline_jobs/json_schema_validation/component_spec.yaml

@@ -0,0 +1,17 @@
+$schema: https://azuremlschemas.azureedge.net/latest/pipelineJob.schema.json
+type: pipeline
+
+settings:
+  default_compute: azureml:cpu-cluster
+
+
+jobs:
+  process_data:
+    type: command
+    component: ./1process_data_component.yaml
+    inputs:
+      raw_data: ${{parent.inputs.raw_data}}
+    outputs:
+      train:
+      validation:
+    compute: azureml:cpu-cluster