Add feature store provision role assignment location (#34646)

* fix role assignment location

* add tests

Author: runhli

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

@@ -65,6 +65,12 @@
                 "description": "Specifies the name of the resource group."
             }
         },
+        "location": {
+            "type": "string",
+            "metadata": {
+                "description": "Specifies the location for all resources."
+            }
+        },
         "update_offline_store_role_assignment": {
             "type": "string",
             "defaultValue": "false",
@@ -144,6 +150,7 @@
                             "apiVersion": "2022-04-01",
                             "name": "[guid(parameters('materialization_identity_resource_id'), parameters('offline_store_target'), 'storage blob data contributor')]",
                             "scope": "[parameters('offline_store_target')]",
+                            "location": "[parameters('location')]",
                             "properties": {
                                 "roleDefinitionId": "[subscriptionResourceId('Microsoft.Authorization/roleDefinitions', 'ba92f5b4-2d11-453d-a403-e96b0029c9fe')]",
                                 "principalId": "[reference(parameters('materialization_identity_resource_id'), '2023-01-31').principalId]",
@@ -172,6 +179,7 @@
                             "apiVersion": "2022-04-01",
                             "name": "[guid(parameters('materialization_identity_resource_id'), parameters('online_store_target'), 'contributor')]",
                             "scope": "[parameters('online_store_target')]",
+                            "location": "[parameters('location')]",
                             "properties": {
                                 "roleDefinitionId": "[subscriptionResourceId('Microsoft.Authorization/roleDefinitions', 'b24988ac-6180-42a0-ab88-20f7382dd24c')]",
                                 "principalId": "[reference(parameters('materialization_identity_resource_id'), '2023-01-31').principalId]",

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

@@ -26,6 +26,9 @@
     "resource_group_name": {
         "value": null
     },
+    "location": {
+        "value": ""
+    },
     "update_offline_store_role_assignment": {
         "value": "false"
     },

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

@@ -1053,6 +1053,7 @@
                             "apiVersion": "2022-04-01",
                             "name": "[guid(variables('materializationIdentity'), resourceId('Microsoft.MachineLearningServices/workspaces', parameters('workspaceName')), 'azureml ds role')]",
                             "scope": "[resourceId(resourceGroup().name, 'Microsoft.MachineLearningServices/workspaces', parameters('workspaceName'))]",
+                            "location": "[parameters('location')]",
                             "properties": {
                                 "roleDefinitionId": "[subscriptionResourceId('Microsoft.Authorization/roleDefinitions', 'f6c7c914-8db3-469d-8ca1-694a8f32e121')]",
                                 "principalId": "[if(not(equals(parameters('materializationIdentityOption'), 'none')), reference(variables('materializationIdentity'), '2018-11-30').principalId, '')]",
@@ -1085,6 +1086,7 @@
                             "apiVersion": "2022-04-01",
                             "name": "[guid(variables('materializationIdentity'), variables('offlineStore'), 'storage blob data contributor')]",
                             "scope": "[variables('offlineStore')]",
+                            "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, '')]",
@@ -1116,6 +1118,7 @@
                             "apiVersion": "2022-04-01",
                             "name": "[guid(variables('materializationIdentity'), parameters('online_store_resource_id'), 'contributor')]",
                             "scope": "[parameters('online_store_resource_id')]",
+                            "location": "[parameters('location')]",
                             "properties": {
                                 "roleDefinitionId": "[subscriptionResourceId('Microsoft.Authorization/roleDefinitions', 'b24988ac-6180-42a0-ab88-20f7382dd24c')]",
                                 "principalId": "[if(not(equals(parameters('materializationIdentityOption'), 'none')), reference(variables('materializationIdentity'), '2023-01-31').principalId, '')]",

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

@@ -783,6 +783,7 @@ def _populate_feature_store_role_assignment_parameters(
         _set_val(param["workspace_name"], workspace.name)
         resource_group = kwargs.get("resource_group", workspace.resource_group)
         _set_val(param["resource_group_name"], resource_group)
+        _set_val(param["location"], workspace.location)
 
         update_workspace_role_assignment = kwargs.get("update_workspace_role_assignment", None)
         if update_workspace_role_assignment:

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

@@ -349,12 +349,13 @@ def test_populate_feature_store_arm_parameters(
         )
 
         # test create feature store
-        feature_store = FeatureStore(name="name", resource_group="rg")
+        feature_store = FeatureStore(name="name", resource_group="rg", location="eastus2euap")
         template, param, _ = mock_workspace_operation_base._populate_arm_parameters(
             workspace=feature_store, grant_materialization_permissions=True
         )
 
         assert param["kind"]["value"] == "featurestore"
+        assert param["location"]["value"] == "eastus2euap"
         assert param["grant_materialization_permissions"]["value"] == "true"
         assert param["materializationIdentityOption"]["value"] == "new"
         assert param["materialization_identity_name"]["value"].startswith("materialization-uai-")
@@ -499,7 +500,7 @@ def test_populate_feature_store_role_assignments_paramaters(
             return_value=("random_id", True),
         )
         template, param, _ = mock_workspace_operation_base._populate_feature_store_role_assignment_parameters(
-            workspace=FeatureStore(name="name"),
+            workspace=FeatureStore(name="name", location="eastus2euap"),
             materialization_identity_id="/subscriptions/sub/resourcegroups/rg/providers/Microsoft.ManagedIdentity/userAssignedIdentities/identity",
             offline_store_target="/subscriptions/sub/resourceGroups/rg/providers/Microsoft.Storage/storageAccounts/test_storage/blobServices/default/containers/offlinestore",
             online_store_target="/subscriptions/sub1/resourceGroups/mdctest/providers/Microsoft.Cache/Redis/onlinestore",
@@ -509,6 +510,7 @@ def test_populate_feature_store_role_assignments_paramaters(
         )
 
         assert template is not None
+        assert param["location"] == {"value": "eastus2euap"}
         assert param["materialization_identity_resource_id"] == {
             "value": "/subscriptions/sub/resourcegroups/rg/providers/Microsoft.ManagedIdentity/userAssignedIdentities/identity"
         }