diff --git a/mkdocs/docs/reference/env.md b/mkdocs/docs/reference/env.md index 88279f089..c54a6b374 100644 --- a/mkdocs/docs/reference/env.md +++ b/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`. diff --git a/src/dstack/_internal/core/backends/azure/compute.py b/src/dstack/_internal/core/backends/azure/compute.py index 3744f7dfb..089ffc122 100644 --- a/src/dstack/_internal/core/backends/azure/compute.py +++ b/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,6 +42,7 @@ 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, @@ -44,10 +50,12 @@ 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, ) @@ -55,6 +63,11 @@ 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, diff --git a/src/tests/_internal/core/backends/azure/test_compute.py b/src/tests/_internal/core/backends/azure/test_compute.py index 039984ef2..7ee8397aa 100644 --- a/src/tests/_internal/core/backends/azure/test_compute.py +++ b/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, + )