Support orphaned resources cleanup for Azure

mkdocs/docs/reference/env.md

@@ -151,7 +151,7 @@ For more details on the options below, refer to the [server deployment](../guide
 - `DSTACK_SERVER_GATEWAY_MAX_REPLICAS`{ #DSTACK_SERVER_GATEWAY_MAX_REPLICAS } - Maximum number of replicas allowed for a gateway. Defaults to `9`.
 - `DSTACK_SERVER_INSTANCE_HEALTH_MIN_COLLECT_INTERVAL_SECONDS`{ #DSTACK_SERVER_INSTANCE_HEALTH_MIN_COLLECT_INTERVAL_SECONDS } – Minimum time interval between consecutive health checks of the same instance.
 - `DSTACK_SERVER_EVENTS_TTL_SECONDS`{ #DSTACK_SERVER_EVENTS_TTL_SECONDS } - Maximum age of event records. Set to `0` to disable event storage. Defaults to 30 days.
-- `DSTACK_SERVER_ORPHAN_CLEANUP_MODE`{ #DSTACK_SERVER_ORPHAN_CLEANUP_MODE } – Controls the cleanup of cloud resources that `dstack` created but no longer tracks, e.g. after the server crashed during provisioning or failed to terminate an instance. `report` emits events about such resources, `delete` also deletes them, and `off` disables the cleanup. Resources are recognized by the tags (or names) `dstack` sets on them, so a resource created by copying the tags of a `dstack` instance is treated as an orphan and deleted in `delete` mode. Defaults to `report`. Currently supported for the `aws`, `gcp`, and `runpod` backends.
+- `DSTACK_SERVER_ORPHAN_CLEANUP_MODE`{ #DSTACK_SERVER_ORPHAN_CLEANUP_MODE } – Controls the cleanup of cloud resources that `dstack` created but no longer tracks, e.g. after the server crashed during provisioning or failed to terminate an instance. `report` emits events about such resources, `delete` also deletes them, and `off` disables the cleanup. Resources are recognized by the tags (or names) `dstack` sets on them, so a resource created by copying the tags of a `dstack` instance is treated as an orphan and deleted in `delete` mode. Defaults to `report`. Currently supported for the `aws`, `azure`, `gcp`, and `runpod` backends.
 - `DSTACK_SERVER_ORPHAN_CLEANUP_INTERVAL_SECONDS`{ #DSTACK_SERVER_ORPHAN_CLEANUP_INTERVAL_SECONDS } – Minimum interval between orphan cleanups of the same backend. A backend is cleaned up only after instances were provisioned or terminated in its project. Defaults to `3600`.
 - `DSTACK_SERVER_DEFAULT_DOCKER_REGISTRY`{ #DSTACK_SERVER_DEFAULT_DOCKER_REGISTRY } – A default Docker registry to use for job images that do not specify an explicit registry. E.g., if set to `registry.example`, then `image: ubuntu` becomes equivalent to `image: registry.example/ubuntu`. **Note**: This setting should only be used for configuring registries that act as a pull-through cache for Docker Hub. The default `dstack` images are also pulled from the configured registry.
 - `DSTACK_SERVER_DEFAULT_DOCKER_REGISTRY_USERNAME`{ #DSTACK_SERVER_DEFAULT_DOCKER_REGISTRY_USERNAME } – Username for authenticating with the default Docker registry. See `DSTACK_SERVER_DEFAULT_DOCKER_REGISTRY_PASSWORD`.

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

@@ -5,7 +5,12 @@
 from typing import Dict, List, Optional, Tuple
 
 from azure.core.credentials import TokenCredential
-from azure.core.exceptions import HttpResponseError, ResourceExistsError, ResourceNotFoundError
+from azure.core.exceptions import (
+    AzureError,
+    HttpResponseError,
+    ResourceExistsError,
+    ResourceNotFoundError,
+)
 from azure.mgmt import compute as compute_mgmt
 from azure.mgmt import network as network_mgmt
 from azure.mgmt.compute.models import (
@@ -37,24 +42,32 @@
 from dstack._internal.core.backends.azure import utils as azure_utils
 from dstack._internal.core.backends.azure.models import AzureConfig
 from dstack._internal.core.backends.base.compute import (
+    PROVIDER_REF_TAG_KEY,
     Compute,
     ComputeWithAllOffersCached,
     ComputeWithCreateInstanceSupport,
     ComputeWithGatewaySupport,
     ComputeWithInstanceVolumesSupport,
     ComputeWithMultinodeSupport,
     ComputeWithPrivilegedSupport,
+    ComputeWithResourceListingSupport,
     generate_unique_gateway_instance_name,
     generate_unique_instance_name,
     get_gateway_user_data,
     get_user_data,
+    is_valid_provider_ref,
     merge_tags,
     requires_nvidia_proprietary_kernel_modules,
 )
 from dstack._internal.core.backends.base.offers import get_catalog_offers
 from dstack._internal.core.consts import DSTACK_OS_IMAGE_WITH_PROPRIETARY_NVIDIA_KERNEL_MODULES
 from dstack._internal.core.errors import ComputeError, NoCapacityError
 from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.backends.orphans import (
+    ListedResourceGroup,
+    OwnerType,
+    ResourceListing,
+)
 from dstack._internal.core.models.gateways import (
     GatewayReplicaConfiguration,
     GatewayReplicaProvisioningData,
@@ -69,6 +82,7 @@
 from dstack._internal.core.models.placement import PlacementGroup
 from dstack._internal.core.models.resources import Memory, Range
 from dstack._internal.core.models.runs import JobProvisioningData
+from dstack._internal.utils.common import get_or_error
 from dstack._internal.utils.logging import get_logger
 
 logger = get_logger(__name__)
@@ -84,6 +98,7 @@ class AzureCompute(
     ComputeWithInstanceVolumesSupport,
     ComputeWithMultinodeSupport,
     ComputeWithGatewaySupport,
+    ComputeWithResourceListingSupport,
     Compute,
 ):
     def __init__(self, config: AzureConfig, credential: TokenCredential):
@@ -160,6 +175,8 @@ def create_instance(
             "dstack_name": instance_config.instance_name,
             "dstack_user": instance_config.user,
         }
+        if instance_config.provider_ref is not None:
+            base_tags[PROVIDER_REF_TAG_KEY] = instance_config.provider_ref
         tags = merge_tags(
             base_tags=base_tags,
             backend_tags=self.config.tags,
@@ -227,6 +244,60 @@ def terminate_instance(
             instance_name=instance_id,
         )
 
+    def list_resource_groups(self) -> ResourceListing:
+        """
+        Lists dstack resource groups (`ListedResourceGroup`) for orphan cleanup.
+        Not related to Azure resource groups.
+
+        Each group is a VM tagged with `provider_ref`, listed across all locations of the
+        configured Azure resource group. Groups contain only VMs since Azure deletes the OS disk,
+        NIC, and public IP created with a VM together with the VM.
+        """
+        try:
+            vms = list(
+                self._compute_client.virtual_machines.list(
+                    resource_group_name=self.config.resource_group
+                )
+            )
+        except AzureError as e:
+            raise ComputeError(f"Failed to list VMs: {e}") from e
+        resource_groups = []
+        for vm in vms:
+            provider_ref = _get_vm_provider_ref(vm)
+            if provider_ref is None:
+                continue
+            if (vm.provisioning_state or "").lower() == "deleting":
+                continue
+            resource_groups.append(
+                ListedResourceGroup(
+                    backend=BackendType.AZURE,
+                    owner_type=OwnerType.INSTANCE,
+                    provider_ref=provider_ref,
+                    primary_id=vm.name,
+                    region=vm.location,
+                    created_at=vm.time_created,
+                    description=f"VM {vm.name} ({_get_vm_size(vm)}, {vm.provisioning_state})",
+                    backend_data=None,
+                )
+            )
+        return ResourceListing(resource_groups=resource_groups)
+
+    def delete_resource_group(self, group: ListedResourceGroup) -> None:
+        """
+        Deletes the VM of a dstack resource group (`ListedResourceGroup`).
+        Not related to Azure resource groups.
+        """
+        try:
+            _terminate_instance(
+                compute_client=self._compute_client,
+                resource_group=self.config.resource_group,
+                instance_name=get_or_error(group.primary_id),
+            )
+        except ResourceNotFoundError:
+            pass
+        except AzureError as e:
+            raise ComputeError(f"Failed to delete resource group {group.provider_ref}: {e}") from e
+
     def create_gateway_replica(
         self,
         configuration: GatewayReplicaConfiguration,
@@ -634,6 +705,8 @@ def _begin_create_instance(
                 ),
                 network_profile=NetworkProfile(
                     network_api_version=network_mgmt.NetworkManagementClient.DEFAULT_API_VERSION,
+                    # Azure deletes the NIC and public IP created from these configurations
+                    # together with the VM.
                     network_interface_configurations=[
                         VirtualMachineNetworkInterfaceConfiguration(
                             name="nic_config",
@@ -792,6 +865,31 @@ def _get_vm_public_private_ips(
     return public_ip.ip_address, private_ip
 
 
+def _get_vm_provider_ref(vm: VirtualMachine) -> Optional[str]:
+    provider_ref = (vm.tags or {}).get(PROVIDER_REF_TAG_KEY)
+    if provider_ref is None:
+        return None
+    if not is_valid_provider_ref(provider_ref):
+        logger.warning(
+            "Skipping Azure VM %s with invalid %s tag value %r",
+            vm.id,
+            PROVIDER_REF_TAG_KEY,
+            provider_ref,
+        )
+        return None
+    return provider_ref
+
+
+def _get_vm_size(vm: VirtualMachine) -> Optional[str]:
+    if vm.hardware_profile is None:
+        return None
+    vm_size = vm.hardware_profile.vm_size
+    # Sizes known to the SDK are `VirtualMachineSizeTypes`, which format as enum names
+    if isinstance(vm_size, enum.Enum):
+        return vm_size.value
+    return vm_size
+
+
 def _terminate_instance(
     compute_client: compute_mgmt.ComputeManagementClient,
     resource_group: str,

src/tests/_internal/core/backends/azure/test_compute.py

@@ -1,18 +1,34 @@
+from datetime import datetime, timezone
+from typing import Optional
 from unittest.mock import Mock
 
 import pytest
-from azure.core.exceptions import ODataV4Format, ResourceNotFoundError
-from azure.mgmt.compute.models import ImageReference
+from azure.core.exceptions import HttpResponseError, ODataV4Format, ResourceNotFoundError
+from azure.mgmt.compute.models import (
+    HardwareProfile,
+    ImageReference,
+    VirtualMachine,
+    VirtualMachineProperties,
+    VirtualMachineSizeTypes,
+)
 
 from dstack._internal import settings
 from dstack._internal.core.backends.azure.compute import (
+    AzureCompute,
     VMImageVariant,
     _begin_create_instance,
     _get_image_ref,
 )
+from dstack._internal.core.backends.azure.models import AzureConfig, AzureDefaultCreds
 from dstack._internal.core.errors import ComputeError
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.backends.orphans import ListedResourceGroup, OwnerType
 from dstack._internal.core.models.instances import Gpu, InstanceType, Resources
 
+RESOURCE_GROUP = "dstack-resources"
+PROVIDER_REF = "abc123def456"
+TIME_CREATED = datetime(2026, 1, 1, 12, tzinfo=timezone.utc)
+
 
 class TestVMImageVariant:
     @pytest.mark.parametrize(
@@ -147,3 +163,127 @@ def _create_instance_kwargs(compute_client: Mock, image_reference: ImageReferenc
         disk_size=100,
         computer_name="runnervm",
     )
+
+
+@pytest.fixture
+def compute_client() -> Mock:
+    return Mock()
+
+
+@pytest.fixture
+def compute(compute_client: Mock) -> AzureCompute:
+    compute = AzureCompute(
+        AzureConfig(
+            tenant_id="tenant",
+            subscription_id="subscription",
+            resource_group=RESOURCE_GROUP,
+            regions=["westeurope"],
+            creds=AzureDefaultCreds(),
+        ),
+        credential=Mock(),
+    )
+    compute._compute_client = compute_client
+    return compute
+
+
+class TestListResourceGroups:
+    def test_lists_tagged_vms(self, compute: AzureCompute, compute_client: Mock):
+        _set_vms(compute_client, [_get_vm("vm-1")])
+        listing = compute.list_resource_groups()
+        compute_client.virtual_machines.list.assert_called_once_with(
+            resource_group_name=RESOURCE_GROUP
+        )
+        assert listing.failures == []
+        assert listing.resource_groups == [
+            ListedResourceGroup(
+                backend=BackendType.AZURE,
+                owner_type=OwnerType.INSTANCE,
+                provider_ref=PROVIDER_REF,
+                primary_id="vm-1",
+                region="westeurope",
+                created_at=TIME_CREATED,
+                description="VM vm-1 (Standard_F2s_v2, Succeeded)",
+                backend_data=None,
+            )
+        ]
+
+    def test_skips_vms_without_provider_ref(self, compute: AzureCompute, compute_client: Mock):
+        _set_vms(compute_client, [_get_vm("vm-1", provider_ref=None)])
+        assert compute.list_resource_groups().resource_groups == []
+
+    def test_skips_vms_with_invalid_provider_ref(
+        self, compute: AzureCompute, compute_client: Mock
+    ):
+        _set_vms(compute_client, [_get_vm("vm-1", provider_ref="Invalid-Ref")])
+        assert compute.list_resource_groups().resource_groups == []
+
+    def test_skips_vms_being_deleted(self, compute: AzureCompute, compute_client: Mock):
+        _set_vms(
+            compute_client,
+            [_get_vm("vm-1", provisioning_state="Deleting"), _get_vm("vm-2")],
+        )
+        groups = compute.list_resource_groups().resource_groups
+        assert [g.primary_id for g in groups] == ["vm-2"]
+
+    def test_raises_compute_error_on_api_error(self, compute: AzureCompute, compute_client: Mock):
+        compute_client.virtual_machines.list.side_effect = HttpResponseError("error")
+        with pytest.raises(ComputeError):
+            compute.list_resource_groups()
+
+
+class TestDeleteResourceGroup:
+    def test_deletes_vm(self, compute: AzureCompute, compute_client: Mock):
+        compute.delete_resource_group(_get_resource_group())
+        compute_client.virtual_machines.begin_delete.assert_called_once_with(
+            resource_group_name=RESOURCE_GROUP, vm_name="vm-1"
+        )
+
+    def test_tolerates_deleted_vm(self, compute: AzureCompute, compute_client: Mock):
+        compute_client.virtual_machines.begin_delete.side_effect = _not_found_error(
+            "ResourceNotFound"
+        )
+        compute.delete_resource_group(_get_resource_group())
+
+    def test_raises_compute_error_on_api_error(self, compute: AzureCompute, compute_client: Mock):
+        compute_client.virtual_machines.begin_delete.side_effect = HttpResponseError("error")
+        with pytest.raises(ComputeError):
+            compute.delete_resource_group(_get_resource_group())
+
+
+def _get_vm(
+    name: str,
+    provider_ref: Optional[str] = PROVIDER_REF,
+    provisioning_state: str = "Succeeded",
+) -> VirtualMachine:
+    tags = {"owner": "dstack"}
+    if provider_ref is not None:
+        tags["dstack_provider_ref"] = provider_ref
+    vm = VirtualMachine(
+        location="westeurope",
+        tags=tags,
+        properties=VirtualMachineProperties(
+            hardware_profile=HardwareProfile(vm_size=VirtualMachineSizeTypes.STANDARD_F2_S_V2)
+        ),
+    )
+    vm.id = f"/subscriptions/subscription/resourceGroups/{RESOURCE_GROUP}/providers/Microsoft.Compute/virtualMachines/{name}"
+    vm.name = name
+    vm.provisioning_state = provisioning_state
+    vm.time_created = TIME_CREATED
+    return vm
+
+
+def _set_vms(compute_client: Mock, vms: list[VirtualMachine]):
+    compute_client.virtual_machines.list.return_value = iter(vms)
+
+
+def _get_resource_group() -> ListedResourceGroup:
+    return ListedResourceGroup(
+        backend=BackendType.AZURE,
+        owner_type=OwnerType.INSTANCE,
+        provider_ref=PROVIDER_REF,
+        primary_id="vm-1",
+        region="westeurope",
+        created_at=TIME_CREATED,
+        description="VM vm-1",
+        backend_data=None,
+    )