Renyix/data import sdkv2 (Azure#29014)

* Add azureml data import verb sdk support

* Create data import pipeline job

* Update doc string

* data import - add unit tests

* data import - skip serverless compute validation

* data import - fix unit tests

* data import - fix DataImportSchema import build error

* data import - remove used JobOperations to fix pylint build check

* add __init__ to fix build checks

* data import - fix line too long build error

* data import - generate recording files

* Fix Run Black build check error

* Refine data import job submission

* Add back some data import parameters

* Mark data import classes as experimental

* Resolve code review comments

* Resolve code review comments

* Remove unused import to fix pyint

* data import - update e2e test recording file

* Add back @experimental for DataImport class and disable asset_name filter

* Update e2e test recording file

* Fix unused 'name' pyint build check

Author: xiongrenyi

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

@@ -22,6 +22,7 @@
     load_component,
     load_compute,
     load_data,
+    load_data_import,
     load_datastore,
     load_environment,
     load_job,
@@ -54,6 +55,7 @@
     "load_component",
     "load_compute",
     "load_data",
+    "load_data_import",
     "load_datastore",
     "load_model",
     "load_environment",

sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py

@@ -472,6 +472,7 @@ def __init__(
             requests_pipeline=self._requests_pipeline,
             **ops_kwargs,
         )
+        self._data._job_operation = self._jobs
         self._operation_container.add(AzureMLResourceType.JOB, self._jobs)
         self._schedules = ScheduleOperations(
             self._operation_scope,

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

@@ -5,6 +5,7 @@
 __path__ = __import__("pkgutil").extend_path(__path__, __name__)
 
 from ._sweep import SweepJobSchema
+from ._data_import import DataImportSchema
 from .assets.code_asset import AnonymousCodeAssetSchema, CodeAssetSchema
 from .assets.data import DataSchema
 from .assets.environment import AnonymousEnvironmentSchema, EnvironmentSchema
@@ -53,4 +54,5 @@
     "ExperimentalField",
     "RegistryStr",
     "WorkspaceAssetReferenceSchema",
+    "DataImportSchema",
 ]

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_data_import/__init__.py

@@ -0,0 +1,9 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)
+
+from .data_import import DataImportSchema
+
+__all__ = ["DataImportSchema"]

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_data_import/data_import.py

@@ -0,0 +1,16 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# pylint: disable=unused-argument,no-self-use
+
+from azure.ai.ml._schema.core.fields import NestedField
+from azure.ai.ml._schema.job.input_output_entry import DatabaseSchema, FileSystemSchema
+from azure.ai.ml._utils._experimental import experimental
+from ..core.fields import UnionField
+from ..assets.data import DataSchema
+
+
+@experimental
+class DataImportSchema(DataSchema):
+    source = UnionField([NestedField(DatabaseSchema), NestedField(FileSystemSchema)], required=True, allow_none=False)

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

@@ -51,6 +51,7 @@
 )
 from ._datastore.adls_gen1 import AzureDataLakeGen1Datastore
 from ._datastore.azure_storage import AzureBlobDatastore, AzureDataLakeGen2Datastore, AzureFileDatastore
+from ._data_import.data_import import DataImport
 from ._datastore.datastore import Datastore
 from ._deployment.batch_deployment import BatchDeployment
 from ._deployment.batch_job import BatchJob
@@ -209,6 +210,7 @@
     "PrivateEndpointDestination",
     "EndpointConnection",
     "CustomerManagedKey",
+    "DataImport",
     "Datastore",
     "AzureDataLakeGen1Datastore",
     "AzureBlobDatastore",

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_data_import/__init__.py

@@ -0,0 +1,5 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_data_import/data_import.py

