Support vm_managed_identity for Azure (#2608)

Author: r4victor

docs/docs/concepts/backends.md

@@ -338,6 +338,8 @@ There are two ways to configure Azure: using a client secret or using the defaul
                     "Microsoft.Compute/disks/write",
                     "Microsoft.Compute/disks/read",
                     "Microsoft.Compute/disks/delete",
+                    "Microsoft.ManagedIdentity/userAssignedIdentities/assign/action",
+                    "Microsoft.ManagedIdentity/userAssignedIdentities/read",
                     "Microsoft.Network/networkSecurityGroups/*",
                     "Microsoft.Network/locations/*",
                     "Microsoft.Network/virtualNetworks/*",

pyproject.toml

@@ -139,6 +139,7 @@ azure = [
     "azure-mgmt-network>=23.0.0,<28.0.0",
     "azure-mgmt-resource>=22.0.0",
     "azure-mgmt-authorization>=3.0.0",
+    "azure-mgmt-msi>=7.0.0",
     "dstack[server]",
 ]
 gcp = [

src/dstack/_internal/core/backends/azure/compute.py

@@ -138,6 +138,10 @@ def create_instance(
             location=location,
         )
 
+        managed_identity_resource_group, managed_identity_name = parse_vm_managed_identity(
+            self.config.vm_managed_identity
+        )
+
         base_tags = {
             "owner": "dstack",
             "dstack_project": instance_config.project_name,
@@ -161,7 +165,8 @@ def create_instance(
             network_security_group=network_security_group,
             network=network,
             subnet=subnet,
-            managed_identity=None,
+            managed_identity_name=managed_identity_name,
+            managed_identity_resource_group=managed_identity_resource_group,
             image_reference=_get_image_ref(
                 compute_client=self._compute_client,
                 location=location,
@@ -257,7 +262,8 @@ def create_gateway(
             network_security_group=network_security_group,
             network=network,
             subnet=subnet,
-            managed_identity=None,
+            managed_identity_name=None,
+            managed_identity_resource_group=None,
             image_reference=_get_gateway_image_ref(),
             vm_size="Standard_B1ms",
             instance_name=instance_name,
@@ -340,6 +346,21 @@ def get_resource_group_network_subnet_or_error(
     return resource_group, network_name, subnet_name
 
 
+def parse_vm_managed_identity(
+    vm_managed_identity: Optional[str],
+) -> Tuple[Optional[str], Optional[str]]:
+    if vm_managed_identity is None:
+        return None, None
+    try:
+        resource_group, managed_identity = vm_managed_identity.split("/")
+        return resource_group, managed_identity
+    except Exception:
+        raise ComputeError(
+            "`vm_managed_identity` specified in incorrect format."
+            " Supported format: 'managedIdentityResourceGroup/managedIdentityName'"
+        )
+
+
 def _parse_config_vpc_id(vpc_id: str) -> Tuple[str, str]:
     resource_group, network_name = vpc_id.split("/")
     return resource_group, network_name
@@ -468,7 +489,8 @@ def _launch_instance(
     network_security_group: str,
     network: str,
     subnet: str,
-    managed_identity: Optional[str],
+    managed_identity_name: Optional[str],
+    managed_identity_resource_group: Optional[str],
     image_reference: ImageReference,
     vm_size: str,
     instance_name: str,
@@ -490,6 +512,20 @@ def _launch_instance(
         public_ip_address_configuration = VirtualMachinePublicIPAddressConfiguration(
             name="public_ip_config",
         )
+    managed_identity = None
+    if managed_identity_name is not None:
+        if managed_identity_resource_group is None:
+            managed_identity_resource_group = resource_group
+        managed_identity = VirtualMachineIdentity(
+            type=ResourceIdentityType.USER_ASSIGNED,
+            user_assigned_identities={
+                azure_utils.get_managed_identity_id(
+                    subscription_id,
+                    managed_identity_resource_group,
+                    managed_identity_name,
+                ): UserAssignedIdentitiesValue(),
+            },
+        )
     try:
         poller = compute_client.virtual_machines.begin_create_or_update(
             resource_group,
@@ -554,16 +590,7 @@ def _launch_instance(
                 ),
                 priority="Spot" if spot else "Regular",
                 eviction_policy="Delete" if spot else None,
-                identity=None
-                if managed_identity is None
-                else VirtualMachineIdentity(
-                    type=ResourceIdentityType.USER_ASSIGNED,
-                    user_assigned_identities={
-                        azure_utils.get_managed_identity_id(
-                            subscription_id, resource_group, managed_identity
-                        ): UserAssignedIdentitiesValue()
-                    },
-                ),
+                identity=managed_identity,
                 user_data=base64.b64encode(user_data.encode()).decode(),
                 tags=tags,
             ),

src/dstack/_internal/core/backends/azure/configurator.py

@@ -4,6 +4,7 @@
 
 import azure.core.exceptions
 from azure.core.credentials import TokenCredential
+from azure.mgmt import msi as msi_mgmt
 from azure.mgmt import network as network_mgmt
 from azure.mgmt import resource as resource_mgmt
 from azure.mgmt import subscription as subscription_mgmt
@@ -97,6 +98,7 @@ def validate_config(self, config: AzureBackendConfigWithCreds, default_creds_ena
         self._check_config_locations(config)
         self._check_config_tags(config)
         self._check_config_resource_group(config=config, credential=credential)
+        self._check_config_vm_managed_identity(config=config, credential=credential)
         self._check_config_vpc(config=config, credential=credential)
 
     def create_backend(
@@ -260,6 +262,25 @@ def _check_config_vpc(
                     except BackendError as e:
                         raise ServerClientError(e.args[0])
 
+    def _check_config_vm_managed_identity(
+        self, config: AzureBackendConfigWithCreds, credential: auth.AzureCredential
+    ):
+        try:
+            resource_group, identity_name = compute.parse_vm_managed_identity(
+                config.vm_managed_identity
+            )
+        except BackendError as e:
+            raise ServerClientError(e.args[0])
+        if resource_group is None or identity_name is None:
+            return
+        msi_client = msi_mgmt.ManagedServiceIdentityClient(credential, config.subscription_id)
+        try:
+            msi_client.user_assigned_identities.get(resource_group, identity_name)
+        except azure.core.exceptions.ResourceNotFoundError:
+            raise ServerClientError(
+                f"Managed identity {identity_name} not found in resource group {resource_group}"
+            )
+
     def _set_client_creds_tenant_id(
         self,
         creds: AzureClientCreds,

src/dstack/_internal/core/backends/azure/models.py

@@ -62,6 +62,15 @@ class AzureBackendConfig(CoreModel):
             )
         ),
     ] = None
+    vm_managed_identity: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "The managed identity to associate with provisioned VMs."
+                " Must have a format `managedIdentityResourceGroup/managedIdentityName`"
+            )
+        ),
+    ] = None
     tags: Annotated[
         Optional[Dict[str, str]],
         Field(description="The tags that will be assigned to resources created by `dstack`"),