[ML-Pipelines] Fix data transfer review feedback (Azure#29384)

* review feedback: update external_data

* remove unnecessary task type

* expose builder classes

* expose component classes

* fix pylint

* fix tests

* fix tests

* update

* fix remaining tests

* update

* fix black

---------

Co-authored-by: Clement Wang <[email protected]>

Author: wangchao1230

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

@@ -205,6 +205,7 @@ class DatabaseSchema(metaclass=PatchedSchemaMeta):
     def make(self, data, **kwargs):
         from azure.ai.ml.data_transfer import Database
 
+        data.pop("type", None)
         return Database(**data)
 
     @pre_dump
@@ -231,6 +232,7 @@ class FileSystemSchema(metaclass=PatchedSchemaMeta):
     def make(self, data, **kwargs):
         from azure.ai.ml.data_transfer import FileSystem
 
+        data.pop("type", None)
         return FileSystem(**data)
 
     @pre_dump

sdk/ml/azure-ai-ml/azure/ai/ml/data_transfer/__init__.py

@@ -2,6 +2,32 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 from azure.ai.ml.entities._inputs_outputs.external_data import Database, FileSystem
-from azure.ai.ml.entities._builders.data_transfer_func import copy_data, import_data, export_data
+from azure.ai.ml.entities._builders.data_transfer_func import (
+    copy_data,
+    import_data,
+    export_data,
+)
+from azure.ai.ml.entities._builders.data_transfer import (
+    DataTransferCopy,
+    DataTransferImport,
+    DataTransferExport,
+)
+from azure.ai.ml.entities._component.datatransfer_component import (
+    DataTransferCopyComponent,
+    DataTransferImportComponent,
+    DataTransferExportComponent,
+)
 
-__all__ = ["import_data", "export_data", "copy_data", "Database", "FileSystem"]
+__all__ = [
+    "import_data",
+    "export_data",
+    "copy_data",
+    "Database",
+    "FileSystem",
+    "DataTransferCopy",
+    "DataTransferImport",
+    "DataTransferExport",
+    "DataTransferCopyComponent",
+    "DataTransferImportComponent",
+    "DataTransferExportComponent",
+]

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

@@ -15,7 +15,11 @@
     DataTransferImportJobSchema,
     DataTransferExportJobSchema,
 )
-from azure.ai.ml.constants._component import NodeType, ExternalDataType, DataTransferTaskType
+from azure.ai.ml.constants._component import (
+    NodeType,
+    ExternalDataType,
+    DataTransferTaskType,
+)
 from azure.ai.ml.entities._component.datatransfer_component import (
     DataTransferCopyComponent,
     DataTransferImportComponent,
@@ -30,7 +34,12 @@
     DataTransferExportJob,
 )
 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.exceptions import (
+    ErrorCategory,
+    ErrorTarget,
+    ValidationErrorType,
+    ValidationException,
+)
 from azure.ai.ml.entities._inputs_outputs.external_data import Database, FileSystem
 
 
@@ -53,17 +62,23 @@ def _build_source_sink(io_dict: Union[Dict, Database, FileSystem]):
         component_io = io_dict
     else:
         if isinstance(io_dict, dict):
