[V2 Datatransfer] Disable reference anoymous component for import/export and update built-in component id for Datatransfer job V2 in Pipeline (Azure#28932)

* add test

* add more invalid test

* update

* update

* update schema

* update builtin component id and test

* format

* format

* delete unnecesssary test

* update test

* disable refer to local component yaml

* rm unused yaml file

* update

* update according to comment and black format

* typo

* pylint

* black

* update recording

---------

Co-authored-by: Ying Chen <[email protected]>

Author: YingChen1996

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

@@ -38,6 +38,15 @@ class DataTransferCopyComponentSchema(DataTransferComponentSchemaMixin):
         values=NestedField(InputPortSchema),
     )
 
+    @validates("outputs")
+    def outputs_key(self, value):
+        outputs_count = len(value)
+        if outputs_count != 1:
+            msg = "Only support single output in {}, but there're {} outputs."
+            raise ValidationError(
+                message=msg.format(DataTransferTaskType.COPY_DATA, outputs_count), field_name="outputs"
+            )
+
 
 class SinkSourceSchema(metaclass=PatchedSchemaMeta):
     type = StringTransformedEnum(

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

@@ -15,15 +15,11 @@
     AnonymousParallelComponentSchema,
     AnonymousSparkComponentSchema,
     AnonymousDataTransferCopyComponentSchema,
-    AnonymousDataTransferImportComponentSchema,
-    AnonymousDataTransferExportComponentSchema,
     ComponentFileRefField,
     ImportComponentFileRefField,
     ParallelComponentFileRefField,
     SparkComponentFileRefField,
     DataTransferCopyComponentFileRefField,
-    DataTransferImportComponentFileRefField,
-    DataTransferExportComponentFileRefField,
 )
 from azure.ai.ml._schema.core.fields import ArmVersionedStr, NestedField, RegistryStr, UnionField
 from azure.ai.ml._schema.core.schema import PathAwareSchema
@@ -418,16 +414,8 @@ def resolve_inputs_outputs(self, job, **kwargs):
 
 class DataTransferImportSchema(BaseNodeSchema):
     # pylint: disable=unused-argument