@@ -0,0 +1,83 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# pylint: disable=no-member
+
+from os import PathLike
+from pathlib import Path
+from typing import Dict, Optional, Union
+from azure.ai.ml._schema import DataImportSchema
+from azure.ai.ml._utils._experimental import experimental
+from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, PARAMS_OVERRIDE_KEY
+from azure.ai.ml.entities._assets import Data
+from azure.ai.ml.entities._inputs_outputs.external_data import Database, FileSystem
+from azure.ai.ml.entities._util import load_from_dict
+
+
+@experimental
+class DataImport(Data):
+    """Data asset with a creating data import job.
+
+    :param name: Name of the asset.
+    :type name: str
+    :param path: The path to the asset being created by data import job.
+    :type path: str
+    :param source: The source of the asset data being copied from.
+    :type source: Union[Database, FileSystem]
+    :param version: Version of the resource.
+    :type version: str
+    :param description: Description of the resource.
+    :type description: str
+    :param tags: Tag dictionary. Tags can be added, removed, and updated.
+    :type tags: dict[str, str]
+    :param properties: The asset property dictionary.
+    :type properties: dict[str, str]
+    :param kwargs: A dictionary of additional configuration parameters.
+    :type kwargs: dict
+    """
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        path: str,
+        source: Union[Database, FileSystem],
+        version: Optional[str] = None,
+        description: Optional[str] = None,
+        tags: Optional[Dict] = None,
+        properties: Optional[Dict] = None,
+        **kwargs,
+    ):
+        super().__init__(
+            name=name,
+            version=version,
+            description=description,
+            tags=tags,
+            properties=properties,
+            path=path,
+            **kwargs,
+        )
+        self.source = source
+
+    @classmethod
+    def _load(
+        cls,
+        data: Optional[Dict] = None,
+        yaml_path: Optional[Union[PathLike, str]] = None,
+        params_override: Optional[list] = None,
+        **kwargs,
+    ) -> "DataImport":
+        data = data or {}
+        params_override = params_override or []
+        context = {
+            BASE_PATH_CONTEXT_KEY: Path(yaml_path).parent if yaml_path else Path("./"),
+            PARAMS_OVERRIDE_KEY: params_override,
+        }
+        data_import = DataImport._load_from_dict(yaml_data=data, context=context, **kwargs)
+
+        return data_import
+
+    @classmethod
+    def _load_from_dict(cls, yaml_data: Dict, context: Dict, **kwargs) -> "DataImport":
+        return DataImport(**load_from_dict(DataImportSchema, yaml_data, context, **kwargs))

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_load_functions.py

@@ -19,6 +19,7 @@
 from azure.ai.ml.entities._component.parallel_component import ParallelComponent
 from azure.ai.ml.entities._component.pipeline_component import PipelineComponent
 from azure.ai.ml.entities._compute.compute import Compute
+from azure.ai.ml.entities._data_import.data_import import DataImport
 from azure.ai.ml.entities._datastore.datastore import Datastore
 from azure.ai.ml.entities._deployment.batch_deployment import BatchDeployment
 from azure.ai.ml.entities._deployment.online_deployment import OnlineDeployment
