[WorkspaceHub] Update enable isolation flag and purge flag support fo…

…r workspaceHub (#31310)

* Update enable isolation flag and purge flag

* Update change log

* reformat code

reformat code

* Address comments and fix tests

run black

revert some changes

* Update docstring

* update docstring

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

@@ -18,6 +18,7 @@
 
 ### Features Added
 - Added support to enable gpu access (local_enable_gpu) for local deployment.
+- Added support for workspaceHub and Project workspace
 
 ### Other Changes
 

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_workspace_hub/workspace_hub.py

@@ -40,3 +40,4 @@ class WorkspaceHubSchema(PathAwareSchema):
     managed_network = ExperimentalField(NestedField(ManagedNetworkSchema))
     existing_workspaces = fields.List(fields.Str())
     workspace_hub_config = ExperimentalField(NestedField(WorkspaceHubConfigSchema))
+    enable_data_isolation = fields.Bool()

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_workspace_hub/workspace_hub_config.py

@@ -5,9 +5,16 @@
 from marshmallow import fields
 from azure.ai.ml._utils._experimental import experimental
 from azure.ai.ml._schema.core.schema_meta import PatchedSchemaMeta
+from marshmallow.decorators import post_load
 
 
 @experimental
 class WorkspaceHubConfigSchema(metaclass=PatchedSchemaMeta):
     additional_workspace_storage_accounts = fields.List(fields.Str())
     default_workspace_resource_group = fields.Str()
+
+    @post_load
+    def make(self, data, **kwargs):
+        from azure.ai.ml.entities import WorkspaceHubConfig
+
+        return WorkspaceHubConfig(**data)

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace_hub/workspace_hub.py

@@ -44,6 +44,7 @@ def __init__(
         public_network_access: Optional[str] = None,
         identity: Optional[IdentityConfiguration] = None,
         primary_user_assigned_identity: Optional[str] = None,
+        enable_data_isolation: bool = False,
         workspace_hub_config: Optional[WorkspaceHubConfig] = None,
         **kwargs,
     ):
@@ -84,6 +85,9 @@ def __init__(
         :type identity: IdentityConfiguration
         :param primary_user_assigned_identity: The workspaceHub's primary user assigned identity
         :type primary_user_assigned_identity: str
+        :param enable_data_isolation: A flag to determine if workspace has data isolation enabled.
+            The flag can only be set at the creation phase, it can't be updated.
+        :type enable_data_isolation: bool
         :param kwargs: A dictionary of additional configuration parameters.
         :type kwargs: dict
         """
@@ -105,6 +109,7 @@ def __init__(
             identity=identity,
             primary_user_assigned_identity=primary_user_assigned_identity,
             managed_network=managed_network,
+            enable_data_isolation=enable_data_isolation,
             **kwargs,
         )
         self.existing_workspaces = existing_workspaces
@@ -142,6 +147,7 @@ def _from_rest_object(cls, rest_obj: RestWorkspace) -> "WorkspaceHub":
             container_registry=rest_obj.container_registry,
             existing_workspaces=rest_obj.existing_workspaces,
             workspace_id=rest_obj.workspace_id,
+            enable_data_isolation=rest_obj.enable_data_isolation,
             workspace_hub_config=workspace_hub_config,
             id=rest_obj.id,
         )

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace_hub/workspace_hub_config.py

@@ -4,11 +4,14 @@
 
 # pylint: disable=too-many-instance-attributes,protected-access
 
-from typing import Dict, List, Optional
+from os import PathLike
+from pathlib import Path
+from typing import Dict, List, Optional, Union
 
 from azure.ai.ml._restclient.v2023_06_01_preview.models import WorkspaceHubConfig as RestWorkspaceHubConfig
 from azure.ai.ml._schema._workspace_hub.workspace_hub import WorkspaceHubConfigSchema
-from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY
+from azure.ai.ml.entities._util import load_from_dict
+from azure.ai.ml.constants._common import BASE_PATH_CONTEXT_KEY, PARAMS_OVERRIDE_KEY
 
 from azure.ai.ml._utils._experimental import experimental
 
@@ -20,7 +23,6 @@ def __init__(
         additional_workspace_storage_accounts: Optional[List[str]] = None,
         default_workspace_resource_group: Optional[str] = None,
     ) -> None:
-
         """WorkspaceHubConfig.
 
         :param additional_workspace_storage_accounts: A list of resource IDs of existing storage accounts that will be utilized in addition to the default one.
@@ -47,3 +49,20 @@ def _from_rest_object(cls, obj: RestWorkspaceHubConfig) -> "WorkspaceHubConfig":
     def _to_dict(self) -> Dict:
         # pylint: disable=no-member
         return WorkspaceHubConfigSchema(context={BASE_PATH_CONTEXT_KEY: "./"}).dump(self)
+
+    @classmethod
+    def _load(
+        cls,
+        data: Optional[Dict] = None,
+        yaml_path: Optional[Union[PathLike, str]] = None,
+        params_override: Optional[list] = None,
+        **kwargs,
+    ) -> "WorkspaceHubConfig":
+        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,
+        }
+        loaded_schema = load_from_dict(WorkspaceHubConfigSchema, data, context, **kwargs)
+        return WorkspaceHubConfig(**loaded_schema)

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

@@ -328,7 +328,6 @@ def _resolve_monitor_schedule_arm_id(  # pylint:disable=too-many-branches,too-ma
                     target=ErrorTarget.SCHEDULE,
                     error_category=ErrorCategory.USER_ERROR,
                 )
-
         # resolve ARM id for each signal and populate any defaults if needed
         for signal_name, signal in schedule.create_monitor.monitoring_signals.items():
             if signal.type == MonitorSignalType.CUSTOM:

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

@@ -164,7 +164,9 @@ def deserialize_callback(rest_obj):
 
     # @monitor_with_activity(logger, "Hub.BeginDelete", ActivityType.PUBLICAPI)
     @distributed_trace
-    def begin_delete(self, name: str, *, delete_dependent_resources: bool, **kwargs: Dict) -> LROPoller:
+    def begin_delete(
+        self, name: str, *, delete_dependent_resources: bool, permanently_delete: bool = False, **kwargs: Dict
+    ) -> LROPoller:
         """Delete a WorkspaceHub.
 
         :param name: Name of the WorkspaceHub
@@ -173,6 +175,9 @@ def begin_delete(self, name: str, *, delete_dependent_resources: bool, **kwargs:
             i.e., container registry, storage account, key vault.
             The default is False. Set to True to delete these resources.
         :type delete_dependent_resources: bool
+        :param permanently_delete: Workspaces are soft-deleted by default to allow recovery of workspace data.
+            Set this flag to true to override the soft-delete behavior and permanently delete your workspace.
+        :type permanently_delete: bool
         :return: A poller to track the operation status.
         :rtype: ~azure.core.polling.LROPoller[None]
         """
@@ -182,7 +187,9 @@ def begin_delete(self, name: str, *, delete_dependent_resources: bool, **kwargs:
             rest_workspace_obj and rest_workspace_obj.kind and rest_workspace_obj.kind.lower() == WORKSPACE_HUB_KIND
         ):
             raise ValidationError("{0} is not a WorkspaceHub".format(name))
-        if hasattr(rest_workspace_obj, 'workspace_hub_config') and hasattr(rest_workspace_obj.workspace_hub_config, 'additional_workspace_storage_accounts'):
+        if hasattr(rest_workspace_obj, "workspace_hub_config") and hasattr(
+            rest_workspace_obj.workspace_hub_config, "additional_workspace_storage_accounts"
+        ):
             for storageaccount in rest_workspace_obj.workspace_hub_config.additional_workspace_storage_accounts:
                 delete_resource_by_arm_id(
                     self._credentials,
@@ -191,4 +198,9 @@ def begin_delete(self, name: str, *, delete_dependent_resources: bool, **kwargs:
                     ArmConstants.AZURE_MGMT_STORAGE_API_VERSION,
                 )
 
-        return super().begin_delete(name=name, delete_dependent_resources=delete_dependent_resources, **kwargs)
+        return super().begin_delete(
+            name=name,
+            delete_dependent_resources=delete_dependent_resources,
+            permanently_delete=permanently_delete,
+            **kwargs,
+        )

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

@@ -199,8 +199,8 @@ def begin_delete(
             i.e., container registry, storage account, key vault, and application insights.
             The default is False. Set to True to delete these resources.
         :type delete_dependent_resources: bool
-        :keyword permanently_delete: Workspaces are soft-deleted state by default to allow recovery of workspace data.
-            Set this flag to override the soft-delete behavior and permanently delete your workspace.
+        :param permanently_delete: Workspaces are soft-deleted by default to allow recovery of workspace data.
+            Set this flag to true to override the soft-delete behavior and permanently delete your workspace.
         :type permanently_delete: bool
         :return: A poller to track the operation status.
         :rtype: ~azure.core.polling.LROPoller[None]

sdk/ml/azure-ai-ml/tests/test_configs/workspace/workspacehub_min.yaml

@@ -10,4 +10,5 @@ workspace_hub_config:
 tags:
   purpose: testing
   team: ws-management
+enable_data_isolation: True
 

sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_hub_operations.py

@@ -3,11 +3,12 @@
 import pytest
 from pytest_mock import MockFixture
 
-from azure.ai.ml import load_workspace
+from azure.ai.ml import load_workspace_hub
 from azure.ai.ml._scope_dependent_operations import OperationScope
 from azure.ai.ml.entities import (
     WorkspaceHub,
     Workspace,
+    WorkspaceHubConfig,
 )
 from azure.ai.ml.operations import WorkspaceHubOperations
 from azure.core.polling import LROPoller
@@ -100,3 +101,11 @@ def outgoing_call(rg, name):
         mocker.patch("azure.ai.ml.operations._workspace_operations_base.delete_resource_by_arm_id", return_value=None)
         with pytest.raises(Exception):
             mock_workspace_hub_operation.begin_delete("randstr", delete_dependent_resources=True)
+
+    def test_load_workspace_hub_from_yaml(self, mock_workspace_hub_operation: WorkspaceHubOperations):
+        params_override = []
+        hub = load_workspace_hub(
+            "./tests/test_configs/workspace/workspacehub_min.yaml", params_override=params_override
+        )
+        assert isinstance(hub.enable_data_isolation, bool)
+        assert isinstance(hub.workspace_hub_config, WorkspaceHubConfig)