-            data_type = io_dict.get("type", None)
+            data_type = io_dict.pop("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),
+                    message=msg.format(
+                        ExternalDataType.DATABASE,
+                        ExternalDataType.FILE_SYSTEM,
+                        data_type,
+                    ),
                     no_personal_data_message=msg.format(
-                        ExternalDataType.DATABASE, ExternalDataType.FILE_SYSTEM, "data_type"
+                        ExternalDataType.DATABASE,
+                        ExternalDataType.FILE_SYSTEM,
+                        "data_type",
                     ),
                     target=ErrorTarget.DATA_TRANSFER_JOB,
                     error_category=ErrorCategory.USER_ERROR,
@@ -157,8 +172,6 @@ class DataTransferCopy(DataTransfer):
     :type experiment_name: str
     :param compute: The compute target the job runs on.
     :type compute: str
-    :param task: task type in data transfer component, possible value is "copy_data".
-    :type task: str
     :param data_copy_mode: data copy mode in copy task, possible value is "merge_with_overwrite", "fail_if_conflict".
     :type data_copy_mode: str
     :raises ~azure.ai.ml.exceptions.ValidationException: Raised if DataTransferCopy cannot be successfully validated.
@@ -173,7 +186,6 @@ def __init__(
         compute: Optional[str] = None,
         inputs: Optional[Dict[str, Union[NodeOutput, Input, str]]] = None,
         outputs: Optional[Dict[str, Union[str, Output]]] = None,
-        task: Optional[str] = DataTransferTaskType.COPY_DATA,
         data_copy_mode: Optional[str] = None,
         **kwargs,
     ):
@@ -188,7 +200,7 @@ def __init__(
         )
         # init mark for _AttrDict
         self._init = True
-        self.task = task
+        self.task = DataTransferTaskType.COPY_DATA
         self.data_copy_mode = data_copy_mode
         is_component = isinstance(component, DataTransferCopyComponent)
         if is_component:
@@ -214,7 +226,10 @@ def _picked_fields_from_dict_to_rest_object(cls) -> List[str]:
 
     def _to_rest_object(self, **kwargs) -> dict:
         rest_obj = super()._to_rest_object(**kwargs)
-        for key, value in {"componentId": self._get_component_id(), "data_copy_mode": self.data_copy_mode}.items():
+        for key, value in {
+            "componentId": self._get_component_id(),
+            "data_copy_mode": self.data_copy_mode,
+        }.items():
             if value is not None:
                 rest_obj[key] = value
         return convert_ordered_dict_to_dict(rest_obj)
@@ -240,7 +255,6 @@ def _to_job(self) -> DataTransferCopyJob:
             outputs=self._job_outputs,
             services=self.services,
             compute=self.compute,
-            task=self.task,
             data_copy_mode=self.data_copy_mode,
         )
 
@@ -301,8 +315,6 @@ class DataTransferImport(DataTransfer):
     :type experiment_name: str
     :param compute: The compute target the job runs on.
     :type compute: str
-    :param task: task type in data transfer component, possible value is "import_data".
-    :type task: str
     :raises ~azure.ai.ml.exceptions.ValidationException: Raised if DataTransferImport cannot be successfully validated.
         Details will be provided in the error message.
     """
@@ -315,7 +327,6 @@ def __init__(
         compute: Optional[str] = None,
         source: Optional[Union[Dict, Database, FileSystem]] = None,
         outputs: Optional[Dict[str, Union[str, Output]]] = None,
-        task: Optional[str] = DataTransferTaskType.IMPORT_DATA,
         **kwargs,
     ):
         # validate init params are valid type
@@ -328,7 +339,7 @@ def __init__(
         )
         # init mark for _AttrDict
         self._init = True
-        self.task = task
+        self.task = DataTransferTaskType.IMPORT_DATA
         is_component = isinstance(component, DataTransferImportComponent)
         if is_component:
             self.task = component.task or self.task
@@ -409,7 +420,6 @@ def _to_job(self) -> DataTransferImportJob:
             outputs=self._job_outputs,
             services=self.services,
             compute=self.compute,
-            task=self.task,
         )
 
 
@@ -438,8 +448,6 @@ class DataTransferExport(DataTransfer):
     :type experiment_name: str
     :param compute: The compute target the job runs on.
     :type compute: str
-    :param task: task type in data transfer component, possible value is "export_data".
-    :type task: str
     :raises ~azure.ai.ml.exceptions.ValidationException: Raised if DataTransferExport cannot be successfully validated.
         Details will be provided in the error message.
     """
@@ -452,7 +460,6 @@ def __init__(
         compute: Optional[str] = None,
         sink: Optional[Union[Dict, Database, FileSystem]] = None,
         inputs: Optional[Dict[str, Union[NodeOutput, Input, str]]] = None,
-        task: Optional[str] = DataTransferTaskType.EXPORT_DATA,
         **kwargs,
     ):
         # validate init params are valid type
@@ -465,7 +472,7 @@ def __init__(
         )
         # init mark for _AttrDict
         self._init = True
-        self.task = task
+        self.task = DataTransferTaskType.EXPORT_DATA
         is_component = isinstance(component, DataTransferExportComponent)
         if is_component:
             self.task = component.task or self.task
@@ -560,5 +567,4 @@ def _to_job(self) -> DataTransferExportJob:
             inputs=self._job_inputs,
             services=self.services,
             compute=self.compute,
-            task=self.task,
         )

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

@@ -13,15 +13,21 @@
     ComponentSource,
     ExternalDataType,
     DataTransferBuiltinComponentUri,
