diff --git a/AGENTS.md b/AGENTS.md index 234464fbd1..c9f86f765b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -2,8 +2,8 @@ Before touching a subsystem, read the relevant notes in `contributing/`: `ARCHITECTURE.md`, `PIPELINES.md`, `LOCKING.md`, `DATABASE.md`, `MIGRATIONS.md`, `RUNS-AND-JOBS.md`, `AUTOSCALING.md`, -`BACKENDS.md`, `GPUHUNT.md`, `PROXY.md`, `RUNNER-AND-SHIM.md`, `FRONTEND.md`, `DOCS.md`, -`DEVELOPMENT.md`, `RELEASE.md`. +`BACKENDS.md`, `ORPHAN-CLEANUP.md`, `GPUHUNT.md`, `PROXY.md`, `RUNNER-AND-SHIM.md`, `FRONTEND.md`, +`DOCS.md`, `DEVELOPMENT.md`, `RELEASE.md`. ## Project Structure & Module Organization - Core Python package lives in `src/dstack`; internal modules (including server) sit under `_internal`, API surfaces under `api`, and plugin integrations under `plugins`. diff --git a/contributing/BACKENDS.md b/contributing/BACKENDS.md index 90b56a55ab..a3ee705c05 100644 --- a/contributing/BACKENDS.md +++ b/contributing/BACKENDS.md @@ -144,6 +144,7 @@ Refer to examples: Go to `compute.py` and implement `Compute` methods. Optionally, extend and implement `ComputeWith*` classes to support additional features such as fleets, volumes, gateways, placement groups, etc. For example, extend `ComputeWithCreateInstanceSupport` to support fleets. +To support the cleanup of orphaned resources, extend `ComputeWithResourceListingSupport` as described in [Orphan cleanup support](#32-orphan-cleanup-support). Refer to examples: [verda](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/verda/compute.py), @@ -228,3 +229,36 @@ To support a container-based backend, `dstack` expects the following: - A way to override the container user to root (as in `docker run --user root ...`) Examples of container-based backends include: `kubernetes`, `vastai`, `runpod`. + +### 3.2. Orphan cleanup support + +`dstack` can find and delete cloud resources it created but no longer tracks, e.g. if the server +crashed after the provider created an instance but before it was recorded, or if instance +termination failed. The `OrphanCleanupPipeline` lists resources of backends that extend +`ComputeWithResourceListingSupport` and matches them with instances by `provider_ref`. +See [ORPHAN-CLEANUP.md](ORPHAN-CLEANUP.md) for the server side. + +To support it in a backend: + +1. Set `provider_ref` on every resource created for an instance or a job. It's + `InstanceConfiguration.provider_ref` in `create_instance()` and `get_job_provider_ref(job)` in + `run_job()`. Internal resources that outlive the instance if left behind (disks, keys, IPs) + need it too, while resources deleted together with the instance don't. + * If the provider supports tags, set the `PROVIDER_REF_TAG_KEY` tag and keep instance names as is. + If the provider does not allow this key, use another one. + * Otherwise, encode `provider_ref` in names with `generate_unique_instance_name(..., with_provider_ref=True)` + and parse it with `parse_provider_ref_from_name()`. + * Do not set it on resources of `run_jobs()` since orphan cleanup does not support compute groups yet. +2. Implement `list_resource_groups()`. It returns a `ListedResourceGroup` per provisioning attempt: the instance + (`primary_id`) together with internal resources created for it. Skip resources without `provider_ref`. + List resources of each region all or nothing: if any listing in a region fails, return no groups + of the region and report it in `ResourceListing.failures`. Otherwise, a tracked instance's resources + may be listed without the instance and deleted. Skip instances whose deletion is in progress + (e.g. shutting down), but list their internal resources. +3. Implement `delete_resource_group()`. Delete everything in the group in the order the provider requires. + The method must be idempotent and never delete resources that don't carry the group's `provider_ref`, + such as attached volumes. If the deletion takes several steps, delete the instance first and raise + `NotYetTerminated` instead of waiting for the next step, e.g. for the instance to be gone before its + disk can be deleted. Next cleanups list the remaining internal resources and call the method again. + +Refer to the [aws](https://github.com/dstackai/dstack/blob/master/src/dstack/_internal/core/backends/aws/compute.py) example. diff --git a/contributing/MIGRATIONS.md b/contributing/MIGRATIONS.md index f494d829d5..3f00e12abb 100644 --- a/contributing/MIGRATIONS.md +++ b/contributing/MIGRATIONS.md @@ -73,3 +73,24 @@ with op.get_context().autocommit_block(): postgresql_concurrently=True, ) ``` + +Keep concurrent index creation in a separate migration. Do not combine `autocommit_block()` +with other schema changes such as `add_column`/`batch_alter_table` in the same migration. +`autocommit_block()` commits the migration's transaction before running the index statements, +so a failure inside the block leaves the earlier changes committed while the migration is not +marked as applied. A retry then re-runs those changes and fails. For example: + +```python +def upgrade() -> None: + with op.batch_alter_table("table", schema=None) as batch_op: + batch_op.add_column(sa.Column("col", sa.String(length=12), nullable=True)) + with op.get_context().autocommit_block(): + # If this fails, "col" is already committed. A retry fails on add_column + # with "column already exists". + op.create_index("ix_table_col", "table", ["col"], postgresql_concurrently=True) +``` + +Instead, add the column in one migration and create the index in the next one. + +Indexes on a table created in the same migration do not need `CONCURRENTLY`, since the table is +empty. Create them without `autocommit_block()` so that the whole migration stays transactional. diff --git a/contributing/ORPHAN-CLEANUP.md b/contributing/ORPHAN-CLEANUP.md new file mode 100644 index 0000000000..5c04d67d4f --- /dev/null +++ b/contributing/ORPHAN-CLEANUP.md @@ -0,0 +1,57 @@ +# Orphan cleanup + +`dstack` can leave cloud resources it no longer tracks: the server crashes or loses its lock between a provider call that creates resources and the commit that records them, the provider returns an error after creating them, or termination gives up after retrying. Orphan cleanup finds such resources by listing them in the provider and matching them with the DB, then reports or deletes them. + +How a backend supports it is described in [BACKENDS.md](BACKENDS.md#32-orphan-cleanup-support). This doc covers the server side. + +## `provider_ref` + +Each record that provisions resources gets a `provider_ref` (12-char string). Backends set it on the resources they create (as a tag or in names) and report it back when listing. + +- Only instances own `provider_ref`s (`InstanceModel.provider_ref`). Fleet instances get it on creation, for `create_instance()`. Placeholder instances, created for job submissions before `run_job()`, get it on creation too, pass it to `run_job()` (`JobSubmission.provider_ref`), and keep it when promoted. +- Every issued `provider_ref` is also recorded in `ProviderRefModel`, committed with its owner before any provider call. Records are not deleted with owners, so resources of a deleted owner are still recognized. This matters for placeholder instances, which are deleted when a `run_job()` fails, possibly after creating resources. +- Compute groups (`run_jobs()`) are not supported: their instances get no `provider_ref`, and backends must not set it on compute group resources. + +## Cleanup unit and enrollment + +An `OrphanCleanupModel` row exists per project backend (`BackendModel`) that supports resource listing. A scheduled task periodically creates rows in batches for backends of projects that have issued `provider_ref`s. Enrolling periodically instead of in a migration covers backends created by any replica, including older replicas during rolling deployments. + +## `OrphanCleanupPipeline` + +The pipeline processes each row at most once per cleanup interval (`DSTACK_SERVER_ORPHAN_CLEANUP_INTERVAL_SECONDS`). Failed cleanups are retried with an increasing delay. + +Processing a row: + +- STEP 1: Check whether the backend may have new orphans: it was never cleaned up, the cleanup mode changed, or the project had instance or job activity since the last clean cleanup (`clean_at`). If not, the row is only marked processed. Checking activity here rather than in the fetch query keeps fetching cheap regardless of project sizes, at the cost of a small query and row update per enrolled backend per interval. +- STEP 2: List the backend's resource groups (`list_resource_groups()`). A resource group is one provisioning attempt as found in the provider: the owner resource (`primary_id`, the instance) and internal resources created with it. +- STEP 3: Load the project's `ProviderRefModel` records and their instances (including soft-deleted ones) by `provider_ref`, in one statement. An instance is committed with its record, so it can't be missed while the record is seen. +- STEP 4: Classify each group: + - Instance found: + - `PENDING`, `TERMINATING`: skip, since the resources may be in flight. An error is logged if a pending instance's resources look stuck. + - `TERMINATED`: orphan. + - `PROVISIONING`, `IDLE`, `BUSY`: tracked if the base backend and `primary_id` match the instance's provisioning data, orphan otherwise. + - Without provisioning data: ignored with an error log. + - Record found without an instance: orphan. The owner was deleted, e.g. a placeholder instance after a failed `run_job()`, so nothing can be creating its resources anymore. + - No record: ignored. Resources of other servers and projects are never touched. +- STEP 5: Owners that were provisioned or terminated around the listing are skipped instead of classified as orphans, since the listing may predate the change. The margin covers clock skew between replicas, the lag between setting timestamps and committing them, and providers listing new resources with a delay. +- STEP 6: Handle orphans according to `DSTACK_SERVER_ORPHAN_CLEANUP_MODE`: + - `report` (default): log, send to Sentry, and emit an event targeting the owner instance. If the instance was deleted, the event still targets it by the id and name kept in `ProviderRefModel`, so it shows next to the instance's earlier events. + - `delete`: call `delete_resource_group()` and emit an event. If the backend raises `NotYetTerminated` for a group with an owner resource, the owner's deletion was requested and the event says the resources are being deleted. The remaining internal resources are listed and deleted by next cleanups. + - `off`: the pipeline and enrollment are disabled. +- STEP 7: Set the outcome: + - Clean (everything listed, nothing skipped, all orphans handled): `clean_at` advances to the listing time minus the margin, and the retry delay is reset. + - Unsettled (a group skipped, a deletion in progress, or some regions in `ResourceListing.failures`): `clean_at` stays, the retry delay is reset, and the backend is retried after the interval. Listed groups are still handled. + - Failed (listing raised or a deletion failed): `clean_at` stays and the retry is backed off. + +Events are emitted in the lock-guarded apply, so a worker that lost the lock emits nothing. + +## Termination give-up + +When instance termination gives up after retrying, the log says that orphan cleanup will retry deleting the instance if the mode is `delete` and the backend supports listing. Otherwise, it asks to terminate the instance manually. + +## Limitations + +- Resources are not cleaned up if the backend credentials change or the backend is deleted. +- Compute groups, placement groups, volumes and gateways are not covered. `OwnerType` only has `INSTANCE`. +- `ProviderRefModel` records are never deleted yet. There is one per fleet instance and per placeholder instance, i.e. per job submission that provisions new capacity. +- On dstack Sky, each project's `dstack` backend lists the shared base backend accounts, so listings repeat per project. Cleaning up each shared account once is a possible optimization. diff --git a/mkdocs/docs/reference/env.md b/mkdocs/docs/reference/env.md index c721de395c..2984ead8ea 100644 --- a/mkdocs/docs/reference/env.md +++ b/mkdocs/docs/reference/env.md @@ -151,6 +151,8 @@ 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. Defaults to `report`. Currently supported for the `aws` backend. +- `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`. - `DSTACK_SERVER_DEFAULT_DOCKER_REGISTRY_PASSWORD`{ #DSTACK_SERVER_DEFAULT_DOCKER_REGISTRY_PASSWORD } – Password for authenticating with the default Docker registry. Applied only when the image has no explicit registry and the run configuration does not specify `registry_auth`. **Note**: The value may be visible to anyone who can SSH into instances managed by `dstack`, which usually includes all users of that `dstack` server. diff --git a/src/dstack/_internal/core/backends/aws/compute.py b/src/dstack/_internal/core/backends/aws/compute.py index c427703710..2631bf58b1 100644 --- a/src/dstack/_internal/core/backends/aws/compute.py +++ b/src/dstack/_internal/core/backends/aws/compute.py @@ -1,4 +1,5 @@ import threading +from collections import defaultdict from collections.abc import Container from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field @@ -20,6 +21,7 @@ AWSOSImageConfig, ) from dstack._internal.core.backends.base.compute import ( + PROVIDER_REF_TAG_KEY, Compute, ComputeCache, ComputeTTLCache, @@ -33,6 +35,7 @@ ComputeWithPrivateGatewaySupport, ComputeWithPrivilegedSupport, ComputeWithReservationSupport, + ComputeWithResourceListingSupport, ComputeWithVolumeSupport, generate_unique_gateway_instance_name, generate_unique_instance_name, @@ -40,6 +43,7 @@ generate_unique_volume_name, get_gateway_user_data, get_user_data, + is_valid_provider_ref, merge_tags, ) from dstack._internal.core.backends.base.offers import get_catalog_offers @@ -51,6 +55,12 @@ ProvisioningError, ) from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.backends.orphans import ( + ListedResourceGroup, + OwnerType, + ResourceListing, + ResourceListingFailure, +) from dstack._internal.core.models.common import CoreModel, validate_json_extra_ignore from dstack._internal.core.models.gateways import ( GatewayLoadBalancerConfiguration, @@ -108,6 +118,11 @@ class AWSInstanceBackendData(CoreModel): """ +class AWSResourceGroupBackendData(CoreModel): + instance_id: Optional[str] = None + eip_allocation_ids: List[str] = [] + + def _ec2client_cache_methodkey(self, ec2_client, *args, **kwargs): return hashkey(*args, **kwargs) @@ -129,6 +144,7 @@ class AWSCompute( ComputeWithGatewayLoadBalancerSupport, ComputeWithPrivateGatewaySupport, ComputeWithVolumeSupport, + ComputeWithResourceListingSupport, Compute, ): def __init__( @@ -253,11 +269,40 @@ def terminate_instance( raise e instance_backend_data = _parse_instance_backend_data(backend_data) if instance_backend_data.eip_allocation_id is not None: - _release_eip( + aws_resources.release_eip( ec2_client=ec2_client, allocation_id=instance_backend_data.eip_allocation_id, ) + def list_resource_groups(self) -> ResourceListing: + regions = get_or_error(self.config.regions) + listing = ResourceListing(resource_groups=[]) + for region in regions: + ec2_client = self.session.client("ec2", region_name=region) + try: + listing.resource_groups.extend(_list_resource_groups(ec2_client, region)) + except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e: + listing.failures.append( + ResourceListingFailure(backend=BackendType.AWS, region=region, message=str(e)) + ) + if len(listing.failures) == len(regions): + raise ComputeError(f"Failed to list resources in all regions: {listing.failures}") + return listing + + def delete_resource_group(self, group: ListedResourceGroup) -> None: + backend_data = validate_json_extra_ignore( + AWSResourceGroupBackendData, get_or_error(group.backend_data) + ) + region = get_or_error(group.region) + try: + if backend_data.instance_id is not None: + self.terminate_instance(instance_id=backend_data.instance_id, region=region) + ec2_client = self.session.client("ec2", region_name=region) + for allocation_id in backend_data.eip_allocation_ids: + aws_resources.release_eip(ec2_client=ec2_client, allocation_id=allocation_id) + except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e: + raise ComputeError(f"Failed to delete resource group {group.provider_ref}: {e}") from e + def create_instance( self, instance_offer: InstanceOfferWithAvailability, @@ -280,6 +325,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, @@ -460,10 +507,13 @@ def update_provisioning_data( # AWS can't auto-assign a public IPv4 to multi-ENI instances (multi-EFA instances). # When `public_ips: true` and no public IP is present after launch, attach an Elastic IP to the primary ENI. # The check relies on running instances always having IP assigned if ever. - public_ip, allocation_id = _allocate_and_associate_eip( + public_ip, allocation_id = aws_resources.allocate_and_associate_eip( ec2_client=ec2_client, instance=instance, - project_name=_get_project_name_from_instance_tags(instance), + project_name=aws_resources.get_tag_value(instance.tags or [], "dstack_project"), + provider_ref=aws_resources.get_tag_value( + instance.tags or [], PROVIDER_REF_TAG_KEY + ), backend_tags=self.config.tags, ) provisioning_data.backend_data = AWSInstanceBackendData( @@ -1490,127 +1540,98 @@ def _parse_instance_backend_data(backend_data: Optional[str]) -> "AWSInstanceBac return AWSInstanceBackendData() -def _get_project_name_from_instance_tags(instance: Any) -> Optional[str]: - for tag in instance.tags or []: - if tag.get("Key") == "dstack_project": - return tag.get("Value") - return None - - -def _allocate_and_associate_eip( - ec2_client: botocore.client.BaseClient, - instance: Any, - project_name: Optional[str], - backend_tags: Optional[Dict[str, str]], -) -> Tuple[str, str]: +def _list_resource_groups( + ec2_client: botocore.client.BaseClient, region: str +) -> List[ListedResourceGroup]: """ - Allocates an Elastic IP and associates it with the primary ENI of `instance`. - Returns `(public_ip, allocation_id)`. + Lists instances and Elastic IPs tagged with `provider_ref`. + An Elastic IP is grouped with its instance by the `dstack_instance` tag. + Elastic IPs whose instance is gone form groups of their own. + Instances in `shutting-down` are not listed since their termination is in progress + and terminating an instance releases its Elastic IPs. """ - primary_nic_id = _get_primary_network_interface_id(instance) - tags = { - "owner": "dstack", - "dstack_instance": instance.instance_id, - } - if project_name is not None: - tags["dstack_project"] = project_name - if backend_tags: - for k, v in backend_tags.items(): - tags.setdefault(k, v) - tags = aws_resources.filter_invalid_tags(tags) - + # Elastic IPs are listed before instances. An Elastic IP is allocated after its instance + # is launched, so the instance of a listed Elastic IP is listed too unless it's gone. try: - allocate_response = ec2_client.allocate_address( - Domain="vpc", - TagSpecifications=[ - { - "ResourceType": "elastic-ip", - "Tags": aws_resources.make_tags(tags), - } - ], - ) + addresses = aws_resources.list_addresses_with_tag_key(ec2_client, PROVIDER_REF_TAG_KEY) except botocore.exceptions.ClientError as e: - code = e.response.get("Error", {}).get("Code", "") - region = ec2_client.meta.region_name - if code == "AddressLimitExceeded": - raise ProvisioningError( - f"Elastic IP quota exceeded in {region}. " - "Raise the EC2 'EC2-VPC Elastic IPs' quota in Service Quotas, " - "or reduce concurrent multi-EFA instances." + if e.response["Error"]["Code"] != "UnauthorizedOperation": + raise + # Elastic IP permissions are optional. They are only needed for multi-EFA instances + # with public IPs, so without them there are no Elastic IPs to clean up. + logger.debug("Skipping Elastic IPs in %s: not authorized to list them", region) + addresses = [] + instances = aws_resources.list_instances_with_tag_key(ec2_client, PROVIDER_REF_TAG_KEY) + + addresses_by_instance_id: Dict[str, List[Dict[str, Any]]] = defaultdict(list) + instance_ids = {instance["InstanceId"] for instance in instances} + resource_groups = [] + for address in addresses: + instance_id = aws_resources.get_tag_value(address.get("Tags", []), "dstack_instance") + if instance_id is not None and instance_id in instance_ids: + addresses_by_instance_id[instance_id].append(address) + continue + # EIP without an instance + provider_ref = _get_resource_provider_ref(address, address["AllocationId"]) + if provider_ref is None: + continue + resource_groups.append( + ListedResourceGroup( + backend=BackendType.AWS, + owner_type=OwnerType.INSTANCE, + provider_ref=provider_ref, + primary_id=None, + region=region, + created_at=None, + description=_get_address_description(address), + backend_data=AWSResourceGroupBackendData( + eip_allocation_ids=[address["AllocationId"]] + ).model_dump_json(), ) - raise ProvisioningError(f"Failed to allocate Elastic IP in {region}: {e}") - - allocation_id = allocate_response["AllocationId"] - public_ip = allocate_response["PublicIp"] - try: - ec2_client.associate_address( - AllocationId=allocation_id, - NetworkInterfaceId=primary_nic_id, - AllowReassociation=False, ) - except botocore.exceptions.ClientError as e: - # Best-effort release; on failure the EIP leaks until manually released. - logger.warning( - "Failed to associate EIP %s to instance %s; releasing.", - allocation_id, - instance.instance_id, + for instance in instances: + provider_ref = _get_resource_provider_ref(instance, instance["InstanceId"]) + if provider_ref is None: + continue + instance_addresses = addresses_by_instance_id[instance["InstanceId"]] + description = ( + f"instance {instance['InstanceId']} ({instance['InstanceType']}," + f" {instance['State']['Name']})" ) - try: - ec2_client.release_address(AllocationId=allocation_id) - except botocore.exceptions.ClientError: - logger.exception( - "Failed to release just-allocated EIP %s; release it manually.", - allocation_id, + for address in instance_addresses: + description += f", {_get_address_description(address)}" + resource_groups.append( + ListedResourceGroup( + backend=BackendType.AWS, + owner_type=OwnerType.INSTANCE, + provider_ref=provider_ref, + primary_id=instance["InstanceId"], + region=region, + created_at=instance["LaunchTime"], + description=description, + backend_data=AWSResourceGroupBackendData( + instance_id=instance["InstanceId"], + eip_allocation_ids=[a["AllocationId"] for a in instance_addresses], + ).model_dump_json(), ) - raise ProvisioningError( - f"Failed to associate Elastic IP {allocation_id} to instance " - f"{instance.instance_id}: {e}" ) - return public_ip, allocation_id + return resource_groups -def _get_primary_network_interface_id(instance: Any) -> str: - for nic in instance.network_interfaces_attribute or []: - attachment = nic.get("Attachment") or {} - if attachment.get("DeviceIndex") == 0: - return nic["NetworkInterfaceId"] - raise ProvisioningError( - f"Instance {instance.instance_id} has no primary network interface (DeviceIndex=0)" - ) +def _get_resource_provider_ref(resource: Dict[str, Any], resource_id: str) -> Optional[str]: + provider_ref = aws_resources.get_tag_value(resource.get("Tags", []), PROVIDER_REF_TAG_KEY) + if provider_ref is None: + return None + if not is_valid_provider_ref(provider_ref): + logger.warning( + "Skipping AWS resource %s with invalid %s tag value %r", + resource_id, + PROVIDER_REF_TAG_KEY, + provider_ref, + ) + return None + return provider_ref -def _release_eip(ec2_client: botocore.client.BaseClient, allocation_id: str) -> None: - """ - Releases an Elastic IP by allocation ID. Disassociates first if the EIP is still - bound to an instance — `TerminateInstances` only initiates shutdown, and AWS - auto-disassociates only once the instance reaches `terminated`. Releasing - explicitly avoids the `InvalidIPAddress.InUse` race and the retry loop. - """ - try: - response = ec2_client.describe_addresses(AllocationIds=[allocation_id]) - except botocore.exceptions.ClientError as e: - code = e.response.get("Error", {}).get("Code", "") - if code in ("InvalidAllocationID.NotFound", "InvalidAddress.NotFound"): - logger.debug("Skipping EIP %s release. Already released.", allocation_id) - return - raise - addresses = response.get("Addresses", []) - if not addresses: - return - association_id = addresses[0].get("AssociationId") - if association_id is not None: - try: - ec2_client.disassociate_address(AssociationId=association_id) - except botocore.exceptions.ClientError as e: - code = e.response.get("Error", {}).get("Code", "") - # AWS may have auto-disassociated between our Describe and Disassociate - # if the instance just reached `terminated`. Tolerated. - if code != "InvalidAssociationID.NotFound": - raise - try: - ec2_client.release_address(AllocationId=allocation_id) - except botocore.exceptions.ClientError as e: - code = e.response.get("Error", {}).get("Code", "") - if code in ("InvalidAllocationID.NotFound", "InvalidAddress.NotFound"): - return - raise +def _get_address_description(address: Dict[str, Any]) -> str: + return f"Elastic IP {address['AllocationId']} ({address.get('PublicIp')})" diff --git a/src/dstack/_internal/core/backends/aws/resources.py b/src/dstack/_internal/core/backends/aws/resources.py index 98a4a61d0e..e3121ac266 100644 --- a/src/dstack/_internal/core/backends/aws/resources.py +++ b/src/dstack/_internal/core/backends/aws/resources.py @@ -1,12 +1,18 @@ import re -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Tuple import botocore.client import botocore.exceptions from dstack._internal import settings from dstack._internal.core.backends.aws.models import AWSOSImageConfig -from dstack._internal.core.errors import BackendError, ComputeError, ComputeResourceNotFoundError +from dstack._internal.core.backends.base.compute import PROVIDER_REF_TAG_KEY +from dstack._internal.core.errors import ( + BackendError, + ComputeError, + ComputeResourceNotFoundError, + ProvisioningError, +) from dstack._internal.utils.logging import get_logger logger = get_logger(__name__) @@ -441,6 +447,168 @@ def list_instance_device_names( return device_names +def list_instances_with_tag_key( + ec2_client: botocore.client.BaseClient, tag_key: str +) -> List[Dict[str, Any]]: + """ + Lists instances that have the `tag_key` tag, except terminated ones and ones being terminated. + """ + instances = [] + paginator = ec2_client.get_paginator("describe_instances") + for page in paginator.paginate( + Filters=[ + {"Name": "tag-key", "Values": [tag_key]}, + { + "Name": "instance-state-name", + "Values": ["pending", "running", "stopping", "stopped"], + }, + ] + ): + # Instances are grouped by the launch request that created them + for reservation in page["Reservations"]: + instances.extend(reservation["Instances"]) + return instances + + +def list_addresses_with_tag_key( + ec2_client: botocore.client.BaseClient, tag_key: str +) -> List[Dict[str, Any]]: + """ + Lists Elastic IPs that have the `tag_key` tag. + """ + response = ec2_client.describe_addresses(Filters=[{"Name": "tag-key", "Values": [tag_key]}]) + return response["Addresses"] + + +def get_tag_value(tags: List[Dict[str, str]], key: str) -> Optional[str]: + for tag in tags: + if tag.get("Key") == key: + return tag.get("Value") + return None + + +def allocate_and_associate_eip( + ec2_client: botocore.client.BaseClient, + instance: Any, + project_name: Optional[str], + provider_ref: Optional[str], + backend_tags: Optional[Dict[str, str]], +) -> Tuple[str, str]: + """ + Allocates an Elastic IP and associates it with the primary ENI of `instance`. + Returns `(public_ip, allocation_id)`. + """ + primary_nic_id = get_primary_network_interface_id(instance) + tags = { + "owner": "dstack", + "dstack_instance": instance.instance_id, + } + if project_name is not None: + tags["dstack_project"] = project_name + if provider_ref is not None: + tags[PROVIDER_REF_TAG_KEY] = provider_ref + if backend_tags: + for k, v in backend_tags.items(): + tags.setdefault(k, v) + tags = filter_invalid_tags(tags) + + try: + allocate_response = ec2_client.allocate_address( + Domain="vpc", + TagSpecifications=[ + { + "ResourceType": "elastic-ip", + "Tags": make_tags(tags), + } + ], + ) + except botocore.exceptions.ClientError as e: + code = e.response.get("Error", {}).get("Code", "") + region = ec2_client.meta.region_name + if code == "AddressLimitExceeded": + raise ProvisioningError( + f"Elastic IP quota exceeded in {region}. " + "Raise the EC2 'EC2-VPC Elastic IPs' quota in Service Quotas, " + "or reduce concurrent multi-EFA instances." + ) + raise ProvisioningError(f"Failed to allocate Elastic IP in {region}: {e}") + + allocation_id = allocate_response["AllocationId"] + public_ip = allocate_response["PublicIp"] + try: + ec2_client.associate_address( + AllocationId=allocation_id, + NetworkInterfaceId=primary_nic_id, + AllowReassociation=False, + ) + except botocore.exceptions.ClientError as e: + # Best-effort release; on failure the EIP leaks until manually released. + logger.warning( + "Failed to associate EIP %s to instance %s; releasing.", + allocation_id, + instance.instance_id, + ) + try: + ec2_client.release_address(AllocationId=allocation_id) + except botocore.exceptions.ClientError: + logger.exception( + "Failed to release just-allocated EIP %s; release it manually.", + allocation_id, + ) + raise ProvisioningError( + f"Failed to associate Elastic IP {allocation_id} to instance " + f"{instance.instance_id}: {e}" + ) + return public_ip, allocation_id + + +def get_primary_network_interface_id(instance: Any) -> str: + for nic in instance.network_interfaces_attribute or []: + attachment = nic.get("Attachment") or {} + if attachment.get("DeviceIndex") == 0: + return nic["NetworkInterfaceId"] + raise ProvisioningError( + f"Instance {instance.instance_id} has no primary network interface (DeviceIndex=0)" + ) + + +def release_eip(ec2_client: botocore.client.BaseClient, allocation_id: str) -> None: + """ + Releases an Elastic IP by allocation ID. Disassociates first if the EIP is still + bound to an instance – `TerminateInstances` only initiates shutdown, and AWS + auto-disassociates only once the instance reaches `terminated`. Releasing + explicitly avoids the `InvalidIPAddress.InUse` race and the retry loop. + """ + try: + response = ec2_client.describe_addresses(AllocationIds=[allocation_id]) + except botocore.exceptions.ClientError as e: + code = e.response.get("Error", {}).get("Code", "") + if code in ("InvalidAllocationID.NotFound", "InvalidAddress.NotFound"): + logger.debug("Skipping EIP %s release. Already released.", allocation_id) + return + raise + addresses = response.get("Addresses", []) + if not addresses: + return + association_id = addresses[0].get("AssociationId") + if association_id is not None: + try: + ec2_client.disassociate_address(AssociationId=association_id) + except botocore.exceptions.ClientError as e: + code = e.response.get("Error", {}).get("Code", "") + # AWS may have auto-disassociated between our Describe and Disassociate + # if the instance just reached `terminated`. Tolerated. + if code != "InvalidAssociationID.NotFound": + raise + try: + ec2_client.release_address(AllocationId=allocation_id) + except botocore.exceptions.ClientError as e: + code = e.response.get("Error", {}).get("Code", "") + if code in ("InvalidAllocationID.NotFound", "InvalidAddress.NotFound"): + return + raise + + def make_tags(tags: Dict[str, str]) -> List[Dict[str, str]]: tags_list = [] for k, v in tags.items(): diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index f6a2cccfe2..b2fb2c8ff9 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/src/dstack/_internal/core/backends/base/compute.py @@ -30,6 +30,7 @@ DSTACK_RUNNER_SSH_PORT, DSTACK_SHIM_HTTP_PORT, ) +from dstack._internal.core.models.backends.orphans import ListedResourceGroup, ResourceListing from dstack._internal.core.models.compute_groups import ComputeGroup, ComputeGroupProvisioningData from dstack._internal.core.models.gateways import ( GatewayLoadBalancerConfiguration, @@ -444,6 +445,7 @@ def run_job( volumes=volumes, reservation=requirements.reservation, tags=run.run_spec.merged_profile.tags, + provider_ref=get_job_provider_ref(job), ) instance_offer = instance_offer.model_copy() self._restrict_instance_offer_az_to_volumes_az(instance_offer, volumes) @@ -472,6 +474,44 @@ def _restrict_instance_offer_az_to_volumes_az( ] +class ComputeWithResourceListingSupport(ABC): + """ + Implement this mixin for automatic cleanup of orphaned resources. + """ + + @abstractmethod + def list_resource_groups(self) -> ResourceListing: + """ + Lists all groups carrying a dstack `provider_ref` across the configured backend regions. + + Resources without `provider_ref` must be skipped. Resources whose `provider_ref` + is present but unparsable must be skipped and logged as a warning with the resource id. + Owner resources whose deletion is in progress (e.g. a VM that is shutting down) + must be skipped, while their internal resources are listed as usual. + + Resources of a backend region are listed all or nothing: if any listing in a region + fails, return no groups of the region and report it in `ResourceListing.failures`. + Otherwise, `primary_id=None` may be reported for resources of a tracked owner + whose owner resource was not listed, and they would be deleted. + Raise if nothing could be listed. + """ + pass + + @abstractmethod + def delete_resource_group(self, group: ListedResourceGroup) -> None: + """ + Deletes everything in the group in the order the provider requires. + Must be idempotent. + + If the deletion takes several steps, request the deletion of the owner resource first + and raise `NotYetTerminated` instead of waiting, e.g. for the VM to be gone before + its disk can be deleted. The server reports the owner resource as being deleted. + Next cleanups list the remaining internal resources without the owner resource and + call the method again until it returns. + """ + pass + + class ComputeWithGroupProvisioningSupport(ABC): @abstractmethod def run_jobs( @@ -779,14 +819,19 @@ def get_job_instance_name(run: Run, job: Job) -> str: def generate_unique_instance_name( instance_configuration: InstanceConfiguration, max_length: int = _DEFAULT_MAX_RESOURCE_NAME_LEN, + with_provider_ref: bool = False, ) -> str: """ Generates a unique instance name valid across all backends. + + Backends that cannot set `PROVIDER_REF_TAG_KEY` on resources pass `with_provider_ref=True` + to encode `provider_ref` in the name. See `generate_unique_backend_name()`. """ return generate_unique_backend_name( resource_name=instance_configuration.instance_name, project_name=instance_configuration.project_name, max_length=max_length, + provider_ref=instance_configuration.provider_ref if with_provider_ref else None, ) @@ -794,14 +839,19 @@ def generate_unique_instance_name_for_job( run: Run, job: Job, max_length: int = _DEFAULT_MAX_RESOURCE_NAME_LEN, + with_provider_ref: bool = False, ) -> str: """ Generates a unique instance name for a job valid across all backends. + + Backends that cannot set `PROVIDER_REF_TAG_KEY` on resources pass `with_provider_ref=True` + to encode `provider_ref` in the name. See `generate_unique_backend_name()`. """ return generate_unique_backend_name( resource_name=get_job_instance_name(run, job), project_name=run.project_name, max_length=max_length, + provider_ref=get_job_provider_ref(job) if with_provider_ref else None, ) @@ -852,12 +902,17 @@ def generate_unique_backend_name( resource_name: str, project_name: Optional[str], max_length: int, + provider_ref: Optional[str] = None, ) -> str: """ Generates a unique resource name valid across all backends. Backend resource names must be unique on every provisioning so that resource re-submission/re-creation doesn't lead to conflicts on backends that require unique names (e.g. Azure, GCP). + + If `provider_ref` is set, the name ends with `provider_ref` followed by fewer random + characters, so that `parse_provider_ref_from_name()` can get it back. Only the prefix + is truncated to fit `max_length`. """ # resource_name is guaranteed to be valid in all backends prefix = f"dstack-{resource_name}" @@ -865,8 +920,15 @@ def generate_unique_backend_name( # project_name is not guaranteed to be valid in all backends, # so we add it only if it passes the validation prefix = f"dstack-{project_name}-{resource_name}" + if provider_ref is None: + return generate_unique_name( + prefix=prefix, + max_length=max_length, + ) return generate_unique_name( prefix=prefix, + fixed_suffix=provider_ref, + random_suffix_length=_PROVIDER_REF_NAME_RANDOM_SUFFIX_LEN, max_length=max_length, ) @@ -882,25 +944,97 @@ def generate_unique_short_backend_name() -> str: def generate_unique_name( *, prefix: Optional[str] = None, - suffix_length: Optional[int] = None, + fixed_suffix: str = "", + random_suffix_length: Optional[int] = None, max_length: Optional[int] = None, ) -> str: - if suffix_length is None: - suffix_length = _CLOUD_RESOURCE_SUFFIX_LEN + """ + Generates a name of the form `{prefix}-{fixed_suffix}{random_suffix}`, + where `random_suffix` has `random_suffix_length` characters. + Only the prefix is truncated to fit `max_length`. + """ + if random_suffix_length is None: + random_suffix_length = _CLOUD_RESOURCE_SUFFIX_LEN + full_suffix_length = len(fixed_suffix) + random_suffix_length if max_length is not None: - assert max_length >= suffix_length + assert max_length >= full_suffix_length if prefix is not None: - prefix_len = max_length - suffix_length - 1 + prefix_len = max_length - full_suffix_length - 1 assert prefix_len > 0 - prefix = prefix[:prefix_len] - suffix = "".join( - random.choice(string.ascii_lowercase + string.digits) for _ in range(suffix_length) + # Avoid a double hyphen if the prefix is cut right after a hyphen + prefix = prefix[:prefix_len].rstrip("-") + suffix = fixed_suffix + "".join( + random.choice(string.ascii_lowercase + string.digits) for _ in range(random_suffix_length) ) if prefix is None: return suffix return f"{prefix}-{suffix}" +PROVIDER_REF_LEN = 12 + +PROVIDER_REF_TAG_KEY = "dstack_provider_ref" +""" +The tag (label) key for `provider_ref` on backends that support tags. +Backends whose tag keys cannot contain `_` define their own key. +""" + +_PROVIDER_REF_NAME_RANDOM_SUFFIX_LEN = 4 +""" +Random characters after `provider_ref` in names. They keep names unique across provisioning +attempts with the same `provider_ref`. +""" + +_PROVIDER_REF_REGEX = re.compile(rf"[0-9a-z]{{{PROVIDER_REF_LEN}}}") + +_PROVIDER_REF_NAME_SUFFIX_REGEX = re.compile( + rf"[0-9a-z]{{{PROVIDER_REF_LEN + _PROVIDER_REF_NAME_RANDOM_SUFFIX_LEN}}}" +) + + +def generate_provider_ref() -> str: + """ + Generates a `provider_ref` that is set on provider resources for orphan cleanup. + Uses lowercase letters and digits only, so it's valid in resource names and tags + on all backends. + """ + return "".join( + random.choice(string.ascii_lowercase + string.digits) for _ in range(PROVIDER_REF_LEN) + ) + + +def is_valid_provider_ref(value: str) -> bool: + """ + Checks a `provider_ref` read from a provider resource, e.g. from a tag. + """ + return _PROVIDER_REF_REGEX.fullmatch(value) is not None + + +def parse_provider_ref_from_name(name: str) -> Optional[str]: + """ + Parses `provider_ref` from a name generated by `generate_unique_backend_name()` + with `provider_ref`. Returns `None` if the name does not contain `provider_ref`. + + Backends that derive names of internal resources from the instance name + (e.g. by appending `-data`) must strip such suffixes before calling it. + """ + if not name.startswith("dstack-"): + return None + suffix = name.rsplit("-", 1)[-1] + if _PROVIDER_REF_NAME_SUFFIX_REGEX.fullmatch(suffix) is None: + return None + return suffix[:PROVIDER_REF_LEN] + + +def get_job_provider_ref(job: Job) -> Optional[str]: + """ + Returns the `provider_ref` of the job's latest submission. + """ + if len(job.job_submissions) == 0: + return None + return job.job_submissions[-1].provider_ref + + def get_cloud_config(**config) -> str: return "#cloud-config\n" + yaml.dump(config, default_flow_style=False) diff --git a/src/dstack/_internal/core/backends/configurators.py b/src/dstack/_internal/core/backends/configurators.py index 826073914c..a277dd9a8a 100644 --- a/src/dstack/_internal/core/backends/configurators.py +++ b/src/dstack/_internal/core/backends/configurators.py @@ -202,13 +202,21 @@ def list_available_backend_types() -> List[BackendType]: return _BACKEND_TYPES -def list_available_configurator_classes() -> List[type[Configurator]]: +def list_builtin_configurator_classes() -> List[type[Configurator]]: """ - Lists all backend configurator classes available on the server. + Lists all built-in backend configurator classes available on the server. """ return _CONFIGURATOR_CLASSES +def list_builtin_and_registered_configurator_classes() -> List[type[Configurator]]: + """ + Lists all built-in backend configurator classes available on the server + and those added with `register_configurator()`. + """ + return list(_BACKEND_TYPE_TO_CONFIGURATOR_CLASS_MAP.values()) + + def register_configurator(configurator: Type[Configurator]): """ A hook to for registering new configurators without importing them. diff --git a/src/dstack/_internal/core/backends/features.py b/src/dstack/_internal/core/backends/features.py index d5c28728dc..8ecc89a2fa 100644 --- a/src/dstack/_internal/core/backends/features.py +++ b/src/dstack/_internal/core/backends/features.py @@ -8,13 +8,23 @@ ComputeWithPrivateGatewaySupport, ComputeWithPrivilegedSupport, ComputeWithReservationSupport, + ComputeWithResourceListingSupport, ComputeWithVolumeSupport, ) from dstack._internal.core.backends.base.configurator import Configurator -from dstack._internal.core.backends.configurators import list_available_configurator_classes +from dstack._internal.core.backends.configurators import ( + list_builtin_and_registered_configurator_classes, + list_builtin_configurator_classes, +) from dstack._internal.core.models.backends.base import BackendType -_configurator_classes = list_available_configurator_classes() +# The `BACKENDS_WITH_*` lists are computed from built-in configurators. They are checked against +# base backend types (offers, volume and gateway configurations), which are always built-in: +# extensions currently only register wrapper backends such as `dstack`, which callers that check +# project backend types special-case or handle with `get_*()` functions below. +# TODO: Compute the lists from built-in and registered configurators and mark wrapper backends +# if extensions ever need to register base backends. +_configurator_classes = list_builtin_configurator_classes() def _get_backends_with_compute_feature( @@ -73,3 +83,21 @@ def _get_backends_with_compute_feature( configurator_classes=_configurator_classes, compute_feature_class=ComputeWithVolumeSupport, ) + + +def get_backends_with_resource_listing_support() -> list[BackendType]: + """ + Unlike the `BACKENDS_WITH_*` lists, includes backends added with `register_configurator()`. + Orphan cleanup checks project backend types, so wrapper backends such as `dstack` + must be included if they support resource listing. + """ + backend_types = [] + for configurator_class in list_builtin_and_registered_configurator_classes(): + # Registered configurators may not declare the classes for introspection + backend_class = getattr(configurator_class, "BACKEND_CLASS", None) + compute_class = getattr(backend_class, "COMPUTE_CLASS", None) + if compute_class is not None and issubclass( + compute_class, ComputeWithResourceListingSupport + ): + backend_types.append(configurator_class.TYPE) + return backend_types diff --git a/src/dstack/_internal/core/backends/runpod/compute.py b/src/dstack/_internal/core/backends/runpod/compute.py index 625c58a134..44dc37032e 100644 --- a/src/dstack/_internal/core/backends/runpod/compute.py +++ b/src/dstack/_internal/core/backends/runpod/compute.py @@ -15,6 +15,7 @@ generate_unique_volume_name, get_docker_commands, get_job_instance_name, + get_job_provider_ref, ) from dstack._internal.core.backends.base.models import JobConfiguration from dstack._internal.core.backends.base.offers import get_catalog_offers @@ -132,6 +133,7 @@ def run_job( instance_name=get_job_instance_name(run, job), ssh_keys=[], user=run.user, + provider_ref=get_job_provider_ref(job), ) pod_name = generate_unique_instance_name(instance_config, max_length=MAX_RESOURCE_NAME_LEN) diff --git a/src/dstack/_internal/core/compatibility/runs.py b/src/dstack/_internal/core/compatibility/runs.py index 66cca5c79c..671522b658 100644 --- a/src/dstack/_internal/core/compatibility/runs.py +++ b/src/dstack/_internal/core/compatibility/runs.py @@ -129,4 +129,5 @@ def get_job_submission_excludes(job_submissions: list[JobSubmission]) -> Include submission_excludes["job_provisioning_data"] = { "instance_type": {"resources": {"description": True}} } + submission_excludes["provider_ref"] = True return submission_excludes diff --git a/src/dstack/_internal/core/models/backends/orphans.py b/src/dstack/_internal/core/models/backends/orphans.py new file mode 100644 index 0000000000..c007bf4170 --- /dev/null +++ b/src/dstack/_internal/core/models/backends/orphans.py @@ -0,0 +1,64 @@ +from datetime import datetime +from enum import Enum +from typing import Optional + +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.common import CoreModel + + +class OwnerType(str, Enum): + INSTANCE = "instance" + # TODO: VOLUME, GATEWAY, PLACEMENT_GROUP + + +class OrphanCleanupMode(str, Enum): + OFF = "off" + REPORT = "report" + DELETE = "delete" + + +class ListedResourceGroup(CoreModel): + """ + Resources provisioned as one backend method call: + the server-level resource (e.g. instance) and the backend-internal resources + created with it (e.g. disks, keys, scripts, NICs). + """ + + backend: BackendType + owner_type: OwnerType + provider_ref: str + """dstack's `provider_ref` of the owner resources read from a tag or a name.""" + primary_id: Optional[str] + """Provider's id of the owner resource (e.g.`JobProvisioningData.instance_id` for the instances). + `None` if the group has only internal resources left. + """ + region: Optional[str] + created_at: Optional[datetime] + """Provider's creation time of the owner resource or, if it's gone, the earliest internal resource. + `None` if the provider does not report it.""" + description: str + """Human-readable summary for events and logs""" + backend_data: Optional[str] + """Everything the backend needs to delete the group. Opaque to the server.""" + + +class ResourceListingFailure(CoreModel): + """ + A part of the resources that could not be listed. It covers all owner types + in `backend` and `region`. + """ + + backend: BackendType + """The backend whose resources could not be listed, as in `ListedResourceGroup.backend`.""" + region: Optional[str] + """The region whose resources could not be listed. `None` if all regions.""" + message: str + + +class ResourceListing(CoreModel): + resource_groups: list[ListedResourceGroup] + failures: list[ResourceListingFailure] = [] + """ + Parts of the resources that could not be listed. Resources of a part are listed + all or nothing, so `resource_groups` has no groups of the parts in `failures`. + """ diff --git a/src/dstack/_internal/core/models/instances.py b/src/dstack/_internal/core/models/instances.py index 68e5c11028..aaeb20cfe9 100644 --- a/src/dstack/_internal/core/models/instances.py +++ b/src/dstack/_internal/core/models/instances.py @@ -166,6 +166,9 @@ class InstanceConfiguration(CoreModel): reservation: Optional[str] = None volumes: Optional[List[Volume]] = None tags: Optional[Dict[str, str]] = None + provider_ref: Optional[str] = None + """`provider_ref` is set on provider resources created for this instance to identify orphaned resources. + """ def get_public_keys(self) -> List[str]: return [ssh_key.public.strip() for ssh_key in self.ssh_keys] diff --git a/src/dstack/_internal/core/models/runs.py b/src/dstack/_internal/core/models/runs.py index a120390079..f02165ae30 100644 --- a/src/dstack/_internal/core/models/runs.py +++ b/src/dstack/_internal/core/models/runs.py @@ -455,6 +455,11 @@ class JobSubmission(CoreModel): error: Optional[str] = None probes: list[Probe] = [] image_pull_progress: Optional[ImagePullProgress] = None + provider_ref: Optional[str] = None + """`provider_ref` to set on provider resources created by `run_job()` for this submission. + The server sets it from the submission's placeholder instance before calling `run_job()`. + It's not stored with the submission. + """ @property def age(self) -> timedelta: diff --git a/src/dstack/_internal/server/background/pipeline_tasks/__init__.py b/src/dstack/_internal/server/background/pipeline_tasks/__init__.py index 2e83e780c0..c2b108ea7a 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/__init__.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/__init__.py @@ -1,5 +1,7 @@ import asyncio +from dstack._internal.core.models.backends.orphans import OrphanCleanupMode +from dstack._internal.server import settings from dstack._internal.server.background.pipeline_tasks.base import Pipeline from dstack._internal.server.background.pipeline_tasks.compute_groups import ComputeGroupPipeline from dstack._internal.server.background.pipeline_tasks.fleets import FleetPipeline @@ -15,6 +17,7 @@ from dstack._internal.server.background.pipeline_tasks.jobs_terminating import ( JobTerminatingPipeline, ) +from dstack._internal.server.background.pipeline_tasks.orphan_cleanup import OrphanCleanupPipeline from dstack._internal.server.background.pipeline_tasks.placement_groups import ( PlacementGroupPipeline, ) @@ -47,6 +50,8 @@ def __init__(self) -> None: VolumePipeline(pipeline_hinter=self._hinter), ]: self.register_pipeline(builtin_pipeline) + if settings.ORPHAN_CLEANUP_MODE != OrphanCleanupMode.OFF: + self.register_pipeline(OrphanCleanupPipeline(pipeline_hinter=self._hinter)) def register_pipeline(self, pipeline: Pipeline): self._pipelines.append(pipeline) diff --git a/src/dstack/_internal/server/background/pipeline_tasks/instances/termination.py b/src/dstack/_internal/server/background/pipeline_tasks/instances/termination.py index c4cce13b27..7b05d7294b 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/instances/termination.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/instances/termination.py @@ -1,6 +1,10 @@ +from dstack._internal.core.backends.features import get_backends_with_resource_listing_support from dstack._internal.core.errors import BackendError, NotYetTerminated from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.backends.orphans import OrphanCleanupMode from dstack._internal.core.models.instances import InstanceStatus +from dstack._internal.core.models.runs import JobProvisioningData +from dstack._internal.server import settings from dstack._internal.server.background.pipeline_tasks.base import NOW_PLACEHOLDER from dstack._internal.server.background.pipeline_tasks.instances.common import ( ProcessResult, @@ -71,11 +75,13 @@ async def terminate_instance(instance_model: InstanceModel) -> ProcessResult: exc_info=not isinstance(exc, BackendError), ) return result + next_step = "Please terminate the instance manually to avoid unexpected charges." + if _is_to_be_deleted_by_orphan_cleanup(instance_model, job_provisioning_data): + next_step = "Orphan cleanup will retry deleting the instance." logger.error( - "Failed all attempts to terminate instance %s." - " Please terminate the instance manually to avoid unexpected charges." - " Error: %r", + "Failed all attempts to terminate instance %s. %s Error: %r", instance_model.name, + next_step, exc, exc_info=not isinstance(exc, BackendError), ) @@ -92,3 +98,13 @@ async def terminate_instance(instance_model: InstanceModel) -> ProcessResult: new_status=InstanceStatus.TERMINATED, ) return result + + +def _is_to_be_deleted_by_orphan_cleanup( + instance_model: InstanceModel, job_provisioning_data: JobProvisioningData +) -> bool: + return ( + settings.ORPHAN_CLEANUP_MODE == OrphanCleanupMode.DELETE + and instance_model.provider_ref is not None + and job_provisioning_data.backend in get_backends_with_resource_listing_support() + ) diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py index 1ab7875d04..fb340d4ddc 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py @@ -21,6 +21,7 @@ BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT, ) from dstack._internal.core.errors import BackendError, ServerClientError, SkipOffer +from dstack._internal.core.models.backends.orphans import OwnerType from dstack._internal.core.models.common import NetworkMode, validate_extra_ignore from dstack._internal.core.models.compute_groups import ( ComputeGroupProvisioningData, @@ -120,6 +121,7 @@ get_instance_offer_with_restricted_az, get_offers_by_requirements, ) +from dstack._internal.server.services.orphan_cleanup import issue_provider_ref from dstack._internal.server.services.pipelines import PipelineHinterProtocol from dstack._internal.server.services.placement import ( can_use_placement_groups, @@ -670,6 +672,7 @@ async def _apply_assignment_result( ) return instance_model = _create_placeholder_instance( + session=session, fleet_model=fleet_model, project=context.project, job_model=job_model, @@ -1179,21 +1182,34 @@ async def _lock_fleet_for_placeholder( def _create_placeholder_instance( + session: AsyncSession, fleet_model: FleetModel, project: ProjectModel, job_model: JobModel, ) -> InstanceModel: taken_instance_nums = {i.instance_num for i in fleet_model.instances} instance_num = get_next_instance_num(taken_instance_nums) + instance_id = uuid.uuid4() + instance_name = f"{fleet_model.name}-{instance_num}" return InstanceModel( - id=uuid.uuid4(), - name=f"{fleet_model.name}-{instance_num}", + id=instance_id, + name=instance_name, instance_num=instance_num, project=project, fleet=fleet_model, status=InstanceStatus.PENDING, unreachable=False, provisioning_job_id=job_model.id, + # Set on resources created by `run_job()` for this placeholder. The placeholder + # keeps it when promoted, and the `ProviderRefModel` record outlives the placeholder + # if it's deleted after a failed `run_job()`. + provider_ref=issue_provider_ref( + session=session, + project=project, + owner_type=OwnerType.INSTANCE, + owner_id=instance_id, + owner_name=instance_name, + ), ) @@ -1852,6 +1868,12 @@ def _promote_placeholder_instance( instance_model.termination_idle_time = termination_idle_time instance_model.total_blocks = 1 instance_model.busy_blocks = 1 + if compute_group_model is not None: + # Compute group instances must not be cleaned up individually. + # TODO: Orphan cleanup for compute groups is not implemented, and backends must not set + # `provider_ref` on compute group resources. Store a `provider_ref` on the compute group, + # set it on all its resources, and clean them up as one unit. + instance_model.provider_ref = None async def _process_volume_attachments( @@ -2511,6 +2533,14 @@ async def _provision_new_capacity( new_placement_group_models=new_placement_group_models, ), ) + placeholder_instance = job_model.instance + if ( + placeholder_instance is not None + and is_placeholder_instance(placeholder_instance) + and len(job.job_submissions) > 0 + ): + # `run_job()` sets the placeholder's `provider_ref` on the resources it creates + job.job_submissions[-1].provider_ref = placeholder_instance.provider_ref job_provisioning_data = await run_async( compute.run_job, run, diff --git a/src/dstack/_internal/server/background/pipeline_tasks/orphan_cleanup.py b/src/dstack/_internal/server/background/pipeline_tasks/orphan_cleanup.py new file mode 100644 index 0000000000..1a401bb5f0 --- /dev/null +++ b/src/dstack/_internal/server/background/pipeline_tasks/orphan_cleanup.py @@ -0,0 +1,800 @@ +import asyncio +import uuid +from dataclasses import dataclass, field +from datetime import datetime, timedelta +from enum import Enum +from typing import Optional, Sequence + +import sentry_sdk +from sqlalchemy import ColumnElement, and_, or_, select, update +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import joinedload, load_only + +from dstack._internal.core.backends.base.compute import ComputeWithResourceListingSupport +from dstack._internal.core.backends.features import get_backends_with_resource_listing_support +from dstack._internal.core.errors import BackendError, NotYetTerminated +from dstack._internal.core.models.backends.orphans import ( + ListedResourceGroup, + OrphanCleanupMode, + OwnerType, +) +from dstack._internal.core.models.events import EventTargetType +from dstack._internal.core.models.instances import InstanceStatus +from dstack._internal.core.models.runs import JobStatus +from dstack._internal.server import settings +from dstack._internal.server.background.pipeline_tasks.base import ( + Fetcher, + Heartbeater, + ItemUpdateMap, + Pipeline, + PipelineItem, + UpdateMapDateTime, + Worker, + log_lock_token_changed_after_processing, + log_lock_token_mismatch, + resolve_now_placeholders, + set_processed_update_map_fields, + set_unlock_update_map_fields, +) +from dstack._internal.server.db import get_db, get_session_ctx +from dstack._internal.server.models import ( + BackendModel, + InstanceModel, + JobModel, + OrphanCleanupModel, + ProjectModel, + ProviderRefModel, +) +from dstack._internal.server.services import backends as backends_services +from dstack._internal.server.services import events +from dstack._internal.server.services.instances import get_instance_provisioning_data +from dstack._internal.server.services.locking import get_locker +from dstack._internal.server.services.pipelines import PipelineHinterProtocol +from dstack._internal.server.utils import tracing +from dstack._internal.utils.common import batched, get_current_datetime, run_async +from dstack._internal.utils.logging import get_logger + +logger = get_logger(__name__) + + +class OrphanCleanupPipeline(Pipeline[PipelineItem]): + def __init__( + self, + workers_num: int = 5, + queue_lower_limit_factor: float = 0.5, + queue_upper_limit_factor: float = 2.0, + min_processing_interval: Optional[timedelta] = None, + lock_timeout: timedelta = timedelta(seconds=60), + heartbeat_trigger: timedelta = timedelta(seconds=20), + *, + pipeline_hinter: PipelineHinterProtocol, + ) -> None: + if min_processing_interval is None: + if settings.ORPHAN_CLEANUP_INTERVAL_SECONDS <= 0: + raise ValueError("DSTACK_SERVER_ORPHAN_CLEANUP_INTERVAL_SECONDS must be positive") + min_processing_interval = timedelta(seconds=settings.ORPHAN_CLEANUP_INTERVAL_SECONDS) + super().__init__( + workers_num=workers_num, + queue_lower_limit_factor=queue_lower_limit_factor, + queue_upper_limit_factor=queue_upper_limit_factor, + min_processing_interval=min_processing_interval, + lock_timeout=lock_timeout, + heartbeat_trigger=heartbeat_trigger, + ) + self.__heartbeater = Heartbeater[PipelineItem]( + model_type=OrphanCleanupModel, + lock_timeout=self._lock_timeout, + heartbeat_trigger=self._heartbeat_trigger, + ) + self.__fetcher = OrphanCleanupFetcher( + queue=self._queue, + queue_desired_minsize=self._queue_desired_minsize, + min_processing_interval=self._min_processing_interval, + lock_timeout=self._lock_timeout, + heartbeater=self._heartbeater, + ) + self.__workers = [ + OrphanCleanupWorker( + queue=self._queue, + heartbeater=self._heartbeater, + pipeline_hinter=pipeline_hinter, + ) + for _ in range(self._workers_num) + ] + + @property + def hint_fetch_model_name(self) -> str: + return OrphanCleanupModel.__name__ + + @property + def _heartbeater(self) -> Heartbeater[PipelineItem]: + return self.__heartbeater + + @property + def _fetcher(self) -> Fetcher[PipelineItem]: + return self.__fetcher + + @property + def _workers(self) -> Sequence["OrphanCleanupWorker"]: + return self.__workers + + +class OrphanCleanupFetcher(Fetcher[PipelineItem]): + def __init__( + self, + queue: asyncio.Queue[PipelineItem], + queue_desired_minsize: int, + min_processing_interval: timedelta, + lock_timeout: timedelta, + heartbeater: Heartbeater[PipelineItem], + queue_check_delay: float = 1.0, + ) -> None: + super().__init__( + queue=queue, + queue_desired_minsize=queue_desired_minsize, + min_processing_interval=min_processing_interval, + lock_timeout=lock_timeout, + heartbeater=heartbeater, + queue_check_delay=queue_check_delay, + ) + + @tracing.instrument_pipeline_task("OrphanCleanupFetcher.fetch") + async def fetch(self, limit: int) -> list[PipelineItem]: + orphan_cleanup_lock, _ = get_locker(get_db().dialect_name).get_lockset( + OrphanCleanupModel.__tablename__ + ) + async with orphan_cleanup_lock: + async with get_session_ctx() as session: + now = get_current_datetime() + res = await session.execute( + select(OrphanCleanupModel) + .join(BackendModel, BackendModel.id == OrphanCleanupModel.backend_id) + .join(ProjectModel, ProjectModel.id == BackendModel.project_id) + .where( + ProjectModel.deleted == False, + BackendModel.type.in_(get_backends_with_resource_listing_support()), + _retry_delay_filter(now), + or_( + OrphanCleanupModel.last_processed_at + <= now - self._min_processing_interval, + OrphanCleanupModel.last_processed_at == OrphanCleanupModel.created_at, + ), + or_( + OrphanCleanupModel.lock_expires_at.is_(None), + OrphanCleanupModel.lock_expires_at < now, + ), + or_( + OrphanCleanupModel.lock_owner.is_(None), + OrphanCleanupModel.lock_owner == OrphanCleanupPipeline.__name__, + ), + ) + .order_by(OrphanCleanupModel.last_processed_at.asc()) + .limit(limit) + .with_for_update(skip_locked=True, key_share=True, of=OrphanCleanupModel) + .options( + load_only( + OrphanCleanupModel.id, + OrphanCleanupModel.lock_token, + OrphanCleanupModel.lock_expires_at, + ) + ) + ) + orphan_cleanup_models = list(res.scalars().all()) + lock_expires_at = get_current_datetime() + self._lock_timeout + lock_token = uuid.uuid4() + items = [] + for orphan_cleanup_model in orphan_cleanup_models: + prev_lock_expired = orphan_cleanup_model.lock_expires_at is not None + orphan_cleanup_model.lock_expires_at = lock_expires_at + orphan_cleanup_model.lock_token = lock_token + orphan_cleanup_model.lock_owner = OrphanCleanupPipeline.__name__ + items.append( + PipelineItem( + __tablename__=OrphanCleanupModel.__tablename__, + id=orphan_cleanup_model.id, + lock_expires_at=lock_expires_at, + lock_token=lock_token, + prev_lock_expired=prev_lock_expired, + ) + ) + await session.commit() + return items + + +class OrphanCleanupWorker(Worker[PipelineItem]): + def __init__( + self, + queue: asyncio.Queue[PipelineItem], + heartbeater: Heartbeater[PipelineItem], + pipeline_hinter: PipelineHinterProtocol, + ) -> None: + super().__init__( + queue=queue, + heartbeater=heartbeater, + pipeline_hinter=pipeline_hinter, + ) + + @tracing.instrument_pipeline_task("OrphanCleanupWorker.process") + async def process(self, item: PipelineItem): + process_context = await _load_process_context(item) + if process_context is None: + return + result = await _process_orphan_cleanup(process_context) + await _apply_process_result(item, result) + + +@dataclass +class _ProcessContext: + orphan_cleanup_model: OrphanCleanupModel + needs_cleanup: bool + + +class _OrphanCleanupUpdateMap(ItemUpdateMap, total=False): + cleanup_attempt: int + clean_at: UpdateMapDateTime + clean_mode: OrphanCleanupMode + + +@dataclass +class _OrphanEvent: + message: str + target: events.Target + + +@dataclass +class _ProcessResult: + update_map: _OrphanCleanupUpdateMap + orphan_events: list[_OrphanEvent] = field(default_factory=list) + + +class _ResourceGroupDecision(str, Enum): + IGNORE = "ignore" + """The group cannot be matched or classified. It's left as is.""" + SKIP = "skip" + """The owner is changing or changed during the listing, or the deletion of the group's + internal resources is in progress. The group is reconsidered by the next cleanup.""" + ORPHAN = "orphan" + """The group is orphaned. After handling, it's been reported or deleted.""" + DELETING = "deleting" + """The group is orphaned and the deletion of its owner resource is requested. + Its internal resources are deleted by next cleanups.""" + DELETE_FAILED = "delete_failed" + """The group is orphaned but failed to be deleted.""" + TRACKED = "tracked" + """The group is the owner's current resources.""" + + +_CLEANUP_RETRY_INTERVAL_MULTIPLIERS = [2, 6, 24] +"""Retry delays after failed cleanups, as multiples of the cleanup interval.""" + + +def _retry_delay_filter(now: datetime) -> ColumnElement[bool]: + """ + SQL-portable way to filter by retry delay depending on `cleanup_attempt`. + """ + interval = timedelta(seconds=settings.ORPHAN_CLEANUP_INTERVAL_SECONDS) + filters = [ + and_( + OrphanCleanupModel.cleanup_attempt == i, + OrphanCleanupModel.last_processed_at < now - interval * multiplier, + ) + for i, multiplier in enumerate(_CLEANUP_RETRY_INTERVAL_MULTIPLIERS[:-1], start=1) + ] + filters.append( + and_( + OrphanCleanupModel.cleanup_attempt >= len(_CLEANUP_RETRY_INTERVAL_MULTIPLIERS), + OrphanCleanupModel.last_processed_at + < now - interval * _CLEANUP_RETRY_INTERVAL_MULTIPLIERS[-1], + ) + ) + return or_(OrphanCleanupModel.cleanup_attempt == 0, *filters) + + +async def _load_process_context(item: PipelineItem) -> Optional[_ProcessContext]: + async with get_session_ctx() as session: + res = await session.execute( + select(OrphanCleanupModel) + .where( + OrphanCleanupModel.id == item.id, + OrphanCleanupModel.lock_token == item.lock_token, + ) + .options( + joinedload(OrphanCleanupModel.backend) + .joinedload(BackendModel.project) + .joinedload(ProjectModel.backends) + ) + ) + orphan_cleanup_model = res.unique().scalar_one_or_none() + if orphan_cleanup_model is None: + log_lock_token_mismatch(logger, item) + return + needs_cleanup = await _needs_cleanup(session, orphan_cleanup_model) + return _ProcessContext(orphan_cleanup_model=orphan_cleanup_model, needs_cleanup=needs_cleanup) + + +async def _needs_cleanup(session: AsyncSession, orphan_cleanup_model: OrphanCleanupModel) -> bool: + """ + Checks if the backend may have new orphans since the last clean cleanup. + + Every enrolled backend is processed once per cleanup interval, and a backend without new + activity is only marked processed without listing its resources. This costs one activity + query and two small row updates (lock and unlock) per enrolled backend per interval, + but keeps the fetch query cheap regardless of project sizes and the number of + idle backends. If the number of enrolled backends grows large, the cost can be cut by: + * Deleting rows of backends that stay idle and re-enrolling backends on recent activity, + so that only recently active backends are processed. + * Cleaning up cloud accounts shared by many backends (e.g. dstack Sky base backends) + once instead of once per backend. + """ + if orphan_cleanup_model.clean_at is None: + return True + if orphan_cleanup_model.clean_mode != settings.ORPHAN_CLEANUP_MODE: + return True + return await _has_activity_since( + session, + project_id=orphan_cleanup_model.backend.project_id, + since=orphan_cleanup_model.clean_at, + ) + + +async def _has_activity_since( + session: AsyncSession, project_id: uuid.UUID, since: datetime +) -> bool: + has_issued_provider_refs = ( + select(ProviderRefModel.provider_ref) + .where( + ProviderRefModel.project_id == project_id, + ProviderRefModel.created_at >= since, + ) + .exists() + ) + has_instance_activity = ( + select(InstanceModel.id) + .where( + InstanceModel.project_id == project_id, + InstanceModel.provider_ref.is_not(None), + or_( + InstanceModel.started_at >= since, + InstanceModel.finished_at >= since, + ), + ) + .exists() + ) + # If `run_job()` fails after creating provider resources, the job is terminated and its + # placeholder instance is deleted, so only the job's activity indicates possible orphans. + # Finished jobs are no longer processed, so `last_processed_at` is the finish time. + has_job_activity = ( + select(JobModel.id) + .where( + JobModel.project_id == project_id, + JobModel.status.in_(JobStatus.finished_statuses()), + JobModel.last_processed_at >= since, + ) + .exists() + ) + res = await session.execute( + select(or_(has_issued_provider_refs, has_instance_activity, has_job_activity)) + ) + return res.scalar_one() + + +async def _process_orphan_cleanup(process_context: _ProcessContext) -> _ProcessResult: + orphan_cleanup_model = process_context.orphan_cleanup_model + if not process_context.needs_cleanup: + return _ProcessResult(update_map=_OrphanCleanupUpdateMap()) + backend = await backends_services.get_project_backend_by_type( + project=orphan_cleanup_model.backend.project, + backend_type=orphan_cleanup_model.backend.type, + ) + if backend is None: + logger.warning( + "Failed to clean up orphaned resources in %s/%s: backend is not available", + orphan_cleanup_model.backend.project.name, + orphan_cleanup_model.backend.type.value, + ) + return _get_backoff_cleanup_process_result(orphan_cleanup_model) + compute = backend.compute() + assert isinstance(compute, ComputeWithResourceListingSupport) + listed_at = get_current_datetime() + try: + resource_listing = await run_async(compute.list_resource_groups) + except BackendError as e: + logger.warning( + "Failed to list resource groups in %s/%s due to backend error: %s", + orphan_cleanup_model.backend.project.name, + orphan_cleanup_model.backend.type.value, + repr(e), + ) + return _get_backoff_cleanup_process_result(orphan_cleanup_model) + except Exception: + logger.exception( + "Failed to list resource groups in %s/%s due to unexpected error", + orphan_cleanup_model.backend.project.name, + orphan_cleanup_model.backend.type.value, + ) + return _get_backoff_cleanup_process_result(orphan_cleanup_model) + for failure in resource_listing.failures: + logger.warning( + "Failed to list resources of %s in %s for orphan cleanup in %s/%s: %s", + failure.backend.value, + failure.region if failure.region is not None else "all regions", + orphan_cleanup_model.backend.project.name, + orphan_cleanup_model.backend.type.value, + failure.message, + ) + resource_groups = resource_listing.resource_groups + owner_records = await _load_owner_records( + project_id=orphan_cleanup_model.backend.project_id, + resource_groups=resource_groups, + ) + decisions = [ + _get_resource_group_decision( + orphan_cleanup_model, resource_group, owner_records, listed_at + ) + for resource_group in resource_groups + ] + final_decisions: list[_ResourceGroupDecision] = [] + orphan_events: list[_OrphanEvent] = [] + for resource_group, decision in zip(resource_groups, decisions): + if decision == _ResourceGroupDecision.ORPHAN: + decision = await _handle_orphaned_resource_group( + compute, orphan_cleanup_model, resource_group + ) + if decision in (_ResourceGroupDecision.ORPHAN, _ResourceGroupDecision.DELETING): + orphan_events.append(_get_orphan_event(resource_group, owner_records, decision)) + final_decisions.append(decision) + if any(d is _ResourceGroupDecision.DELETE_FAILED for d in final_decisions): + result = _get_backoff_cleanup_process_result(orphan_cleanup_model) + elif ( + any( + d in (_ResourceGroupDecision.SKIP, _ResourceGroupDecision.DELETING) + for d in final_decisions + ) + or len(resource_listing.failures) > 0 + ): + result = _get_unsettled_cleanup_process_result() + else: + result = _get_success_cleanup_process_result(listed_at) + result.orphan_events = orphan_events + return result + + +_PROVIDER_REF_BATCH_SIZE = 100 + + +@dataclass +class _OwnerRecord: + provider_ref_model: ProviderRefModel + instance_model: Optional[InstanceModel] + """Including soft-deleted instances. `None` if the instance was deleted, + e.g. a placeholder instance of a failed `run_job()`.""" + + +_OwnerRecords = dict[str, _OwnerRecord] +"""Owner records by `provider_ref`. `provider_ref`s missing from the dict were not issued.""" + + +async def _load_owner_records( + project_id: uuid.UUID, resource_groups: list[ListedResourceGroup] +) -> _OwnerRecords: + instance_refs = { + resource_group.provider_ref + for resource_group in resource_groups + if resource_group.owner_type == OwnerType.INSTANCE + } + owner_records: _OwnerRecords = {} + async with get_session_ctx() as session: + for batch in batched(instance_refs, _PROVIDER_REF_BATCH_SIZE): + # `provider_ref`s and instances are read in one statement, so that an instance + # committed with its `provider_ref` record is never missed while the record is seen. + res = await session.execute( + select(ProviderRefModel, InstanceModel) + .outerjoin( + InstanceModel, InstanceModel.provider_ref == ProviderRefModel.provider_ref + ) + .where( + ProviderRefModel.project_id == project_id, + ProviderRefModel.owner_type == OwnerType.INSTANCE, + ProviderRefModel.provider_ref.in_(batch), + ) + .options( + load_only( + InstanceModel.id, + InstanceModel.project_id, + InstanceModel.fleet_id, + InstanceModel.name, + InstanceModel.provider_ref, + InstanceModel.status, + InstanceModel.started_at, + InstanceModel.finished_at, + InstanceModel.job_provisioning_data, + ) + ) + ) + for provider_ref_model, instance_model in res.all(): + owner_records[provider_ref_model.provider_ref] = _OwnerRecord( + provider_ref_model=provider_ref_model, instance_model=instance_model + ) + return owner_records + + +def _get_resource_group_decision( + orphan_cleanup_model: OrphanCleanupModel, + resource_group: ListedResourceGroup, + owner_records: _OwnerRecords, + listed_at: datetime, +) -> _ResourceGroupDecision: + if resource_group.owner_type == OwnerType.INSTANCE: + if resource_group.provider_ref not in owner_records: + return _ResourceGroupDecision.IGNORE + instance_model = owner_records[resource_group.provider_ref].instance_model + if instance_model is None: + # Placeholder instances are deleted when their jobs are terminated, + # so `run_job()` can no longer be creating resources for them. + return _ResourceGroupDecision.ORPHAN + return _get_instance_group_decision( + orphan_cleanup_model, resource_group, instance_model, listed_at + ) + # Backends may list owner types that orphan cleanup does not support yet + logger.warning( + "Ignoring resource group %s in %s/%s with unsupported owner type %s", + resource_group.provider_ref, + orphan_cleanup_model.backend.project.name, + orphan_cleanup_model.backend.type.value, + resource_group.owner_type.value, + ) + return _ResourceGroupDecision.IGNORE + + +async def _handle_orphaned_resource_group( + compute: ComputeWithResourceListingSupport, + orphan_cleanup_model: OrphanCleanupModel, + resource_group: ListedResourceGroup, +) -> _ResourceGroupDecision: + if settings.ORPHAN_CLEANUP_MODE == OrphanCleanupMode.REPORT: + logger.info( + "Found orphaned %s %s in %s/%s", + resource_group.owner_type.value, + resource_group.provider_ref, + orphan_cleanup_model.backend.project.name, + orphan_cleanup_model.backend.type.value, + ) + sentry_sdk.capture_message( + "Detected orphaned {} {} in {}/{}".format( + resource_group.owner_type.value, + resource_group.provider_ref, + orphan_cleanup_model.backend.project.name, + orphan_cleanup_model.backend.type.value, + ) + ) + return _ResourceGroupDecision.ORPHAN + try: + await run_async(compute.delete_resource_group, resource_group) + except NotYetTerminated: + # Backends do not list owner resources whose deletion is in progress, + # so a listed owner resource has just been requested to be deleted. + if resource_group.primary_id is not None: + logger.info( + "Started cleaning up orphaned %s %s in %s/%s", + resource_group.owner_type.value, + resource_group.provider_ref, + orphan_cleanup_model.backend.project.name, + orphan_cleanup_model.backend.type.value, + ) + return _ResourceGroupDecision.DELETING + logger.info( + "Cleanup of orphaned %s %s in %s/%s is skipped because the termination is in progress", + resource_group.owner_type.value, + resource_group.provider_ref, + orphan_cleanup_model.backend.project.name, + orphan_cleanup_model.backend.type.value, + ) + return _ResourceGroupDecision.SKIP + except Exception as e: + if isinstance(e, BackendError): + logger.warning( + "Cleanup of orphaned %s %s in %s/%s is skipped due to error: %s", + resource_group.owner_type.value, + resource_group.provider_ref, + orphan_cleanup_model.backend.project.name, + orphan_cleanup_model.backend.type.value, + repr(e), + ) + else: + logger.exception( + "Cleanup of orphaned %s %s in %s/%s is skipped because of unexpected exception", + resource_group.owner_type.value, + resource_group.provider_ref, + orphan_cleanup_model.backend.project.name, + orphan_cleanup_model.backend.type.value, + ) + return _ResourceGroupDecision.DELETE_FAILED + logger.info( + "Cleaned up orphaned %s %s in %s/%s", + resource_group.owner_type.value, + resource_group.provider_ref, + orphan_cleanup_model.backend.project.name, + orphan_cleanup_model.backend.type.value, + ) + return _ResourceGroupDecision.ORPHAN + + +def _get_orphan_event( + resource_group: ListedResourceGroup, + owner_records: _OwnerRecords, + decision: _ResourceGroupDecision, +) -> _OrphanEvent: + assert resource_group.owner_type == OwnerType.INSTANCE + owner_record = owner_records[resource_group.provider_ref] + if owner_record.instance_model is not None: + target = events.Target.from_model(owner_record.instance_model) + else: + # Targeting the deleted instance puts the event next to the instance's earlier events, + # e.g. the job it was created for. + provider_ref_model = owner_record.provider_ref_model + target = events.Target( + type=EventTargetType.INSTANCE, + project_id=provider_ref_model.project_id, + id=provider_ref_model.owner_id, + name=provider_ref_model.owner_name, + ) + return _OrphanEvent( + message=_get_orphan_event_message(resource_group, decision), + target=target, + ) + + +def _get_orphan_event_message( + resource_group: ListedResourceGroup, decision: _ResourceGroupDecision +) -> str: + # Built from the required fields so that the message is useful + # even if the backend fills `description` poorly. + if settings.ORPHAN_CLEANUP_MODE == OrphanCleanupMode.REPORT: + action = "found" + elif decision == _ResourceGroupDecision.DELETING: + action = "being deleted" + else: + action = "deleted" + location = resource_group.backend.value + if resource_group.region is not None: + location += f"/{resource_group.region}" + message = ( + f"Orphaned resources {action} in {location}, provider_ref {resource_group.provider_ref}" + ) + if resource_group.primary_id is not None: + message += f", {resource_group.owner_type.value} id {resource_group.primary_id}" + if resource_group.description: + message += f": {resource_group.description}" + return message + + +_STUCK_THRESHOLD = timedelta(hours=1) + +_LISTING_MARGIN = timedelta(minutes=10) +""" +Owners changed later than `listed_at - _LISTING_MARGIN` are not judged by the listing, +and `clean_at` is set to `listed_at - _LISTING_MARGIN`. The margin covers clock skew +between replicas, the lag between setting timestamps and committing them, +and providers listing new resources with a delay. +""" + + +def _get_instance_group_decision( + orphan_cleanup_model: OrphanCleanupModel, + resource_group: ListedResourceGroup, + instance_model: InstanceModel, + listed_at: datetime, +) -> _ResourceGroupDecision: + if instance_model.status == InstanceStatus.PENDING: + if _is_stuck(resource_group, listed_at): + # Stuck pending instances are unexpected. If this changes, + # we may treat them as orphans instead of error log and skip. + logger.error( + "Resource group %s in %s/%s was created at %s for instance %s" + " that is still pending. The instance may be stuck in pending.", + resource_group.provider_ref, + orphan_cleanup_model.backend.project.name, + orphan_cleanup_model.backend.type.value, + resource_group.created_at, + instance_model.name, + ) + return _ResourceGroupDecision.SKIP + if instance_model.status == InstanceStatus.TERMINATING: + return _ResourceGroupDecision.SKIP + if instance_model.status == InstanceStatus.TERMINATED: + # The listing may predate the termination + if _changed_after(instance_model.finished_at, listed_at - _LISTING_MARGIN): + return _ResourceGroupDecision.SKIP + return _ResourceGroupDecision.ORPHAN + provisioning_data = get_instance_provisioning_data(instance_model) + if provisioning_data is None: + logger.error( + "Failed to classify resource group %s in %s/%s:" + " provisioning_data is None for active instance %s", + resource_group.provider_ref, + orphan_cleanup_model.backend.project.name, + orphan_cleanup_model.backend.type.value, + instance_model.name, + ) + return _ResourceGroupDecision.IGNORE + if ( + provisioning_data.get_base_backend() != resource_group.backend + or provisioning_data.instance_id != resource_group.primary_id + ): + # The listing may predate the provisioning: the instance may be missing from it + # while its internal resources created earlier are listed. + # `started_at` is set together with the provisioning data. + if _changed_after(instance_model.started_at, listed_at - _LISTING_MARGIN): + return _ResourceGroupDecision.SKIP + return _ResourceGroupDecision.ORPHAN + return _ResourceGroupDecision.TRACKED + + +def _changed_after(changed_at: Optional[datetime], at: datetime) -> bool: + return changed_at is not None and changed_at >= at + + +def _is_stuck(resource_group: ListedResourceGroup, listed_at: datetime) -> bool: + return ( + resource_group.created_at is not None + and resource_group.created_at < listed_at - _STUCK_THRESHOLD + ) + + +def _get_backoff_cleanup_process_result( + orphan_cleanup_model: OrphanCleanupModel, +) -> _ProcessResult: + return _ProcessResult( + update_map=_OrphanCleanupUpdateMap( + cleanup_attempt=orphan_cleanup_model.cleanup_attempt + 1 + ) + ) + + +def _get_unsettled_cleanup_process_result() -> _ProcessResult: + # Some groups were skipped or some parts could not be listed. + # The backend stays eligible since `clean_at` is not updated. + # It's retried after the cleanup interval rather than backed off, + # since skipped groups are expected to settle and the listed parts + # should be cleaned up as usual while other parts fail. + return _ProcessResult(update_map=_OrphanCleanupUpdateMap(cleanup_attempt=0)) + + +def _get_success_cleanup_process_result(listed_at: datetime) -> _ProcessResult: + return _ProcessResult( + update_map=_OrphanCleanupUpdateMap( + cleanup_attempt=0, + clean_at=listed_at - _LISTING_MARGIN, + clean_mode=settings.ORPHAN_CLEANUP_MODE, + ) + ) + + +async def _apply_process_result(item: PipelineItem, result: _ProcessResult): + update_map = _OrphanCleanupUpdateMap() + update_map.update(result.update_map) + set_processed_update_map_fields(update_map) + set_unlock_update_map_fields(update_map) + async with get_session_ctx() as session: + now = get_current_datetime() + resolve_now_placeholders(update_map, now=now) + res = await session.execute( + update(OrphanCleanupModel) + .where( + OrphanCleanupModel.id == item.id, + OrphanCleanupModel.lock_token == item.lock_token, + ) + .values(**update_map) + .returning(OrphanCleanupModel.id) + ) + updated_ids = list(res.scalars().all()) + if len(updated_ids) == 0: + log_lock_token_changed_after_processing(logger, item) + return + for orphan_event in result.orphan_events: + events.emit( + session, + orphan_event.message, + actor=events.SystemActor(), + targets=[orphan_event.target], + ) diff --git a/src/dstack/_internal/server/background/scheduled_tasks/__init__.py b/src/dstack/_internal/server/background/scheduled_tasks/__init__.py index 0e929811e6..46408d2602 100644 --- a/src/dstack/_internal/server/background/scheduled_tasks/__init__.py +++ b/src/dstack/_internal/server/background/scheduled_tasks/__init__.py @@ -2,6 +2,7 @@ from apscheduler.triggers.date import DateTrigger from apscheduler.triggers.interval import IntervalTrigger +from dstack._internal.core.models.backends.orphans import OrphanCleanupMode from dstack._internal.server import settings from dstack._internal.server.background.scheduled_tasks.events import delete_events from dstack._internal.server.background.scheduled_tasks.gateways import ( @@ -21,6 +22,9 @@ from dstack._internal.server.background.scheduled_tasks.offers_catalog import ( preload_offers_catalog, ) +from dstack._internal.server.background.scheduled_tasks.orphan_cleanup_enrollment import ( + enroll_backends_in_orphan_cleanup, +) from dstack._internal.server.background.scheduled_tasks.probes import process_probes from dstack._internal.server.background.scheduled_tasks.prometheus_metrics import ( collect_prometheus_metrics, @@ -53,6 +57,10 @@ def start_scheduled_tasks() -> AsyncIOScheduler: process_idle_volumes, IntervalTrigger(seconds=60, jitter=10), max_instances=1 ) _scheduler.add_job(delete_instance_healthchecks, IntervalTrigger(minutes=5), max_instances=1) + if settings.ORPHAN_CLEANUP_MODE != OrphanCleanupMode.OFF: + _scheduler.add_job( + enroll_backends_in_orphan_cleanup, IntervalTrigger(minutes=10), max_instances=1 + ) if settings.ENABLE_PROMETHEUS_METRICS: _scheduler.add_job( collect_prometheus_metrics, IntervalTrigger(seconds=10), max_instances=1 diff --git a/src/dstack/_internal/server/background/scheduled_tasks/orphan_cleanup_enrollment.py b/src/dstack/_internal/server/background/scheduled_tasks/orphan_cleanup_enrollment.py new file mode 100644 index 0000000000..523a136bb8 --- /dev/null +++ b/src/dstack/_internal/server/background/scheduled_tasks/orphan_cleanup_enrollment.py @@ -0,0 +1,86 @@ +import uuid + +from sqlalchemy import select +from sqlalchemy.dialects import postgresql, sqlite +from sqlalchemy.ext.asyncio import AsyncSession + +from dstack._internal.core.backends.features import get_backends_with_resource_listing_support +from dstack._internal.server.db import get_db, get_session_ctx +from dstack._internal.server.models import ( + BackendModel, + OrphanCleanupModel, + ProjectModel, + ProviderRefModel, +) +from dstack._internal.server.utils import tracing +from dstack._internal.utils.common import get_current_datetime + +_ENROLLMENT_BATCH_SIZE = 100 + + +@tracing.instrument_scheduled_task +async def enroll_backends_in_orphan_cleanup(): + """ + Creates `OrphanCleanupModel` rows for backends that support resource listing, + so that `OrphanCleanupPipeline` starts processing them. + + Only backends of projects that have issued `provider_ref`s are enrolled, + since other backends cannot have orphaned resources that can be matched. + + At most `_ENROLLMENT_BATCH_SIZE` backends are enrolled per run. This spreads out + the first cleanups when many backends become eligible at once, e.g. after an upgrade. + """ + async with get_session_ctx() as session: + backend_ids = await _select_backend_ids_to_enroll(session) + if len(backend_ids) == 0: + return + now = get_current_datetime() + await session.execute( + # Another replica may enroll the same backends concurrently + _insert(OrphanCleanupModel) + .values( + [ + { + "id": uuid.uuid4(), + "backend_id": backend_id, + "created_at": now, + "last_processed_at": now, + } + for backend_id in backend_ids + ] + ) + .on_conflict_do_nothing(index_elements=[OrphanCleanupModel.backend_id]) + ) + + +async def _select_backend_ids_to_enroll(session: AsyncSession) -> list[uuid.UUID]: + has_provider_refs = ( + select(ProviderRefModel.provider_ref) + .where(ProviderRefModel.project_id == BackendModel.project_id) + .exists() + ) + is_enrolled = ( + select(OrphanCleanupModel.id) + .where(OrphanCleanupModel.backend_id == BackendModel.id) + .exists() + ) + res = await session.execute( + select(BackendModel.id) + .join(ProjectModel, ProjectModel.id == BackendModel.project_id) + .where( + ProjectModel.deleted == False, + BackendModel.type.in_(get_backends_with_resource_listing_support()), + ~is_enrolled, + has_provider_refs, + ) + .limit(_ENROLLMENT_BATCH_SIZE) + ) + return list(res.scalars().all()) + + +def _insert(model: type[OrphanCleanupModel]): + # The generic `sqlalchemy.insert()` does not support `ON CONFLICT`. + # Postgres and SQLite inserts both provide `on_conflict_do_nothing()`. + if get_db().dialect_name == "postgresql": + return postgresql.insert(model) + return sqlite.insert(model) diff --git a/src/dstack/_internal/server/migrations/versions/2026/10_19_1003_aff11cf0ff72_add_instancemodel_provider_ref.py b/src/dstack/_internal/server/migrations/versions/2026/10_19_1003_aff11cf0ff72_add_instancemodel_provider_ref.py new file mode 100644 index 0000000000..4733c3469b --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/10_19_1003_aff11cf0ff72_add_instancemodel_provider_ref.py @@ -0,0 +1,26 @@ +"""Add InstanceModel.provider_ref + +Revision ID: aff11cf0ff72 +Revises: 40f19bea9d2a +Create Date: 2026-10-19 10:03:45.638994+00:00 + +""" + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "aff11cf0ff72" +down_revision = "40f19bea9d2a" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + with op.batch_alter_table("instances", schema=None) as batch_op: + batch_op.add_column(sa.Column("provider_ref", sa.String(length=12), nullable=True)) + + +def downgrade() -> None: + with op.batch_alter_table("instances", schema=None) as batch_op: + batch_op.drop_column("provider_ref") diff --git a/src/dstack/_internal/server/migrations/versions/2026/10_19_1004_9caf3d6a1d09_add_instances_provider_ref_and_project_id_indexes.py b/src/dstack/_internal/server/migrations/versions/2026/10_19_1004_9caf3d6a1d09_add_instances_provider_ref_and_project_id_indexes.py new file mode 100644 index 0000000000..d2d208f5e4 --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/10_19_1004_9caf3d6a1d09_add_instances_provider_ref_and_project_id_indexes.py @@ -0,0 +1,50 @@ +"""Add ix_instances_provider_ref and ix_instances_project_id + +Revision ID: 9caf3d6a1d09 +Revises: aff11cf0ff72 +Create Date: 2026-10-19 10:04:00.000000+00:00 + +""" + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "9caf3d6a1d09" +down_revision = "aff11cf0ff72" +branch_labels = None +depends_on = None + + +TABLE_NAME = "instances" +INDEXES = [ + ("ix_instances_provider_ref", ["provider_ref"], True), + ("ix_instances_project_id", ["project_id"], False), +] + + +def _drop_index(name: str) -> None: + op.drop_index( + name, + table_name=TABLE_NAME, + if_exists=True, + postgresql_concurrently=True, + ) + + +def upgrade() -> None: + with op.get_context().autocommit_block(): + for name, columns, unique in INDEXES: + _drop_index(name) + op.create_index( + name, + TABLE_NAME, + columns, + unique=unique, + postgresql_concurrently=True, + ) + + +def downgrade() -> None: + with op.get_context().autocommit_block(): + for name, _, _ in INDEXES: + _drop_index(name) diff --git a/src/dstack/_internal/server/migrations/versions/2026/10_19_1005_d0f36d5050ea_add_orphancleanupmodel.py b/src/dstack/_internal/server/migrations/versions/2026/10_19_1005_d0f36d5050ea_add_orphancleanupmodel.py new file mode 100644 index 0000000000..2ee92a29b9 --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/10_19_1005_d0f36d5050ea_add_orphancleanupmodel.py @@ -0,0 +1,67 @@ +"""Add OrphanCleanupModel + +Revision ID: d0f36d5050ea +Revises: 9caf3d6a1d09 +Create Date: 2026-10-19 10:05:33.775566+00:00 + +""" + +import sqlalchemy as sa +import sqlalchemy_utils +from alembic import op + +import dstack._internal.server.models + +# revision identifiers, used by Alembic. +revision = "d0f36d5050ea" +down_revision = "9caf3d6a1d09" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "orphan_cleanups", + sa.Column("id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False), + sa.Column( + "backend_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False + ), + sa.Column("created_at", dstack._internal.server.models.NaiveDateTime(), nullable=False), + sa.Column( + "last_processed_at", dstack._internal.server.models.NaiveDateTime(), nullable=False + ), + sa.Column("cleanup_attempt", sa.Integer(), server_default="0", nullable=False), + sa.Column("clean_at", dstack._internal.server.models.NaiveDateTime(), nullable=True), + sa.Column("clean_mode", sa.String(length=100), nullable=True), + sa.Column( + "lock_expires_at", dstack._internal.server.models.NaiveDateTime(), nullable=True + ), + sa.Column("lock_token", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=True), + sa.Column("lock_owner", sa.String(length=100), nullable=True), + sa.ForeignKeyConstraint( + ["backend_id"], + ["backends.id"], + name=op.f("fk_orphan_cleanups_backend_id_backends"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_orphan_cleanups")), + sa.UniqueConstraint("backend_id", name=op.f("uq_orphan_cleanups_backend_id")), + ) + with op.batch_alter_table("orphan_cleanups", schema=None) as batch_op: + batch_op.create_index( + "ix_orphan_cleanups_pipeline_fetch_q", + [sa.literal_column("last_processed_at ASC")], + unique=False, + ) + + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("orphan_cleanups", schema=None) as batch_op: + batch_op.drop_index("ix_orphan_cleanups_pipeline_fetch_q") + + op.drop_table("orphan_cleanups") + # ### end Alembic commands ### diff --git a/src/dstack/_internal/server/migrations/versions/2026/10_22_0901_2e7479a16e36_add_providerrefmodel.py b/src/dstack/_internal/server/migrations/versions/2026/10_22_0901_2e7479a16e36_add_providerrefmodel.py new file mode 100644 index 0000000000..a83b8f1a02 --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/10_22_0901_2e7479a16e36_add_providerrefmodel.py @@ -0,0 +1,56 @@ +"""Add ProviderRefModel + +Revision ID: 2e7479a16e36 +Revises: d0f36d5050ea +Create Date: 2026-10-22 09:01:29.598613+00:00 + +""" + +import sqlalchemy as sa +import sqlalchemy_utils +from alembic import op + +import dstack._internal.server.models + +# revision identifiers, used by Alembic. +revision = "2e7479a16e36" +down_revision = "d0f36d5050ea" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "provider_refs", + sa.Column("provider_ref", sa.String(length=12), nullable=False), + sa.Column( + "project_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False + ), + sa.Column("owner_type", sa.String(length=100), nullable=False), + sa.Column("owner_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False), + sa.Column("owner_name", sa.String(length=50), nullable=False), + sa.Column("created_at", dstack._internal.server.models.NaiveDateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["project_id"], + ["projects.id"], + name=op.f("fk_provider_refs_project_id_projects"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("provider_ref", name=op.f("pk_provider_refs")), + ) + with op.batch_alter_table("provider_refs", schema=None) as batch_op: + batch_op.create_index( + "ix_provider_refs_project_id_created_at", ["project_id", "created_at"], unique=False + ) + + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("provider_refs", schema=None) as batch_op: + batch_op.drop_index("ix_provider_refs_project_id_created_at") + + op.drop_table("provider_refs") + # ### end Alembic commands ### diff --git a/src/dstack/_internal/server/migrations/versions/2026/10_22_0901_ca5fb194134e_add_ix_jobs_project_id.py b/src/dstack/_internal/server/migrations/versions/2026/10_22_0901_ca5fb194134e_add_ix_jobs_project_id.py new file mode 100644 index 0000000000..c433c1c030 --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/10_22_0901_ca5fb194134e_add_ix_jobs_project_id.py @@ -0,0 +1,42 @@ +"""Add ix_jobs_project_id + +Revision ID: ca5fb194134e +Revises: 2e7479a16e36 +Create Date: 2026-10-22 09:01:45.369057+00:00 + +""" + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "ca5fb194134e" +down_revision = "2e7479a16e36" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + with op.get_context().autocommit_block(): + op.drop_index( + "ix_jobs_project_id", + table_name="jobs", + if_exists=True, + postgresql_concurrently=True, + ) + op.create_index( + "ix_jobs_project_id", + "jobs", + ["project_id"], + unique=False, + postgresql_concurrently=True, + ) + + +def downgrade() -> None: + with op.get_context().autocommit_block(): + op.drop_index( + "ix_jobs_project_id", + table_name="jobs", + if_exists=True, + postgresql_concurrently=True, + ) diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py index a80f4922a4..71db4aeab1 100644 --- a/src/dstack/_internal/server/models.py +++ b/src/dstack/_internal/server/models.py @@ -25,6 +25,7 @@ from dstack._internal.core.errors import DstackError from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.backends.orphans import OrphanCleanupMode, OwnerType from dstack._internal.core.models.common import CoreModel from dstack._internal.core.models.compute_groups import ComputeGroupStatus from dstack._internal.core.models.events import EventTargetType @@ -516,7 +517,9 @@ class JobModel(PipelineModelMixin, BaseModel): UUIDType(binary=False), primary_key=True, default=uuid.uuid4 ) - project_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("projects.id", ondelete="CASCADE")) + project_id: Mapped[uuid.UUID] = mapped_column( + ForeignKey("projects.id", ondelete="CASCADE"), index=True + ) project: Mapped["ProjectModel"] = relationship() run_id: Mapped[uuid.UUID] = mapped_column( @@ -935,7 +938,9 @@ class InstanceModel(PipelineModelMixin, BaseModel): deleted: Mapped[bool] = mapped_column(Boolean, default=False) deleted_at: Mapped[Optional[datetime]] = mapped_column(NaiveDateTime) - project_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("projects.id", ondelete="CASCADE")) + project_id: Mapped[uuid.UUID] = mapped_column( + ForeignKey("projects.id", ondelete="CASCADE"), index=True + ) project: Mapped["ProjectModel"] = relationship(foreign_keys=[project_id]) # TODO: drop `pool_id` after the release without pools. @@ -1038,6 +1043,11 @@ class InstanceModel(PipelineModelMixin, BaseModel): cascade=CASCADE_DEFAULT_WITH_DELETE_ORPHAN, ) + provider_ref: Mapped[Optional[str]] = mapped_column(String(12), unique=True, index=True) + """Set on provider resources for orphans matching. Recorded in `ProviderRefModel`. + Placeholder instances get it on creation and keep it when promoted, so `run_job()` + resources carry it too. `None` for older instances and compute group instances.""" + __table_args__ = ( Index( "ix_instances_pipeline_fetch_q", @@ -1439,3 +1449,56 @@ class UserPublicKeyModel(BaseModel): """`fingerprint` stores a key digest in the format used by OpenSSH: `SHA256:`.""" key: Mapped[str] = mapped_column(Text) """`key` stores a public key in the OpenSSH disk (ASCII-armored) format.""" + + +class OrphanCleanupModel(PipelineModelMixin, BaseModel): + __tablename__ = "orphan_cleanups" + + id: Mapped[uuid.UUID] = mapped_column( + UUIDType(binary=False), primary_key=True, default=uuid.uuid4 + ) + backend_id: Mapped[uuid.UUID] = mapped_column( + ForeignKey("backends.id", ondelete="CASCADE"), unique=True + ) + backend: Mapped["BackendModel"] = relationship() + created_at: Mapped[datetime] = mapped_column(NaiveDateTime, default=get_current_datetime) + last_processed_at: Mapped[datetime] = mapped_column( + NaiveDateTime, default=get_current_datetime + ) + cleanup_attempt: Mapped[int] = mapped_column(Integer, server_default="0") + """Consecutive failed cleanups. Determines the retry delay.""" + clean_at: Mapped[Optional[datetime]] = mapped_column(NaiveDateTime) + """Start of the last clean cleanup minus a margin for timing inaccuracies. + Later instance or job activity makes the backend eligible for cleanup.""" + clean_mode: Mapped[Optional[OrphanCleanupMode]] = mapped_column( + EnumAsString(OrphanCleanupMode, 100) + ) + """Mode `clean_at` was reached in. The backend is eligible for cleanup again if the current mode differs.""" + + __table_args__ = (Index("ix_orphan_cleanups_pipeline_fetch_q", last_processed_at.asc()),) + + +class ProviderRefModel(BaseModel): + """ + A `provider_ref` issued by the server, committed with its owner record before any provider + call that uses it. Rows are not deleted together with owners, so that resources of deleted + owners (e.g. placeholder instances of failed `run_job()` calls) are still matched as orphans. + Resources with a `provider_ref` that has no row are not owned by the project. + """ + + # TODO: Delete rows of deleted or terminated owners after a retention period. + # A row is added per fleet instance and per placeholder instance, i.e. per job submission + # that provisions new capacity, including no-capacity retries. + + __tablename__ = "provider_refs" + + provider_ref: Mapped[str] = mapped_column(String(12), primary_key=True) + project_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("projects.id", ondelete="CASCADE")) + owner_type: Mapped[OwnerType] = mapped_column(EnumAsString(OwnerType, 100)) + owner_id: Mapped[uuid.UUID] = mapped_column(UUIDType(binary=False)) + """Not a foreign key, since owners may be deleted. Kept to report orphans of deleted owners + in the owner's events.""" + owner_name: Mapped[str] = mapped_column(String(50)) + created_at: Mapped[datetime] = mapped_column(NaiveDateTime, default=get_current_datetime) + + __table_args__ = (Index("ix_provider_refs_project_id_created_at", project_id, created_at),) diff --git a/src/dstack/_internal/server/services/instances.py b/src/dstack/_internal/server/services/instances.py index cdcab7ce34..ef3f15efb6 100644 --- a/src/dstack/_internal/server/services/instances.py +++ b/src/dstack/_internal/server/services/instances.py @@ -16,6 +16,7 @@ from dstack._internal.core.backends.features import BACKENDS_WITH_MULTINODE_SUPPORT from dstack._internal.core.errors import ResourceNotExistsError from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.backends.orphans import OwnerType from dstack._internal.core.models.common import EntityReference, validate_json_extra_ignore from dstack._internal.core.models.envs import Env from dstack._internal.core.models.health import HealthCheck, HealthEvent, HealthStatus @@ -60,6 +61,7 @@ from dstack._internal.server.services import events from dstack._internal.server.services.logging import fmt from dstack._internal.server.services.offers import generate_shared_offer +from dstack._internal.server.services.orphan_cleanup import issue_provider_ref from dstack._internal.server.services.projects import list_user_project_models from dstack._internal.server.services.runner.client import ShimClient from dstack._internal.utils import common as common_utils @@ -902,6 +904,13 @@ def create_instance_model( ) if instance_id is None: instance_id = uuid.uuid4() + provider_ref = issue_provider_ref( + session=session, + project=project, + owner_type=OwnerType.INSTANCE, + owner_id=instance_id, + owner_name=instance_name, + ) project_ssh_key = SSHKey( public=project.ssh_public_key.strip(), private=project.ssh_private_key.strip(), @@ -914,6 +923,7 @@ def create_instance_model( instance_id=str(instance_id), reservation=reservation, tags=tags, + provider_ref=provider_ref, ) now = common_utils.get_current_datetime() instance = InstanceModel( @@ -932,6 +942,7 @@ def create_instance_model( termination_idle_time=termination_idle_time, total_blocks=None if blocks == "auto" else blocks, busy_blocks=0, + provider_ref=provider_ref, ) session.add(instance) return instance diff --git a/src/dstack/_internal/server/services/orphan_cleanup.py b/src/dstack/_internal/server/services/orphan_cleanup.py new file mode 100644 index 0000000000..674662f63c --- /dev/null +++ b/src/dstack/_internal/server/services/orphan_cleanup.py @@ -0,0 +1,32 @@ +import uuid + +from sqlalchemy.ext.asyncio import AsyncSession + +from dstack._internal.core.backends.base.compute import generate_provider_ref +from dstack._internal.core.models.backends.orphans import OwnerType +from dstack._internal.server.models import ProjectModel, ProviderRefModel + + +def issue_provider_ref( + session: AsyncSession, + project: ProjectModel, + owner_type: OwnerType, + owner_id: uuid.UUID, + owner_name: str, +) -> str: + """ + Generates a `provider_ref` for a new owner record and records it in `ProviderRefModel`. + The record must be committed in the same transaction as the owner record, + before any provider call that uses the `provider_ref`. + """ + provider_ref = generate_provider_ref() + session.add( + ProviderRefModel( + provider_ref=provider_ref, + project_id=project.id, + owner_type=owner_type, + owner_id=owner_id, + owner_name=owner_name, + ) + ) + return provider_ref diff --git a/src/dstack/_internal/server/settings.py b/src/dstack/_internal/server/settings.py index c077feb4ea..806c45b53b 100644 --- a/src/dstack/_internal/server/settings.py +++ b/src/dstack/_internal/server/settings.py @@ -7,6 +7,7 @@ from pathlib import Path from typing import Optional +from dstack._internal.core.models.backends.orphans import OrphanCleanupMode from dstack._internal.server.utils.settings import parse_hostname_port from dstack._internal.utils.env import environ from dstack._internal.utils.logging import get_logger @@ -235,3 +236,13 @@ class JobNetworkMode(Enum): value_type=int, default=DEFAULT_JOB_NETWORK_MODE, ) + +ORPHAN_CLEANUP_INTERVAL_SECONDS = environ.get_int( + "DSTACK_SERVER_ORPHAN_CLEANUP_INTERVAL_SECONDS", default=3600 +) + + +DEFAULT_ORPHAN_CLEANUP_MODE = OrphanCleanupMode.REPORT +ORPHAN_CLEANUP_MODE = environ.get_enum( + "DSTACK_SERVER_ORPHAN_CLEANUP_MODE", OrphanCleanupMode, default=DEFAULT_ORPHAN_CLEANUP_MODE +) diff --git a/src/tests/_internal/core/backends/aws/test_compute.py b/src/tests/_internal/core/backends/aws/test_compute.py new file mode 100644 index 0000000000..e2d11e18d8 --- /dev/null +++ b/src/tests/_internal/core/backends/aws/test_compute.py @@ -0,0 +1,279 @@ +import json +from datetime import datetime, timezone +from typing import Iterator +from unittest.mock import Mock + +import boto3 +import pytest +from botocore.stub import Stubber + +from dstack._internal.core.backends.aws.compute import AWSCompute +from dstack._internal.core.backends.aws.models import AWSAccessKeyCreds, AWSConfig +from dstack._internal.core.backends.aws.resources import allocate_and_associate_eip +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 + +REGION = "us-east-1" +PROVIDER_REF = "abc123def456" +LAUNCH_TIME = datetime(2026, 1, 1, 12, tzinfo=timezone.utc) +PROVIDER_REF_FILTER = {"Name": "tag-key", "Values": ["dstack_provider_ref"]} + + +@pytest.fixture +def ec2_client() -> Iterator[Stubber]: + client = boto3.client( + "ec2", + region_name=REGION, + aws_access_key_id="test", + aws_secret_access_key="test", + ) + with Stubber(client) as stubber: + yield stubber + stubber.assert_no_pending_responses() + + +@pytest.fixture +def compute(ec2_client: Stubber) -> AWSCompute: + compute = AWSCompute( + AWSConfig( + creds=AWSAccessKeyCreds(access_key="test", secret_key="test"), + regions=[REGION], + ) + ) + compute.session = Mock() + compute.session.client.return_value = ec2_client.client + return compute + + +def _get_instance(instance_id: str, provider_ref: str = PROVIDER_REF) -> dict: + return { + "InstanceId": instance_id, + "InstanceType": "p5.48xlarge", + "LaunchTime": LAUNCH_TIME, + "State": {"Code": 16, "Name": "running"}, + "Tags": [ + {"Key": "owner", "Value": "dstack"}, + {"Key": "dstack_provider_ref", "Value": provider_ref}, + ], + } + + +def _get_address(allocation_id: str, instance_id: str) -> dict: + return { + "AllocationId": allocation_id, + "PublicIp": "203.0.113.10", + "Tags": [ + {"Key": "dstack_instance", "Value": instance_id}, + {"Key": "dstack_provider_ref", "Value": PROVIDER_REF}, + ], + } + + +def _add_list_responses(ec2_client: Stubber, instances: list[dict], addresses: list[dict]): + ec2_client.add_response( + "describe_addresses", + {"Addresses": addresses}, + {"Filters": [PROVIDER_REF_FILTER]}, + ) + ec2_client.add_response( + "describe_instances", + {"Reservations": [{"Instances": instances}]}, + { + "Filters": [ + PROVIDER_REF_FILTER, + { + "Name": "instance-state-name", + "Values": ["pending", "running", "stopping", "stopped"], + }, + ] + }, + ) + + +def _get_resource_group(eip_allocation_ids: list[str]) -> ListedResourceGroup: + return ListedResourceGroup( + backend=BackendType.AWS, + owner_type=OwnerType.INSTANCE, + provider_ref=PROVIDER_REF, + primary_id="i-1", + region=REGION, + created_at=LAUNCH_TIME, + description="instance i-1", + backend_data=json.dumps({"instance_id": "i-1", "eip_allocation_ids": eip_allocation_ids}), + ) + + +class TestListResourceGroups: + def test_groups_elastic_ip_with_its_instance(self, compute: AWSCompute, ec2_client: Stubber): + _add_list_responses( + ec2_client, + instances=[_get_instance("i-1")], + addresses=[_get_address("eipalloc-1", "i-1")], + ) + groups = compute.list_resource_groups().resource_groups + assert len(groups) == 1 + group = groups[0] + assert group.backend == BackendType.AWS + assert group.owner_type == OwnerType.INSTANCE + assert group.provider_ref == PROVIDER_REF + assert group.primary_id == "i-1" + assert group.region == REGION + assert group.created_at == LAUNCH_TIME + assert "eipalloc-1" in group.description + assert json.loads(group.backend_data or "") == { + "instance_id": "i-1", + "eip_allocation_ids": ["eipalloc-1"], + } + + def test_lists_elastic_ip_of_gone_instance_as_separate_group( + self, compute: AWSCompute, ec2_client: Stubber + ): + _add_list_responses( + ec2_client, instances=[], addresses=[_get_address("eipalloc-1", "i-gone")] + ) + groups = compute.list_resource_groups().resource_groups + assert len(groups) == 1 + group = groups[0] + assert group.provider_ref == PROVIDER_REF + assert group.primary_id is None + assert group.created_at is None + assert json.loads(group.backend_data or "") == { + "instance_id": None, + "eip_allocation_ids": ["eipalloc-1"], + } + + def test_skips_resources_with_invalid_provider_ref( + self, compute: AWSCompute, ec2_client: Stubber + ): + _add_list_responses( + ec2_client, instances=[_get_instance("i-1", provider_ref="Invalid-Ref")], addresses=[] + ) + assert compute.list_resource_groups().resource_groups == [] + + def test_lists_instances_without_elastic_ip_permissions( + self, compute: AWSCompute, ec2_client: Stubber + ): + ec2_client.add_client_error( + "describe_addresses", service_error_code="UnauthorizedOperation" + ) + ec2_client.add_response( + "describe_instances", + {"Reservations": [{"Instances": [_get_instance("i-1")]}]}, + ) + groups = compute.list_resource_groups().resource_groups + assert [g.primary_id for g in groups] == ["i-1"] + + def test_reports_failed_region_and_lists_other_regions(self, ec2_client: Stubber): + other_ec2_client = boto3.client( + "ec2", + region_name="eu-west-1", + aws_access_key_id="test", + aws_secret_access_key="test", + ) + compute = AWSCompute( + AWSConfig( + creds=AWSAccessKeyCreds(access_key="test", secret_key="test"), + regions=[REGION, "eu-west-1"], + ) + ) + compute.session = Mock() + compute.session.client.side_effect = lambda service, region_name: { + REGION: ec2_client.client, + "eu-west-1": other_ec2_client, + }[region_name] + _add_list_responses(ec2_client, instances=[_get_instance("i-1")], addresses=[]) + with Stubber(other_ec2_client) as other_stubber: + other_stubber.add_client_error( + "describe_addresses", service_error_code="RequestLimitExceeded" + ) + listing = compute.list_resource_groups() + assert [g.primary_id for g in listing.resource_groups] == ["i-1"] + assert [(f.backend, f.region) for f in listing.failures] == [ + (BackendType.AWS, "eu-west-1") + ] + + def test_raises_if_all_regions_fail(self, compute: AWSCompute, ec2_client: Stubber): + ec2_client.add_client_error( + "describe_addresses", service_error_code="RequestLimitExceeded" + ) + with pytest.raises(ComputeError): + compute.list_resource_groups() + + +class TestDeleteResourceGroup: + def test_terminates_instance_and_releases_elastic_ip( + self, compute: AWSCompute, ec2_client: Stubber + ): + group = _get_resource_group(eip_allocation_ids=["eipalloc-1"]) + ec2_client.add_response( + "terminate_instances", {"TerminatingInstances": []}, {"InstanceIds": ["i-1"]} + ) + ec2_client.add_response( + "describe_addresses", + {"Addresses": [{"AllocationId": "eipalloc-1", "AssociationId": "eipassoc-1"}]}, + {"AllocationIds": ["eipalloc-1"]}, + ) + ec2_client.add_response("disassociate_address", {}, {"AssociationId": "eipassoc-1"}) + ec2_client.add_response("release_address", {}, {"AllocationId": "eipalloc-1"}) + compute.delete_resource_group(group) + + def test_tolerates_already_terminated_instance(self, compute: AWSCompute, ec2_client: Stubber): + group = _get_resource_group(eip_allocation_ids=[]) + ec2_client.add_client_error( + "terminate_instances", + service_error_code="InvalidInstanceID.NotFound", + expected_params={"InstanceIds": ["i-1"]}, + ) + compute.delete_resource_group(group) + + def test_raises_compute_error_on_api_error(self, compute: AWSCompute, ec2_client: Stubber): + ec2_client.add_client_error( + "terminate_instances", service_error_code="UnauthorizedOperation" + ) + with pytest.raises(ComputeError): + compute.delete_resource_group(_get_resource_group(eip_allocation_ids=[])) + + +class TestAllocateAndAssociateEip: + def test_tags_elastic_ip_with_provider_ref(self, ec2_client: Stubber): + instance = Mock( + instance_id="i-1", + network_interfaces_attribute=[ + {"NetworkInterfaceId": "eni-1", "Attachment": {"DeviceIndex": 0}} + ], + ) + ec2_client.add_response( + "allocate_address", + {"AllocationId": "eipalloc-1", "PublicIp": "203.0.113.10"}, + { + "Domain": "vpc", + "TagSpecifications": [ + { + "ResourceType": "elastic-ip", + "Tags": [ + {"Key": "owner", "Value": "dstack"}, + {"Key": "dstack_instance", "Value": "i-1"}, + {"Key": "dstack_project", "Value": "main"}, + {"Key": "dstack_provider_ref", "Value": PROVIDER_REF}, + ], + } + ], + }, + ) + ec2_client.add_response( + "associate_address", + {"AssociationId": "eipassoc-1"}, + { + "AllocationId": "eipalloc-1", + "NetworkInterfaceId": "eni-1", + "AllowReassociation": False, + }, + ) + assert allocate_and_associate_eip( + ec2_client=ec2_client.client, + instance=instance, + project_name="main", + provider_ref=PROVIDER_REF, + backend_tags=None, + ) == ("203.0.113.10", "eipalloc-1") diff --git a/src/tests/_internal/core/backends/base/test_compute.py b/src/tests/_internal/core/backends/base/test_compute.py index c885bafcc3..f6c52f44bd 100644 --- a/src/tests/_internal/core/backends/base/test_compute.py +++ b/src/tests/_internal/core/backends/base/test_compute.py @@ -12,6 +12,7 @@ generate_unique_instance_name, generate_unique_volume_name, normalize_arch, + parse_provider_ref_from_name, ) from dstack._internal.core.models.instances import Disk, InstanceOfferWithAvailability from dstack._internal.core.models.resources import Memory, Range, ResourcesSpec @@ -91,6 +92,29 @@ def test_generates_name(self): name = generate_unique_instance_name(configuration, 60) assert re.match(r"^dstack-project1-my-instance-[a-z0-9]{8}$", name) + def test_generates_name_with_provider_ref(self): + configuration = get_instance_configuration( + project_name="project1", instance_name="my-instance" + ) + configuration.provider_ref = "abc123def456" + name = generate_unique_instance_name(configuration, 60, with_provider_ref=True) + assert re.match(r"^dstack-project1-my-instance-abc123def456[a-z0-9]{4}$", name) + + def test_ignores_provider_ref_by_default(self): + configuration = get_instance_configuration( + project_name="project1", instance_name="my-instance" + ) + configuration.provider_ref = "abc123def456" + name = generate_unique_instance_name(configuration, 60) + assert re.match(r"^dstack-project1-my-instance-[a-z0-9]{8}$", name) + + def test_falls_back_to_random_suffix_without_provider_ref(self): + configuration = get_instance_configuration( + project_name="project1", instance_name="my-instance" + ) + name = generate_unique_instance_name(configuration, 60, with_provider_ref=True) + assert re.match(r"^dstack-project1-my-instance-[a-z0-9]{8}$", name) + class TestGenerateUniqueGatewayInstanceName: def test_generates_name(self): @@ -121,10 +145,46 @@ def test_truncates_long_names(self): name = generate_unique_backend_name("a" * 100, "project", 30) assert re.match(r"^dstack-project-aaaaaa-[a-z0-9]{8}$", name) + def test_strips_trailing_hyphens_of_truncated_prefix(self): + name = generate_unique_backend_name("batch-inference", "acme-corp-production-us-east", 40) + assert re.match(r"^dstack-acme-corp-production-us-[a-z0-9]{8}$", name) + def test_validates_project_name(self): name = generate_unique_backend_name("instance", "invalid_project!@", 60) assert re.match(r"^dstack-instance-[a-z0-9]{8}$", name) + def test_generates_name_with_provider_ref(self): + name = generate_unique_backend_name("instance", "project", 60, "abc123def456") + assert re.match(r"^dstack-project-instance-abc123def456[a-z0-9]{4}$", name) + + def test_truncates_prefix_but_keeps_provider_ref(self): + name = generate_unique_backend_name("a" * 100, "project", 40, "abc123def456") + assert len(name) == 40 + assert re.match(r"^dstack-project-aaaaaaaa-abc123def456[a-z0-9]{4}$", name) + + +class TestParseProviderRefFromName: + @pytest.mark.parametrize( + ("project_name", "max_length"), + [("project", 60), (None, 60), ("project", 24)], + ) + def test_parses_generated_name(self, project_name: Optional[str], max_length: int): + name = generate_unique_backend_name("instance", project_name, max_length, "abc123def456") + assert parse_provider_ref_from_name(name) == "abc123def456" + + @pytest.mark.parametrize( + "name", + [ + "dstack-project-instance-1a2b3c4d", # name without provider_ref + "dstack-project-instance-abc123def456", # no random suffix + "dstack-project-instance-ABC123DEF456abcd", # uppercase + "other-project-instance-abc123def456abcd", # not a dstack name + "dstack-project-instance-abc123def456abcd-data", # unstripped internal suffix + ], + ) + def test_returns_none_for_names_without_provider_ref(self, name: str): + assert parse_provider_ref_from_name(name) is None + class TestNormalizeArch: @pytest.mark.parametrize( diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_termination.py b/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_termination.py index 3077fb8e75..d0d273cae4 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_termination.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_termination.py @@ -9,12 +9,19 @@ from dstack._internal.core.errors import BackendError, NotYetTerminated from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.backends.orphans import OrphanCleanupMode from dstack._internal.core.models.instances import InstanceStatus, InstanceTerminationReason +from dstack._internal.server import settings from dstack._internal.server.background.pipeline_tasks.instances import InstanceWorker from dstack._internal.server.background.pipeline_tasks.instances import ( termination as instances_termination, ) -from dstack._internal.server.testing.common import create_instance, create_project +from dstack._internal.server.models import InstanceModel +from dstack._internal.server.testing.common import ( + create_instance, + create_project, + get_job_provisioning_data, +) from tests._internal.server.background.pipeline_tasks.test_instances.helpers import ( instance_to_pipeline_item, lock_instance, @@ -217,3 +224,32 @@ async def test_terminate_on_termination_deadline( mock.assert_called_once() await session.refresh(instance) assert instance.status == InstanceStatus.TERMINATED + + +class TestIsDeletedByOrphanCleanup: + @pytest.mark.parametrize( + ("mode", "provider_ref", "backend", "expected"), + [ + (OrphanCleanupMode.DELETE, "abc123def456", BackendType.AWS, True), + (OrphanCleanupMode.REPORT, "abc123def456", BackendType.AWS, False), + (OrphanCleanupMode.OFF, "abc123def456", BackendType.AWS, False), + (OrphanCleanupMode.DELETE, None, BackendType.AWS, False), + (OrphanCleanupMode.DELETE, "abc123def456", BackendType.LAMBDA, False), + ], + ) + def test_is_deleted_by_orphan_cleanup( + self, + mode: OrphanCleanupMode, + provider_ref: Optional[str], + backend: BackendType, + expected: bool, + ): + instance = InstanceModel(provider_ref=provider_ref) + job_provisioning_data = get_job_provisioning_data(backend=backend) + with patch.object(settings, "ORPHAN_CLEANUP_MODE", mode): + assert ( + instances_termination._is_to_be_deleted_by_orphan_cleanup( + instance, job_provisioning_data + ) + == expected + ) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_orphan_cleanup.py b/src/tests/_internal/server/background/pipeline_tasks/test_orphan_cleanup.py new file mode 100644 index 0000000000..ccd2513230 --- /dev/null +++ b/src/tests/_internal/server/background/pipeline_tasks/test_orphan_cleanup.py @@ -0,0 +1,757 @@ +import asyncio +import uuid +from datetime import datetime, timedelta, timezone +from typing import Optional +from unittest.mock import AsyncMock, Mock, patch + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession + +from dstack._internal.core.backends.base.compute import ComputeWithResourceListingSupport +from dstack._internal.core.errors import BackendError, NotYetTerminated +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.backends.orphans import ( + ListedResourceGroup, + OrphanCleanupMode, + OwnerType, + ResourceListing, + ResourceListingFailure, +) +from dstack._internal.core.models.instances import InstanceStatus +from dstack._internal.core.models.runs import JobStatus +from dstack._internal.server import settings +from dstack._internal.server.background.pipeline_tasks import ( + orphan_cleanup as orphan_cleanup_pipeline, +) +from dstack._internal.server.background.pipeline_tasks.base import PipelineItem +from dstack._internal.server.background.pipeline_tasks.orphan_cleanup import ( + OrphanCleanupFetcher, + OrphanCleanupPipeline, + OrphanCleanupWorker, + _get_resource_group_decision, + _OwnerRecord, + _OwnerRecords, + _ResourceGroupDecision, +) +from dstack._internal.server.models import ( + BackendModel, + InstanceModel, + OrphanCleanupModel, + ProjectModel, + ProviderRefModel, +) +from dstack._internal.server.testing.common import ( + create_backend, + create_instance, + create_job, + create_project, + create_repo, + create_run, + create_user, + get_job_provisioning_data, + list_events, +) +from dstack._internal.utils.common import get_current_datetime + +LISTED_AT = datetime(2026, 1, 1, 12, tzinfo=timezone.utc) +PROVIDER_REF = "abc123def456" + + +def _get_orphan_cleanup_model() -> OrphanCleanupModel: + return OrphanCleanupModel( + backend=BackendModel(type=BackendType.AWS, project=ProjectModel(name="test-project")) + ) + + +def _get_resource_group( + primary_id: Optional[str] = "instance_id", + backend: BackendType = BackendType.AWS, + created_at: datetime = LISTED_AT - timedelta(minutes=5), +) -> ListedResourceGroup: + return ListedResourceGroup( + backend=backend, + owner_type=OwnerType.INSTANCE, + provider_ref=PROVIDER_REF, + primary_id=primary_id, + region="us-east-1", + created_at=created_at, + description="VM", + backend_data=None, + ) + + +def _get_instance_model( + status: InstanceStatus, + provisioned: bool = True, + started_at: Optional[datetime] = None, + finished_at: Optional[datetime] = None, +) -> InstanceModel: + job_provisioning_data = None + if provisioned: + # `instance_id` is "instance_id" + job_provisioning_data = get_job_provisioning_data(backend=BackendType.AWS) + return InstanceModel( + name="test-instance", + status=status, + provider_ref=PROVIDER_REF, + started_at=started_at, + finished_at=finished_at, + job_provisioning_data=( + job_provisioning_data.model_dump_json() if job_provisioning_data else None + ), + ) + + +def _get_decision( + resource_group: ListedResourceGroup, + instance_model: Optional[InstanceModel] = None, + issued: bool = True, +) -> _ResourceGroupDecision: + owner_records: _OwnerRecords = {} + if issued: + owner_records[PROVIDER_REF] = _OwnerRecord( + provider_ref_model=ProviderRefModel(provider_ref=PROVIDER_REF), + instance_model=instance_model, + ) + return _get_resource_group_decision( + _get_orphan_cleanup_model(), resource_group, owner_records, LISTED_AT + ) + + +class TestGetResourceGroupDecision: + def test_ignores_group_with_unknown_provider_ref(self): + assert _get_decision(_get_resource_group(), issued=False) == _ResourceGroupDecision.IGNORE + + def test_orphans_group_of_deleted_instance(self): + assert _get_decision(_get_resource_group()) == _ResourceGroupDecision.ORPHAN + + @pytest.mark.parametrize("status", [InstanceStatus.PENDING, InstanceStatus.TERMINATING]) + def test_skips_group_of_transitional_instance(self, status: InstanceStatus): + instance_model = _get_instance_model(status, provisioned=False) + decision = _get_decision(_get_resource_group(), instance_model=instance_model) + assert decision == _ResourceGroupDecision.SKIP + + def test_skips_old_group_of_pending_instance(self): + instance_model = _get_instance_model(InstanceStatus.PENDING, provisioned=False) + resource_group = _get_resource_group(created_at=LISTED_AT - timedelta(hours=2)) + decision = _get_decision(resource_group, instance_model=instance_model) + assert decision == _ResourceGroupDecision.SKIP + + @pytest.mark.parametrize("provisioned", [True, False]) + def test_orphans_group_of_terminated_instance(self, provisioned: bool): + instance_model = _get_instance_model(InstanceStatus.TERMINATED, provisioned=provisioned) + decision = _get_decision(_get_resource_group(), instance_model=instance_model) + assert decision == _ResourceGroupDecision.ORPHAN + + @pytest.mark.parametrize( + ("minutes_before_listing", "expected_decision"), + [(5, _ResourceGroupDecision.SKIP), (60, _ResourceGroupDecision.ORPHAN)], + ) + def test_skips_group_of_instance_terminated_around_listing( + self, minutes_before_listing: int, expected_decision: _ResourceGroupDecision + ): + instance_model = _get_instance_model( + InstanceStatus.TERMINATED, + finished_at=LISTED_AT - timedelta(minutes=minutes_before_listing), + ) + decision = _get_decision(_get_resource_group(), instance_model=instance_model) + assert decision == expected_decision + + def test_tracks_group_matching_active_instance(self): + instance_model = _get_instance_model(InstanceStatus.BUSY) + decision = _get_decision(_get_resource_group(), instance_model=instance_model) + assert decision == _ResourceGroupDecision.TRACKED + + def test_orphans_group_with_other_primary_id_of_active_instance(self): + instance_model = _get_instance_model(InstanceStatus.BUSY) + resource_group = _get_resource_group(primary_id="other_instance_id") + decision = _get_decision(resource_group, instance_model=instance_model) + assert decision == _ResourceGroupDecision.ORPHAN + + def test_orphans_group_without_primary_resource_of_active_instance(self): + instance_model = _get_instance_model(InstanceStatus.BUSY) + resource_group = _get_resource_group(primary_id=None) + decision = _get_decision(resource_group, instance_model=instance_model) + assert decision == _ResourceGroupDecision.ORPHAN + + @pytest.mark.parametrize( + ("minutes_before_listing", "expected_decision"), + [(5, _ResourceGroupDecision.SKIP), (60, _ResourceGroupDecision.ORPHAN)], + ) + def test_skips_other_group_of_instance_provisioned_around_listing( + self, minutes_before_listing: int, expected_decision: _ResourceGroupDecision + ): + instance_model = _get_instance_model( + InstanceStatus.PROVISIONING, + started_at=LISTED_AT - timedelta(minutes=minutes_before_listing), + ) + resource_group = _get_resource_group(primary_id=None) + decision = _get_decision(resource_group, instance_model=instance_model) + assert decision == expected_decision + + def test_orphans_group_in_other_backend_of_active_instance(self): + instance_model = _get_instance_model(InstanceStatus.BUSY) + resource_group = _get_resource_group(backend=BackendType.GCP) + decision = _get_decision(resource_group, instance_model=instance_model) + assert decision == _ResourceGroupDecision.ORPHAN + + def test_ignores_active_instance_without_provisioning_data(self): + instance_model = _get_instance_model(InstanceStatus.BUSY, provisioned=False) + decision = _get_decision(_get_resource_group(), instance_model=instance_model) + assert decision == _ResourceGroupDecision.IGNORE + + +CLEANUP_INTERVAL = timedelta(seconds=settings.ORPHAN_CLEANUP_INTERVAL_SECONDS) + + +class _ListingCompute(ComputeWithResourceListingSupport): + def __init__( + self, + listing: Optional[ResourceListing] = None, + error: Optional[Exception] = None, + delete_error: Optional[Exception] = None, + ): + self.listing = listing or ResourceListing(resource_groups=[]) + self.error = error + self.delete_error = delete_error + self.listed = False + self.deleted: list[ListedResourceGroup] = [] + + def list_resource_groups(self) -> ResourceListing: + self.listed = True + if self.error is not None: + raise self.error + return self.listing + + def delete_resource_group(self, group: ListedResourceGroup) -> None: + if self.delete_error is not None: + raise self.delete_error + self.deleted.append(group) + + +async def _create_orphan_cleanup( + session: AsyncSession, + project: ProjectModel, + *, + last_processed_at: datetime, + created_at: Optional[datetime] = None, + clean_at: Optional[datetime] = None, + clean_mode: Optional[OrphanCleanupMode] = None, + cleanup_attempt: int = 0, + backend_type: BackendType = BackendType.AWS, +) -> OrphanCleanupModel: + backend = await create_backend( + session=session, project_id=project.id, backend_type=backend_type + ) + orphan_cleanup = OrphanCleanupModel( + backend_id=backend.id, + created_at=created_at or last_processed_at - timedelta(days=1), + last_processed_at=last_processed_at, + clean_at=clean_at, + clean_mode=clean_mode, + cleanup_attempt=cleanup_attempt, + ) + session.add(orphan_cleanup) + await session.commit() + return orphan_cleanup + + +async def _create_instance_with_provider_ref( + session: AsyncSession, + project: ProjectModel, + status: InstanceStatus, + created_at: datetime = LISTED_AT, +) -> InstanceModel: + instance = await create_instance( + session=session, project=project, status=status, created_at=created_at + ) + instance.provider_ref = PROVIDER_REF + session.add( + ProviderRefModel( + provider_ref=PROVIDER_REF, + project_id=project.id, + owner_type=OwnerType.INSTANCE, + owner_id=instance.id, + owner_name=instance.name, + created_at=created_at, + ) + ) + await session.commit() + return instance + + +async def _fetch_ids(fetcher: OrphanCleanupFetcher) -> set[uuid.UUID]: + return {item.id for item in await fetcher.fetch(limit=10)} + + +@pytest.fixture +def fetcher() -> OrphanCleanupFetcher: + return OrphanCleanupFetcher( + queue=asyncio.Queue(), + queue_desired_minsize=1, + min_processing_interval=CLEANUP_INTERVAL, + lock_timeout=timedelta(seconds=60), + heartbeater=Mock(), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) +class TestOrphanCleanupFetcher: + async def test_fetches_new_rows(self, test_db, session: AsyncSession, fetcher): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + now = get_current_datetime() + orphan_cleanup = await _create_orphan_cleanup( + session, project, last_processed_at=now, created_at=now + ) + assert await _fetch_ids(fetcher) == {orphan_cleanup.id} + + async def test_fetches_clean_rows_after_cleanup_interval( + self, test_db, session: AsyncSession, fetcher + ): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + clean_at = get_current_datetime() - 2 * CLEANUP_INTERVAL + orphan_cleanup = await _create_orphan_cleanup( + session, + project, + last_processed_at=clean_at, + clean_at=clean_at, + clean_mode=settings.ORPHAN_CLEANUP_MODE, + ) + assert await _fetch_ids(fetcher) == {orphan_cleanup.id} + + async def test_skips_rows_processed_within_cleanup_interval( + self, test_db, session: AsyncSession, fetcher + ): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + await _create_orphan_cleanup( + session, project, last_processed_at=get_current_datetime() - CLEANUP_INTERVAL / 2 + ) + assert await _fetch_ids(fetcher) == set() + + async def test_skips_rows_of_deleted_projects_and_unsupported_backends( + self, test_db, session: AsyncSession, fetcher + ): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + now = get_current_datetime() + deleted_project = await create_project( + session=session, owner=user, name="deleted-project", deleted=True + ) + await _create_orphan_cleanup( + session, deleted_project, last_processed_at=now, created_at=now + ) + await _create_orphan_cleanup( + session, + project, + last_processed_at=now, + created_at=now, + backend_type=BackendType.LAMBDA, + ) + assert await _fetch_ids(fetcher) == set() + + @pytest.mark.parametrize( + ("cleanup_attempt", "intervals_since_processed", "fetched"), + [ + (1, 1.5, False), + (1, 2.5, True), + (5, 23, False), + (5, 25, True), + ], + ) + async def test_applies_retry_delay_after_failed_cleanups( + self, + test_db, + session: AsyncSession, + fetcher, + cleanup_attempt: int, + intervals_since_processed: float, + fetched: bool, + ): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + orphan_cleanup = await _create_orphan_cleanup( + session, + project, + last_processed_at=get_current_datetime() + - CLEANUP_INTERVAL * intervals_since_processed, + cleanup_attempt=cleanup_attempt, + ) + assert await _fetch_ids(fetcher) == ({orphan_cleanup.id} if fetched else set()) + + +async def _lock(session: AsyncSession, orphan_cleanup: OrphanCleanupModel) -> PipelineItem: + lock_token = uuid.uuid4() + lock_expires_at = get_current_datetime() + timedelta(seconds=60) + orphan_cleanup.lock_token = lock_token + orphan_cleanup.lock_expires_at = lock_expires_at + orphan_cleanup.lock_owner = OrphanCleanupPipeline.__name__ + await session.commit() + return PipelineItem( + __tablename__=OrphanCleanupModel.__tablename__, + id=orphan_cleanup.id, + lock_token=lock_token, + lock_expires_at=lock_expires_at, + prev_lock_expired=False, + ) + + +def _patch_compute(compute: _ListingCompute): + backend = Mock() + backend.compute.return_value = compute + return patch.object( + orphan_cleanup_pipeline.backends_services, + "get_project_backend_by_type", + AsyncMock(return_value=backend), + ) + + +@pytest.fixture +def worker() -> OrphanCleanupWorker: + return OrphanCleanupWorker(queue=Mock(), heartbeater=Mock(), pipeline_hinter=Mock()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) +class TestOrphanCleanupWorker: + async def test_reports_orphans_and_marks_backend_clean( + self, test_db, session: AsyncSession, worker + ): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + orphan_cleanup = await _create_orphan_cleanup( + session, project, last_processed_at=get_current_datetime(), cleanup_attempt=2 + ) + await _create_instance_with_provider_ref(session, project, InstanceStatus.TERMINATED) + compute = _ListingCompute(ResourceListing(resource_groups=[_get_resource_group()])) + item = await _lock(session, orphan_cleanup) + with ( + patch.object(settings, "ORPHAN_CLEANUP_MODE", OrphanCleanupMode.REPORT), + _patch_compute(compute), + ): + await worker.process(item) + await session.refresh(orphan_cleanup) + assert orphan_cleanup.clean_at is not None + assert ( + orphan_cleanup.clean_at + <= get_current_datetime() - orphan_cleanup_pipeline._LISTING_MARGIN + ) + assert orphan_cleanup.clean_mode == OrphanCleanupMode.REPORT + assert orphan_cleanup.cleanup_attempt == 0 + assert orphan_cleanup.lock_token is None + assert compute.deleted == [] + events = await list_events(session) + assert len(events) == 1 + assert events[0].message.startswith("Orphaned resources found in aws/us-east-1") + + @pytest.mark.parametrize( + "activity", + ["provider_ref_issued", "instance_started", "instance_finished", "job_finished"], + ) + async def test_cleans_up_backend_with_activity_after_clean_at( + self, test_db, session: AsyncSession, worker, activity: str + ): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + clean_at = get_current_datetime() - 2 * CLEANUP_INTERVAL + before_clean_at = clean_at - timedelta(minutes=10) + after_clean_at = clean_at + timedelta(minutes=10) + orphan_cleanup = await _create_orphan_cleanup( + session, + project, + last_processed_at=clean_at, + clean_at=clean_at, + clean_mode=settings.ORPHAN_CLEANUP_MODE, + ) + if activity == "job_finished": + repo = await create_repo(session=session, project_id=project.id) + run = await create_run(session=session, project=project, repo=repo, user=user) + await create_job( + session=session, + run=run, + status=JobStatus.FAILED, + last_processed_at=after_clean_at, + ) + else: + instance = await _create_instance_with_provider_ref( + session, + project, + InstanceStatus.TERMINATED, + created_at=after_clean_at + if activity == "provider_ref_issued" + else before_clean_at, + ) + if activity == "instance_started": + instance.started_at = after_clean_at + if activity == "instance_finished": + instance.finished_at = after_clean_at + await session.commit() + compute = _ListingCompute() + item = await _lock(session, orphan_cleanup) + with _patch_compute(compute): + await worker.process(item) + await session.refresh(orphan_cleanup) + assert compute.listed + assert orphan_cleanup.clean_at is not None + assert orphan_cleanup.clean_at > clean_at + + async def test_skips_listing_without_activity_after_clean_at( + self, test_db, session: AsyncSession, worker + ): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + clean_at = get_current_datetime() - 2 * CLEANUP_INTERVAL + orphan_cleanup = await _create_orphan_cleanup( + session, + project, + last_processed_at=clean_at, + clean_at=clean_at, + clean_mode=settings.ORPHAN_CLEANUP_MODE, + ) + await _create_instance_with_provider_ref( + session, + project, + InstanceStatus.IDLE, + created_at=clean_at - timedelta(minutes=10), + ) + compute = _ListingCompute() + item = await _lock(session, orphan_cleanup) + with _patch_compute(compute): + await worker.process(item) + await session.refresh(orphan_cleanup) + assert not compute.listed + assert orphan_cleanup.clean_at == clean_at + assert orphan_cleanup.last_processed_at > clean_at + assert orphan_cleanup.lock_token is None + + async def test_cleans_up_backend_cleaned_in_other_mode( + self, test_db, session: AsyncSession, worker + ): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + clean_at = get_current_datetime() - 2 * CLEANUP_INTERVAL + orphan_cleanup = await _create_orphan_cleanup( + session, + project, + last_processed_at=clean_at, + clean_at=clean_at, + clean_mode=OrphanCleanupMode.DELETE, + ) + compute = _ListingCompute() + item = await _lock(session, orphan_cleanup) + with ( + patch.object(settings, "ORPHAN_CLEANUP_MODE", OrphanCleanupMode.REPORT), + _patch_compute(compute), + ): + await worker.process(item) + await session.refresh(orphan_cleanup) + assert compute.listed + assert orphan_cleanup.clean_mode == OrphanCleanupMode.REPORT + + async def test_deletes_orphans_in_delete_mode(self, test_db, session: AsyncSession, worker): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + orphan_cleanup = await _create_orphan_cleanup( + session, project, last_processed_at=get_current_datetime() + ) + await _create_instance_with_provider_ref(session, project, InstanceStatus.TERMINATED) + resource_group = _get_resource_group() + compute = _ListingCompute(ResourceListing(resource_groups=[resource_group])) + item = await _lock(session, orphan_cleanup) + with ( + patch.object(settings, "ORPHAN_CLEANUP_MODE", OrphanCleanupMode.DELETE), + _patch_compute(compute), + ): + await worker.process(item) + await session.refresh(orphan_cleanup) + assert orphan_cleanup.clean_mode == OrphanCleanupMode.DELETE + assert compute.deleted == [resource_group] + events = await list_events(session) + assert [e.message.split(" in ")[0] for e in events] == ["Orphaned resources deleted"] + + async def test_keeps_backend_due_without_backoff_after_partial_listing( + self, test_db, session: AsyncSession, worker + ): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + orphan_cleanup = await _create_orphan_cleanup( + session, project, last_processed_at=get_current_datetime(), cleanup_attempt=2 + ) + compute = _ListingCompute( + ResourceListing( + resource_groups=[], + failures=[ + ResourceListingFailure( + backend=BackendType.AWS, region="us-east-1", message="throttled" + ) + ], + ) + ) + item = await _lock(session, orphan_cleanup) + with _patch_compute(compute): + await worker.process(item) + await session.refresh(orphan_cleanup) + assert orphan_cleanup.clean_at is None + assert orphan_cleanup.cleanup_attempt == 0 + + async def test_backs_off_after_listing_error(self, test_db, session: AsyncSession, worker): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + orphan_cleanup = await _create_orphan_cleanup( + session, project, last_processed_at=get_current_datetime(), cleanup_attempt=1 + ) + compute = _ListingCompute(error=BackendError("unauthorized")) + item = await _lock(session, orphan_cleanup) + with _patch_compute(compute): + await worker.process(item) + await session.refresh(orphan_cleanup) + assert orphan_cleanup.clean_at is None + assert orphan_cleanup.cleanup_attempt == 2 + + async def test_retries_without_backoff_when_owner_is_in_flight( + self, test_db, session: AsyncSession, worker + ): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + orphan_cleanup = await _create_orphan_cleanup( + session, project, last_processed_at=get_current_datetime(), cleanup_attempt=2 + ) + await _create_instance_with_provider_ref(session, project, InstanceStatus.PENDING) + compute = _ListingCompute(ResourceListing(resource_groups=[_get_resource_group()])) + item = await _lock(session, orphan_cleanup) + with _patch_compute(compute): + await worker.process(item) + await session.refresh(orphan_cleanup) + assert orphan_cleanup.clean_at is None + assert orphan_cleanup.cleanup_attempt == 0 + assert compute.deleted == [] + + async def test_backs_off_after_delete_error(self, test_db, session: AsyncSession, worker): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + orphan_cleanup = await _create_orphan_cleanup( + session, project, last_processed_at=get_current_datetime(), cleanup_attempt=1 + ) + await _create_instance_with_provider_ref(session, project, InstanceStatus.TERMINATED) + compute = _ListingCompute( + ResourceListing(resource_groups=[_get_resource_group()]), + delete_error=BackendError("rate limited"), + ) + item = await _lock(session, orphan_cleanup) + with ( + patch.object(settings, "ORPHAN_CLEANUP_MODE", OrphanCleanupMode.DELETE), + _patch_compute(compute), + ): + await worker.process(item) + await session.refresh(orphan_cleanup) + assert orphan_cleanup.clean_at is None + assert orphan_cleanup.cleanup_attempt == 2 + assert await list_events(session) == [] + + @pytest.mark.parametrize( + ("primary_id", "expected_messages"), + [ + ("instance_id", ["Orphaned resources being deleted"]), + (None, []), + ], + ) + async def test_retries_without_backoff_when_deletion_is_in_progress( + self, + test_db, + session: AsyncSession, + worker, + primary_id: Optional[str], + expected_messages: list[str], + ): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + orphan_cleanup = await _create_orphan_cleanup( + session, project, last_processed_at=get_current_datetime(), cleanup_attempt=1 + ) + await _create_instance_with_provider_ref(session, project, InstanceStatus.TERMINATED) + compute = _ListingCompute( + ResourceListing(resource_groups=[_get_resource_group(primary_id=primary_id)]), + delete_error=NotYetTerminated("waiting for the instance to be gone"), + ) + item = await _lock(session, orphan_cleanup) + with ( + patch.object(settings, "ORPHAN_CLEANUP_MODE", OrphanCleanupMode.DELETE), + _patch_compute(compute), + ): + await worker.process(item) + await session.refresh(orphan_cleanup) + assert orphan_cleanup.clean_at is None + assert orphan_cleanup.cleanup_attempt == 0 + events = await list_events(session) + assert [e.message.split(" in ")[0] for e in events] == expected_messages + + async def test_reports_orphans_of_deleted_instance_to_the_instance( + self, test_db, session: AsyncSession, worker + ): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + orphan_cleanup = await _create_orphan_cleanup( + session, project, last_processed_at=get_current_datetime() + ) + # A placeholder instance deleted after a failed `run_job()` leaves only the record + deleted_instance_id = uuid.uuid4() + session.add( + ProviderRefModel( + provider_ref=PROVIDER_REF, + project_id=project.id, + owner_type=OwnerType.INSTANCE, + owner_id=deleted_instance_id, + owner_name="fleet-0", + ) + ) + await session.commit() + compute = _ListingCompute(ResourceListing(resource_groups=[_get_resource_group()])) + item = await _lock(session, orphan_cleanup) + with ( + patch.object(settings, "ORPHAN_CLEANUP_MODE", OrphanCleanupMode.REPORT), + _patch_compute(compute), + ): + await worker.process(item) + events = await list_events(session) + assert len(events) == 1 + assert events[0].message.startswith("Orphaned resources found") + assert [(t.entity_id, t.entity_name) for t in events[0].targets] == [ + (deleted_instance_id, "fleet-0") + ] + + async def test_does_not_apply_result_if_lock_token_changes( + self, test_db, session: AsyncSession, worker + ): + user = await create_user(session=session) + project = await create_project(session=session, owner=user) + orphan_cleanup = await _create_orphan_cleanup( + session, project, last_processed_at=get_current_datetime() + ) + await _create_instance_with_provider_ref(session, project, InstanceStatus.TERMINATED) + compute = _ListingCompute(ResourceListing(resource_groups=[_get_resource_group()])) + item = await _lock(session, orphan_cleanup) + new_lock_token = uuid.uuid4() + load_owner_records = orphan_cleanup_pipeline._load_owner_records + + async def load_owner_records_after_lock_takeover(**kwargs): + # Another worker takes over the row while it's being processed + orphan_cleanup.lock_token = new_lock_token + await session.commit() + return await load_owner_records(**kwargs) + + with ( + _patch_compute(compute), + patch.object( + orphan_cleanup_pipeline, + "_load_owner_records", + side_effect=load_owner_records_after_lock_takeover, + ), + ): + await worker.process(item) + await session.refresh(orphan_cleanup) + assert orphan_cleanup.lock_token == new_lock_token + assert orphan_cleanup.clean_at is None + assert await list_events(session) == [] diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py index 536df52cc1..8294e86238 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py @@ -61,6 +61,7 @@ InstanceModel, JobModel, PlacementGroupModel, + ProviderRefModel, VolumeAttachmentModel, ) from dstack._internal.server.services.docker import ImageConfig @@ -446,6 +447,10 @@ async def test_provisions_new_capacity_for_assigned_job_with_placeholder( placeholder_id = job.instance.id assert job.used_instance_id == placeholder_id assert job.instance.status == InstanceStatus.PENDING + provider_ref = job.instance.provider_ref + assert provider_ref is not None + res = await session.execute(select(ProviderRefModel.provider_ref)) + assert res.scalars().all() == [provider_ref] await _process_job(session=session, worker=worker, job_model=job) @@ -458,6 +463,9 @@ async def test_provisions_new_capacity_for_assigned_job_with_placeholder( assert job.instance.fleet_id == fleet.id assert job.instance.offer is not None assert job.instance.provisioning_job_id == job.id # never cleared + assert job.instance.provider_ref == provider_ref + run_job_job = backend_mock.compute.return_value.run_job.call_args[0][1] + assert run_job_job.job_submissions[-1].provider_ref == provider_ref res = await session.execute( select(InstanceModel).where( InstanceModel.fleet_id == fleet.id, @@ -2124,7 +2132,11 @@ async def test_provisions_compute_group( assert job1.status == JobStatus.PROVISIONING assert job2.status == JobStatus.PROVISIONING res = await session.execute(select(ComputeGroupModel)) - assert res.scalar_one_or_none() is not None + compute_group = res.scalar_one() + res = await session.execute( + select(InstanceModel).where(InstanceModel.compute_group_id == compute_group.id) + ) + assert [i.provider_ref for i in res.scalars().all()] == [None, None] async def test_provisions_one_job_node_group_via_run_job( self, test_db, session: AsyncSession, worker: JobSubmittedWorker diff --git a/src/tests/_internal/server/background/scheduled_tasks/test_orphan_cleanup_enrollment.py b/src/tests/_internal/server/background/scheduled_tasks/test_orphan_cleanup_enrollment.py new file mode 100644 index 0000000000..1c7471e84e --- /dev/null +++ b/src/tests/_internal/server/background/scheduled_tasks/test_orphan_cleanup_enrollment.py @@ -0,0 +1,132 @@ +import uuid +from unittest.mock import AsyncMock, patch + +import pytest +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.backends.orphans import OwnerType +from dstack._internal.server.background.scheduled_tasks import orphan_cleanup_enrollment +from dstack._internal.server.background.scheduled_tasks.orphan_cleanup_enrollment import ( + enroll_backends_in_orphan_cleanup, +) +from dstack._internal.server.models import ( + OrphanCleanupModel, + ProjectModel, + ProviderRefModel, + UserModel, +) +from dstack._internal.server.testing.common import ( + create_backend, + create_instance, + create_project, + create_user, +) + + +async def _list_enrolled_backend_ids(session: AsyncSession) -> set: + res = await session.execute(select(OrphanCleanupModel.backend_id)) + return set(res.scalars().all()) + + +async def _create_project_with_provider_ref( + session: AsyncSession, + owner: UserModel, + provider_ref: str, + name: str = "test_project", + deleted: bool = False, +) -> ProjectModel: + project = await create_project(session=session, owner=owner, name=name, deleted=deleted) + session.add( + ProviderRefModel( + provider_ref=provider_ref, + project_id=project.id, + owner_type=OwnerType.INSTANCE, + owner_id=uuid.uuid4(), + owner_name="instance", + ) + ) + await session.commit() + return project + + +@pytest.mark.asyncio +@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) +class TestEnrollBackendsInOrphanCleanup: + async def test_enrolls_backend_of_project_with_provider_refs( + self, test_db, session: AsyncSession + ): + user = await create_user(session=session) + project = await _create_project_with_provider_ref(session, user, "instref00001") + backend = await create_backend(session=session, project_id=project.id) + await enroll_backends_in_orphan_cleanup() + res = await session.execute(select(OrphanCleanupModel)) + orphan_cleanup = res.scalar_one() + assert orphan_cleanup.backend_id == backend.id + assert orphan_cleanup.cleanup_attempt == 0 + assert orphan_cleanup.clean_at is None + assert orphan_cleanup.last_processed_at == orphan_cleanup.created_at + + async def test_skips_backends_that_cannot_be_cleaned_up(self, test_db, session: AsyncSession): + user = await create_user(session=session) + project_without_refs = await create_project(session=session, owner=user, name="no-refs") + await create_instance(session=session, project=project_without_refs) + await create_backend(session=session, project_id=project_without_refs.id) + deleted_project = await _create_project_with_provider_ref( + session, user, "instref00001", name="deleted", deleted=True + ) + await create_backend(session=session, project_id=deleted_project.id) + # Backend type without resource listing support + project = await _create_project_with_provider_ref(session, user, "instref00002") + await create_backend( + session=session, project_id=project.id, backend_type=BackendType.LAMBDA + ) + await enroll_backends_in_orphan_cleanup() + assert await _list_enrolled_backend_ids(session) == set() + + async def test_enrolls_limited_number_of_backends_per_run( + self, test_db, session: AsyncSession + ): + user = await create_user(session=session) + project = await _create_project_with_provider_ref(session, user, "instref00001") + backend = await create_backend(session=session, project_id=project.id) + other_project = await _create_project_with_provider_ref( + session, user, "instref00002", name="other" + ) + other_backend = await create_backend(session=session, project_id=other_project.id) + with patch.object(orphan_cleanup_enrollment, "_ENROLLMENT_BATCH_SIZE", 1): + await enroll_backends_in_orphan_cleanup() + assert len(await _list_enrolled_backend_ids(session)) == 1 + await enroll_backends_in_orphan_cleanup() + assert await _list_enrolled_backend_ids(session) == {backend.id, other_backend.id} + + async def test_is_idempotent(self, test_db, session: AsyncSession): + user = await create_user(session=session) + project = await _create_project_with_provider_ref(session, user, "instref00001") + await create_backend(session=session, project_id=project.id) + await enroll_backends_in_orphan_cleanup() + res = await session.execute(select(OrphanCleanupModel)) + orphan_cleanup = res.scalar_one() + await enroll_backends_in_orphan_cleanup() + res = await session.execute(select(OrphanCleanupModel)) + assert [r.id for r in res.scalars().all()] == [orphan_cleanup.id] + + async def test_tolerates_backends_enrolled_concurrently(self, test_db, session: AsyncSession): + user = await create_user(session=session) + project = await _create_project_with_provider_ref(session, user, "instref00001") + backend = await create_backend(session=session, project_id=project.id) + other_project = await _create_project_with_provider_ref( + session, user, "instref00002", name="other" + ) + other_backend = await create_backend(session=session, project_id=other_project.id) + session.add(OrphanCleanupModel(backend_id=backend.id)) + await session.commit() + # Another replica enrolls `backend` between the query and the insert + with patch.object( + orphan_cleanup_enrollment, + "_select_backend_ids_to_enroll", + AsyncMock(return_value=[backend.id, other_backend.id]), + ): + await enroll_backends_in_orphan_cleanup() + assert await _list_enrolled_backend_ids(session) == {backend.id, other_backend.id} diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index d88c802373..8d64936c19 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -623,6 +623,7 @@ def get_dev_env_run_dict( "job_runtime_data": None, "probes": [], "image_pull_progress": None, + "provider_ref": None, } ], "job_connection_info": None, @@ -646,6 +647,7 @@ def get_dev_env_run_dict( "job_runtime_data": None, "probes": [], "image_pull_progress": None, + "provider_ref": None, }, "cost": 0.0, "service": None, @@ -791,6 +793,7 @@ async def test_lists_runs(self, test_db, session: AsyncSession, client: AsyncCli "job_runtime_data": None, "probes": [], "image_pull_progress": None, + "provider_ref": None, } ], "job_connection_info": None, @@ -814,6 +817,7 @@ async def test_lists_runs(self, test_db, session: AsyncSession, client: AsyncCli "job_runtime_data": None, "probes": [], "image_pull_progress": None, + "provider_ref": None, }, "cost": 0.0, "service": None, @@ -985,6 +989,7 @@ async def test_limits_job_submissions( "job_runtime_data": None, "probes": [], "image_pull_progress": None, + "provider_ref": None, } ], "job_connection_info": None, @@ -1008,6 +1013,7 @@ async def test_limits_job_submissions( "job_runtime_data": None, "probes": [], "image_pull_progress": None, + "provider_ref": None, }, "cost": 0.0, "service": None,