From b1e8264c140773ffbe72e64ef945c0b9ee4eaac3 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Wed, 28 Oct 2026 17:03:14 +0100 Subject: [PATCH 1/4] Use Nebius Capacity Advisor for GPU offers --- .../_internal/core/backends/nebius/compute.py | 144 ++++++- .../_internal/core/backends/nebius/models.py | 3 + .../core/backends/nebius/resources.py | 39 +- .../core/backends/nebius/test_compute.py | 396 +++++++++++++++++- .../core/backends/nebius/test_resources.py | 46 ++ 5 files changed, 614 insertions(+), 14 deletions(-) diff --git a/src/dstack/_internal/core/backends/nebius/compute.py b/src/dstack/_internal/core/backends/nebius/compute.py index 89b5795a69..2fa1ade386 100644 --- a/src/dstack/_internal/core/backends/nebius/compute.py +++ b/src/dstack/_internal/core/backends/nebius/compute.py @@ -2,14 +2,18 @@ import random import shlex import time +from collections import defaultdict +from collections.abc import Iterable from dataclasses import dataclass from decimal import ROUND_CEILING, Decimal from functools import cached_property from typing import List, Optional +from cachetools import TTLCache, cachedmethod from grpc import StatusCode from nebius.aio.service_error import RequestError from nebius.api.nebius.billing.v1 import PricingPolicy, PricingPolicyStatus +from nebius.api.nebius.capacity.v1 import ResourceAdvice, ResourceAdviceStatus from nebius.api.nebius.common.v1 import ResourceMetadata from nebius.api.nebius.compute.v1 import Disk, DiskStatus, Instance, InstanceStatus from nebius.sdk import SDK @@ -17,6 +21,7 @@ from dstack._internal.core.backends.base.backend import Compute from dstack._internal.core.backends.base.compute import ( PROVIDER_REF_TAG_KEY, + ComputeTTLCache, ComputeWithAllOffersCached, ComputeWithCreateInstanceSupport, ComputeWithInstanceVolumesSupport, @@ -40,6 +45,7 @@ from dstack._internal.core.errors import ( BackendError, ComputeError, + NoCapacityError, NotYetTerminated, ProvisioningError, ) @@ -76,6 +82,8 @@ min=Memory.parse("40GB"), # min for the ubuntu24.04-cuda13.0 image max=Memory.parse("8192GB"), # max for the NETWORK_SSD disk type ) +RESOURCE_ADVICE_CACHE_TTL = 600 +"""Reuse capacity advice across offer-cache refreshes, like AWS's quota cache.""" WAIT_FOR_DISK_TIMEOUT = 20 WAIT_FOR_PRICING_POLICY_TIMEOUT = 10 WAIT_FOR_INSTANCE_TIMEOUT = 30 @@ -126,18 +134,26 @@ def __init__(self, config: NebiusConfig): super().__init__() self.config = config self._subnet_id_cache: dict[str, str] = {} + self._resource_advice_cache = ComputeTTLCache( + cache=TTLCache(maxsize=1, ttl=RESOURCE_ADVICE_CACHE_TTL) + ) @cached_property def _sdk(self) -> SDK: assert isinstance(self.config.creds, NebiusServiceAccountCreds) return resources.make_sdk(self.config.creds) + @cached_property + def _tenant_id(self) -> str: + return resources.get_tenant_id(self._sdk) + @cached_property def _region_to_project_id(self) -> dict[str, str]: return resources.get_region_to_project_id_map( self._sdk, configured_regions=self.config.regions, configured_project_ids=self.config.projects, + tenant_id=self._tenant_id, ) def _get_subnet_id(self, region: str) -> str: @@ -155,14 +171,46 @@ def get_all_offers_with_availability( locations=list(self._region_to_project_id), extra_filter=_supported_instances, ) + # Capacity Advisor batches GPU advice; prices stay in the offline catalog. + advice_by_offer: dict[tuple[str, str, bool], dict[str, InstanceAvailability]] = ( + defaultdict(dict) + ) + if any(offer.instance.resources.gpus for offer in offers): + try: + advice = self._get_resource_advice() + except Exception: + # Missing permissions or a Capacity Advisor outage must not hide catalog offers. + logger.warning("Failed to get Nebius capacity advice", exc_info=True) + else: + for item in advice: + instance = item.spec.compute_instance + if instance is None: + continue + name = f"{instance.platform} {instance.preset.name}" + for spot, status in ( + (False, item.status.on_demand), + (True, item.status.preemptible), + ): + key = (item.spec.region, name, spot) + advice_by_offer[key][item.spec.fabric] = _get_capacity_availability(status) return [ - offer.with_availability( - availability=InstanceAvailability.UNKNOWN, - price=_get_price_adjusted_for_pricing_policy(offer), + _apply_resource_advice( + offer, + advice_by_offer.get( + (offer.region, offer.instance.name, offer.instance.resources.spot), {} + ), ) for offer in offers ] + @cachedmethod( + cache=lambda self: self._resource_advice_cache.cache, + lock=lambda self: self._resource_advice_cache.lock, + ) + def _get_resource_advice(self) -> list[ResourceAdvice]: + """Supply cached advice to offer collection; propagate failures for UNKNOWN fallback.""" + return resources.list_resource_advice(self._sdk, self._tenant_id) + def get_offer_disk_size_range( self, offer: InstanceOfferWithAvailability ) -> Optional[Range[Memory]]: @@ -387,14 +435,14 @@ def create_placement_group( master_instance_offer_backend_data = validate_extra_ignore( NebiusOfferBackendData, master_instance_offer.backend_data ) - fabrics = list(master_instance_offer_backend_data.fabrics) - if self.config.fabrics is not None: - fabrics = [f for f in fabrics if f in self.config.fabrics] + fabrics = self._get_eligible_fabrics(master_instance_offer_backend_data) + if master_instance_offer_backend_data.fabrics and not fabrics: + raise NoCapacityError("No eligible InfiniBand fabrics for the instance offer") placement_group_backend_data = NebiusPlacementGroupBackendData(cluster=None) # Only create a Nebius cluster if the instance supports it. # For other instances, return dummy PlacementGroupProvisioningData. if fabrics: - fabric = random.choice(fabrics) + fabric = random.choice(list(fabrics)) op = resources.create_cluster( self._sdk, name=placement_group.name, @@ -435,9 +483,19 @@ def is_suitable_placement_group( ) return ( placement_group_backend_data.cluster is None - or placement_group_backend_data.cluster.fabric in instance_offer_backend_data.fabrics + or placement_group_backend_data.cluster.fabric + in self._get_eligible_fabrics(instance_offer_backend_data) ) + def _get_eligible_fabrics(self, data: NebiusOfferBackendData) -> set[str]: + """Constrain cluster creation/reuse by catalog compatibility, config, and cached advice.""" + return { + fabric + for fabric in data.fabrics + if (self.config.fabrics is None or fabric in self.config.fabrics) + and data.fabric_availability.get(fabric, InstanceAvailability.UNKNOWN).is_available() + } + class NebiusInstanceBackendData(CoreModel): boot_disk_id: str @@ -471,6 +529,76 @@ def load(cls, raw: Optional[str]) -> "NebiusPlacementGroupBackendData": return validate_json_extra_ignore(cls, raw) +def _get_capacity_availability( + status: ResourceAdviceStatus.Availability, +) -> InstanceAvailability: + """Map advice for offer collection; stale or unsupported advice remains UNKNOWN.""" + if status.data_state != ResourceAdviceStatus.Availability.DataState.DATA_STATE_FRESH: + return InstanceAvailability.UNKNOWN + level = status.availability_level + levels = ResourceAdviceStatus.Availability.AvailabilityLevel + # Nebius defines LIMIT_REACHED as a customer quota or capacity-block limit. + # `limit` is the maximum quota, not the remaining quota; zero `available` alone + # does not mean quota exhaustion. + if level == levels.AVAILABILITY_LEVEL_LIMIT_REACHED: + return InstanceAvailability.NO_QUOTA + if level not in ( + levels.AVAILABILITY_LEVEL_HIGH, + levels.AVAILABILITY_LEVEL_MEDIUM, + levels.AVAILABILITY_LEVEL_LOW, + ): + return InstanceAvailability.UNKNOWN + if status.available > 0: + return InstanceAvailability.AVAILABLE + return InstanceAvailability.NOT_AVAILABLE + + +def _apply_resource_advice( + offer: InstanceOffer, + advice_by_fabric: dict[str, InstanceAvailability], +) -> InstanceOfferWithAvailability: + """Annotate catalog offers, retaining per-fabric advice for cluster creation and reuse.""" + availability = InstanceAvailability.UNKNOWN + backend_data = offer.backend_data + if offer.instance.resources.gpus: + data = validate_extra_ignore(NebiusOfferBackendData, offer.backend_data) + if data.fabrics: + fabrics = data.fabrics + fabric_availability = { + f: advice_by_fabric.get(f, InstanceAvailability.UNKNOWN) for f in fabrics + } + availability = _get_best_availability(fabric_availability.values()) + backend_data = { + **offer.backend_data, + "fabric_availability": { + fabric: status.value for fabric, status in fabric_availability.items() + }, + } + else: + # Presets without GPU clustering can use any fabric in the region. + availability = _get_best_availability(advice_by_fabric.values()) + return offer.with_availability( + availability=availability, + price=_get_price_adjusted_for_pricing_policy(offer), + backend_data=backend_data, + ) + + +def _get_best_availability(availabilities: Iterable[InstanceAvailability]) -> InstanceAvailability: + """Combine fabric statuses for _apply_resource_advice(); UNKNOWN keeps an offer eligible.""" + statuses = set(availabilities) + # Report NO_QUOTA only when every compatible fabric explicitly reports a quota limit. + for availability in ( + InstanceAvailability.AVAILABLE, + InstanceAvailability.UNKNOWN, + InstanceAvailability.NOT_AVAILABLE, + InstanceAvailability.NO_QUOTA, + ): + if availability in statuses: + return availability + return InstanceAvailability.UNKNOWN + + def _wait_for_instance(sdk: SDK, op: resources.NebiusOperation) -> None: start = time.monotonic() while True: diff --git a/src/dstack/_internal/core/backends/nebius/models.py b/src/dstack/_internal/core/backends/nebius/models.py index 34af0e6fa1..5a92408f5a 100644 --- a/src/dstack/_internal/core/backends/nebius/models.py +++ b/src/dstack/_internal/core/backends/nebius/models.py @@ -7,6 +7,7 @@ from dstack._internal.core.backends.base.models import fill_data from dstack._internal.core.models.common import CoreModel +from dstack._internal.core.models.instances import InstanceAvailability DEFAULT_PROJECT_NAME_PREFIX = "default" @@ -182,6 +183,8 @@ class NebiusConfig(NebiusStoredConfig): class NebiusOfferBackendData(CoreModel): fabrics: set[str] = set() + fabric_availability: dict[str, InstanceAvailability] = {} + """Cached capacity advice for cluster creation and placement-group suitability checks.""" is_preemptible_flat_rate: bool = False """ True if the price of this spot instance does not change and pricing policies are not supported. diff --git a/src/dstack/_internal/core/backends/nebius/resources.py b/src/dstack/_internal/core/backends/nebius/resources.py index 2928cfc7ab..539c62df3c 100644 --- a/src/dstack/_internal/core/backends/nebius/resources.py +++ b/src/dstack/_internal/core/backends/nebius/resources.py @@ -25,6 +25,11 @@ PricingPolicyServiceClient, PricingPolicySpec, ) +from nebius.api.nebius.capacity.v1 import ( + ListResourceAdviceRequest, + ResourceAdvice, + ResourceAdviceServiceClient, +) from nebius.api.nebius.common.v1 import Operation, ResourceMetadata from nebius.api.nebius.compute.v1 import ( AttachedDiskSpec, @@ -158,11 +163,14 @@ def get_operation_error(op: NebiusOperation) -> str: def get_region_to_project_id_map( - sdk: SDK, configured_regions: Optional[list[str]], configured_project_ids: Optional[list[str]] + sdk: SDK, + configured_regions: Optional[list[str]], + configured_project_ids: Optional[list[str]], + tenant_id: Optional[str] = None, ) -> dict[str, str]: - """Validate backend settings and build region->project_id map""" + """Resolve projects for backend validation and compute; reuse tenant_id when supplied.""" - projects = list_tenant_projects(sdk) + projects = list_tenant_projects(sdk, tenant_id=tenant_id) if configured_regions: validate_regions( configured=set(configured_regions), available={p.status.region for p in projects} @@ -187,7 +195,8 @@ def validate_regions(configured: set[str], available: set[str]) -> None: ) -def list_tenant_projects(sdk: SDK) -> Sequence[Container]: +def get_tenant_id(sdk: SDK) -> str: + """Resolve the sole tenant for project discovery and Capacity Advisor.""" tenants = LOOP.await_( TenantServiceClient(sdk).list( ListTenantsRequest(), @@ -197,7 +206,13 @@ def list_tenant_projects(sdk: SDK) -> Sequence[Container]: ) if len(tenants.items) != 1: raise ValueError(f"Expected to find 1 tenant, found {(len(tenants.items))}") - tenant_id = tenants.items[0].metadata.id + return tenants.items[0].metadata.id + + +def list_tenant_projects(sdk: SDK, tenant_id: Optional[str] = None) -> Sequence[Container]: + """List projects for backend discovery, resolving the tenant only if its ID is not supplied.""" + if tenant_id is None: + tenant_id = get_tenant_id(sdk) projects = LOOP.await_( ProjectServiceClient(sdk).list( ListProjectsRequest(parent_id=tenant_id, page_size=999), @@ -208,6 +223,20 @@ def list_tenant_projects(sdk: SDK) -> Sequence[Container]: return projects.items +def list_resource_advice(sdk: SDK, tenant_id: str) -> list[ResourceAdvice]: + """Fetch all advice for NebiusCompute offer collection; propagate failures for UNKNOWN fallback.""" + client = ResourceAdviceServiceClient(sdk) + return _list_all_pages( + lambda page_token: client.list( + ListResourceAdviceRequest( + parent_id=tenant_id, page_size=LIST_PAGE_SIZE, page_token=page_token + ), + timeout=REQUEST_TIMEOUT, + auth_options=REQUEST_AUTH_OPTIONS, + ) + ) + + def _get_region_to_default_project_id_map( all_tenant_projects: Iterable[Container], configured_regions: Optional[list[str]] ) -> dict[str, str]: diff --git a/src/tests/_internal/core/backends/nebius/test_compute.py b/src/tests/_internal/core/backends/nebius/test_compute.py index ab1441434a..9373a3559e 100644 --- a/src/tests/_internal/core/backends/nebius/test_compute.py +++ b/src/tests/_internal/core/backends/nebius/test_compute.py @@ -5,9 +5,16 @@ from unittest.mock import Mock import pytest +from cachetools import TTLCache from grpc import StatusCode from nebius.aio.service_error import RequestError, RequestStatusExtended from nebius.api.nebius.billing.v1 import PricingPolicy, PricingPolicyStatus +from nebius.api.nebius.capacity.v1 import ( + ComputeInstanceDetails, + ResourceAdvice, + ResourceAdviceSpec, + ResourceAdviceStatus, +) from nebius.api.nebius.common.v1 import ResourceMetadata from nebius.api.nebius.compute.v1 import ( Disk, @@ -24,7 +31,7 @@ NebiusConfig, NebiusServiceAccountCreds, ) -from dstack._internal.core.errors import NotYetTerminated +from dstack._internal.core.errors import NoCapacityError, 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 ( @@ -34,6 +41,16 @@ InstanceType, Resources, ) +from dstack._internal.core.models.placement import ( + PlacementGroup, + PlacementGroupConfiguration, + PlacementGroupProvisioningData, + PlacementStrategy, +) + +CapacityAvailability = ResourceAdviceStatus.Availability +CAPACITY_LEVELS = CapacityAvailability.AvailabilityLevel +CAPACITY_DATA_STATES = CapacityAvailability.DataState def make_compute() -> NebiusCompute: @@ -71,10 +88,102 @@ def make_offer( ) +def make_capacity_status( + available: int = 0, + level: CapacityAvailability.AvailabilityLevel = CAPACITY_LEVELS.AVAILABILITY_LEVEL_LOW, + data_state: CapacityAvailability.DataState = CAPACITY_DATA_STATES.DATA_STATE_FRESH, +) -> CapacityAvailability: + return CapacityAvailability( + available=available, availability_level=level, data_state=data_state + ) + + +def make_resource_advice( + on_demand: Optional[CapacityAvailability] = None, + preemptible: Optional[CapacityAvailability] = None, + fabric: str = "fabric-2", + region: str = "eu-north1", + platform: str = "gpu-h100-sxm", + gpu_count: int = 8, +) -> ResourceAdvice: + return ResourceAdvice( + spec=ResourceAdviceSpec( + region=region, + fabric=fabric, + compute_instance=ComputeInstanceDetails( + platform=platform, + preset=ComputeInstanceDetails.Preset(name=f"{gpu_count}gpu-128vcpu-1600gb"), + ), + ), + status=ResourceAdviceStatus(on_demand=on_demand, preemptible=preemptible), + ) + + +class TestOfferCache: + def test_reuses_tenant_and_capacity_across_offer_refreshes(self, mocker): + mocker.patch.object(NebiusCompute, "_sdk", Mock()) + get_tenant_id = mocker.patch.object( + compute_module.resources, "get_tenant_id", return_value="tenant-id" + ) + get_projects = mocker.patch.object( + compute_module.resources, + "get_region_to_project_id_map", + return_value={"eu-north1": "project-id"}, + ) + list_advice = mocker.patch.object( + compute_module.resources, + "list_resource_advice", + return_value=[make_resource_advice(on_demand=make_capacity_status(1))], + ) + get_catalog_offers = mocker.patch.object( + compute_module, + "get_catalog_offers", + return_value=[make_offer(1, spot=False, is_preemptible_flat_rate=False)], + ) + compute = make_compute() + timer = Mock(return_value=0) + compute._offers_cache = TTLCache(maxsize=1, ttl=compute._offers_cache.ttl, timer=timer) + compute._resource_advice_cache.cache = TTLCache( + maxsize=1, ttl=compute._resource_advice_cache.cache.ttl, timer=timer + ) + + first = compute._get_all_offers_with_availability_cached(False) + assert first[0].availability == InstanceAvailability.AVAILABLE + list_advice.return_value = [make_resource_advice(on_demand=make_capacity_status(0))] + + timer.return_value = 179 + assert compute._get_all_offers_with_availability_cached(False) is first + get_catalog_offers.assert_called_once() + + for now in (180, 360, 540): + timer.return_value = now + offers = compute._get_all_offers_with_availability_cached(False) + assert offers[0].availability == InstanceAvailability.AVAILABLE + list_advice.assert_called_once() + assert get_catalog_offers.call_count == 4 + + timer.return_value = 720 + offers = compute._get_all_offers_with_availability_cached(False) + assert offers[0].availability == InstanceAvailability.NOT_AVAILABLE + + get_tenant_id.assert_called_once() + assert get_projects.call_args.kwargs["tenant_id"] == "tenant-id" + assert list_advice.call_count == 2 + assert all(call.args[1] == "tenant-id" for call in list_advice.call_args_list) + + class TestGetAllOffersWithAvailability: @pytest.fixture(autouse=True) def _mock_region_to_project_id(self, mocker): mocker.patch.object(NebiusCompute, "_region_to_project_id", {"eu-north1": "project-id"}) + mocker.patch.object(NebiusCompute, "_tenant_id", "tenant-id") + + @pytest.fixture(autouse=True) + def mock_capacity_advice(self, mocker): + mocker.patch.object(NebiusCompute, "_sdk", Mock()) + return mocker.patch.object( + compute_module.resources, "list_resource_advice", return_value=[] + ) @pytest.mark.parametrize( ("price", "gpu_count", "expected_price"), @@ -130,6 +239,291 @@ def test_keeps_flat_rate_spot_price_as_is(self, mocker): assert offers[0].price == 1.2341 + @pytest.mark.parametrize( + ("status", "expected"), + [ + (make_capacity_status(1), InstanceAvailability.AVAILABLE), + ( + make_capacity_status(4, level=CAPACITY_LEVELS.AVAILABILITY_LEVEL_HIGH), + InstanceAvailability.AVAILABLE, + ), + ( + make_capacity_status(2, level=CAPACITY_LEVELS.AVAILABILITY_LEVEL_MEDIUM), + InstanceAvailability.AVAILABLE, + ), + (make_capacity_status(0), InstanceAvailability.NOT_AVAILABLE), + ( + make_capacity_status(level=CAPACITY_LEVELS.AVAILABILITY_LEVEL_LIMIT_REACHED), + InstanceAvailability.NO_QUOTA, + ), + (CapacityAvailability(), InstanceAvailability.UNKNOWN), + ( + make_capacity_status(1, level=CAPACITY_LEVELS.AVAILABILITY_LEVEL_UNKNOWN), + InstanceAvailability.UNKNOWN, + ), + ( + make_capacity_status(level=CAPACITY_LEVELS.AVAILABILITY_LEVEL_UNSPECIFIED), + InstanceAvailability.UNKNOWN, + ), + ( + make_capacity_status(1, data_state=CAPACITY_DATA_STATES.DATA_STATE_STALE), + InstanceAvailability.UNKNOWN, + ), + ( + make_capacity_status( + level=CAPACITY_LEVELS.AVAILABILITY_LEVEL_LIMIT_REACHED, + data_state=CAPACITY_DATA_STATES.DATA_STATE_UNKNOWN, + ), + InstanceAvailability.UNKNOWN, + ), + ], + ) + def test_capacity_status(self, mocker, mock_capacity_advice, status, expected): + mocker.patch.object( + compute_module, + "get_catalog_offers", + return_value=[make_offer(1, spot=False, is_preemptible_flat_rate=False)], + ) + mock_capacity_advice.return_value = [make_resource_advice(on_demand=status)] + + offers = make_compute().get_all_offers_with_availability(unallocated_resources=False) + + assert offers[0].availability == expected + assert offers[0].backend_data["fabrics"] == ["fabric-2"] + + def test_distinguishes_on_demand_and_preemptible(self, mocker, mock_capacity_advice): + mocker.patch.object( + compute_module, + "get_catalog_offers", + return_value=[ + make_offer(1, spot=False, is_preemptible_flat_rate=False), + make_offer(0.5, spot=True, is_preemptible_flat_rate=False), + ], + ) + mock_capacity_advice.return_value = [ + make_resource_advice( + on_demand=make_capacity_status( + level=CAPACITY_LEVELS.AVAILABILITY_LEVEL_LIMIT_REACHED + ), + preemptible=make_capacity_status(1), + ) + ] + + offers = make_compute().get_all_offers_with_availability(unallocated_resources=False) + + assert [o.availability for o in offers] == [ + InstanceAvailability.NO_QUOTA, + InstanceAvailability.AVAILABLE, + ] + mock_capacity_advice.assert_called_once() + + def test_matches_exact_region_platform_and_preset(self, mocker, mock_capacity_advice): + one_gpu = make_offer(1, spot=False, is_preemptible_flat_rate=False, gpu_count=1) + one_gpu.instance.name = "gpu-h200-sxm 1gpu-128vcpu-1600gb" + one_gpu.backend_data["fabrics"] = [] + eight_gpus = make_offer(8, spot=False, is_preemptible_flat_rate=False) + eight_gpus.instance.name = "gpu-h200-sxm 8gpu-128vcpu-1600gb" + mocker.patch.object( + compute_module, "get_catalog_offers", return_value=[one_gpu, eight_gpus] + ) + mock_capacity_advice.return_value = [ + ResourceAdvice(), + make_resource_advice( + on_demand=make_capacity_status(13), platform="gpu-h200-sxm", gpu_count=1 + ), + make_resource_advice(on_demand=make_capacity_status(0), platform="gpu-h200-sxm"), + # Available H100s or H200s in another region do not make H200:8 available here. + make_resource_advice(on_demand=make_capacity_status(4)), + make_resource_advice( + on_demand=make_capacity_status(2), platform="gpu-h200-sxm", region="eu-west1" + ), + ] + + offers = make_compute().get_all_offers_with_availability(unallocated_resources=False) + + assert [o.availability for o in offers] == [ + InstanceAvailability.AVAILABLE, + InstanceAvailability.NOT_AVAILABLE, + ] + + @pytest.mark.parametrize( + ("fabrics", "expected"), + [ + (["fabric-2"], InstanceAvailability.NOT_AVAILABLE), + (["fabric-2", "fabric-3"], InstanceAvailability.AVAILABLE), + ], + ) + def test_matches_compatible_fabrics(self, mocker, mock_capacity_advice, fabrics, expected): + original = make_offer(1, spot=False, is_preemptible_flat_rate=False) + original.backend_data["fabrics"] = fabrics + mocker.patch.object(compute_module, "get_catalog_offers", return_value=[original]) + mock_capacity_advice.return_value = [ + make_resource_advice(on_demand=make_capacity_status(0), fabric="fabric-2"), + make_resource_advice(on_demand=make_capacity_status(4), fabric="fabric-3"), + ] + compute = make_compute() + compute.config.fabrics = ["fabric-2"] + + offers = compute.get_all_offers_with_availability(unallocated_resources=False) + + assert offers[0].availability == expected + assert offers[0].backend_data["fabrics"] == original.backend_data["fabrics"] + assert original.backend_data["fabrics"] == fabrics + assert "fabric_availability" not in original.backend_data + + @pytest.mark.parametrize( + ("allowed_fabrics", "status", "expected_fabric", "existing_suitable"), + [ + (None, make_capacity_status(0), "fabric-3", False), + (["fabric-2"], make_capacity_status(0), None, False), + ( + ["fabric-2"], + make_capacity_status(level=CAPACITY_LEVELS.AVAILABILITY_LEVEL_LIMIT_REACHED), + None, + False, + ), + (["fabric-2"], None, "fabric-2", True), + (["fabric-3"], make_capacity_status(1), "fabric-3", False), + ], + ) + def test_capacity_advice_respects_cluster_fabric( + self, + mocker, + mock_capacity_advice, + allowed_fabrics, + status, + expected_fabric, + existing_suitable, + ): + original = make_offer(1, spot=False, is_preemptible_flat_rate=False) + original.backend_data["fabrics"] = ["fabric-2", "fabric-3"] + mocker.patch.object(compute_module, "get_catalog_offers", return_value=[original]) + mock_capacity_advice.return_value = [ + make_resource_advice(on_demand=make_capacity_status(4), fabric="fabric-3"), + ] + if status is not None: + mock_capacity_advice.return_value.append( + make_resource_advice(on_demand=status, fabric="fabric-2") + ) + create_cluster = mocker.patch.object( + compute_module.resources, + "create_cluster", + return_value=SimpleNamespace(resource_id="cluster-id"), + ) + compute = make_compute() + compute.config.fabrics = allowed_fabrics + offer = compute.get_all_offers_with_availability(unallocated_resources=False)[0] + group = PlacementGroup( + name="cluster", + project_name="main", + configuration=PlacementGroupConfiguration( + backend=BackendType.NEBIUS, + region="eu-north1", + placement_strategy=PlacementStrategy.CLUSTER, + ), + provisioning_data=PlacementGroupProvisioningData( + backend=BackendType.NEBIUS, + backend_data=json.dumps({"cluster": {"id": "cluster-id", "fabric": "fabric-2"}}), + ), + ) + + assert offer.backend_data["fabrics"] == ["fabric-2", "fabric-3"] + assert "fabric_availability" not in original.backend_data + assert compute.is_suitable_placement_group(group, offer) is existing_suitable + if expected_fabric is None: + with pytest.raises(NoCapacityError): + compute.create_placement_group(group, offer) + create_cluster.assert_not_called() + else: + group.provisioning_data = compute.create_placement_group(group, offer) + assert create_cluster.call_args.kwargs["fabric"] == expected_fabric + assert compute.is_suitable_placement_group(group, offer) + + @pytest.mark.parametrize( + ("second_status", "expected"), + [ + (None, InstanceAvailability.UNKNOWN), + ( + make_capacity_status(0), + InstanceAvailability.NOT_AVAILABLE, + ), + ( + make_capacity_status(level=CAPACITY_LEVELS.AVAILABILITY_LEVEL_LIMIT_REACHED), + InstanceAvailability.NO_QUOTA, + ), + ], + ) + def test_aggregates_fabric_statuses( + self, mocker, mock_capacity_advice, second_status, expected + ): + original = make_offer(1, spot=False, is_preemptible_flat_rate=False) + original.backend_data["fabrics"] = ["fabric-2", "fabric-3"] + mocker.patch.object(compute_module, "get_catalog_offers", return_value=[original]) + mock_capacity_advice.return_value = [ + make_resource_advice( + on_demand=make_capacity_status( + level=CAPACITY_LEVELS.AVAILABILITY_LEVEL_LIMIT_REACHED + ), + fabric="fabric-2", + ) + ] + if second_status is not None: + mock_capacity_advice.return_value.append( + make_resource_advice(on_demand=second_status, fabric="fabric-3") + ) + + offers = make_compute().get_all_offers_with_availability(unallocated_resources=False) + + assert offers[0].availability == expected + assert offers[0].backend_data["fabrics"] == original.backend_data["fabrics"] + assert "fabric_availability" not in original.backend_data + + def test_preserves_offers_when_advice_fails(self, mocker, mock_capacity_advice): + mocker.patch.object( + compute_module, + "get_catalog_offers", + return_value=[make_offer(1, spot=False, is_preemptible_flat_rate=False)], + ) + mock_capacity_advice.side_effect = make_request_error(StatusCode.PERMISSION_DENIED) + + offers = make_compute().get_all_offers_with_availability(unallocated_resources=False) + + assert len(offers) == 1 + assert offers[0].availability == InstanceAvailability.UNKNOWN + + def test_skips_advice_for_cpu_offers(self, mocker, mock_capacity_advice): + mocker.patch.object( + compute_module, + "get_catalog_offers", + return_value=[make_offer(1, spot=False, is_preemptible_flat_rate=False, gpu_count=0)], + ) + + offers = make_compute().get_all_offers_with_availability(unallocated_resources=False) + + assert offers[0].availability == InstanceAvailability.UNKNOWN + mock_capacity_advice.assert_not_called() + + def test_nonclustered_preset_can_use_any_fabric(self, mocker, mock_capacity_advice): + original = make_offer(1, spot=False, is_preemptible_flat_rate=False, gpu_count=1) + original.backend_data["fabrics"] = [] + mocker.patch.object(compute_module, "get_catalog_offers", return_value=[original]) + mock_capacity_advice.return_value = [ + make_resource_advice( + on_demand=make_capacity_status(0), fabric="fabric-2", gpu_count=1 + ), + make_resource_advice( + on_demand=make_capacity_status(1), fabric="fabric-3", gpu_count=1 + ), + ] + compute = make_compute() + # The fabric setting restricts InfiniBand clusters, not nonclustered VMs. + compute.config.fabrics = ["fabric-2"] + + offers = compute.get_all_offers_with_availability(unallocated_resources=False) + + assert offers[0].availability == InstanceAvailability.AVAILABLE + assert offers[0].backend_data["fabrics"] == [] + PROVIDER_REF = "abcdef123456" CREATED_AT = datetime(2026, 10, 1, tzinfo=timezone.utc) diff --git a/src/tests/_internal/core/backends/nebius/test_resources.py b/src/tests/_internal/core/backends/nebius/test_resources.py index 2764e867c0..96b9849b3c 100644 --- a/src/tests/_internal/core/backends/nebius/test_resources.py +++ b/src/tests/_internal/core/backends/nebius/test_resources.py @@ -1,3 +1,9 @@ +from unittest.mock import AsyncMock, Mock + +import pytest +from nebius.api.nebius.capacity.v1 import ListResourceAdviceResponse, ResourceAdvice +from nebius.api.nebius.common.v1 import ResourceMetadata + from dstack._internal.core.backends.nebius import resources from dstack._internal.core.backends.nebius.models import NebiusServiceAccountCreds @@ -14,3 +20,43 @@ def test_make_sdk_sets_user_agent_prefix(mocker): resources.make_sdk(creds) assert sdk_cls.call_args.kwargs["user_agent_prefix"] == "dstack/1.2.3" + + +class TestListResourceAdvice: + def test_lists_all_pages(self, mocker): + sdk = Mock() + client_cls = mocker.patch.object(resources, "ResourceAdviceServiceClient") + first = ResourceAdvice(metadata=ResourceMetadata(id="advice-1")) + second = ResourceAdvice(metadata=ResourceMetadata(id="advice-2")) + client_cls.return_value.list = AsyncMock( + side_effect=[ + ListResourceAdviceResponse(items=[first], next_page_token="page-2"), + ListResourceAdviceResponse(items=[second]), + ] + ) + + advice = resources.list_resource_advice(sdk, "tenant-id") + + assert [item.metadata.id for item in advice] == ["advice-1", "advice-2"] + client_cls.assert_called_once_with(sdk) + calls = client_cls.return_value.list.call_args_list + assert [call.args[0].page_token for call in calls] == ["", "page-2"] + for call in calls: + assert call.args[0].parent_id == "tenant-id" + assert call.args[0].page_size == resources.LIST_PAGE_SIZE + assert call.kwargs == { + "timeout": resources.REQUEST_TIMEOUT, + "auth_options": resources.REQUEST_AUTH_OPTIONS, + } + + def test_does_not_return_partial_advice_on_failure(self, mocker): + client = mocker.patch.object(resources, "ResourceAdviceServiceClient").return_value + client.list = AsyncMock( + side_effect=[ + ListResourceAdviceResponse(items=[ResourceAdvice()], next_page_token="page-2"), + TimeoutError("capacity request timed out"), + ] + ) + + with pytest.raises(TimeoutError): + resources.list_resource_advice(Mock(), "tenant-id") From e070a2324fb63aed3c343d1f7b9706308bdbfa65 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Wed, 28 Oct 2026 17:08:07 +0100 Subject: [PATCH 2/4] Clarify Nebius helper contracts --- .../_internal/core/backends/nebius/compute.py | 12 ++++++------ .../_internal/core/backends/nebius/models.py | 2 +- .../_internal/core/backends/nebius/resources.py | 14 ++++++++++---- 3 files changed, 17 insertions(+), 11 deletions(-) diff --git a/src/dstack/_internal/core/backends/nebius/compute.py b/src/dstack/_internal/core/backends/nebius/compute.py index 2fa1ade386..34abcc92ad 100644 --- a/src/dstack/_internal/core/backends/nebius/compute.py +++ b/src/dstack/_internal/core/backends/nebius/compute.py @@ -83,7 +83,7 @@ max=Memory.parse("8192GB"), # max for the NETWORK_SSD disk type ) RESOURCE_ADVICE_CACHE_TTL = 600 -"""Reuse capacity advice across offer-cache refreshes, like AWS's quota cache.""" +"""Seconds to cache tenant advice across offer refreshes; matches AWS's quota cache.""" WAIT_FOR_DISK_TIMEOUT = 20 WAIT_FOR_PRICING_POLICY_TIMEOUT = 10 WAIT_FOR_INSTANCE_TIMEOUT = 30 @@ -208,7 +208,7 @@ def get_all_offers_with_availability( lock=lambda self: self._resource_advice_cache.lock, ) def _get_resource_advice(self) -> list[ResourceAdvice]: - """Supply cached advice to offer collection; propagate failures for UNKNOWN fallback.""" + """Return advice for offer collection; fetch on cache miss and propagate API errors.""" return resources.list_resource_advice(self._sdk, self._tenant_id) def get_offer_disk_size_range( @@ -488,7 +488,7 @@ def is_suitable_placement_group( ) def _get_eligible_fabrics(self, data: NebiusOfferBackendData) -> set[str]: - """Constrain cluster creation/reuse by catalog compatibility, config, and cached advice.""" + """Cluster fabrics that are compatible, allowed, and AVAILABLE or UNKNOWN.""" return { fabric for fabric in data.fabrics @@ -532,7 +532,7 @@ def load(cls, raw: Optional[str]) -> "NebiusPlacementGroupBackendData": def _get_capacity_availability( status: ResourceAdviceStatus.Availability, ) -> InstanceAvailability: - """Map advice for offer collection; stale or unsupported advice remains UNKNOWN.""" + """Convert fresh Nebius advice to offer status; stale or unknown data returns UNKNOWN.""" if status.data_state != ResourceAdviceStatus.Availability.DataState.DATA_STATE_FRESH: return InstanceAvailability.UNKNOWN level = status.availability_level @@ -557,7 +557,7 @@ def _apply_resource_advice( offer: InstanceOffer, advice_by_fabric: dict[str, InstanceAvailability], ) -> InstanceOfferWithAvailability: - """Annotate catalog offers, retaining per-fabric advice for cluster creation and reuse.""" + """Set offer status across compatible fabrics; keep per-fabric advice for cluster placement.""" availability = InstanceAvailability.UNKNOWN backend_data = offer.backend_data if offer.instance.resources.gpus: @@ -585,7 +585,7 @@ def _apply_resource_advice( def _get_best_availability(availabilities: Iterable[InstanceAvailability]) -> InstanceAvailability: - """Combine fabric statuses for _apply_resource_advice(); UNKNOWN keeps an offer eligible.""" + """Prefer AVAILABLE, UNKNOWN, NOT_AVAILABLE, then NO_QUOTA; no advice means UNKNOWN.""" statuses = set(availabilities) # Report NO_QUOTA only when every compatible fabric explicitly reports a quota limit. for availability in ( diff --git a/src/dstack/_internal/core/backends/nebius/models.py b/src/dstack/_internal/core/backends/nebius/models.py index 5a92408f5a..3d5ea4384d 100644 --- a/src/dstack/_internal/core/backends/nebius/models.py +++ b/src/dstack/_internal/core/backends/nebius/models.py @@ -184,7 +184,7 @@ class NebiusConfig(NebiusStoredConfig): class NebiusOfferBackendData(CoreModel): fabrics: set[str] = set() fabric_availability: dict[str, InstanceAvailability] = {} - """Cached capacity advice for cluster creation and placement-group suitability checks.""" + """Per-fabric offer status for cluster placement; missing entries mean UNKNOWN.""" is_preemptible_flat_rate: bool = False """ True if the price of this spot instance does not change and pricing policies are not supported. diff --git a/src/dstack/_internal/core/backends/nebius/resources.py b/src/dstack/_internal/core/backends/nebius/resources.py index 539c62df3c..137f41837b 100644 --- a/src/dstack/_internal/core/backends/nebius/resources.py +++ b/src/dstack/_internal/core/backends/nebius/resources.py @@ -168,7 +168,10 @@ def get_region_to_project_id_map( configured_project_ids: Optional[list[str]], tenant_id: Optional[str] = None, ) -> dict[str, str]: - """Resolve projects for backend validation and compute; reuse tenant_id when supplied.""" + """Select one project per region for NebiusConfigurator and NebiusCompute. + + Use configured projects when supplied, otherwise per-region defaults. + """ projects = list_tenant_projects(sdk, tenant_id=tenant_id) if configured_regions: @@ -196,7 +199,7 @@ def validate_regions(configured: set[str], available: set[str]) -> None: def get_tenant_id(sdk: SDK) -> str: - """Resolve the sole tenant for project discovery and Capacity Advisor.""" + """Require one visible tenant; used by project selection and Capacity Advisor.""" tenants = LOOP.await_( TenantServiceClient(sdk).list( ListTenantsRequest(), @@ -210,7 +213,10 @@ def get_tenant_id(sdk: SDK) -> str: def list_tenant_projects(sdk: SDK, tenant_id: Optional[str] = None) -> Sequence[Container]: - """List projects for backend discovery, resolving the tenant only if its ID is not supplied.""" + """Return projects for get_region_to_project_id_map() to validate and select. + + Without tenant_id, require exactly one visible tenant. + """ if tenant_id is None: tenant_id = get_tenant_id(sdk) projects = LOOP.await_( @@ -224,7 +230,7 @@ def list_tenant_projects(sdk: SDK, tenant_id: Optional[str] = None) -> Sequence[ def list_resource_advice(sdk: SDK, tenant_id: str) -> list[ResourceAdvice]: - """Fetch all advice for NebiusCompute offer collection; propagate failures for UNKNOWN fallback.""" + """Fetch every advice page for NebiusCompute's cache; raise if any page fails.""" client = ResourceAdviceServiceClient(sdk) return _list_all_pages( lambda page_token: client.list( From 81bae08d032a79fd12836674abd7e3948c13cf3e Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Thu, 29 Oct 2026 10:42:09 +0100 Subject: [PATCH 3/4] Keep Nebius reserved capacity eligible for provisioning --- .../_internal/core/backends/nebius/compute.py | 19 +++-- .../core/backends/nebius/test_compute.py | 70 +++++++++++++++++-- 2 files changed, 77 insertions(+), 12 deletions(-) diff --git a/src/dstack/_internal/core/backends/nebius/compute.py b/src/dstack/_internal/core/backends/nebius/compute.py index 34abcc92ad..3021b60e5d 100644 --- a/src/dstack/_internal/core/backends/nebius/compute.py +++ b/src/dstack/_internal/core/backends/nebius/compute.py @@ -187,12 +187,21 @@ def get_all_offers_with_availability( if instance is None: continue name = f"{instance.platform} {instance.preset.name}" - for spot, status in ( - (False, item.status.on_demand), - (True, item.status.preemptible), + # AUTO reservation policy can use reserved capacity before the shared pool. + # Advice is tenant-wide; project reservation limits remain unchecked, so + # reserved capacity keeps offers eligible as UNKNOWN, not AVAILABLE. + reserved = _get_capacity_availability(item.status.reserved) + if reserved == InstanceAvailability.AVAILABLE: + reserved = InstanceAvailability.UNKNOWN + on_demand = _get_best_availability( + (_get_capacity_availability(item.status.on_demand), reserved) + ) + for spot, availability in ( + (False, on_demand), + (True, _get_capacity_availability(item.status.preemptible)), ): key = (item.spec.region, name, spot) - advice_by_offer[key][item.spec.fabric] = _get_capacity_availability(status) + advice_by_offer[key][item.spec.fabric] = availability return [ _apply_resource_advice( offer, @@ -587,7 +596,7 @@ def _apply_resource_advice( def _get_best_availability(availabilities: Iterable[InstanceAvailability]) -> InstanceAvailability: """Prefer AVAILABLE, UNKNOWN, NOT_AVAILABLE, then NO_QUOTA; no advice means UNKNOWN.""" statuses = set(availabilities) - # Report NO_QUOTA only when every compatible fabric explicitly reports a quota limit. + # Report NO_QUOTA only when every allocation option explicitly reports a limit. for availability in ( InstanceAvailability.AVAILABLE, InstanceAvailability.UNKNOWN, diff --git a/src/tests/_internal/core/backends/nebius/test_compute.py b/src/tests/_internal/core/backends/nebius/test_compute.py index 9373a3559e..cb2ff3220a 100644 --- a/src/tests/_internal/core/backends/nebius/test_compute.py +++ b/src/tests/_internal/core/backends/nebius/test_compute.py @@ -105,6 +105,7 @@ def make_resource_advice( region: str = "eu-north1", platform: str = "gpu-h100-sxm", gpu_count: int = 8, + reserved: Optional[CapacityAvailability] = None, ) -> ResourceAdvice: return ResourceAdvice( spec=ResourceAdviceSpec( @@ -115,7 +116,13 @@ def make_resource_advice( preset=ComputeInstanceDetails.Preset(name=f"{gpu_count}gpu-128vcpu-1600gb"), ), ), - status=ResourceAdviceStatus(on_demand=on_demand, preemptible=preemptible), + status=ResourceAdviceStatus( + on_demand=on_demand, + preemptible=preemptible, + reserved=reserved + if reserved is not None + else make_capacity_status(level=CAPACITY_LEVELS.AVAILABILITY_LEVEL_LIMIT_REACHED), + ), ) @@ -317,6 +324,51 @@ def test_distinguishes_on_demand_and_preemptible(self, mocker, mock_capacity_adv ] mock_capacity_advice.assert_called_once() + @pytest.mark.parametrize( + ("shared", "reserved", "expected"), + [ + ( + make_capacity_status(level=CAPACITY_LEVELS.AVAILABILITY_LEVEL_LIMIT_REACHED), + make_capacity_status(1), + InstanceAvailability.UNKNOWN, + ), + (make_capacity_status(0), make_capacity_status(1), InstanceAvailability.UNKNOWN), + (make_capacity_status(1), make_capacity_status(1), InstanceAvailability.AVAILABLE), + (make_capacity_status(0), CapacityAvailability(), InstanceAvailability.UNKNOWN), + ( + make_capacity_status(0), + make_capacity_status(1, data_state=CAPACITY_DATA_STATES.DATA_STATE_STALE), + InstanceAvailability.UNKNOWN, + ), + ], + ) + def test_reservations_keep_only_on_demand_eligible( + self, mocker, mock_capacity_advice, shared, reserved, expected + ): + mocker.patch.object( + compute_module, + "get_catalog_offers", + return_value=[ + make_offer(1, spot=False, is_preemptible_flat_rate=False), + make_offer(0.5, spot=True, is_preemptible_flat_rate=False), + ], + ) + mock_capacity_advice.return_value = [ + make_resource_advice( + on_demand=shared, preemptible=make_capacity_status(0), reserved=reserved + ) + ] + + offers = make_compute().get_all_offers_with_availability(unallocated_resources=False) + + assert [offer.availability for offer in offers] == [ + expected, + InstanceAvailability.NOT_AVAILABLE, + ] + assert offers[0].backend_data["fabric_availability"] == {"fabric-2": expected.value} + assert offers[1].backend_data["fabric_availability"] == {"fabric-2": "not_available"} + mock_capacity_advice.assert_called_once() + def test_matches_exact_region_platform_and_preset(self, mocker, mock_capacity_advice): one_gpu = make_offer(1, spot=False, is_preemptible_flat_rate=False, gpu_count=1) one_gpu.instance.name = "gpu-h200-sxm 1gpu-128vcpu-1600gb" @@ -372,18 +424,21 @@ def test_matches_compatible_fabrics(self, mocker, mock_capacity_advice, fabrics, assert "fabric_availability" not in original.backend_data @pytest.mark.parametrize( - ("allowed_fabrics", "status", "expected_fabric", "existing_suitable"), + ("allowed_fabrics", "status", "reserved", "expected_fabric", "existing_suitable"), [ - (None, make_capacity_status(0), "fabric-3", False), - (["fabric-2"], make_capacity_status(0), None, False), + (None, make_capacity_status(0), None, "fabric-3", False), + (["fabric-2"], make_capacity_status(0), None, None, False), ( ["fabric-2"], make_capacity_status(level=CAPACITY_LEVELS.AVAILABILITY_LEVEL_LIMIT_REACHED), None, + None, False, ), - (["fabric-2"], None, "fabric-2", True), - (["fabric-3"], make_capacity_status(1), "fabric-3", False), + (["fabric-2"], None, None, "fabric-2", True), + (["fabric-3"], make_capacity_status(1), None, "fabric-3", False), + (["fabric-2"], make_capacity_status(0), make_capacity_status(1), "fabric-2", True), + (["fabric-3"], make_capacity_status(0), make_capacity_status(1), "fabric-3", False), ], ) def test_capacity_advice_respects_cluster_fabric( @@ -392,6 +447,7 @@ def test_capacity_advice_respects_cluster_fabric( mock_capacity_advice, allowed_fabrics, status, + reserved, expected_fabric, existing_suitable, ): @@ -403,7 +459,7 @@ def test_capacity_advice_respects_cluster_fabric( ] if status is not None: mock_capacity_advice.return_value.append( - make_resource_advice(on_demand=status, fabric="fabric-2") + make_resource_advice(on_demand=status, reserved=reserved, fabric="fabric-2") ) create_cluster = mocker.patch.object( compute_module.resources, From 1893a350231c0e7ef697cbfdb283c04358b0b840 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Thu, 29 Oct 2026 11:18:18 +0100 Subject: [PATCH 4/4] Allow sharing Nebius capacity advice across backends --- .../_internal/core/backends/nebius/backend.py | 6 ++-- .../_internal/core/backends/nebius/compute.py | 21 ++++++++--- .../core/backends/nebius/test_compute.py | 35 +++++++++++++++++-- 3 files changed, 54 insertions(+), 8 deletions(-) diff --git a/src/dstack/_internal/core/backends/nebius/backend.py b/src/dstack/_internal/core/backends/nebius/backend.py index c79709b8b7..afa1963c13 100644 --- a/src/dstack/_internal/core/backends/nebius/backend.py +++ b/src/dstack/_internal/core/backends/nebius/backend.py @@ -1,3 +1,5 @@ +from typing import Optional + from dstack._internal.core.backends.base.backend import Backend from dstack._internal.core.backends.nebius.compute import NebiusCompute from dstack._internal.core.backends.nebius.models import NebiusConfig @@ -8,9 +10,9 @@ class NebiusBackend(Backend): TYPE = BackendType.NEBIUS COMPUTE_CLASS = NebiusCompute - def __init__(self, config: NebiusConfig): + def __init__(self, config: NebiusConfig, compute: Optional[NebiusCompute] = None): self.config = config - self._compute = NebiusCompute(self.config) + self._compute = compute if compute is not None else NebiusCompute(self.config) def compute(self) -> NebiusCompute: return self._compute diff --git a/src/dstack/_internal/core/backends/nebius/compute.py b/src/dstack/_internal/core/backends/nebius/compute.py index 3021b60e5d..0056fa754d 100644 --- a/src/dstack/_internal/core/backends/nebius/compute.py +++ b/src/dstack/_internal/core/backends/nebius/compute.py @@ -10,6 +10,7 @@ from typing import List, Optional from cachetools import TTLCache, cachedmethod +from cachetools.keys import hashkey from grpc import StatusCode from nebius.aio.service_error import RequestError from nebius.api.nebius.billing.v1 import PricingPolicy, PricingPolicyStatus @@ -130,13 +131,17 @@ class NebiusCompute( ComputeWithResourceListingSupport, Compute, ): - def __init__(self, config: NebiusConfig): + def __init__( + self, config: NebiusConfig, resource_advice_cache: Optional[ComputeTTLCache] = None + ): super().__init__() self.config = config self._subnet_id_cache: dict[str, str] = {} - self._resource_advice_cache = ComputeTTLCache( - cache=TTLCache(maxsize=1, ttl=RESOURCE_ADVICE_CACHE_TTL) - ) + if resource_advice_cache is None: + resource_advice_cache = ComputeTTLCache( + cache=TTLCache(maxsize=1, ttl=RESOURCE_ADVICE_CACHE_TTL) + ) + self._resource_advice_cache = resource_advice_cache @cached_property def _sdk(self) -> SDK: @@ -212,8 +217,16 @@ def get_all_offers_with_availability( for offer in offers ] + def _get_resource_advice_key(self) -> tuple: + return hashkey( + self._tenant_id, + self.config.creds.service_account_id, + self.config.creds.public_key_id, + ) + @cachedmethod( cache=lambda self: self._resource_advice_cache.cache, + key=_get_resource_advice_key, lock=lambda self: self._resource_advice_cache.lock, ) def _get_resource_advice(self) -> list[ResourceAdvice]: diff --git a/src/tests/_internal/core/backends/nebius/test_compute.py b/src/tests/_internal/core/backends/nebius/test_compute.py index cb2ff3220a..6483d6cf2d 100644 --- a/src/tests/_internal/core/backends/nebius/test_compute.py +++ b/src/tests/_internal/core/backends/nebius/test_compute.py @@ -25,6 +25,7 @@ ResourcesSpec, ) +from dstack._internal.core.backends.base.compute import ComputeTTLCache from dstack._internal.core.backends.nebius import compute as compute_module from dstack._internal.core.backends.nebius.compute import NebiusCompute from dstack._internal.core.backends.nebius.models import ( @@ -53,7 +54,7 @@ CAPACITY_DATA_STATES = CapacityAvailability.DataState -def make_compute() -> NebiusCompute: +def make_compute(resource_advice_cache: Optional[ComputeTTLCache] = None) -> NebiusCompute: return NebiusCompute( NebiusConfig( creds=NebiusServiceAccountCreds( @@ -61,7 +62,8 @@ def make_compute() -> NebiusCompute: public_key_id="public-key-id", private_key_content="private-key", ) - ) + ), + resource_advice_cache=resource_advice_cache, ) @@ -127,6 +129,35 @@ def make_resource_advice( class TestOfferCache: + @pytest.mark.parametrize("identity", ["tenant_id", "service_account_id", "public_key_id"]) + def test_shared_advice_cache_is_scoped_to_tenant_and_credentials(self, mocker, identity): + mocker.patch.object(NebiusCompute, "_sdk", Mock()) + mocker.patch.object(NebiusCompute, "_tenant_id", "tenant-id") + list_advice = mocker.patch.object( + compute_module.resources, + "list_resource_advice", + return_value=[make_resource_advice(on_demand=make_capacity_status(1))], + ) + cache = ComputeTTLCache( + cache=TTLCache(maxsize=10, ttl=compute_module.RESOURCE_ADVICE_CACHE_TTL) + ) + first = make_compute(resource_advice_cache=cache) + second = make_compute(resource_advice_cache=cache) + + advice = first._get_resource_advice() + assert second._get_resource_advice() is advice + list_advice.assert_called_once() + + if identity == "tenant_id": + mocker.patch.object(second, "_tenant_id", "other-tenant-id") + else: + setattr(second.config.creds, identity, "other-id") + list_advice.return_value = [] + + assert second._get_resource_advice() == [] + assert first._get_resource_advice() is advice + assert list_advice.call_count == 2 + def test_reuses_tenant_and_capacity_across_offer_refreshes(self, mocker): mocker.patch.object(NebiusCompute, "_sdk", Mock()) get_tenant_id = mocker.patch.object(