-    component = TypeSensitiveUnionField(
-        {
-            NodeType.DATA_TRANSFER: [
-                # inline component or component file reference starting with FILE prefix
-                NestedField(AnonymousDataTransferImportComponentSchema, unknown=INCLUDE),
-                # component file reference
-                DataTransferImportComponentFileRefField(),
-            ],
-        },
-        plain_union_fields=[
+    component = UnionField(
+        [
             # for registry type assets
             RegistryStr(),
             # existing component
@@ -477,16 +465,8 @@ def resolve_inputs_outputs(self, job, **kwargs):
 
 class DataTransferExportSchema(BaseNodeSchema):
     # pylint: disable=unused-argument
-    component = TypeSensitiveUnionField(
-        {
-            NodeType.DATA_TRANSFER: [
-                # inline component or component file reference starting with FILE prefix
-                NestedField(AnonymousDataTransferExportComponentSchema, unknown=INCLUDE),
-                # component file reference
-                DataTransferExportComponentFileRefField(),
-            ],
-        },
-        plain_union_fields=[
+    component = UnionField(
+        [
             # for registry type assets
             RegistryStr(),
             # existing component

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

@@ -84,14 +84,6 @@ def PipelineJobsField():
             NestedField(SparkSchema, unknown=INCLUDE),
             NestedField(PipelineSparkJobSchema),
         ],
-        NodeType.DATA_TRANSFER: [
-            NestedField(DataTransferCopySchema, unknown=INCLUDE),
-            NestedField(DataTransferImportSchema, unknown=INCLUDE),
-            NestedField(DataTransferExportSchema, unknown=INCLUDE),
-            NestedField(PipelineDataTransferCopyJobSchema),
-            NestedField(PipelineDataTransferImportJobSchema),
-            NestedField(PipelineDataTransferExportJobSchema),
-        ],
     }
 
     # Note: the private node types only available when private preview flag opened before init of pipeline job
@@ -101,6 +93,29 @@ def PipelineJobsField():
         pipeline_enable_job_type[ControlFlowType.IF_ELSE] = [NestedField(ConditionNodeSchema, unknown=INCLUDE)]
         pipeline_enable_job_type[ControlFlowType.PARALLEL_FOR] = [NestedField(ParallelForSchema, unknown=INCLUDE)]
 
+    # Todo: Put data_transfer logic to the last to avoid error message conflict, open a item to track:
+    #  https://msdata.visualstudio.com/Vienna/_workitems/edit/2244262/
+    pipeline_enable_job_type[NodeType.DATA_TRANSFER] = [
+        TypeSensitiveUnionField(
+            {
+                DataTransferTaskType.COPY_DATA: [
+                    NestedField(DataTransferCopySchema, unknown=INCLUDE),
+                    NestedField(PipelineDataTransferCopyJobSchema),
+                ],
+                DataTransferTaskType.IMPORT_DATA: [
+                    NestedField(DataTransferImportSchema, unknown=INCLUDE),
+                    NestedField(PipelineDataTransferImportJobSchema),
+                ],
+                DataTransferTaskType.EXPORT_DATA: [
+                    NestedField(DataTransferExportSchema, unknown=INCLUDE),
+                    NestedField(PipelineDataTransferExportJobSchema),
+                ],
+            },
+            type_field_name="task",
+            unknown=INCLUDE,
+        )
+    ]
+
     pipeline_job_field = fields.Dict(
         keys=NodeNameStr(),
         values=TypeSensitiveUnionField(pipeline_enable_job_type),
@@ -157,11 +172,11 @@ def _post_load_pipeline_jobs(context, data: dict) -> dict:
                 context=context,
                 pipeline_job_dict=data,
             )
-            if not (
-                job_instance.type == NodeType.DATA_TRANSFER and job_instance.task != DataTransferTaskType.COPY_DATA
-            ):
+            if job_instance.type == NodeType.DATA_TRANSFER and job_instance.task != DataTransferTaskType.COPY_DATA:
+                job_instance._source = ComponentSource.BUILTIN
+            else:
                 job_instance.component._source = ComponentSource.YAML_JOB
-            job_instance._source = job_instance.component._source
+                job_instance._source = job_instance.component._source
             jobs[key] = job_instance
         # update job instance name to key
         job_instance.name = key

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

@@ -11,6 +11,7 @@
 
 from azure.ai.ml._schema.core.fields import NestedField, UnionField
 from azure.ai.ml._schema.job.input_output_entry import OutputSchema
+from azure.ai.ml._schema.pipeline.pipeline_job_io import OutputBindingStr
 from azure.ai.ml._schema.job.data_transfer_job import (
     DataTransferCopyJobSchema,
     DataTransferImportJobSchema,
@@ -23,7 +24,7 @@
 class PipelineDataTransferCopyJobSchema(DataTransferCopyJobSchema):
     outputs = fields.Dict(
         keys=fields.Str(),
-        values=UnionField([NestedField(OutputSchema), fields.Str()], allow_none=True),
+        values=UnionField([NestedField(OutputSchema), OutputBindingStr], allow_none=True),
     )
 
     @post_load
@@ -36,7 +37,7 @@ def make(self, data: Any, **kwargs: Any):
 class PipelineDataTransferImportJobSchema(DataTransferImportJobSchema):
     outputs = fields.Dict(
         keys=fields.Str(),
-        values=UnionField([NestedField(OutputSchema)], allow_none=True),
+        values=UnionField([NestedField(OutputSchema), OutputBindingStr], allow_none=True),
     )
 
     @post_load

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

@@ -54,11 +54,9 @@ class ExternalDataType(object):
 
 
 class DataTransferBuiltinComponentUri(object):
-    # todo: need update
-    IMPORT_DATABASE = "azureml://registries/import_database"
-    IMPORT_FILE_SYSTEM = "azureml://registries/import_file_system"
-    EXPORT_DATABASE = "azureml://registries/export_database"
-    EXPORT_FILE_SYSTEM = "azureml://registries/export_file_system"
+    IMPORT_DATABASE = "azureml://registries/azureml-dev/components/import_data_database/versions/0.0.1"
+    IMPORT_FILE_SYSTEM = "azureml://registries/azureml-dev/components/import_data_file_system/versions/0.0.1"
+    EXPORT_DATABASE = "azureml://registries/azureml-dev/components/export_data_database/versions/0.0.1"
 
 
 class ComponentSource:

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

@@ -29,7 +29,7 @@
     DataTransferImportJob,
     DataTransferExportJob,
 )
-from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY
+from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, AssetTypes
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, ValidationErrorType, ValidationException
 from azure.ai.ml.entities._inputs_outputs.external_data import Database, FileSystem
 
@@ -47,25 +47,34 @@
 
 
 def _build_source_sink(io_dict: Union[Dict, Database, FileSystem]):
-    if not io_dict:
+    if io_dict is None:
         return io_dict
-    io_dict = io_dict or {}
     if isinstance(io_dict, (Database, FileSystem)):
         component_io = io_dict
     else:
-        data_type = io_dict.get("type", None)
-        if data_type == ExternalDataType.DATABASE:
-            component_io = Database(**io_dict)
-        elif data_type == ExternalDataType.FILE_SYSTEM:
-            component_io = FileSystem(**io_dict)
+        if isinstance(io_dict, dict):
+            data_type = io_dict.get("type", None)
+            if data_type == ExternalDataType.DATABASE:
+                component_io = Database(**io_dict)
+            elif data_type == ExternalDataType.FILE_SYSTEM:
+                component_io = FileSystem(**io_dict)
+            else:
+                msg = "Type in source or sink only support {} and {}, currently got {}."
+                raise ValidationException(
+                    message=msg.format(ExternalDataType.DATABASE, ExternalDataType.FILE_SYSTEM, data_type),
+                    no_personal_data_message=msg.format(
+                        ExternalDataType.DATABASE, ExternalDataType.FILE_SYSTEM, "data_type"
+                    ),
+                    target=ErrorTarget.DATA_TRANSFER_JOB,
+                    error_category=ErrorCategory.USER_ERROR,
+                    error_type=ValidationErrorType.INVALID_VALUE,
+                )
         else:
-            msg = "Source or sink only support type {} and {}, currently got {}."
+            msg = "Source or sink only support dict, Database and FileSystem"
             raise ValidationException(
-                message=msg.format(ExternalDataType.DATABASE, ExternalDataType.FILE_SYSTEM, data_type),
-                no_personal_data_message=msg.format(
-                    ExternalDataType.DATABASE, ExternalDataType.FILE_SYSTEM, "data_type"
-                ),
-                target=ErrorTarget.COMPONENT,
+                message=msg,
+                no_personal_data_message=msg,
+                target=ErrorTarget.DATA_TRANSFER_JOB,
                 error_category=ErrorCategory.USER_ERROR,
                 error_type=ValidationErrorType.INVALID_VALUE,
             )
@@ -356,6 +365,20 @@ def _customized_validate(self):
                 yaml_path="outputs.sink",
                 message="Outputs field only support one output called sink in import task",
             )
+        if "sink" in self.outputs and isinstance(self.outputs["sink"]._data, Output):
+            sink_output = self.outputs["sink"]._data
+            if (self.source.type == ExternalDataType.DATABASE and sink_output.type != AssetTypes.MLTABLE) or (
+                self.source.type == ExternalDataType.FILE_SYSTEM and sink_output.type != AssetTypes.URI_FOLDER
+            ):
+                result.append_error(
+                    yaml_path="outputs.sink.type",
+                    message="Outputs field only support type {} for {} and {} for {}".format(
+                        AssetTypes.MLTABLE,
+                        ExternalDataType.DATABASE,
+                        AssetTypes.URI_FOLDER,
+                        ExternalDataType.FILE_SYSTEM,
+                    ),
+                )
         return result
 
     def _to_rest_object(self, **kwargs) -> dict:
@@ -493,6 +516,21 @@ def _customized_validate(self):
                 yaml_path="inputs.source",
                 message="Inputs field only support one input called source in export task",
             )
+        if "source" in self.inputs and isinstance(self.inputs["source"]._data, Input):
+            source_input = self.inputs["source"]._data
+            if (self.sink.type == ExternalDataType.DATABASE and source_input.type != AssetTypes.URI_FILE) or (
+                self.sink.type == ExternalDataType.FILE_SYSTEM and source_input.type != AssetTypes.URI_FOLDER
+            ):
+                result.append_error(
+                    yaml_path="inputs.source.type",
+                    message="Inputs field only support type {} for {} and {} for {}".format(
+                        AssetTypes.URI_FILE,
+                        ExternalDataType.DATABASE,
+                        AssetTypes.URI_FOLDER,
+                        ExternalDataType.FILE_SYSTEM,
+                    ),
+                )
+
         return result
 
     def _to_rest_object(self, **kwargs) -> dict: