set storage account access to identity-based for feature store creation (#35941)

* set storage account access to identity-based for feature store creation

* revert

* add workspace default storage role assignment

Author: runhli

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

@@ -599,6 +599,13 @@
                 "description": "Serverless compute settings to be used for the workspace."
             }
         },
+        "system_datastore_auth_mode": {
+            "type": "string",
+            "defaultValue": "",
+            "metadata": {
+                "description": "Settings to control workspace storage account access auth type"
+            }
+        },
         "endpoint_resource_id": {
             "type": "string",
             "defaultValue": "null",
@@ -831,6 +838,7 @@
                     "SearchAccountArmId": "[parameters('encryption_search_resourceid')]"
                 },
                 "primaryUserAssignedIdentity": "[parameters('primaryUserAssignedIdentity')]",
+                "systemDatastoresAuthMode": "[parameters('system_datastore_auth_mode')]",
                 "managedNetwork": "[parameters('managedNetwork')]",
                 "featureStoreSettings": {
                     "computeruntime": {
@@ -895,7 +903,7 @@
         {
             "condition":"[equals(parameters('kind'), 'featurestore')]",
 			"type": "Microsoft.Resources/deployments",
-			"apiVersion": "2022-05-01",
+			"apiVersion": "2024-03-01",
 			"name": "[concat(parameters('workspaceName'), '-deploy-feature-store')]",
 			"dependsOn": [
                 "[resourceId('Microsoft.MachineLearningServices/workspaces', parameters('workspaceName'))]",
@@ -945,6 +953,7 @@
                                     "SearchAccountArmId": "[parameters('encryption_search_resourceid')]"
                                 },
                                 "primaryUserAssignedIdentity": "[parameters('primaryUserAssignedIdentity')]",
+                                "systemDatastoresAuthMode": "[parameters('system_datastore_auth_mode')]",
                                 "managedNetwork": "[parameters('managedNetwork')]",
                                 "featureStoreSettings": {
                                     "computeruntime": {
@@ -1056,6 +1065,38 @@
                 }
             }
         },
+        {
+            "condition": "[and(equals(parameters('kind'), 'featurestore'), equals(parameters('grant_materialization_permissions'), 'true'), not(equals(parameters('materializationIdentityOption'), 'none')))]",
+            "type": "Microsoft.Resources/deployments",
+            "name": "[concat('ws-storage-role-assign-', guid(variables('materializationIdentity'), variables('storageAccount'), 'storage blob data contributor'))]",
+            "apiVersion": "2020-06-01",
+            "dependsOn": [
+                "[resourceId('Microsoft.MachineLearningServices/workspaces', parameters('workspaceName'))]",
+                "[resourceId('Microsoft.ManagedIdentity/userAssignedIdentities', parameters('materialization_identity_name'))]"
+            ],
+            "resourceGroup": "[parameters('storageAccountResourceGroupName')]",
+            "properties": {
+                "mode": "Incremental",
+                "template": {
+                    "$schema": "https://schema.management.azure.com/schemas/2015-01-01/deploymentTemplate.json#",
+                    "contentVersion": "1.0.0.0",
+                    "resources": [
+                        {
+                            "type": "Microsoft.Authorization/roleAssignments",
+                            "apiVersion": "2022-04-01",
+                            "name": "[guid(variables('materializationIdentity'), variables('storageAccount'), 'storage blob data contributor')]",
+                            "scope": "[variables('storageAccount')]",
+                            "location": "[parameters('location')]",
+                            "properties": {
+                                "roleDefinitionId": "[subscriptionResourceId('Microsoft.Authorization/roleDefinitions', 'ba92f5b4-2d11-453d-a403-e96b0029c9fe')]",
+                                "principalId": "[if(not(equals(parameters('materializationIdentityOption'), 'none')), reference(variables('materializationIdentity'), '2023-01-31').principalId, '')]",
+                                "principalType": "ServicePrincipal"
+                            }
+                        }
+                    ]
+                }
+            }
+        },
         {
             "condition": "[and(variables('enablePE'), not(equals(parameters('privateEndpointType'), 'none')))]",
             "type": "Microsoft.Resources/deployments",

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

@@ -197,6 +197,9 @@
     "serverless_compute_settings": {
         "value": {}
     },
+    "system_datastore_auth_mode": {
+        "value": ""
+    },
     "endpoint_resource_id": {
         "value": "null"
     },

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

@@ -684,6 +684,9 @@ def _populate_arm_parameters(self, workspace: Workspace, **kwargs: Any) -> Tuple
 
             from azure.ai.ml._utils._arm_id_utils import AzureResourceId, AzureStorageContainerResourceId
 
+            # set workspace storage account access auth type to identity-based
+            _set_val(param["system_datastore_auth_mode"], "identity")
+
             if offline_store_target:
                 arm_id = AzureStorageContainerResourceId(offline_store_target)
                 _set_val(param["offlineStoreStorageAccountOption"], "existing")

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

@@ -380,6 +380,7 @@ def test_populate_feature_store_arm_parameters(
         assert param["online_store_resource_group_name"]["value"] is None
         assert param["online_store_subscription_id"]["value"] is None
         assert param["online_store_connection_name"]["value"] is None
+        assert param["system_datastore_auth_mode"]["value"] == "identity"
 
         # test create feature store with materialization identity
         mock_materialization_identity_resource_id = (