-    DataTransferTaskType,
 )
 from azure.ai.ml.entities._inputs_outputs.external_data import Database, FileSystem
 from azure.ai.ml.entities._inputs_outputs import Output, Input
 from azure.ai.ml.entities._job.pipeline._io import PipelineInput, NodeOutput
 from azure.ai.ml.entities._builders.base_node import pipeline_node_decorator
-from azure.ai.ml.entities._job.pipeline._component_translatable import ComponentTranslatableMixin
+from azure.ai.ml.entities._job.pipeline._component_translatable import (
+    ComponentTranslatableMixin,
+)
 from azure.ai.ml.exceptions import ErrorTarget, ValidationErrorType, ValidationException
-from .data_transfer import DataTransferCopy, DataTransferImport, DataTransferExport, _build_source_sink
+from .data_transfer import (
+    DataTransferCopy,
+    DataTransferImport,
+    DataTransferExport,
+    _build_source_sink,
+)
 
 
 SUPPORTED_INPUTS = [
@@ -122,7 +128,6 @@ def copy_data(
     inputs: Optional[Dict] = None,
     outputs: Optional[Dict] = None,
     is_deterministic: bool = True,
-    task: Optional[str] = DataTransferTaskType.COPY_DATA,
     data_copy_mode: Optional[str] = None,
     **kwargs,
 ) -> DataTransferCopy:
@@ -150,8 +155,6 @@ def copy_data(
         In this case, this step will not use any compute resource.
         Default to be True, specify is_deterministic=False if you would like to avoid such reuse behavior.
     :type is_deterministic: bool
-    :param task: task type in data transfer component, possible value is "copy_data".
-    :type task: str
     :param data_copy_mode: data copy mode in copy task, possible value is "merge_with_overwrite", "fail_if_conflict".
     :type data_copy_mode: str
     """
@@ -170,7 +173,6 @@ def copy_data(
             description=description,
             inputs=component_inputs,
             outputs=component_outputs,
-            task=task,
             data_copy_mode=data_copy_mode,
             _source=ComponentSource.BUILDER,
             is_deterministic=is_deterministic,
@@ -186,7 +188,6 @@ def copy_data(
         compute=compute,
         inputs=job_inputs,
         outputs=job_outputs,
-        task=task,
         data_copy_mode=data_copy_mode,
         **kwargs,
     )
@@ -204,7 +205,6 @@ def import_data(
     compute: Optional[str] = None,
     source: Optional[Union[Dict, Database, FileSystem]] = None,
     outputs: Optional[Dict] = None,
-    task: Optional[str] = DataTransferTaskType.IMPORT_DATA,
     **kwargs,
 ) -> DataTransferImport:
     """Create a DataTransferImport object which can be used inside dsl.pipeline.
@@ -226,8 +226,6 @@ def import_data(
     :param outputs: Mapping of outputs data bindings used in the job, default will be an output port with key "sink"
     and type "mltable".
     :type outputs: dict
-    :param task: task type in data transfer component, possible value is "copy_data".
-    :type task: str
     """
     source = _build_source_sink(source)
     outputs = outputs or {"sink": Output(type=AssetTypes.MLTABLE)}
@@ -253,7 +251,6 @@ def import_data(
         compute=compute,
         source=source,
         outputs=job_outputs,
-        task=task,
         **kwargs,
     )
     if update_source:
@@ -273,7 +270,6 @@ def export_data(
     compute: Optional[str] = None,
     sink: Optional[Union[Dict, Database, FileSystem]] = None,
     inputs: Optional[Dict] = None,
-    task: Optional[str] = DataTransferTaskType.EXPORT_DATA,
     **kwargs,
 ) -> DataTransferExport:
     """Create a DataTransferExport object which can be used inside dsl.pipeline.
@@ -294,8 +290,6 @@ def export_data(
     :type sink: Union[Dict, Database, FileSystem]
     :param inputs: Mapping of inputs data bindings used in the job.
     :type inputs: dict
-    :param task: task type in data transfer component, possible value is "copy_data".
-    :type task: str
     """
     sink = _build_source_sink(sink)
     _, job_inputs = _parse_inputs_outputs(inputs, parse_func=_parse_input)
@@ -326,7 +320,6 @@ def export_data(
         compute=compute,
         sink=sink,
         inputs=job_inputs,
-        task=task,
         **kwargs,
     )
     if update_source: