From e977864590c53800fe1846b2d62bfa7c5835c1a7 Mon Sep 17 00:00:00 2001 From: Matthew Ferrari Date: Tue, 4 Jun 2024 00:13:58 -0700 Subject: [PATCH 1/4] Python SDK System Datastore Auth AML --- .../ml/_arm_deployments/arm_templates/workspace_base.json | 7 +++++++ .../ml/_arm_deployments/arm_templates/workspace_param.json | 3 +++ .../azure-ai-ml/azure/ai/ml/_schema/workspace/workspace.py | 1 + .../azure/ai/ml/entities/_workspace/workspace.py | 6 ++++++ .../azure/ai/ml/operations/_workspace_operations_base.py | 6 +++++- .../workspace/unittests/test_workspace_operations_base.py | 5 +++++ 6 files changed, 27 insertions(+), 1 deletion(-) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_base.json b/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_base.json index 0b58b9975947..728b67b44622 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_base.json +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_base.json @@ -416,6 +416,13 @@ "description": "Managed network settings to be used for the workspace. If not specified, isolation mode Disabled is the default" } }, + "system_datastores_auth_mode": { + "type": "string", + "defaultValue": "accesskey", + "metadata": { + "description": "The authentication mode for system datastores" + } + }, "spark_runtime_version": { "type": "string", "defaultValue": "", diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_param.json b/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_param.json index b3051d5ebd73..fad9a157fa11 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_param.json +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_param.json @@ -128,6 +128,9 @@ "managedNetwork": { "value": {} }, + "system_datastores_auth_mode": { + "value": "accesskey" + }, "spark_runtime_version" : { "value": null }, diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/workspace.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/workspace.py index 764ca3dc57e1..850a5b5db02b 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/workspace.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/workspace.py @@ -37,6 +37,7 @@ class WorkspaceSchema(PathAwareSchema): allowed_values=[PublicNetworkAccess.DISABLED, PublicNetworkAccess.ENABLED], casing_transform=snake_to_pascal, ) + system_datastores_auth_mode = fields.Str() identity = NestedField(IdentitySchema) primary_user_assigned_identity = fields.Str() workspace_hub = fields.Str(validate=validate_arm_str) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py index ee556e7b0ba2..dc5b3c0e2631 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py @@ -81,6 +81,8 @@ class Workspace(Resource): :type primary_user_assigned_identity: str :param managed_network: workspace's Managed Network configuration :type managed_network: ~azure.ai.ml.entities.ManagedNetwork + :param system_datastores_auth_mode: The authentication mode for system datastores. + :type system_datastores_auth_mode: 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 @@ -120,6 +122,7 @@ def __init__( identity: Optional[IdentityConfiguration] = None, primary_user_assigned_identity: Optional[str] = None, managed_network: Optional[ManagedNetwork] = None, + system_datastores_auth_mode: Optional[str] = None, enable_data_isolation: bool = False, hub_id: Optional[str] = None, # Hidden input, surfaced by Project workspace_hub: Optional[str] = None, # Deprecated input maintained for backwards compat. @@ -159,6 +162,7 @@ def __init__( self.identity = identity self.primary_user_assigned_identity = primary_user_assigned_identity self.managed_network = managed_network + self.system_datastores_auth_mode = system_datastores_auth_mode self.enable_data_isolation = enable_data_isolation if workspace_hub and not hub_id: hub_id = workspace_hub @@ -378,6 +382,7 @@ def _from_rest_object(cls, rest_obj: RestWorkspace) -> Optional["Workspace"]: identity=identity, primary_user_assigned_identity=rest_obj.primary_user_assigned_identity, managed_network=managed_network, + system_datastores_auth_mode=rest_obj.system_datastores_auth_mode, feature_store_settings=feature_store_settings, enable_data_isolation=rest_obj.enable_data_isolation, hub_id=rest_obj.hub_resource_id, @@ -423,6 +428,7 @@ def _to_rest_object(self) -> RestWorkspace: if self.managed_network else None ), # pylint: disable=protected-access + system_datastores_auth_mode=self.system_datastores_auth_mode, feature_store_settings=feature_store_settings, enable_data_isolation=self.enable_data_isolation, hub_resource_id=self._hub_id, diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py index a74d426a3c39..3efadd5f77a7 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py @@ -106,7 +106,7 @@ def begin_create( :rtype: ~azure.core.polling.LROPoller[~azure.ai.ml.entities.Workspace] :raises ~azure.ai.ml.ValidationException: Raised if workspace is Project workspace and user specifies any of the following in workspace object: storage_account, container_registry, key_vault, - public_network_access, managed_network, customer_managed_key. + public_network_access, managed_network, customer_managed_key, system_datastores_auth_mode. """ existing_workspace = None resource_group = kwargs.get("resource_group") or workspace.resource_group or self._resource_group_name @@ -338,6 +338,7 @@ def begin_update( description=kwargs.get("description", workspace.description), friendly_name=kwargs.get("display_name", workspace.display_name), public_network_access=kwargs.get("public_network_access", workspace.public_network_access), + system_datastores_auth_mode=kwargs.get("system_datastores_auth_mode", workspace.system_datastores_auth_mode), image_build_compute=kwargs.get("image_build_compute", workspace.image_build_compute), identity=identity, primary_user_assigned_identity=kwargs.get( @@ -634,6 +635,9 @@ def _populate_arm_parameters(self, workspace: Workspace, **kwargs: Any) -> Tuple if workspace.public_network_access: _set_val(param["publicNetworkAccess"], workspace.public_network_access) + if workspace.system_datastores_auth_mode: + _set_val(param["system_datastores_auth_mode"], workspace.system_datastores_auth_mode) + if workspace.image_build_compute: _set_val(param["imageBuildCompute"], workspace.image_build_compute) diff --git a/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations_base.py b/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations_base.py index 38ecaf8f6a02..05ac748452b6 100644 --- a/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations_base.py +++ b/sdk/ml/azure-ai-ml/tests/workspace/unittests/test_workspace_operations_base.py @@ -152,12 +152,15 @@ def outgoing_get_call(rg, name): ), ], ) + ws.system_datastores_auth_mode = "identity" return ws._to_rest_object() mock_workspace_operation_base._operation.get.side_effect = outgoing_get_call ws = mock_workspace_operation_base.get(name="random_name", resource_group="rg") mock_workspace_operation_base._operation.get.assert_called_once() + assert ws.system_datastores_auth_mode == "identity" + assert ws.managed_network is not None assert ws.managed_network.isolation_mode == IsolationMode.ALLOW_ONLY_APPROVED_OUTBOUND rules = ws.managed_network.outbound_rules @@ -217,6 +220,7 @@ def test_update(self, mock_workspace_operation_base: WorkspaceOperationsBase) -> ], ), managed_network=ManagedNetwork(), + system_datastores_auth_mode="identity", primary_user_assigned_identity="resource2", customer_managed_key=CustomerManagedKey(key_uri="new_cmk_uri"), ) @@ -240,6 +244,7 @@ def outgoing_call(rg, name, params, polling, cls): ) ) assert params.managed_network.isolation_mode == "Disabled" + assert params.system_datastores_auth_mode == "identity" assert params.managed_network.outbound_rules == {} assert polling is True assert callable(cls) From 11296b71fb0aa54414b317b256ff98c21efb7956 Mon Sep 17 00:00:00 2001 From: Matthew Ferrari Date: Tue, 4 Jun 2024 00:29:42 -0700 Subject: [PATCH 2/4] Python SDK System Datastore Auth AML --- .../ai/ml/_arm_deployments/arm_templates/workspace_base.json | 2 +- .../ai/ml/_arm_deployments/arm_templates/workspace_param.json | 2 +- .../azure/ai/ml/operations/_workspace_operations_base.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_base.json b/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_base.json index 728b67b44622..f892ca887d3d 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_base.json +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_base.json @@ -416,7 +416,7 @@ "description": "Managed network settings to be used for the workspace. If not specified, isolation mode Disabled is the default" } }, - "system_datastores_auth_mode": { + "systemDatastoresAuthMode": { "type": "string", "defaultValue": "accesskey", "metadata": { diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_param.json b/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_param.json index fad9a157fa11..57e608d96310 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_param.json +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_param.json @@ -128,7 +128,7 @@ "managedNetwork": { "value": {} }, - "system_datastores_auth_mode": { + "systemDatastoresAuthMode": { "value": "accesskey" }, "spark_runtime_version" : { diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py index 3efadd5f77a7..60fecf625bec 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py @@ -636,7 +636,7 @@ def _populate_arm_parameters(self, workspace: Workspace, **kwargs: Any) -> Tuple _set_val(param["publicNetworkAccess"], workspace.public_network_access) if workspace.system_datastores_auth_mode: - _set_val(param["system_datastores_auth_mode"], workspace.system_datastores_auth_mode) + _set_val(param["systemDatastoresAuthMode"], workspace.system_datastores_auth_mode) if workspace.image_build_compute: _set_val(param["imageBuildCompute"], workspace.image_build_compute) From f1aaaed1e34fbb5e47825901e2abdd40a85e6033 Mon Sep 17 00:00:00 2001 From: Matthew Ferrari Date: Tue, 4 Jun 2024 06:11:27 -0700 Subject: [PATCH 3/4] Python SDK System Datastore Auth AML --- .../ml/_restclient/v2023_08_01_preview/models/_models.py | 8 ++++++++ .../v2023_08_01_preview/models/_models_py3.py | 9 +++++++++ 2 files changed, 17 insertions(+) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/v2023_08_01_preview/models/_models.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/v2023_08_01_preview/models/_models.py index 0c838eec6bc5..81c9aff5d7a6 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/v2023_08_01_preview/models/_models.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/v2023_08_01_preview/models/_models.py @@ -32348,6 +32348,9 @@ class WorkspaceUpdateParameters(msrest.serialization.Model): :ivar v1_legacy_mode: Enabling v1_legacy_mode may prevent you from using features provided by the v2 API. :vartype v1_legacy_mode: bool + :ivar system_datastores_auth_mode: The auth mode used for accessing the system datastores of + the workspace. + :vartype system_datastores_auth_mode: str """ _attribute_map = { @@ -32369,6 +32372,7 @@ class WorkspaceUpdateParameters(msrest.serialization.Model): 'service_managed_resources_settings': {'key': 'properties.serviceManagedResourcesSettings', 'type': 'ServiceManagedResourcesSettings'}, 'soft_delete_retention_in_days': {'key': 'properties.softDeleteRetentionInDays', 'type': 'int'}, 'v1_legacy_mode': {'key': 'properties.v1LegacyMode', 'type': 'bool'}, + 'system_datastores_auth_mode': {'key': 'properties.systemDatastoresAuthMode', 'type': 'str'}, } def __init__( @@ -32422,6 +32426,9 @@ def __init__( :keyword v1_legacy_mode: Enabling v1_legacy_mode may prevent you from using features provided by the v2 API. :paramtype v1_legacy_mode: bool + :keyword system_datastores_auth_mode: The auth mode used for accessing the system datastores of + the workspace. + :paramtype system_datastores_auth_mode: str """ super(WorkspaceUpdateParameters, self).__init__(**kwargs) self.identity = kwargs.get('identity', None) @@ -32442,3 +32449,4 @@ def __init__( self.service_managed_resources_settings = kwargs.get('service_managed_resources_settings', None) self.soft_delete_retention_in_days = kwargs.get('soft_delete_retention_in_days', None) self.v1_legacy_mode = kwargs.get('v1_legacy_mode', None) + self.system_datastores_auth_mode = kwargs.get('system_datastores_auth_mode', None) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/v2023_08_01_preview/models/_models_py3.py b/sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/v2023_08_01_preview/models/_models_py3.py index 034b9137668b..32435ecfb832 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/v2023_08_01_preview/models/_models_py3.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/v2023_08_01_preview/models/_models_py3.py @@ -35117,6 +35117,9 @@ class WorkspaceUpdateParameters(msrest.serialization.Model): :ivar v1_legacy_mode: Enabling v1_legacy_mode may prevent you from using features provided by the v2 API. :vartype v1_legacy_mode: bool + :ivar system_datastores_auth_mode: The auth mode used for accessing the system datastores of + the workspace. + :vartype system_datastores_auth_mode: str """ _attribute_map = { @@ -35138,6 +35141,7 @@ class WorkspaceUpdateParameters(msrest.serialization.Model): 'service_managed_resources_settings': {'key': 'properties.serviceManagedResourcesSettings', 'type': 'ServiceManagedResourcesSettings'}, 'soft_delete_retention_in_days': {'key': 'properties.softDeleteRetentionInDays', 'type': 'int'}, 'v1_legacy_mode': {'key': 'properties.v1LegacyMode', 'type': 'bool'}, + 'system_datastores_auth_mode': {'key': 'properties.systemDatastoresAuthMode', 'type': 'str'}, } def __init__( @@ -35161,6 +35165,7 @@ def __init__( service_managed_resources_settings: Optional["ServiceManagedResourcesSettings"] = None, soft_delete_retention_in_days: Optional[int] = None, v1_legacy_mode: Optional[bool] = None, + system_datastores_auth_mode: Optional[str] = None, **kwargs ): """ @@ -35210,6 +35215,9 @@ def __init__( :keyword v1_legacy_mode: Enabling v1_legacy_mode may prevent you from using features provided by the v2 API. :paramtype v1_legacy_mode: bool + :keyword system_datastores_auth_mode: The auth mode used for accessing the system datastores of + the workspace. + :paramtype system_datastores_auth_mode: str """ super(WorkspaceUpdateParameters, self).__init__(**kwargs) self.identity = identity @@ -35230,3 +35238,4 @@ def __init__( self.service_managed_resources_settings = service_managed_resources_settings self.soft_delete_retention_in_days = soft_delete_retention_in_days self.v1_legacy_mode = v1_legacy_mode + self.system_datastores_auth_mode = system_datastores_auth_mode From e971bb4d5273fe365210cebe5121966a9d8e6111 Mon Sep 17 00:00:00 2001 From: Matthew Ferrari Date: Tue, 4 Jun 2024 13:24:41 -0700 Subject: [PATCH 4/4] SDK Updates --- .../azure/ai/ml/operations/_workspace_operations_base.py | 4 +++- .../preview/2023-08-01-preview/workspaceRP.json | 4 ++++ .../preview/2024-01-01-preview/workspaceRP.json | 4 ++++ .../preview/2024-04-01-preview/workspaceRP.json | 4 ++++ .../stable/2023-10-01/machineLearningServices.json | 4 ++++ 5 files changed, 19 insertions(+), 1 deletion(-) diff --git a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py index 60fecf625bec..4c7605543c00 100644 --- a/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py +++ b/sdk/ml/azure-ai-ml/azure/ai/ml/operations/_workspace_operations_base.py @@ -338,7 +338,9 @@ def begin_update( description=kwargs.get("description", workspace.description), friendly_name=kwargs.get("display_name", workspace.display_name), public_network_access=kwargs.get("public_network_access", workspace.public_network_access), - system_datastores_auth_mode=kwargs.get("system_datastores_auth_mode", workspace.system_datastores_auth_mode), + system_datastores_auth_mode=kwargs.get( + "system_datastores_auth_mode", workspace.system_datastores_auth_mode + ), image_build_compute=kwargs.get("image_build_compute", workspace.image_build_compute), identity=identity, primary_user_assigned_identity=kwargs.get( diff --git a/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/preview/2023-08-01-preview/workspaceRP.json b/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/preview/2023-08-01-preview/workspaceRP.json index 7b038fbb8080..b3b6e721daf6 100644 --- a/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/preview/2023-08-01-preview/workspaceRP.json +++ b/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/preview/2023-08-01-preview/workspaceRP.json @@ -3571,6 +3571,10 @@ "description": "The parameters for updating a machine learning workspace.", "type": "object", "properties": { + "systemDatastoresAuthMode": { + "description": "The auth mode used for accessing the system datastores of the workspace.", + "type": "string" + }, "applicationInsights": { "description": "ARM id of the application insights associated with this workspace.", "type": "string" diff --git a/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/preview/2024-01-01-preview/workspaceRP.json b/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/preview/2024-01-01-preview/workspaceRP.json index 0aee201cdd37..42ab1adc3b7a 100644 --- a/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/preview/2024-01-01-preview/workspaceRP.json +++ b/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/preview/2024-01-01-preview/workspaceRP.json @@ -5020,6 +5020,10 @@ "description": "The parameters for updating a machine learning workspace.", "type": "object", "properties": { + "systemDatastoresAuthMode": { + "description": "The auth mode used for accessing the system datastores of the workspace.", + "type": "string" + }, "applicationInsights": { "description": "ARM id of the application insights associated with this workspace.", "type": "string" diff --git a/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/preview/2024-04-01-preview/workspaceRP.json b/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/preview/2024-04-01-preview/workspaceRP.json index 3358be323c8b..635c140070fc 100644 --- a/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/preview/2024-04-01-preview/workspaceRP.json +++ b/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/preview/2024-04-01-preview/workspaceRP.json @@ -6588,6 +6588,10 @@ "description": "The parameters for updating a machine learning workspace.", "type": "object", "properties": { + "systemDatastoresAuthMode": { + "description": "The auth mode used for accessing the system datastores of the workspace.", + "type": "string" + }, "applicationInsights": { "description": "ARM id of the application insights associated with this workspace.", "type": "string" diff --git a/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/stable/2023-10-01/machineLearningServices.json b/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/stable/2023-10-01/machineLearningServices.json index 468587358b00..1b3937eb2d6e 100644 --- a/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/stable/2023-10-01/machineLearningServices.json +++ b/sdk/ml/azure-ai-ml/swagger/machinelearningservices/resource-manager/Microsoft.MachineLearningServices/stable/2023-10-01/machineLearningServices.json @@ -2580,6 +2580,10 @@ "type": "object", "description": "The parameters for updating the properties of a machine learning workspace.", "properties": { + "systemDatastoresAuthMode": { + "description": "The auth mode used for accessing the system datastores of the workspace.", + "type": "string" + }, "description": { "description": "The description of this workspace.", "type": "string"