fix role assignment when update feature store (#35751)

Azure · May 24, 2024 · 8e7522f · 8e7522f
1 parent ca6b8ff
commit 8e7522f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
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
@@ -383,7 +383,7 @@ def callback(_: Any, deserialized: Any, args: Any) -> Workspace:
             ):
                 module_logger.info("updating feature store materialization identity role assignments..")
                 template, param, resources_being_deployed = self._populate_feature_store_role_assignment_parameters(
-                    workspace, resource_group=resource_group, **kwargs
+                    workspace, resource_group=resource_group, location=existing_workspace.location, **kwargs
                 )
 
                 arm_submit = ArmDeploymentExecutor(
@@ -780,7 +780,8 @@ 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)
+        location = kwargs.get("location", workspace.location)
+        _set_val(param["location"], location)
 
         update_workspace_role_assignment = kwargs.get("update_workspace_role_assignment", None)
         if update_workspace_role_assignment:

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
@@ -500,13 +500,14 @@ 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", location="eastus2euap"),
+            workspace=FeatureStore(name="name"),
             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",
             update_workspace_role_assignment=True,
             update_offline_store_role_assignment=True,
             update_online_store_role_assignment=True,
+            location="eastus2euap",
         )
 
         assert template is not None