update role assignment removal logic

skypilot-org · Jul 31, 2024 · 15b4d75 · 15b4d75
1 parent ee77e4e
commit 15b4d75
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 20 deletions.
diff --git a/sky/clouds/azure.py b/sky/clouds/azure.py
@@ -384,9 +384,6 @@ def _failover_disk_tier() -> Optional[resources_utils.DiskTier]:
             'azure_subscription_id': self.get_project_id(dryrun),
             'resource_group': resource_group_name,
             'use_external_resource_group': use_external_resource_group,
-            'role_assignment_name': (
-                provision_constants.ROLE_ASSIGNMENT_NAME.format(
-                    cluster_name_on_cloud=cluster_name.name_on_cloud))
         }
 
     def _get_feasible_launchable_resources(

diff --git a/sky/provision/azure/azure-config-template.json b/sky/provision/azure/azure-config-template.json
@@ -13,19 +13,13 @@
             "metadata": {
                 "description": "Subnet parameters."
             }
-        },
-        "roleAssignmentName": {
-            "type": "string",
-            "metadata": {
-                "description": "Raw name of the Role Assignment created before it is passed to guid."
-            }
         }
     },
     "variables": {
         "contributor": "[subscriptionResourceId('Microsoft.Authorization/roleDefinitions', 'b24988ac-6180-42a0-ab88-20f7382dd24c')]",
         "location": "[resourceGroup().location]",
         "msiName": "[concat('sky-', parameters('clusterId'), '-msi')]",
-        "roleAssignmentName": "[parameters('roleAssignmentName')]",
+        "roleAssignmentName": "[concat('sky-', parameters('clusterId'), '-ra')]",
         "nsgName": "[concat('sky-', parameters('clusterId'), '-nsg')]",
         "nsg": "[resourceId('Microsoft.Network/networkSecurityGroups', variables('nsgName'))]",
         "vnetName": "[concat('sky-', parameters('clusterId'), '-vnet')]",

diff --git a/sky/provision/azure/config.py b/sky/provision/azure/config.py
@@ -127,10 +127,6 @@ def bootstrap_instances(
                     # as we have already appended the user hash to the cluster
                     # name.
                     'value': cluster_name_on_cloud
-                },
-                'roleAssignmentName': {
-                    'value': constants.ROLE_ASSIGNMENT_NAME.format(
-                        cluster_name_on_cloud=cluster_name_on_cloud)
                 }
             },
         }

diff --git a/sky/provision/azure/instance.py b/sky/provision/azure/instance.py
@@ -620,7 +620,6 @@ def terminate_instances(
 
     use_external_resource_group = provider_config.get(
         'use_external_resource_group', False)
-
     # When user specified resource group through config.yaml to create a VM, we
     # cannot remove the entire resource group as it may contain other resources
     # unrelated to this VM being removed.
@@ -798,15 +797,23 @@ def delete_vm_and_attached_resources(
                 resource_group_name=resource_group))
         for identity in user_assigned_identities:
             if msi_name == identity.name:
+                # We use the principal_id to find the correct guid converted
+                # role assignment name because each managed identity has a
+                # unique principal_id, and role assignments are associated
+                # with security principals (like managed identities) via this
+                # principal_id.
                 target_principal_id = identity.principal_id
-                scope = f'/subscriptions/{subscription_id}/resourceGroups/{resource_group}'
-                # List role assignments for the specified scope
-                role_assignments = auth_client.role_assignments.list_for_scope(scope)
+                scope = (f'/subscriptions/{subscription_id}'
+                         f'/resourceGroups/{resource_group}')
+                role_assignments = auth_client.role_assignments.list_for_scope(
+                    scope)
                 for assignment in role_assignments:
                     if target_principal_id == assignment.principal_id:
                         guid_role_assignment_name = assignment.name
                         delete_role_assignment(
-                            scope=scope, role_assignment_name=guid_role_assignment_name)
+                            scope=scope,
+                            role_assignment_name=guid_role_assignment_name)
+                        break
         delete_managed_identity(resource_group_name=resource_group,
                                 resource_name=msi_name)
 

diff --git a/sky/provision/constants.py b/sky/provision/constants.py
@@ -22,4 +22,3 @@
 LEGACY_DEPLOYMENT_NAME = 'ray-config'
 EXTERNAL_RG_BOOTSTRAP_DEPLOYMENT_NAME = 'skypilot-bootstrap-config-{cluster_name_on_cloud}'
 EXTERNAL_RG_VM_DEPLOYMENT_NAME = 'skypilot-vm-config-{cluster_name_on_cloud}'
-ROLE_ASSIGNMENT_NAME = 'sky-{cluster_name_on_cloud}-ra'