add a new parameter allow_roleassignment_on_rg to allow/disallow role assignment on resource group level (#35939)

* Python SDK System Datastore Auth AML

* Python SDK System Datastore Auth AML

* Python SDK System Datastore Auth AML

* SDK Updates

* Add allow_roleassignment_on_rg for customer to determine whether to allow workspace role assignment on RG level

Add allow_roleassignment_on_rg for customer to determine whether to allow workspace role assignment on resource group level

* update the type of allow_roleassignment_on_rg

* update allow_roleassignment_on_rg to be boolean

* address lint check errors

* run black to fix the formatting issue

---------

Co-authored-by: Matthew Ferrari <[email protected]>
Co-authored-by: jiewe <[email protected]>

Author: 3 people

sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_base.json

@@ -416,6 +416,24 @@
                 "description": "Managed network settings to be used for the workspace. If not specified, isolation mode Disabled is the default"
             }
         },
+        "systemDatastoresAuthMode": {
+            "type": "string",
+            "defaultValue": "accesskey",
+            "metadata": {
+                "description": "The authentication mode for system datastores"
+            }
+        },
+        "allowRoleAssignmentOnRG": {
+            "type": "string",
+            "defaultValue": "true",
+            "allowedValues": [
+                "true",
+                "false"
+            ],
+            "metadata": {
+                "description": "A flag to determine if workspace should be created/update with allow role assignment on resource group level."
+            }
+        },
         "spark_runtime_version": {
             "type": "string",
             "defaultValue": "",

sdk/ml/azure-ai-ml/azure/ai/ml/_arm_deployments/arm_templates/workspace_param.json

@@ -128,6 +128,9 @@
     "managedNetwork": {
         "value": {}
     },
+    "systemDatastoresAuthMode": {
+        "value": "accesskey"
+    },
     "spark_runtime_version" : {
         "value": null
     },
@@ -188,6 +191,9 @@
     "enable_data_isolation": {
         "value": "false"
     },
+    "allowRoleAssignmentOnRG": {
+        "value": "true"
+    },
     "serverless_compute_settings": {
         "value": {}
     },

sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/v2023_08_01_preview/models/_models.py

@@ -31621,6 +31621,8 @@ class Workspace(Resource):
     :vartype discovery_url: str
     :ivar enable_data_isolation:
     :vartype enable_data_isolation: bool
+    :ivar allow_roleassignment_on_rg: Determine whether allow workspace role assignment on resource group level.  
+    :vartype allow_roleassignment_on_rg: bool
     :ivar encryption:
     :vartype encryption: ~azure.mgmt.machinelearningservices.models.EncryptionProperty
     :ivar existing_workspaces:
@@ -31737,6 +31739,7 @@ class Workspace(Resource):
         'description': {'key': 'properties.description', 'type': 'str'},
         'discovery_url': {'key': 'properties.discoveryUrl', 'type': 'str'},
         'enable_data_isolation': {'key': 'properties.enableDataIsolation', 'type': 'bool'},
+        'allow_roleassignment_on_rg': {'key': 'properties.allowRoleAssignmentOnRG', 'type': 'bool'},
         'encryption': {'key': 'properties.encryption', 'type': 'EncryptionProperty'},
         'existing_workspaces': {'key': 'properties.existingWorkspaces', 'type': '[str]'},
         'feature_store_settings': {'key': 'properties.featureStoreSettings', 'type': 'FeatureStoreSettings'},
@@ -31804,6 +31807,8 @@ def __init__(
         :paramtype discovery_url: str
         :keyword enable_data_isolation:
         :paramtype enable_data_isolation: bool
+        :keyword allow_roleassignment_on_rg: Determine whether allow workspace role assignment on resource group level.
+        :paramtype allow_roleassignment_on_rg: bool
         :keyword encryption:
         :paramtype encryption: ~azure.mgmt.machinelearningservices.models.EncryptionProperty
         :keyword existing_workspaces:
@@ -31875,6 +31880,7 @@ def __init__(
         self.description = kwargs.get('description', None)
         self.discovery_url = kwargs.get('discovery_url', None)
         self.enable_data_isolation = kwargs.get('enable_data_isolation', None)
+        self.allow_roleassignment_on_rg = kwargs.get('allow_roleassignment_on_rg', None)
         self.encryption = kwargs.get('encryption', None)
         self.existing_workspaces = kwargs.get('existing_workspaces', None)
         self.feature_store_settings = kwargs.get('feature_store_settings', None)
@@ -32319,6 +32325,8 @@ class WorkspaceUpdateParameters(msrest.serialization.Model):
     :vartype description: str
     :ivar enable_data_isolation:
     :vartype enable_data_isolation: bool
+    :ivar allow_roleassignment_on_rg: Determine whether allow workspace role assignment on resource group level.  
+    :vartype allow_roleassignment_on_rg: bool
     :ivar encryption:
     :vartype encryption: ~azure.mgmt.machinelearningservices.models.EncryptionUpdateProperties
     :ivar feature_store_settings: Settings for feature store type workspace.
@@ -32348,6 +32356,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 +32380,8 @@ 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'},
+        'allow_roleassignment_on_rg' : {'Key': 'properties.allowRoleAssignmentOnRG', 'type': 'bool'}
     }
 
     def __init__(
@@ -32392,6 +32405,8 @@ def __init__(
         :paramtype description: str
         :keyword enable_data_isolation:
         :paramtype enable_data_isolation: bool
+        :keyword allow_roleassignment_on_rg: Determine whether allow workspace role assignment on resource group level.
+        :paramtype allow_roleassignment_on_rg: bool
         :keyword encryption:
         :paramtype encryption: ~azure.mgmt.machinelearningservices.models.EncryptionUpdateProperties
         :keyword feature_store_settings: Settings for feature store type workspace.
@@ -32422,6 +32437,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)
@@ -32431,6 +32449,7 @@ def __init__(
         self.container_registry = kwargs.get('container_registry', None)
         self.description = kwargs.get('description', None)
         self.enable_data_isolation = kwargs.get('enable_data_isolation', None)
+        self.allow_roleassignment_on_rg = kwargs.get('allow_roleassignment_on_rg', None)
         self.encryption = kwargs.get('encryption', None)
         self.feature_store_settings = kwargs.get('feature_store_settings', None)
         self.friendly_name = kwargs.get('friendly_name', None)
@@ -32442,3 +32461,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)

sdk/ml/azure-ai-ml/azure/ai/ml/_restclient/v2023_08_01_preview/models/_models_py3.py

@@ -34323,6 +34323,8 @@ class Workspace(Resource):
     :vartype discovery_url: str
     :ivar enable_data_isolation:
     :vartype enable_data_isolation: bool
+    :ivar allow_roleassignment_on_rg: Determine whether allow workspace role assignment on resource group level.
+    :vartype allow_roleassignment_on_rg: bool
     :ivar encryption:
     :vartype encryption: ~azure.mgmt.machinelearningservices.models.EncryptionProperty
     :ivar existing_workspaces:
@@ -34439,6 +34441,7 @@ class Workspace(Resource):
         'description': {'key': 'properties.description', 'type': 'str'},
         'discovery_url': {'key': 'properties.discoveryUrl', 'type': 'str'},
         'enable_data_isolation': {'key': 'properties.enableDataIsolation', 'type': 'bool'},
+        'allow_roleassignment_on_rg': {'key': 'properties.allowRoleAssignmentOnRG', 'type': 'bool'},
         'encryption': {'key': 'properties.encryption', 'type': 'EncryptionProperty'},
         'existing_workspaces': {'key': 'properties.existingWorkspaces', 'type': '[str]'},
         'feature_store_settings': {'key': 'properties.featureStoreSettings', 'type': 'FeatureStoreSettings'},
@@ -34487,6 +34490,7 @@ def __init__(
         description: Optional[str] = None,
         discovery_url: Optional[str] = None,
         enable_data_isolation: Optional[bool] = None,
+        allow_roleassignment_on_rg: Optional[bool] = None,
         encryption: Optional["EncryptionProperty"] = None,
         existing_workspaces: Optional[List[str]] = None,
         feature_store_settings: Optional["FeatureStoreSettings"] = None,
@@ -34541,6 +34545,8 @@ def __init__(
         :paramtype discovery_url: str
         :keyword enable_data_isolation:
         :paramtype enable_data_isolation: bool
+        :keyword allow_roleassignment_on_rg: Determine whether allow workspace role assignment on resource group level.
+        :paramtype allow_roleassignment_on_rg: bool
         :keyword encryption:
         :paramtype encryption: ~azure.mgmt.machinelearningservices.models.EncryptionProperty
         :keyword existing_workspaces:
@@ -34612,6 +34618,7 @@ def __init__(
         self.description = description
         self.discovery_url = discovery_url
         self.enable_data_isolation = enable_data_isolation
+        self.allow_roleassignment_on_rg = allow_roleassignment_on_rg
         self.encryption = encryption
         self.existing_workspaces = existing_workspaces
         self.feature_store_settings = feature_store_settings
@@ -35088,6 +35095,8 @@ class WorkspaceUpdateParameters(msrest.serialization.Model):
     :vartype description: str
     :ivar enable_data_isolation:
     :vartype enable_data_isolation: bool
+    :ivar allow_roleassignment_on_rg: Determine whether allow workspace role assignment on resource group level.
+    :vartype allow_roleassignment_on_rg: bool
     :ivar encryption:
     :vartype encryption: ~azure.mgmt.machinelearningservices.models.EncryptionUpdateProperties
     :ivar feature_store_settings: Settings for feature store type workspace.
@@ -35117,6 +35126,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 +35150,8 @@ 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'},
+        'allow_roleassignment_on_rg' : {'Key': 'properties.allowRoleAssignmentOnRG', 'type': 'bool'}
     }
 
     def __init__(
@@ -35150,6 +35164,7 @@ def __init__(
         container_registry: Optional[str] = None,
         description: Optional[str] = None,
         enable_data_isolation: Optional[bool] = None,
+        allow_roleassignment_on_rg: Optional[bool] = None,
         encryption: Optional["EncryptionUpdateProperties"] = None,
         feature_store_settings: Optional["FeatureStoreSettings"] = None,
         friendly_name: Optional[str] = None,
@@ -35161,6 +35176,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
     ):
         """
@@ -35180,6 +35196,8 @@ def __init__(
         :paramtype description: str
         :keyword enable_data_isolation:
         :paramtype enable_data_isolation: bool
+        :keyword allow_roleassignment_on_rg: Determine whether allow workspace role assignment on resource group level.
+        :paramtype allow_roleassignment_on_rg: bool
         :keyword encryption:
         :paramtype encryption: ~azure.mgmt.machinelearningservices.models.EncryptionUpdateProperties
         :keyword feature_store_settings: Settings for feature store type workspace.
@@ -35210,6 +35228,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
@@ -35219,6 +35240,7 @@ def __init__(
         self.container_registry = container_registry
         self.description = description
         self.enable_data_isolation = enable_data_isolation
+        self.allow_roleassignment_on_rg = allow_roleassignment_on_rg
         self.encryption = encryption
         self.feature_store_settings = feature_store_settings
         self.friendly_name = friendly_name
@@ -35230,3 +35252,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

sdk/ml/azure-ai-ml/azure/ai/ml/_schema/workspace/workspace.py

@@ -37,9 +37,11 @@ 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)
     managed_network = ExperimentalField(NestedField(ManagedNetworkSchema, unknown=EXCLUDE))
     enable_data_isolation = fields.Bool()
+    allow_roleassignment_on_rg = fields.Bool()
     serverless_compute = NestedField(ServerlessComputeSettingsSchema)

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_workspace/workspace.py

@@ -81,9 +81,13 @@ 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
+    :param allow_roleassignment_on_rg: Determine whether allow workspace role assignment on resource group level.
+    :type allow_roleassignment_on_rg: bool
     :param serverless_compute: The serverless compute settings for the workspace.
     :type: ~azure.ai.ml.entities.ServerlessComputeSettings
     :param workspace_hub: Deprecated resource ID of an existing workspace hub to help create project workspace.
@@ -120,7 +124,9 @@ 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,
+        allow_roleassignment_on_rg: bool = True,
         hub_id: Optional[str] = None,  # Hidden input, surfaced by Project
         workspace_hub: Optional[str] = None,  # Deprecated input maintained for backwards compat.
         serverless_compute: Optional[ServerlessComputeSettings] = None,
@@ -159,7 +165,9 @@ 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
+        self.allow_roleassignment_on_rg = allow_roleassignment_on_rg
         if workspace_hub and not hub_id:
             hub_id = workspace_hub
         self.__hub_id = hub_id
@@ -378,8 +386,10 @@ 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,
+            allow_roleassignment_on_rg=rest_obj.allow_roleassignment_on_rg,
             hub_id=rest_obj.hub_resource_id,
             workspace_id=rest_obj.workspace_id,
             serverless_compute=serverless_compute,
@@ -423,8 +433,10 @@ 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,
+            allow_roleassignment_on_rg=self.allow_roleassignment_on_rg,
             hub_resource_id=self._hub_id,
             serverless_compute_settings=serverless_compute_settings,
         )

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
@@ -336,6 +336,10 @@ 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
+            ),
+            allow_roleassignment_on_rg=kwargs.get("allow_roleassignment_on_rg", workspace.allow_roleassignment_on_rg),
             image_build_compute=kwargs.get("image_build_compute", workspace.image_build_compute),
             identity=identity,
             primary_user_assigned_identity=kwargs.get(
@@ -632,6 +636,12 @@ 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["systemDatastoresAuthMode"], workspace.system_datastores_auth_mode)
+
+        if not workspace.allow_roleassignment_on_rg:
+            _set_val(param["allowRoleAssignmentOnRG"], "false")
+
         if workspace.image_build_compute:
             _set_val(param["imageBuildCompute"], workspace.image_build_compute)