Feature set and Feature store entity archive and restore (Azure#29356)

* Updating feature set and feature store entity operations to use get_entity

* Remove archive and restore for feature set and feature store entity

* fix lint

* Remove tests

* Adding archive and restore for version

* Adding line

Author: maykumar-msft

sdk/ml/azure-ai-ml/azure/ai/ml/_utils/_feature_set_utils.py

@@ -0,0 +1,146 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+
+# pylint: disable=protected-access
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Dict, Optional, Union
+
+
+from azure.ai.ml._restclient.v2023_02_01_preview.operations import (  # pylint: disable = unused-import
+    FeaturesetContainersOperations,
+    FeaturesetVersionsOperations,
+    FeaturestoreEntityContainersOperations,
+    FeaturestoreEntityVersionsOperations,
+)
+from azure.ai.ml.entities._assets.asset import Asset
+from azure.ai.ml.exceptions import (
+    ErrorCategory,
+    ErrorTarget,
+    ValidationErrorType,
+    ValidationException,
+)
+from azure.core.exceptions import ResourceNotFoundError
+
+from .utils import load_yaml
+
+if TYPE_CHECKING:
+    from azure.ai.ml.operations._feature_set_operations import _FeatureSetOperations
+    from azure.ai.ml.operations._feature_store_entity_operations import _FeatureStoreEntityOperations
+
+
+def read_feature_set_metadata_contents(*, path: str) -> Dict:
+    metadata_path = str(Path(path, "FeaturesetSpec.yaml"))
+    return load_yaml(metadata_path)
+
+
+def _get_latest_version_from_container(
+    asset_name: str,
+    container_operation: Union[
+        "FeaturesetContainersOperations",
+        "FeaturestoreEntityContainersOperations",
+    ],
+    resource_group_name: str,
+    workspace_name: str,
+    **kwargs,
+) -> str:
+    try:
+        container = container_operation.get_entity(
+            name=asset_name,
+            resource_group_name=resource_group_name,
+            workspace_name=workspace_name,
+            **kwargs,
+        )
+        version = container.properties.latest_version
+
+    except ResourceNotFoundError:
+        message = f"Asset {asset_name} does not exist in feature store {workspace_name}."
+        no_personal_data_message = "Asset {asset_name} does not exist in feature store {workspace_name}."
+        raise ValidationException(
+            message=message,
+            no_personal_data_message=no_personal_data_message,
+            target=ErrorTarget.ASSET,
+            error_category=ErrorCategory.USER_ERROR,
+            error_type=ValidationErrorType.RESOURCE_NOT_FOUND,
+        )
+    return version
+
+
+def _archive_or_restore(
+    asset_operations: Union["_FeatureSetOperations", "_FeatureStoreEntityOperations"],
+    version_operation: Union[
+        "FeaturesetVersionsOperations",
+        "FeaturestoreEntityVersionsOperations",
+    ],
+    is_archived: bool,
+    name: str,
+    version: Optional[str] = None,
+    label: Optional[str] = None,
+) -> None:
+    resource_group_name = asset_operations._operation_scope._resource_group_name
+    workspace_name = asset_operations._workspace_name
+    if not version and not label:
+        msg = "Must provide either version or label."
+        raise ValidationException(
+            message=msg,
+            target=ErrorTarget.ASSET,
+            no_personal_data_message=msg,
+            error_category=ErrorCategory.USER_ERROR,
+            error_type=ValidationErrorType.MISSING_FIELD,
+        )
+
+    if version and label:
+        msg = "Cannot specify both version and label."
+        raise ValidationException(
+            message=msg,
+            no_personal_data_message=msg,
+            target=ErrorTarget.ASSET,
+            error_category=ErrorCategory.USER_ERROR,
+            error_type=ValidationErrorType.RESOURCE_NOT_FOUND,
+        )
+
+    if label:
+        version = _resolve_label_to_asset(asset_operations, name, label).version
+
+    if version:
+        version_resource = version_operation.get(
+            name=name,
+            version=version,
+            resource_group_name=resource_group_name,
+            workspace_name=workspace_name,
+        )
+        version_resource.properties.is_archived = is_archived
+        version_resource.properties.stage = "Archived" if is_archived else "Development"
+        version_operation.begin_create_or_update(
+            name=name,
+            version=version,
+            resource_group_name=resource_group_name,
+            workspace_name=workspace_name,
+            body=version_resource,
+        )
+
+
+def _resolve_label_to_asset(
+    assetOperations: Union[
+        "_FeatureSetOperations",
+        "_FeatureStoreEntityOperations",
+    ],
+    name: str,
+    label: str,
+) -> Asset:
+    """Returns the asset referred to by the given label.
+
+    Throws if label does not refer to a version of the named asset
+    """
+
+    resolver = assetOperations._managed_label_resolver.get(label, None)
+    if not resolver:
+        msg = "Asset {} with version label {} does not exist in workspace."
+        raise ValidationException(
+            message=msg.format(name, label),
+            no_personal_data_message=msg.format("[name]", "[label]"),
+            target=ErrorTarget.ASSET,
+            error_type=ValidationErrorType.RESOURCE_NOT_FOUND,
+        )
+    return resolver(name)

sdk/ml/azure-ai-ml/azure/ai/ml/_utils/_feature_store_utils.py

@@ -144,6 +144,7 @@
 from ._feature_store_entity.feature_store_entity import _FeatureStoreEntity
 from ._feature_store_entity.data_column import _DataColumn
 from ._feature_store_entity.data_column_type import _DataColumnType
+from ._feature_set.feature import _Feature
 from ._feature_set.feature_set_specification import _FeatureSetSpecification
 from ._feature_set.materialization_compute_resource import _MaterializationComputeResource
 from ._feature_set.materialization_settings import _MaterializationSettings
@@ -287,6 +288,7 @@
     "AutoScaleSettings",
     "AutoPauseSettings",
     "WorkspaceModelReference",
+    "_Feature",
     "_FeatureSet",
     "_ComputeRuntime",
     "_FeatureStoreSettings",

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

@@ -27,12 +27,12 @@
 from azure.ai.ml.operations._datastore_operations import DatastoreOperations
 
 # from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
-from azure.ai.ml._utils._asset_utils import (
+from azure.ai.ml._utils._feature_store_utils import (
     _archive_or_restore,
     _get_latest_version_from_container,
     _resolve_label_to_asset,
+    read_feature_set_metadata_contents,
 )
-from azure.ai.ml._utils._feature_set_utils import read_feature_set_metadata_contents
 from azure.ai.ml._utils._logger_utils import OpsLogger
 from azure.ai.ml.entities._assets._artifacts.feature_set import _FeatureSet
 from azure.ai.ml.entities._feature_set.featureset_spec_metadata import FeaturesetSpecMetadata
@@ -365,7 +365,6 @@ def archive(
         _archive_or_restore(
             asset_operations=self,
             version_operation=self._operation,
-            container_operation=self._container_operation,
             is_archived=True,
             name=name,
             version=version,
@@ -395,7 +394,6 @@ def restore(
         _archive_or_restore(
             asset_operations=self,
             version_operation=self._operation,
-            container_operation=self._container_operation,
             is_archived=False,
             name=name,
             version=version,

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

@@ -16,7 +16,7 @@
 
 
 # from azure.ai.ml._telemetry import ActivityType, monitor_with_activity
-from azure.ai.ml._utils._asset_utils import (
+from azure.ai.ml._utils._feature_store_utils import (
     _archive_or_restore,
     _get_latest_version_from_container,
     _resolve_label_to_asset,
@@ -184,7 +184,6 @@ def archive(
         _archive_or_restore(
             asset_operations=self,
             version_operation=self._operation,
-            container_operation=self._container_operation,
             is_archived=True,
             name=name,
             version=version,
@@ -214,7 +213,6 @@ def restore(
         _archive_or_restore(
             asset_operations=self,
             version_operation=self._operation,
-            container_operation=self._container_operation,
             is_archived=False,
             name=name,
             version=version,

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

@@ -137,52 +137,24 @@ def test_archive_version(self, mock_feature_set_operations: _FeatureSetOperation
         version = "1"
         mock_feature_set_operations._operation.get.return_value = featureset_version
         mock_feature_set_operations.archive(name=name, version=version)
-        mock_feature_set_operations._operation.create_or_update.assert_called_once_with(
+        mock_feature_set_operations._operation.begin_create_or_update.assert_called_once_with(
             name=name,
             version=version,
             workspace_name=mock_feature_set_operations._workspace_name,
             body=featureset_version,
             resource_group_name=mock_feature_set_operations._resource_group_name,
         )
 
-    def test_archive_container(self, mock_feature_set_operations: _FeatureSetOperations):
-        name = "random_name"
-        featureset_container = Mock(
-            FeaturesetContainer(properties=Mock(FeaturesetContainerProperties(description="test")))
-        )
-        mock_feature_set_operations._container_operation.get.return_value = featureset_container
-        mock_feature_set_operations.archive(name=name)
-        mock_feature_set_operations._container_operation.create_or_update.assert_called_once_with(
-            name=name,
-            workspace_name=mock_feature_set_operations._workspace_name,
-            body=featureset_container,
-            resource_group_name=mock_feature_set_operations._resource_group_name,
-        )
-
     def test_restore_version(self, mock_feature_set_operations: _FeatureSetOperations):
         name = "random_name"
         featureset_version = Mock(FeaturesetVersion(properties=Mock(FeaturesetVersionProperties(entities=["test"]))))
         version = "1"
         mock_feature_set_operations._operation.get.return_value = featureset_version
         mock_feature_set_operations.restore(name=name, version=version)
-        mock_feature_set_operations._operation.create_or_update.assert_called_once_with(
+        mock_feature_set_operations._operation.begin_create_or_update.assert_called_once_with(
             name=name,
             version=version,
             workspace_name=mock_feature_set_operations._workspace_name,
             body=featureset_version,
             resource_group_name=mock_feature_set_operations._resource_group_name,
         )
-
-    def test_restore_container(self, mock_feature_set_operations: _FeatureSetOperations):
-        name = "random_name"
-        featureset_container = Mock(
-            FeaturesetContainer(properties=Mock(FeaturesetContainerProperties(entities=["test"])))
-        )
-        mock_feature_set_operations._container_operation.get.return_value = featureset_container
-        mock_feature_set_operations.restore(name=name)
-        mock_feature_set_operations._container_operation.create_or_update.assert_called_once_with(
-            name=name,
-            workspace_name=mock_feature_set_operations._workspace_name,
-            body=featureset_container,
-            resource_group_name=mock_feature_set_operations._resource_group_name,
-        )

sdk/ml/azure-ai-ml/tests/feature_set/unittests/test_feature_set_operations.py

@@ -75,28 +75,14 @@ def test_archive_version(self, mock_feature_store_entity_operations: _FeatureSto
         version = "1"
         mock_feature_store_entity_operations._operation.get.return_value = featureStoreEntity_version
         mock_feature_store_entity_operations.archive(name=name, version=version)
-        mock_feature_store_entity_operations._operation.create_or_update.assert_called_once_with(
+        mock_feature_store_entity_operations._operation.begin_create_or_update.assert_called_once_with(
             name=name,
             version=version,
             workspace_name=mock_feature_store_entity_operations._workspace_name,
             body=featureStoreEntity_version,
             resource_group_name=mock_feature_store_entity_operations._resource_group_name,
         )
 
-    def test_archive_container(self, mock_feature_store_entity_operations: _FeatureStoreEntityOperations):
-        name = "random_name"
-        featureStoreEntity_container = Mock(
-            FeaturestoreEntityContainer(properties=Mock(FeaturestoreEntityContainerProperties(description="test")))
-        )
-        mock_feature_store_entity_operations._container_operation.get.return_value = featureStoreEntity_container
-        mock_feature_store_entity_operations.archive(name=name)
-        mock_feature_store_entity_operations._container_operation.create_or_update.assert_called_once_with(
-            name=name,
-            workspace_name=mock_feature_store_entity_operations._workspace_name,
-            body=featureStoreEntity_container,
-            resource_group_name=mock_feature_store_entity_operations._resource_group_name,
-        )
-
     def test_restore_version(self, mock_feature_store_entity_operations: _FeatureStoreEntityOperations):
         name = "random_name"
         featureStoreEntity_version = Mock(
@@ -105,24 +91,10 @@ def test_restore_version(self, mock_feature_store_entity_operations: _FeatureSto
         version = "1"
         mock_feature_store_entity_operations._operation.get.return_value = featureStoreEntity_version
         mock_feature_store_entity_operations.restore(name=name, version=version)
-        mock_feature_store_entity_operations._operation.create_or_update.assert_called_once_with(
+        mock_feature_store_entity_operations._operation.begin_create_or_update.assert_called_once_with(
             name=name,
             version=version,
             workspace_name=mock_feature_store_entity_operations._workspace_name,
             body=featureStoreEntity_version,
             resource_group_name=mock_feature_store_entity_operations._resource_group_name,
         )
-
-    def test_restore_container(self, mock_feature_store_entity_operations: _FeatureStoreEntityOperations):
-        name = "random_name"
-        featureStoreEntity_container = Mock(
-            FeaturestoreEntityContainer(properties=Mock(FeaturestoreEntityContainerProperties()))
-        )
-        mock_feature_store_entity_operations._container_operation.get.return_value = featureStoreEntity_container
-        mock_feature_store_entity_operations.restore(name=name)
-        mock_feature_store_entity_operations._container_operation.create_or_update.assert_called_once_with(
-            name=name,
-            workspace_name=mock_feature_store_entity_operations._workspace_name,
-            body=featureStoreEntity_container,
-            resource_group_name=mock_feature_store_entity_operations._resource_group_name,
-        )