diff --git a/mkdocs/docs/reference/env.md b/mkdocs/docs/reference/env.md index c54a6b3740..4cc395a5fa 100644 --- a/mkdocs/docs/reference/env.md +++ b/mkdocs/docs/reference/env.md @@ -151,7 +151,7 @@ For more details on the options below, refer to the [server deployment](../guide - `DSTACK_SERVER_GATEWAY_MAX_REPLICAS`{ #DSTACK_SERVER_GATEWAY_MAX_REPLICAS } - Maximum number of replicas allowed for a gateway. Defaults to `9`. - `DSTACK_SERVER_INSTANCE_HEALTH_MIN_COLLECT_INTERVAL_SECONDS`{ #DSTACK_SERVER_INSTANCE_HEALTH_MIN_COLLECT_INTERVAL_SECONDS } – Minimum time interval between consecutive health checks of the same instance. - `DSTACK_SERVER_EVENTS_TTL_SECONDS`{ #DSTACK_SERVER_EVENTS_TTL_SECONDS } - Maximum age of event records. Set to `0` to disable event storage. Defaults to 30 days. -- `DSTACK_SERVER_ORPHAN_CLEANUP_MODE`{ #DSTACK_SERVER_ORPHAN_CLEANUP_MODE } – Controls the cleanup of cloud resources that `dstack` created but no longer tracks, e.g. after the server crashed during provisioning or failed to terminate an instance. `report` emits events about such resources, `delete` also deletes them, and `off` disables the cleanup. Resources are recognized by the tags (or names) `dstack` sets on them, so a resource created by copying the tags of a `dstack` instance is treated as an orphan and deleted in `delete` mode. Defaults to `report`. Currently supported for the `aws`, `azure`, `gcp`, and `runpod` backends. +- `DSTACK_SERVER_ORPHAN_CLEANUP_MODE`{ #DSTACK_SERVER_ORPHAN_CLEANUP_MODE } – Controls the cleanup of cloud resources that `dstack` created but no longer tracks, e.g. after the server crashed during provisioning or failed to terminate an instance. `report` emits events about such resources, `delete` also deletes them, and `off` disables the cleanup. Resources are recognized by the tags (or names) `dstack` sets on them, so a resource created by copying the tags of a `dstack` instance is treated as an orphan and deleted in `delete` mode. Defaults to `report`. Currently supported for the `aws`, `azure`, `gcp`, `nebius`, and `runpod` backends. - `DSTACK_SERVER_ORPHAN_CLEANUP_INTERVAL_SECONDS`{ #DSTACK_SERVER_ORPHAN_CLEANUP_INTERVAL_SECONDS } – Minimum interval between orphan cleanups of the same backend. A backend is cleaned up only after instances were provisioned or terminated in its project. Defaults to `3600`. - `DSTACK_SERVER_DEFAULT_DOCKER_REGISTRY`{ #DSTACK_SERVER_DEFAULT_DOCKER_REGISTRY } – A default Docker registry to use for job images that do not specify an explicit registry. E.g., if set to `registry.example`, then `image: ubuntu` becomes equivalent to `image: registry.example/ubuntu`. **Note**: This setting should only be used for configuring registries that act as a pull-through cache for Docker Hub. The default `dstack` images are also pulled from the configured registry. - `DSTACK_SERVER_DEFAULT_DOCKER_REGISTRY_USERNAME`{ #DSTACK_SERVER_DEFAULT_DOCKER_REGISTRY_USERNAME } – Username for authenticating with the default Docker registry. See `DSTACK_SERVER_DEFAULT_DOCKER_REGISTRY_PASSWORD`. diff --git a/src/dstack/_internal/core/backends/nebius/compute.py b/src/dstack/_internal/core/backends/nebius/compute.py index 97cafb6171..7acd3a5707 100644 --- a/src/dstack/_internal/core/backends/nebius/compute.py +++ b/src/dstack/_internal/core/backends/nebius/compute.py @@ -2,26 +2,33 @@ import random import shlex import time +from dataclasses import dataclass from decimal import ROUND_CEILING, Decimal from functools import cached_property from typing import List, Optional from nebius.aio.operation import Operation as SDKOperation from nebius.aio.service_error import RequestError, StatusCode -from nebius.api.nebius.common.v1 import Operation +from nebius.api.nebius.billing.v1 import PricingPolicy, PricingPolicyStatus +from nebius.api.nebius.common.v1 import Operation, ResourceMetadata +from nebius.api.nebius.compute.v1 import Disk, DiskStatus, Instance, InstanceStatus from nebius.sdk import SDK from dstack._internal.core.backends.base.backend import Compute from dstack._internal.core.backends.base.compute import ( + PROVIDER_REF_TAG_KEY, ComputeWithAllOffersCached, ComputeWithCreateInstanceSupport, ComputeWithInstanceVolumesSupport, ComputeWithMultinodeSupport, ComputeWithPlacementGroupSupport, ComputeWithPrivilegedSupport, + ComputeWithResourceListingSupport, generate_unique_instance_name, get_user_data, + is_valid_provider_ref, merge_tags, + parse_provider_ref_from_name, ) from dstack._internal.core.backends.base.offers import get_catalog_offers from dstack._internal.core.backends.nebius import resources @@ -32,10 +39,17 @@ ) from dstack._internal.core.errors import ( BackendError, + ComputeError, NotYetTerminated, 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_extra_ignore, @@ -54,6 +68,7 @@ ) from dstack._internal.core.models.resources import Memory, Range from dstack._internal.core.models.runs import JobProvisioningData +from dstack._internal.utils.common import get_or_error from dstack._internal.utils.logging import get_logger logger = get_logger(__name__) @@ -88,6 +103,13 @@ "cpu-d3", "cpu-e2", ] +PROVIDER_NAME_LABEL_KEY = "dstack_provider_name" +""" +The label that groups resources created by one `create_instance()` call for orphan cleanup. +Its value is the provider resource name that dstack generated for the call and that +the disk, the pricing policy, and the instance share. Unlike the name itself, it survives +renaming the resource. +""" class NebiusCompute( @@ -97,6 +119,7 @@ class NebiusCompute( ComputeWithInstanceVolumesSupport, ComputeWithMultinodeSupport, ComputeWithPlacementGroupSupport, + ComputeWithResourceListingSupport, Compute, ): def __init__(self, config: NebiusConfig): @@ -155,7 +178,9 @@ def create_instance( # and the instance to enter the STARTING state. This has to be done in create_instance so # that we can handle quota and availability errors that may occur even after creating an # instance. - instance_name = generate_unique_instance_name(instance_config) + # The disk, the pricing policy, and the instance share this name. Pricing policies + # don't keep labels, so orphan cleanup finds them by `provider_ref` in the name. + instance_name = generate_unique_instance_name(instance_config, with_provider_ref=True) platform, preset = instance_offer.instance.name.split() offer_backend_data = validate_extra_ignore( NebiusOfferBackendData, instance_offer.backend_data @@ -181,6 +206,9 @@ def create_instance( resource_tags=instance_config.tags, ) labels = resources.filter_invalid_labels(labels) + if instance_config.provider_ref is not None: + labels[PROVIDER_REF_TAG_KEY] = instance_config.provider_ref + labels[PROVIDER_NAME_LABEL_KEY] = instance_name create_disk_op = resources.create_disk( sdk=self._sdk, name=instance_name, @@ -244,15 +272,18 @@ def create_instance( ) _wait_for_instance(self._sdk, create_instance_op) except BaseException: + # Best-effort rollback. It may leave resources behind, e.g. the boot disk if + # the instance is not deleted within DELETE_INSTANCE_TIMEOUT. This is accepted: + # leftovers carry provider_ref in labels or names and are handled by orphan cleanup. if create_instance_op is not None: try: with resources.ignore_errors([StatusCode.NOT_FOUND]): delete_instance_op = resources.delete_instance( self._sdk, create_instance_op.resource_id ) - resources.wait_for_operation( - delete_instance_op, timeout=DELETE_INSTANCE_TIMEOUT - ) + resources.wait_for_operation( + delete_instance_op, timeout=DELETE_INSTANCE_TIMEOUT + ) except Exception as e: logger.exception( "Could not delete instance %s: %s", create_instance_op.resource_id, e @@ -311,30 +342,43 @@ def terminate_instance( self, instance_id: str, region: str, backend_data: Optional[str] = None ): backend_data_parsed = NebiusInstanceBackendData.load(backend_data) - try: - instance = resources.get_instance(self._sdk, instance_id) - except RequestError as e: - if e.status.code != StatusCode.NOT_FOUND: - raise - instance = None - if instance is not None: - if instance.status.state != instance.status.InstanceState.DELETING: - resources.delete_instance(self._sdk, instance_id) - raise NotYetTerminated( - "Requested instance deletion." - " Will wait for deletion before deleting the boot disk." - f" Instance state was: {instance.status.state.name}" + _delete_resources( + self._sdk, + instance_id=instance_id, + disk_id=backend_data_parsed.boot_disk_id, + pricing_policy_id=backend_data_parsed.pricing_policy_id, + ) + + def list_resource_groups(self) -> ResourceListing: + listing = ResourceListing(resource_groups=[]) + for region, project_id in self._region_to_project_id.items(): + try: + listing.resource_groups.extend( + _list_resource_groups(self._sdk, region=region, project_id=project_id) ) - else: - raise NotYetTerminated( - "Waiting for instance deletion before deleting the boot disk." - f" Instance state: {instance.status.state.name}" + except RequestError as e: + listing.failures.append( + ResourceListingFailure( + backend=BackendType.NEBIUS, region=region, message=str(e) + ) ) - with resources.ignore_errors([StatusCode.NOT_FOUND]): - resources.delete_disk(self._sdk, backend_data_parsed.boot_disk_id) - if backend_data_parsed.pricing_policy_id is not None: - with resources.ignore_errors([StatusCode.NOT_FOUND]): - resources.delete_pricing_policy(self._sdk, backend_data_parsed.pricing_policy_id) + if len(listing.failures) == len(self._region_to_project_id): + 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( + NebiusResourceGroupBackendData, get_or_error(group.backend_data) + ) + try: + _delete_resources( + self._sdk, + instance_id=backend_data.instance_id, + disk_id=backend_data.disk_id, + pricing_policy_id=backend_data.pricing_policy_id, + ) + except RequestError as e: + raise ComputeError(f"Failed to delete resource group {group.provider_ref}: {e}") from e def create_placement_group( self, @@ -407,6 +451,14 @@ def load(cls, raw: Optional[str]) -> "NebiusInstanceBackendData": return validate_json_extra_ignore(cls, raw) +class NebiusResourceGroupBackendData(CoreModel): + instance_id: Optional[str] = None + """The group's instance. It's set without `ListedResourceGroup.primary_id` + if the instance is being deleted.""" + disk_id: Optional[str] = None + pricing_policy_id: Optional[str] = None + + class NebiusClusterBackendData(CoreModel): id: str fabric: str @@ -455,6 +507,258 @@ def _wait_for_instance(sdk: SDK, op: SDKOperation[Operation]) -> None: ) +def _delete_resources( + sdk: SDK, + instance_id: Optional[str], + disk_id: Optional[str], + pricing_policy_id: Optional[str], +) -> None: + """ + Deletes the instance, then its boot disk and pricing policy. + Raises `NotYetTerminated` while the instance exists, since an attached disk can't be deleted, + and while a pricing policy is being created, since Nebius rejects its deletion. + """ + if instance_id is not None: + try: + instance = resources.get_instance(sdk, instance_id) + except RequestError as e: + if e.status.code != StatusCode.NOT_FOUND: + raise + instance = None + if instance is not None: + if instance.status.state != InstanceStatus.InstanceState.DELETING: + resources.delete_instance(sdk, instance_id) + raise NotYetTerminated( + "Requested instance deletion." + " Will wait for deletion before deleting the boot disk." + f" Instance state was: {instance.status.state.name}" + ) + raise NotYetTerminated( + "Waiting for instance deletion before deleting the boot disk." + f" Instance state: {instance.status.state.name}" + ) + if disk_id is not None: + with resources.ignore_errors([StatusCode.NOT_FOUND]): + resources.delete_disk(sdk, disk_id) + if pricing_policy_id is not None: + _delete_pricing_policy(sdk, pricing_policy_id) + + +def _delete_pricing_policy(sdk: SDK, pricing_policy_id: str) -> None: + try: + pricing_policy = resources.get_pricing_policy(sdk, pricing_policy_id) + except RequestError as e: + if e.status.code == StatusCode.NOT_FOUND: + return + raise + if pricing_policy.status.state == PricingPolicyStatus.State.STATE_DELETING: + return + if pricing_policy.status.state == PricingPolicyStatus.State.STATE_CREATING: + raise NotYetTerminated( + f"Waiting for pricing policy {pricing_policy_id} to be created before deleting it" + ) + with resources.ignore_errors([StatusCode.NOT_FOUND]): + resources.delete_pricing_policy(sdk, pricing_policy_id) + + +def _list_resource_groups(sdk: SDK, region: str, project_id: str) -> list[ListedResourceGroup]: + """ + Lists the resources of `create_instance()` calls, grouped by the name they share. + Instances being deleted are not listed, but their ids are kept in the groups + of their disks and pricing policies, so that these are deleted after the instances. + """ + # Disks and pricing policies are created before their instance and deleted after it, + # so listing them first ensures the instance of a listed disk or pricing policy + # is listed too unless it's gone or created after the instances are listed. + # The latter is a provisioning in progress, which the server skips. + disks = resources.list_disks(sdk, project_id) + pricing_policies = resources.list_pricing_policies(sdk, project_id) + instances = resources.list_instances(sdk, project_id) + + collector = _ProvisioningAttemptCollector() + for disk in disks: + if disk.status.state != DiskStatus.State.DELETING: + collector.add_disk(disk) + for pricing_policy in pricing_policies: + if pricing_policy.status.state != PricingPolicyStatus.State.STATE_DELETING: + collector.add_pricing_policy(pricing_policy) + for instance in instances: + collector.add_instance(instance) + + resource_groups = [] + for attempt in collector.attempts(): + resource_group = _get_resource_group(attempt, region) + if resource_group is not None: + resource_groups.append(resource_group) + return resource_groups + + +@dataclass +class _ProvisioningAttempt: + """ + Resources created by one `create_instance()` call. + """ + + provider_ref: str + instance: Optional[Instance] = None + disk: Optional[Disk] = None + pricing_policy: Optional[PricingPolicy] = None + + +class _ProvisioningAttemptCollector: + """ + Groups resources into `create_instance()` calls by the name these resources share. + Resources without a valid `provider_ref` are ignored. If labels were copied to another + disk or instance, only the first one listed is kept in the attempt. + """ + + def __init__(self) -> None: + self._attempts: dict[str, _ProvisioningAttempt] = {} + + def attempts(self) -> list[_ProvisioningAttempt]: + return list(self._attempts.values()) + + def add_disk(self, disk: Disk) -> None: + attempt = self._get_or_add(disk.metadata, _get_attempt_key_from_labels(disk.metadata)) + if attempt is None: + return + if attempt.disk is not None: + _log_duplicate_resource(disk.metadata, attempt.disk.metadata) + return + attempt.disk = disk + + def add_pricing_policy(self, pricing_policy: PricingPolicy) -> None: + # Pricing policies don't keep labels. Their names are unique within the project, + # so there are no duplicates. + attempt = self._get_or_add( + pricing_policy.metadata, _get_attempt_key_from_name(pricing_policy.metadata) + ) + if attempt is not None: + attempt.pricing_policy = pricing_policy + + def add_instance(self, instance: Instance) -> None: + attempt = self._get_or_add( + instance.metadata, _get_attempt_key_from_labels(instance.metadata) + ) + if attempt is None: + return + if attempt.instance is not None: + _log_duplicate_resource(instance.metadata, attempt.instance.metadata) + return + attempt.instance = instance + + def _get_or_add( + self, metadata: ResourceMetadata, key: Optional[tuple[str, str]] + ) -> Optional[_ProvisioningAttempt]: + if key is None: + return None + name, provider_ref = key + attempt = self._attempts.setdefault(name, _ProvisioningAttempt(provider_ref=provider_ref)) + if attempt.provider_ref != provider_ref: + logger.warning( + "Skipping Nebius resource %s of %s with provider_ref %s, expected %s", + metadata.id, + name, + provider_ref, + attempt.provider_ref, + ) + return None + return attempt + + +def _log_duplicate_resource(metadata: ResourceMetadata, kept_metadata: ResourceMetadata) -> None: + logger.warning( + "Skipping Nebius resource %s labeled with the same %s as %s", + metadata.id, + PROVIDER_NAME_LABEL_KEY, + kept_metadata.id, + ) + + +def _get_attempt_key_from_labels(metadata: ResourceMetadata) -> Optional[tuple[str, str]]: + """ + Returns the attempt's name and `provider_ref` from the resource's labels, + or `None` if the resource has no valid `provider_ref` label. + """ + provider_ref = metadata.labels.get(PROVIDER_REF_TAG_KEY) + if provider_ref is None: + return None + if not is_valid_provider_ref(provider_ref): + logger.warning( + "Skipping Nebius resource %s with invalid %s label value %r", + metadata.id, + PROVIDER_REF_TAG_KEY, + provider_ref, + ) + return None + name = metadata.labels.get(PROVIDER_NAME_LABEL_KEY) + if name is None: + logger.warning( + "Skipping Nebius resource %s with %s label but without %s label", + metadata.id, + PROVIDER_REF_TAG_KEY, + PROVIDER_NAME_LABEL_KEY, + ) + return None + return name, provider_ref + + +def _get_attempt_key_from_name(metadata: ResourceMetadata) -> Optional[tuple[str, str]]: + """ + Returns the attempt's name and `provider_ref` from the resource's name, + or `None` if the name has no `provider_ref`. + """ + provider_ref = parse_provider_ref_from_name(metadata.name) + if provider_ref is None: + return None + return metadata.name, provider_ref + + +def _get_resource_group( + attempt: _ProvisioningAttempt, region: str +) -> Optional[ListedResourceGroup]: + instance = attempt.instance + owner = None + if instance is not None and instance.status.state != InstanceStatus.InstanceState.DELETING: + owner = instance + internal_resources = [r for r in (attempt.disk, attempt.pricing_policy) if r is not None] + if owner is None and not internal_resources: + return None + descriptions = [] + if owner is not None: + descriptions.append( + f"instance {owner.metadata.id} ({owner.spec.resources.platform}" + f" {owner.spec.resources.preset}, {owner.status.state.name})" + ) + if attempt.disk is not None: + descriptions.append(f"boot disk {attempt.disk.metadata.id}") + if attempt.pricing_policy is not None: + descriptions.append(f"pricing policy {attempt.pricing_policy.metadata.id}") + if owner is not None: + created_at = owner.metadata.created_at + else: + created_at = min( + (r.metadata.created_at for r in internal_resources if r.metadata.created_at), + default=None, + ) + return ListedResourceGroup( + backend=BackendType.NEBIUS, + owner_type=OwnerType.INSTANCE, + provider_ref=attempt.provider_ref, + primary_id=owner.metadata.id if owner is not None else None, + region=region, + created_at=created_at, + description=", ".join(descriptions), + backend_data=NebiusResourceGroupBackendData( + instance_id=instance.metadata.id if instance is not None else None, + disk_id=attempt.disk.metadata.id if attempt.disk is not None else None, + pricing_policy_id=( + attempt.pricing_policy.metadata.id if attempt.pricing_policy is not None else None + ), + ).model_dump_json(), + ) + + def _supported_instances(offer: InstanceOffer) -> bool: platform, _ = offer.instance.name.split() return platform in SUPPORTED_PLATFORMS diff --git a/src/dstack/_internal/core/backends/nebius/resources.py b/src/dstack/_internal/core/backends/nebius/resources.py index 4b04cdb29c..4fe3417126 100644 --- a/src/dstack/_internal/core/backends/nebius/resources.py +++ b/src/dstack/_internal/core/backends/nebius/resources.py @@ -2,11 +2,11 @@ import re import time from collections import defaultdict +from collections.abc import Awaitable, Callable, Generator, Iterable, Sequence from collections.abc import Container as ContainerT -from collections.abc import Generator, Iterable, Sequence from contextlib import contextmanager from tempfile import NamedTemporaryFile -from typing import Dict, Optional +from typing import Any, Dict, Optional from nebius.aio.operation import Operation as SDKOperation from nebius.aio.service_error import RequestError, StatusCode @@ -16,8 +16,11 @@ ComputeInstanceSpecV1, CreatePricingPolicyRequest, DeletePricingPolicyRequest, + GetPricingPolicyRequest, + ListPricingPoliciesRequest, MaxPriceV1, PricingMethod, + PricingPolicy, PricingPolicyServiceClient, PricingPolicySpec, ) @@ -30,6 +33,7 @@ DeleteDiskRequest, DeleteGpuClusterRequest, DeleteInstanceRequest, + Disk, DiskServiceClient, DiskSpec, ExistingDisk, @@ -43,6 +47,8 @@ InstanceServiceClient, InstanceSpec, IPAddress, + ListDisksRequest, + ListInstancesRequest, NetworkInterfaceSpec, PreemptibleSpec, PublicIPAddress, @@ -87,6 +93,7 @@ OPTION_RENEW_SYNCHRONOUS: "true", OPTION_RENEW_REQUEST_TIMEOUT: "5", } +LIST_PAGE_SIZE = 100 # disables log messages about errors such as invalid creds or expired timeouts logging.getLogger("nebius").setLevel(logging.CRITICAL) @@ -311,6 +318,19 @@ def delete_disk(sdk: SDK, disk_id: str) -> None: ) +def list_disks(sdk: SDK, project_id: str) -> list[Disk]: + client = DiskServiceClient(sdk) + return _list_all_pages( + lambda page_token: client.list( + ListDisksRequest( + parent_id=project_id, page_size=LIST_PAGE_SIZE, page_token=page_token + ), + per_retry_timeout=REQUEST_TIMEOUT, + auth_options=REQUEST_AUTH_OPTIONS, + ) + ) + + def create_pricing_policy( sdk: SDK, name: str, project_id: str, platform: str, max_price: str ) -> SDKOperation[Operation]: @@ -329,6 +349,16 @@ def create_pricing_policy( ) +def get_pricing_policy(sdk: SDK, pricing_policy_id: str) -> PricingPolicy: + return LOOP.await_( + PricingPolicyServiceClient(sdk).get( + GetPricingPolicyRequest(id=pricing_policy_id), + per_retry_timeout=REQUEST_TIMEOUT, + auth_options=REQUEST_AUTH_OPTIONS, + ) + ) + + def delete_pricing_policy(sdk: SDK, pricing_policy_id: str) -> None: LOOP.await_( PricingPolicyServiceClient(sdk).delete( @@ -339,6 +369,19 @@ def delete_pricing_policy(sdk: SDK, pricing_policy_id: str) -> None: ) +def list_pricing_policies(sdk: SDK, project_id: str) -> list[PricingPolicy]: + client = PricingPolicyServiceClient(sdk) + return _list_all_pages( + lambda page_token: client.list( + ListPricingPoliciesRequest( + parent_id=project_id, page_size=LIST_PAGE_SIZE, page_token=page_token + ), + per_retry_timeout=REQUEST_TIMEOUT, + auth_options=REQUEST_AUTH_OPTIONS, + ) + ) + + def create_instance( sdk: SDK, name: str, @@ -420,6 +463,19 @@ def delete_instance(sdk: SDK, instance_id: str) -> SDKOperation[Operation]: ) +def list_instances(sdk: SDK, project_id: str) -> list[Instance]: + client = InstanceServiceClient(sdk) + return _list_all_pages( + lambda page_token: client.list( + ListInstancesRequest( + parent_id=project_id, page_size=LIST_PAGE_SIZE, page_token=page_token + ), + per_retry_timeout=REQUEST_TIMEOUT, + auth_options=REQUEST_AUTH_OPTIONS, + ) + ) + + def create_cluster(sdk: SDK, name: str, project_id: str, fabric: str) -> SDKOperation[Operation]: with wrap_capacity_errors(): return LOOP.await_( @@ -481,3 +537,17 @@ def is_valid_resource_name(name: str) -> bool: def is_valid_label_value(value: str) -> bool: match = re.match(LABEL_VALUE_PATTERN, value) return match is not None + + +def _list_all_pages(list_page: Callable[[str], Awaitable[Any]]) -> list: + """ + Calls `list_page` with the next page token until all pages are fetched. + """ + items = [] + page_token = "" + while True: + response = LOOP.await_(list_page(page_token)) + items.extend(response.items) + page_token = response.next_page_token + if not page_token: + return items diff --git a/src/tests/_internal/core/backends/nebius/test_compute.py b/src/tests/_internal/core/backends/nebius/test_compute.py index e40416b225..4e70b3e30f 100644 --- a/src/tests/_internal/core/backends/nebius/test_compute.py +++ b/src/tests/_internal/core/backends/nebius/test_compute.py @@ -1,4 +1,20 @@ +import json +from datetime import datetime, timezone +from typing import Optional +from unittest.mock import Mock + import pytest +from nebius.aio.service_error import RequestError, RequestStatusExtended, StatusCode +from nebius.api.nebius.billing.v1 import PricingPolicy, PricingPolicyStatus +from nebius.api.nebius.common.v1 import ResourceMetadata +from nebius.api.nebius.compute.v1 import ( + Disk, + DiskStatus, + Instance, + InstanceSpec, + InstanceStatus, + ResourcesSpec, +) from dstack._internal.core.backends.nebius import compute as compute_module from dstack._internal.core.backends.nebius.compute import NebiusCompute @@ -6,7 +22,9 @@ NebiusConfig, NebiusServiceAccountCreds, ) +from dstack._internal.core.errors import NotYetTerminated from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.backends.orphans import ListedResourceGroup, OwnerType from dstack._internal.core.models.instances import ( Gpu, InstanceAvailability, @@ -109,3 +127,349 @@ def test_keeps_flat_rate_spot_price_as_is(self, mocker): offers = make_compute().get_all_offers_with_availability(unallocated_resources=False) assert offers[0].price == 1.2341 + + +PROVIDER_REF = "abcdef123456" +CREATED_AT = datetime(2026, 10, 1, tzinfo=timezone.utc) + + +def make_name(provider_ref: str = PROVIDER_REF, suffix: str = "a1b2") -> str: + return f"dstack-main-fleet-0-{provider_ref}{suffix}" + + +def make_labels(name: str, provider_ref: str = PROVIDER_REF) -> dict[str, str]: + return { + "owner": "dstack", + "dstack_provider_ref": provider_ref, + "dstack_provider_name": name, + } + + +def make_metadata( + resource_id: str, name: str, labels: Optional[dict[str, str]], created_at: datetime +) -> ResourceMetadata: + return ResourceMetadata(id=resource_id, name=name, labels=labels or {}, created_at=created_at) + + +def make_disk( + disk_id: str = "disk-1", + name: Optional[str] = None, + labels: Optional[dict[str, str]] = None, + state: DiskStatus.State = DiskStatus.State.READY, + created_at: datetime = CREATED_AT, +) -> Disk: + name = name or make_name() + return Disk( + metadata=make_metadata(disk_id, name, labels, created_at), + status=DiskStatus(state=state), + ) + + +def make_pricing_policy( + pricing_policy_id: str = "pricingpolicy-1", + name: Optional[str] = None, + state: PricingPolicyStatus.State = PricingPolicyStatus.State.STATE_ACTIVE, + created_at: datetime = CREATED_AT, +) -> PricingPolicy: + # Nebius drops labels of pricing policies + return PricingPolicy( + metadata=make_metadata(pricing_policy_id, name or make_name(), None, created_at), + status=PricingPolicyStatus(state=state), + ) + + +def make_instance( + instance_id: str = "instance-1", + name: Optional[str] = None, + labels: Optional[dict[str, str]] = None, + state: InstanceStatus.InstanceState = InstanceStatus.InstanceState.RUNNING, + created_at: datetime = CREATED_AT, +) -> Instance: + name = name or make_name() + return Instance( + metadata=make_metadata(instance_id, name, labels, created_at), + spec=InstanceSpec(resources=ResourcesSpec(platform="cpu-d3", preset="2vcpu-8gb")), + status=InstanceStatus(state=state), + ) + + +def make_request_error(code: StatusCode) -> RequestError: + return RequestError( + RequestStatusExtended( + code=code, + message="error", + details=[], + service_errors=[], + request_id="", + trace_id="", + ) + ) + + +def make_resource_group( + instance_id: Optional[str] = None, + disk_id: Optional[str] = None, + pricing_policy_id: Optional[str] = None, + instance_is_owner: bool = True, +) -> ListedResourceGroup: + return ListedResourceGroup( + backend=BackendType.NEBIUS, + owner_type=OwnerType.INSTANCE, + provider_ref=PROVIDER_REF, + primary_id=instance_id if instance_is_owner else None, + region="eu-north1", + created_at=CREATED_AT, + description="", + backend_data=json.dumps( + { + "instance_id": instance_id, + "disk_id": disk_id, + "pricing_policy_id": pricing_policy_id, + } + ), + ) + + +class TestListResourceGroups: + @pytest.fixture(autouse=True) + def _mock_sdk(self, mocker): + mocker.patch.object(NebiusCompute, "_sdk", Mock()) + mocker.patch.object(NebiusCompute, "_region_to_project_id", {"eu-north1": "project-id"}) + + @pytest.fixture + def mock_listings(self, mocker): + def mock( + disks: Optional[list[Disk]] = None, + pricing_policies: Optional[list[PricingPolicy]] = None, + instances: Optional[list[Instance]] = None, + ) -> None: + resources = compute_module.resources + mocker.patch.object(resources, "list_disks", return_value=disks or []) + mocker.patch.object( + resources, "list_pricing_policies", return_value=pricing_policies or [] + ) + mocker.patch.object(resources, "list_instances", return_value=instances or []) + + return mock + + def test_groups_instance_with_its_disk_and_pricing_policy(self, mock_listings): + name = make_name() + mock_listings( + disks=[make_disk(name=name, labels=make_labels(name))], + pricing_policies=[make_pricing_policy(name=name)], + instances=[make_instance(name=name, labels=make_labels(name))], + ) + + listing = make_compute().list_resource_groups() + + assert listing.failures == [] + assert len(listing.resource_groups) == 1 + group = listing.resource_groups[0] + assert group.provider_ref == PROVIDER_REF + assert group.primary_id == "instance-1" + assert group.region == "eu-north1" + assert group.created_at == CREATED_AT + assert json.loads(group.backend_data or "") == { + "instance_id": "instance-1", + "disk_id": "disk-1", + "pricing_policy_id": "pricingpolicy-1", + } + + def test_lists_leftovers_without_owner(self, mock_listings): + name = make_name() + mock_listings( + disks=[make_disk(name=name, labels=make_labels(name), created_at=CREATED_AT)], + pricing_policies=[ + make_pricing_policy( + name=name, created_at=datetime(2026, 10, 2, tzinfo=timezone.utc) + ) + ], + ) + + listing = make_compute().list_resource_groups() + + assert len(listing.resource_groups) == 1 + group = listing.resource_groups[0] + assert group.primary_id is None + assert group.created_at == CREATED_AT + assert json.loads(group.backend_data or "") == { + "instance_id": None, + "disk_id": "disk-1", + "pricing_policy_id": "pricingpolicy-1", + } + + def test_keeps_id_of_instance_being_deleted_without_owner(self, mock_listings): + name = make_name() + mock_listings( + disks=[make_disk(name=name, labels=make_labels(name))], + instances=[ + make_instance( + name=name, + labels=make_labels(name), + state=InstanceStatus.InstanceState.DELETING, + ) + ], + ) + + listing = make_compute().list_resource_groups() + + assert len(listing.resource_groups) == 1 + group = listing.resource_groups[0] + assert group.primary_id is None + assert json.loads(group.backend_data or "")["instance_id"] == "instance-1" + + def test_skips_resources_being_deleted(self, mock_listings): + name = make_name() + mock_listings( + disks=[ + make_disk(name=name, labels=make_labels(name), state=DiskStatus.State.DELETING) + ], + pricing_policies=[ + make_pricing_policy(name=name, state=PricingPolicyStatus.State.STATE_DELETING) + ], + instances=[ + make_instance( + name=name, + labels=make_labels(name), + state=InstanceStatus.InstanceState.DELETING, + ) + ], + ) + + listing = make_compute().list_resource_groups() + + assert listing.resource_groups == [] + + def test_groups_attempts_with_same_provider_ref_separately(self, mock_listings): + name_1 = make_name(suffix="aaaa") + name_2 = make_name(suffix="bbbb") + mock_listings( + disks=[ + make_disk("disk-1", name=name_1, labels=make_labels(name_1)), + make_disk("disk-2", name=name_2, labels=make_labels(name_2)), + ], + instances=[make_instance(name=name_2, labels=make_labels(name_2))], + ) + + listing = make_compute().list_resource_groups() + + groups = {g.primary_id: g for g in listing.resource_groups} + assert set(groups) == {None, "instance-1"} + assert json.loads(groups[None].backend_data or "")["disk_id"] == "disk-1" + assert json.loads(groups["instance-1"].backend_data or "")["disk_id"] == "disk-2" + + def test_skips_resources_without_valid_provider_ref(self, mock_listings): + name = make_name() + mock_listings( + disks=[ + make_disk("disk-unlabeled", name=name, labels={"owner": "dstack"}), + make_disk( + "disk-invalid", + name=name, + labels=make_labels(name, provider_ref="INVALID"), + ), + make_disk( + "disk-without-name", + name=name, + labels={"dstack_provider_ref": PROVIDER_REF}, + ), + ], + pricing_policies=[ + make_pricing_policy("pricingpolicy-old", name="dstack-main-fleet-0-ab12cd34"), + make_pricing_policy("pricingpolicy-foreign", name="my-policy"), + ], + ) + + listing = make_compute().list_resource_groups() + + assert listing.resource_groups == [] + + def test_reports_failed_region(self, mocker, mock_listings): + mocker.patch.object( + NebiusCompute, + "_region_to_project_id", + {"eu-north1": "project-1", "us-central1": "project-2"}, + ) + name = make_name() + mock_listings(disks=[make_disk(name=name, labels=make_labels(name))]) + + def list_instances(sdk, project_id: str) -> list[Instance]: + if project_id == "project-2": + raise make_request_error(StatusCode.UNAVAILABLE) + return [] + + compute_module.resources.list_instances.side_effect = list_instances + + listing = make_compute().list_resource_groups() + + assert [g.region for g in listing.resource_groups] == ["eu-north1"] + assert [f.region for f in listing.failures] == ["us-central1"] + + +class TestDeleteResourceGroup: + @pytest.fixture(autouse=True) + def _mock_sdk(self, mocker): + mocker.patch.object(NebiusCompute, "_sdk", Mock()) + + @pytest.fixture + def resources(self, mocker): + resources = compute_module.resources + for name in [ + "get_instance", + "delete_instance", + "delete_disk", + "get_pricing_policy", + "delete_pricing_policy", + ]: + mocker.patch.object(resources, name) + resources.get_instance.side_effect = make_request_error(StatusCode.NOT_FOUND) + resources.get_pricing_policy.return_value = make_pricing_policy() + return resources + + def test_deletes_instance_first(self, resources): + resources.get_instance.side_effect = None + resources.get_instance.return_value = make_instance() + group = make_resource_group("instance-1", "disk-1", "pricingpolicy-1") + + with pytest.raises(NotYetTerminated): + make_compute().delete_resource_group(group) + + resources.delete_instance.assert_called_once() + resources.delete_disk.assert_not_called() + resources.delete_pricing_policy.assert_not_called() + + def test_waits_for_instance_being_deleted(self, resources): + resources.get_instance.side_effect = None + resources.get_instance.return_value = make_instance( + state=InstanceStatus.InstanceState.DELETING + ) + group = make_resource_group( + "instance-1", "disk-1", "pricingpolicy-1", instance_is_owner=False + ) + + with pytest.raises(NotYetTerminated): + make_compute().delete_resource_group(group) + + resources.delete_instance.assert_not_called() + resources.delete_disk.assert_not_called() + + def test_deletes_disk_and_pricing_policy_once_instance_is_gone(self, resources): + group = make_resource_group("instance-1", "disk-1", "pricingpolicy-1") + + make_compute().delete_resource_group(group) + + resources.delete_disk.assert_called_once_with(NebiusCompute._sdk, "disk-1") + resources.delete_pricing_policy.assert_called_once_with( + NebiusCompute._sdk, "pricingpolicy-1" + ) + + def test_waits_for_pricing_policy_being_created(self, resources): + resources.get_pricing_policy.return_value = make_pricing_policy( + state=PricingPolicyStatus.State.STATE_CREATING + ) + group = make_resource_group(None, None, "pricingpolicy-1") + + with pytest.raises(NotYetTerminated): + make_compute().delete_resource_group(group) + + resources.delete_pricing_policy.assert_not_called()