@@ -452,6 +453,37 @@ def load_data(
     return load_common(Data, source, relative_origin, **kwargs)
 
 
+def load_data_import(
+    source: Union[str, PathLike, IO[AnyStr]],
+    *,
+    relative_origin: Optional[str] = None,
+    **kwargs,
+) -> DataImport:
+    """Construct a data object from yaml file.
+
+    :param source: The local yaml source of a data object. Must be either a
+        path to a local file, or an already-open file.
+        If the source is a path, it will be open and read.
+        An exception is raised if the file does not exist.
+        If the source is an open file, the file will be read directly,
+        and an exception is raised if the file is not readable.
+    :type source: Union[PathLike, str, io.TextIOWrapper]
+    :param relative_origin: The origin to be used when deducing
+        the relative locations of files referenced in the parsed yaml.
+        Defaults to the inputted source's directory if it is a file or file path input.
+        Defaults to "./" if the source is a stream input with no name value.
+    :type relative_origin: str
+    :param params_override: Fields to overwrite on top of the yaml file.
+        Format is [{"field1": "value1"}, {"field2": "value2"}]
+    :type params_override: List[Dict]
+    :raises ~azure.ai.ml.exceptions.ValidationException: Raised if DataImport cannot be successfully validated.
+        Details will be provided in the error message.
+    :return: Constructed data_import object.
+    :rtype: DataImport
+    """
+    return load_common(DataImport, source, relative_origin, **kwargs)
+
+
 def load_environment(
     source: Union[str, PathLike, IO[AnyStr]],
     *,

sdk/ml/azure-ai-ml/azure/ai/ml/operations/_data_operations.py

@@ -6,10 +6,15 @@
 
 import os
 from pathlib import Path
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union, Iterable
 
 from marshmallow.exceptions import ValidationError as SchemaValidationError
 
+from azure.ai.ml._utils._experimental import experimental
+from azure.ai.ml.entities import Job, PipelineJob, PipelineJobSettings
+from azure.ai.ml.data_transfer import import_data as import_data_func
+from azure.ai.ml.entities._inputs_outputs import Output
+from azure.ai.ml.entities._inputs_outputs.external_data import Database
 from azure.ai.ml._artifacts._artifact_utilities import _check_and_upload_path
 from azure.ai.ml._artifacts._constants import (
     ASSET_PATH_ERROR,
@@ -52,6 +57,7 @@
     AzureMLResourceType,
 )
 from azure.ai.ml.entities._assets import Data, WorkspaceAssetReference
+from azure.ai.ml.entities._data_import.data_import import DataImport
 from azure.ai.ml.entities._data.mltable_metadata import MLTableMetadata
 from azure.ai.ml.exceptions import (
     AssetPathException,
@@ -347,6 +353,64 @@ def create_or_update(self, data: Data) -> Data:
                     )
             raise ex
 
+    @experimental
+    def import_data(self, data_import: DataImport) -> Job:
+        """Returns the data import job that is creating the data asset.
+
+        :param data_import: DataImport object.
+        :type data_import: azure.ai.ml.entities.DataImport
+        :return: data import job object.
+        :rtype: ~azure.ai.ml.entities.Job
+        """
+
+        experiment_name = "data_import_" + data_import.name
+        data_import.type = AssetTypes.MLTABLE if isinstance(data_import.source, Database) else AssetTypes.URI_FOLDER
+        if "{name}" not in data_import.path:
+            data_import.path = data_import.path.rstrip("/") + "/{name}"
+        import_job = import_data_func(
+            description=data_import.description or experiment_name,
+            display_name=experiment_name,
+            experiment_name=experiment_name,
+            compute="serverless",
+            source=data_import.source,
+            outputs={
+                "sink": Output(
+                    type=data_import.type, path=data_import.path, name=data_import.name, version=data_import.version
+                )
+            },
+        )
+        import_pipeline = PipelineJob(
+            description=data_import.description or experiment_name,
+            tags=data_import.tags,
+            display_name=experiment_name,
+            experiment_name=experiment_name,
+            properties=data_import.properties or {},
+            settings=PipelineJobSettings(force_rerun=True),
+            jobs={experiment_name: import_job},
+        )
+        import_pipeline.properties["azureml.materializationAssetName"] = data_import.name
+        return self._job_operation.create_or_update(job=import_pipeline, skip_validation=True)
+
+    @experimental
+    def show_materialization_status(
+        self,
+        name: str,
+        *,
+        list_view_type: ListViewType = ListViewType.ACTIVE_ONLY,
+    ) -> Iterable[Job]:
+        """List materialization jobs of the asset.
+
+        :param name: name of asset being created by the materialization jobs.
+        :type name: str
+        :param list_view_type: View type for including/excluding (for example) archived jobs. Default: ACTIVE_ONLY.
+        :type list_view_type: Optional[ListViewType]
+        :return: An iterator like instance of Job objects.
+        :rtype: ~azure.core.paging.ItemPaged[Job]
+        """
+
+        # TODO: Add back 'asset_name=name' filter once client switches to mfe 2023-02-01-preview and above
+        return self._job_operation.list(job_type="Pipeline", tag=name, list_view_type=list_view_type)
+
     # @monitor_with_activity(logger, "Data.Validate", ActivityType.INTERNALCALL)
     def _validate(self, data: Data) -> Union[List[str], None]:
         if not data.path: