Add support for promoting data asset from a workspace to a registry (Azure#28866)

* Support list/show and create data asset in a registry
---------

Co-authored-by: Valerie Pham <[email protected]>

Author: vmpham1012

sdk/ml/azure-ai-ml/CHANGELOG.md

@@ -4,6 +4,7 @@
 
 ### Features Added
 - Added support for `tags` on Compute Resources.
+- Added support for promoting data asset from a workspace to a registry
 
 ### Bugs Fixed
 

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

@@ -102,6 +102,7 @@
 )
 ASSET_ID_RESOURCE_REGEX_FORMAT = "azureml://resource[gG]roups/([^/]+)/workspaces/([^/]+)/([^/]+)/([^/]+)/versions/(.+)"
 MODEL_ID_REGEX_FORMAT = "azureml://models/([^/]+)/versions/(.+)"
+DATA_ID_REGEX_FORMAT = "azureml://data/([^/]+)/versions/(.+)"
 ASSET_ID_URI_REGEX_FORMAT = "azureml://locations/([^/]+)/workspaces/([^/]+)/([^/]+)/([^/]+)/versions/(.+)"
 AZUREML_CLI_SYSTEM_EXECUTED_ENV_VAR = "AZUREML_CLI_SYSTEM_EXECUTED"
 DOCSTRING_TEMPLATE = ".. note::    {0} {1}\n\n"

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

@@ -45,8 +45,13 @@
     get_sas_uri_for_registry_asset,
 )
 from azure.ai.ml._utils.utils import is_url
-from azure.ai.ml.constants._common import MLTABLE_METADATA_SCHEMA_URL_FALLBACK, AssetTypes
-from azure.ai.ml.entities._assets import Data
+from azure.ai.ml.constants._common import (
+    MLTABLE_METADATA_SCHEMA_URL_FALLBACK,
+    AssetTypes,
+    ASSET_ID_FORMAT,
+    AzureMLResourceType,
+)
+from azure.ai.ml.entities._assets import Data, WorkspaceAssetReference
 from azure.ai.ml.entities._data.mltable_metadata import MLTableMetadata
 from azure.ai.ml.exceptions import (
     AssetPathException,
@@ -58,6 +63,7 @@
 from azure.ai.ml.operations._datastore_operations import DatastoreOperations
 from azure.core.exceptions import HttpResponseError
 from azure.core.paging import ItemPaged
+from azure.core.exceptions import ResourceNotFoundError
 
 ops_logger = OpsLogger(__name__)
 module_logger = ops_logger.module_logger
@@ -230,7 +236,6 @@ def create_or_update(self, data: Data) -> Data:
         :return: Data asset object.
         :rtype: ~azure.ai.ml.entities.Data
         """
-
         try:
             name = data.name
             if not data.version and self._registry_name:
@@ -246,6 +251,34 @@ def create_or_update(self, data: Data) -> Data:
 
             sas_uri = None
             if self._registry_name:
+                # If the data asset is a workspace asset, promote to registry
+                if isinstance(data, WorkspaceAssetReference):
+                    try:
+                        self._operation.get(
+                            name=data.name,
+                            version=data.version,
+                            resource_group_name=self._resource_group_name,
+                            registry_name=self._registry_name,
+                        )
+                    except Exception as err:  # pylint: disable=broad-except
+                        if isinstance(err, ResourceNotFoundError):
+                            pass
+                        else:
+                            raise err
+                    else:
+                        msg = "An data asset with this name and version already exists in registry"
+                        raise ValidationException(
+                            message=msg,
+                            no_personal_data_message=msg,
+                            target=ErrorTarget.DATA,
+                            error_category=ErrorCategory.USER_ERROR,
+                        )
+                    data = data._to_rest_object()
+                    result = self._service_client.resource_management_asset_reference.begin_import_method(
+                        resource_group_name=self._resource_group_name, registry_name=self._registry_name, body=data
+                    )
+                    return result
+
                 sas_uri = get_sas_uri_for_registry_asset(
                     service_client=self._service_client,
                     name=name,
@@ -445,6 +478,43 @@ def _get_latest_version(self, name: str) -> Data:
         )
         return self.get(name, version=latest_version)
 
+    # pylint: disable=no-self-use
+    def _prepare_to_copy(
+        self, data: Data, name: Optional[str] = None, version: Optional[str] = None
+    ) -> WorkspaceAssetReference:
+
+        """Returns WorkspaceAssetReference
+        to copy a registered data to registry given the asset id
+
+        :param data: Registered data
+        :type data: Data
+        :param name: Destination name
+        :type name: str
+        :param version: Destination version
+        :type version: str
+        """
+        #  Get workspace info to get workspace GUID
+        workspace = self._service_client.workspaces.get(
+            resource_group_name=self._resource_group_name, workspace_name=self._workspace_name
+        )
+        workspace_guid = workspace.workspace_id
+        workspace_location = workspace.location
+
+        # Get data asset ID
+        asset_id = ASSET_ID_FORMAT.format(
+            workspace_location,
+            workspace_guid,
+            AzureMLResourceType.DATA,
+            data.name,
+            data.version,
+        )
+
+        return WorkspaceAssetReference(
+            name=name if name else data.name,
+            version=version if version else data.version,
+            asset_id=asset_id,
+        )
+
 
 def _assert_local_path_matches_asset_type(
     local_path: str,

sdk/ml/azure-ai-ml/tests/dataset/unittests/test_data_operations.py

@@ -24,6 +24,7 @@
 from azure.ai.ml.exceptions import ErrorTarget
 from azure.ai.ml.operations import DataOperations, DatastoreOperations
 from azure.core.paging import ItemPaged
+from azure.core.exceptions import ResourceNotFoundError
 from unittest.mock import ANY
 
 
@@ -562,3 +563,40 @@ def test_create_with_datastore(
                 show_progress=True,
                 ignore_file=None,
             )
+
+    def test_promote_data_from_workspace(
+        self, mock_data_operations_in_registry: DataOperations, mock_data_operations: DataOperations, tmp_path: Path
+    ) -> None:
+        data_asset_name = f"data_random_string"
+        p = tmp_path / "data_full.yml"
+        data_path = tmp_path / "data.pkl"
+        data_path.write_text("hello world")
+        p.write_text(
+            f"""
+    name: {data_asset_name}
+    path: ./data.pkl
+    version: 3"""
+        )
+
+        with patch(
+            "azure.ai.ml._artifacts._artifact_utilities._upload_to_datastore",
+            return_value=ArtifactStorageInfo(
+                name=data_asset_name,
+                version="3",
+                relative_path="path",
+                datastore_arm_id="/subscriptions/mock/resourceGroups/mock/providers/Microsoft.MachineLearningServices/workspaces/mock/datastores/datastore_id",
+                container_name="containerName",
+            ),
+        ) as mock_upload, patch(
+            "azure.ai.ml.operations._data_operations.Data._from_rest_object",
+            return_value=Data(),
+        ):
+            data = load_data(source=p)
+            data_to_promote = mock_data_operations._prepare_to_copy(data, "new_name", "new_version")
+            assert data_to_promote.name == "new_name"
+            assert data_to_promote.version == "new_version"
+            mock_data_operations_in_registry._operation.get.side_effect = Mock(
+                side_effect=ResourceNotFoundError("Test")
+            )
+            mock_data_operations_in_registry.create_or_update(data_to_promote)
+            mock_data_operations_in_registry._service_client.resource_management_asset_reference.begin_import_method.assert_called_once()