diff --git a/sdk/ml/azure-ai-ml/CHANGELOG.md b/sdk/ml/azure-ai-ml/CHANGELOG.md index 5a11803454da..4f67b79be25e 100644 --- a/sdk/ml/azure-ai-ml/CHANGELOG.md +++ b/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 diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_workspace_hub/workspace_hub.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_workspace_hub/workspace_hub.py index 483c6cd3f323..bbc8c5671b1a 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_workspace_hub/workspace_hub.py +++ b/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() diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_workspace_hub/workspace_hub_config.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_workspace_hub/workspace_hub_config.py index db3c8db12f3c..d01969cdcb2d 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/_workspace_hub/workspace_hub_config.py +++ b/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) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace_hub/workspace_hub.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace_hub/workspace_hub.py index 36e4baeb6970..120f7f09302a 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace_hub/workspace_hub.py +++ b/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, ) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace_hub/workspace_hub_config.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace_hub/workspace_hub_config.py index 101c1d1f3110..de2d7d0400c6 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace_hub/workspace_hub_config.py +++ b/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) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_schedule_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_schedule_operations.py index 41b046b34e0d..442adbc70963 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_schedule_operations.py +++ b/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: diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_hub_operation.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_hub_operation.py index beafbb3dd8b0..c7883f68a7e6 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_hub_operation.py +++ b/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, + ) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations.py index c25d231db1b1..82e9050d5eca 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations.py +++ b/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] diff --git a/sdk/ml/azure-ai-ml/tests/test_configs/workspace/workspacehub_min.yaml b/sdk/ml/azure-ai-ml/tests/test_configs/workspace/workspacehub_min.yaml index c569014f7b02..39a751404b75 100644 --- a/sdk/ml/azure-ai-ml/tests/test_configs/workspace/workspacehub_min.yaml +++ b/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 diff --git a/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_hub_operations.py b/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_hub_operations.py index bff5ebc6fbf4..04af19c0096c 100644 --- a/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_hub_operations.py +++ b/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)