diff --git a/src/dstack/_internal/core/backends/aws/compute.py b/src/dstack/_internal/core/backends/aws/compute.py index b5e9af0bdc..c427703710 100644 --- a/src/dstack/_internal/core/backends/aws/compute.py +++ b/src/dstack/_internal/core/backends/aws/compute.py @@ -1,5 +1,5 @@ import threading -from collections.abc import Container, Iterable +from collections.abc import Container from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field from functools import partial @@ -42,11 +42,7 @@ get_user_data, merge_tags, ) -from dstack._internal.core.backends.base.offers import ( - OfferModifier, - get_catalog_offers, - get_offers_disk_modifier, -) +from dstack._internal.core.backends.base.offers import get_catalog_offers from dstack._internal.core.errors import ( ComputeError, NoCapacityError, @@ -196,10 +192,10 @@ def get_all_offers_with_availability( ) return availability_offers - def get_offers_modifiers( - self, requirements: Requirements, full_offers: bool - ) -> Iterable[OfferModifier]: - return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)] + def get_offer_disk_size_range( + self, offer: InstanceOfferWithAvailability + ) -> Optional[Range[Memory]]: + return CONFIGURABLE_DISK_SIZE def get_offers_post_filter( self, requirements: Requirements diff --git a/src/dstack/_internal/core/backends/azure/compute.py b/src/dstack/_internal/core/backends/azure/compute.py index 0a8036fdcc..3744f7dfb9 100644 --- a/src/dstack/_internal/core/backends/azure/compute.py +++ b/src/dstack/_internal/core/backends/azure/compute.py @@ -1,7 +1,6 @@ import base64 import enum import re -from collections.abc import Iterable from concurrent.futures import ThreadPoolExecutor, as_completed from typing import Dict, List, Optional, Tuple @@ -52,11 +51,7 @@ merge_tags, requires_nvidia_proprietary_kernel_modules, ) -from dstack._internal.core.backends.base.offers import ( - OfferModifier, - get_catalog_offers, - get_offers_disk_modifier, -) +from dstack._internal.core.backends.base.offers import get_catalog_offers from dstack._internal.core.consts import DSTACK_OS_IMAGE_WITH_PROPRIETARY_NVIDIA_KERNEL_MODULES from dstack._internal.core.errors import ComputeError, NoCapacityError from dstack._internal.core.models.backends.base import BackendType @@ -73,7 +68,7 @@ ) from dstack._internal.core.models.placement import PlacementGroup from dstack._internal.core.models.resources import Memory, Range -from dstack._internal.core.models.runs import JobProvisioningData, Requirements +from dstack._internal.core.models.runs import JobProvisioningData from dstack._internal.utils.logging import get_logger logger = get_logger(__name__) @@ -117,10 +112,10 @@ def get_all_offers_with_availability( ) return offers_with_availability - def get_offers_modifiers( - self, requirements: Requirements, full_offers: bool - ) -> Iterable[OfferModifier]: - return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)] + def get_offer_disk_size_range( + self, offer: InstanceOfferWithAvailability + ) -> Optional[Range[Memory]]: + return CONFIGURABLE_DISK_SIZE def create_instance( self, diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index 2a952aff7e..f6a2cccfe2 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/src/dstack/_internal/core/backends/base/compute.py @@ -20,7 +20,11 @@ from dstack._internal import settings from dstack._internal.core.backends.base.models import JobConfiguration -from dstack._internal.core.backends.base.offers import OfferModifier, filter_offers_by_requirements +from dstack._internal.core.backends.base.offers import ( + OfferDiskFitter, + OfferModifier, + filter_offers_by_requirements, +) from dstack._internal.core.consts import ( DSTACK_RUNNER_HTTP_PORT, DSTACK_RUNNER_SSH_PORT, @@ -40,6 +44,7 @@ SSHKey, ) from dstack._internal.core.models.placement import PlacementGroup, PlacementGroupProvisioningData +from dstack._internal.core.models.resources import Memory, Range from dstack._internal.core.models.runs import Job, JobProvisioningData, Requirements, Run from dstack._internal.core.models.volumes import ( Volume, @@ -237,12 +242,25 @@ def get_offers_modifiers( """ Returns functions that modify offers before they are filtered by requirements. A modifier function can return `None` to exclude the offer. - E.g. can be used to set appropriate disk size based on requirements. + E.g. can be used to fit node resources to the requested resources. + A modifier must not modify the offer passed to it, since it is a cached offer. + + Modifiers run on every cached offer, so prefer `get_offer_disk_size_range()` + and `get_offers_post_filter()` when they are enough. See `Compute.get_offers()` for the `full_offers` argument description. """ return [] + def get_offer_disk_size_range( + self, offer: InstanceOfferWithAvailability + ) -> Optional[Range[Memory]]: + """ + Returns disk sizes the offer can be provisioned with, or `None` if the offer's disk is fixed. + `get_offers()` sets the offer's disk to the smallest size that satisfies requirements. + """ + return None + def get_offers_post_filter( self, requirements: Requirements ) -> Optional[Callable[[InstanceOfferWithAvailability], bool]]: @@ -262,7 +280,8 @@ def get_offers( offers = self.__apply_modifiers( cached_offers, self.get_offers_modifiers(requirements, full_offers) ) - offers = filter_offers_by_requirements(offers, requirements) + offers = filter_offers_by_requirements(offers, requirements, ignore_disk=True) + offers = self.__fit_disk_size(offers, requirements) post_filter = self.get_offers_post_filter(requirements) if post_filter is not None: offers = (o for o in offers if post_filter(o)) @@ -295,6 +314,15 @@ def __apply_modifiers( else: yield offer + def __fit_disk_size( + self, offers: Iterable[InstanceOfferWithAvailability], requirements: Requirements + ) -> Iterator[InstanceOfferWithAvailability]: + disk_fitter = OfferDiskFitter(requirements) + for offer in offers: + fitted_offer = disk_fitter.fit(offer, self.get_offer_disk_size_range(offer)) + if fitted_offer is not None: + yield fitted_offer + class ComputeWithFilteredOffersCached(ABC): """ diff --git a/src/dstack/_internal/core/backends/base/offers.py b/src/dstack/_internal/core/backends/base/offers.py index 43d6cb484d..f82df795a4 100644 --- a/src/dstack/_internal/core/backends/base/offers.py +++ b/src/dstack/_internal/core/backends/base/offers.py @@ -161,7 +161,9 @@ def offer_to_catalog_item(offer: InstanceOffer) -> gpuhunt.CatalogItem: ) -def requirements_to_query_filter(req: Optional[Requirements]) -> gpuhunt.QueryFilter: +def requirements_to_query_filter( + req: Optional[Requirements], ignore_disk: bool = False +) -> gpuhunt.QueryFilter: q = gpuhunt.QueryFilter(allowed_flags=SUPPORTED_GPUHUNT_FLAGS) if req is None: return q @@ -178,7 +180,7 @@ def requirements_to_query_filter(req: Optional[Requirements]) -> gpuhunt.QueryFi if res.memory: q.min_memory = res.memory.min q.max_memory = res.memory.max - if res.disk: + if res.disk and not ignore_disk: q.min_disk_size = res.disk.size.min q.max_disk_size = res.disk.size.max @@ -206,8 +208,9 @@ def requirements_to_query_filter(req: Optional[Requirements]) -> gpuhunt.QueryFi def filter_offers_by_requirements( offers: Iterable[InstanceOfferT], requirements: Optional[Requirements], + ignore_disk: bool = False, ) -> Iterator[InstanceOfferT]: - query_filter = requirements_to_query_filter(requirements) + query_filter = requirements_to_query_filter(requirements, ignore_disk=ignore_disk) for offer in offers: catalog_item = offer_to_catalog_item(offer) if gpuhunt.matches(catalog_item, q=query_filter): @@ -234,28 +237,71 @@ def choose_disk_size_mib( OfferModifier = Callable[[InstanceOfferWithAvailability], Optional[InstanceOfferWithAvailability]] -def get_offers_disk_modifier( - configurable_disk_size: Range[Memory], requirements: Requirements -) -> OfferModifier: +class OfferDiskFitter: """ - Returns a func that modifies offers disk by setting min value that satisfies both - `configurable_disk_size` and `requirements`. + Fits offers' disks to the `requirements` disk. + + Create one per `requirements` and reuse it for all offers: the fitted disk is computed + once per distinct configurable disk size range and shared by all offers it is set on. """ - def modifier(offer: InstanceOfferWithAvailability) -> Optional[InstanceOfferWithAvailability]: - requirements_disk_range = DEFAULT_DISK.size - if requirements.resources.disk is not None: - requirements_disk_range = requirements.resources.disk.size - disk_size_range = requirements_disk_range.intersect(configurable_disk_size) - if disk_size_range is None: + def __init__(self, requirements: Requirements) -> None: + disk_spec = requirements.resources.disk + # The range fixed disks are checked against, `None` if fixed disks are not checked. + self._fixed_disk_range = disk_spec.size if disk_spec is not None else None + # The range configurable disks are fitted to. + self._target_disk_range = disk_spec.size if disk_spec is not None else DEFAULT_DISK.size + self._disk_by_range: dict[tuple[Optional[float], Optional[float]], Optional[Disk]] = {} + + def fit( + self, + offer: InstanceOfferWithAvailability, + configurable_disk_size: Optional[Range[Memory]], + ) -> Optional[InstanceOfferWithAvailability]: + """ + Returns the offer with the fitted disk, or `None` if the offer doesn't fit. + + If `configurable_disk_size` is `None`, the offer's disk is fixed and is only checked + against `requirements`, the same way `filter_offers_by_requirements()` checks it. + Otherwise, the disk is set to the smallest size within both `configurable_disk_size` + and the required range (`DEFAULT_DISK` if `requirements` don't specify disk). + + Never modifies `offer`. A changed offer shares all objects with `offer` except those on + the path to the disk. + """ + disk_size_mib = offer.instance.resources.disk.size_mib + if configurable_disk_size is None: + if self._fixed_disk_fits(disk_size_mib): + return offer return None - offer_copy = offer.model_copy(deep=True) - offer_copy.instance.resources.disk = Disk( - size_mib=get_or_error(disk_size_range.min) * 1024 - ) - return offer_copy - - return modifier + disk = self._get_disk(configurable_disk_size) + if disk is None: + return None + if disk.size_mib == disk_size_mib: + return offer + resources = offer.instance.resources.model_copy(update={"disk": disk}) + instance = offer.instance.model_copy(update={"resources": resources}) + return offer.model_copy(update={"instance": instance}) + + def _fixed_disk_fits(self, disk_size_mib: int) -> bool: + if self._fixed_disk_range is None or disk_size_mib == 0: + return True + disk_size = disk_size_mib / 1024 + if self._fixed_disk_range.min is not None and disk_size < self._fixed_disk_range.min: + return False + if self._fixed_disk_range.max is not None and disk_size > self._fixed_disk_range.max: + return False + return True + + def _get_disk(self, configurable_disk_size: Range[Memory]) -> Optional[Disk]: + key = (configurable_disk_size.min, configurable_disk_size.max) + if key not in self._disk_by_range: + disk_size_range = self._target_disk_range.intersect(configurable_disk_size) + disk = None + if disk_size_range is not None: + disk = Disk(size_mib=round(get_or_error(disk_size_range.min) * 1024)) + self._disk_by_range[key] = disk + return self._disk_by_range[key] def gpu_matches_gpu_spec(gpu: Gpu, gpu_spec: GPUSpec) -> bool: diff --git a/src/dstack/_internal/core/backends/crusoe/compute.py b/src/dstack/_internal/core/backends/crusoe/compute.py index bbe4db95ed..f73dfaed6c 100644 --- a/src/dstack/_internal/core/backends/crusoe/compute.py +++ b/src/dstack/_internal/core/backends/crusoe/compute.py @@ -1,4 +1,3 @@ -from collections.abc import Iterable from typing import List, Optional import gpuhunt @@ -15,11 +14,7 @@ generate_unique_instance_name, get_shim_commands, ) -from dstack._internal.core.backends.base.offers import ( - OfferModifier, - get_catalog_offers, - get_offers_disk_modifier, -) +from dstack._internal.core.backends.base.offers import get_catalog_offers from dstack._internal.core.backends.crusoe.models import CrusoeConfig from dstack._internal.core.backends.crusoe.resources import CrusoeClient from dstack._internal.core.errors import BackendError, NotYetTerminated @@ -37,7 +32,7 @@ PlacementStrategy, ) from dstack._internal.core.models.resources import Memory, Range -from dstack._internal.core.models.runs import JobProvisioningData, Requirements +from dstack._internal.core.models.runs import JobProvisioningData from dstack._internal.utils.logging import get_logger logger = get_logger(__name__) @@ -183,21 +178,13 @@ def _get_quota_map(self) -> dict[str, int]: result[prog_name] = available return result - def get_offers_modifiers( - self, requirements: Requirements, full_offers: bool - ) -> Iterable[OfferModifier]: - # Only adjust disk size for types without ephemeral NVMe (disk_gb == 0). - # Types with ephemeral NVMe already have their disk_size set by gpuhunt. - base_modifier = get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements) - - def modifier( - offer: InstanceOfferWithAvailability, - ) -> Optional[InstanceOfferWithAvailability]: - if _has_ephemeral_disk(offer): - return offer - return base_modifier(offer) - - return [modifier] + def get_offer_disk_size_range( + self, offer: InstanceOfferWithAvailability + ) -> Optional[Range[Memory]]: + # Types with ephemeral NVMe (disk_gb > 0) already have their disk_size set by gpuhunt. + if _has_ephemeral_disk(offer): + return None + return CONFIGURABLE_DISK_SIZE def create_instance( self, diff --git a/src/dstack/_internal/core/backends/gcp/compute.py b/src/dstack/_internal/core/backends/gcp/compute.py index 662a5455c3..243bc13caa 100644 --- a/src/dstack/_internal/core/backends/gcp/compute.py +++ b/src/dstack/_internal/core/backends/gcp/compute.py @@ -45,7 +45,6 @@ from dstack._internal.core.backends.base.offers import ( OfferModifier, get_catalog_offers, - get_offers_disk_modifier, ) from dstack._internal.core.backends.gcp.features import tcpx as tcpx_features from dstack._internal.core.backends.gcp.models import GCPConfig @@ -242,9 +241,13 @@ def reservation_modifier( modifiers.append(reservation_modifier) - modifiers.append(get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)) return modifiers + def get_offer_disk_size_range( + self, offer: InstanceOfferWithAvailability + ) -> Optional[Range[Memory]]: + return CONFIGURABLE_DISK_SIZE + def get_offers_post_filter( self, requirements: Requirements ) -> Optional[Callable[[InstanceOfferWithAvailability], bool]]: diff --git a/src/dstack/_internal/core/backends/jarvislabs/compute.py b/src/dstack/_internal/core/backends/jarvislabs/compute.py index 7cb094d741..01062ca051 100644 --- a/src/dstack/_internal/core/backends/jarvislabs/compute.py +++ b/src/dstack/_internal/core/backends/jarvislabs/compute.py @@ -1,7 +1,6 @@ import shlex import subprocess import tempfile -from collections.abc import Iterable from typing import Callable, List, Optional, cast import gpuhunt @@ -18,11 +17,7 @@ generate_unique_instance_name, get_shim_commands, ) -from dstack._internal.core.backends.base.offers import ( - OfferModifier, - get_catalog_offers, - get_offers_disk_modifier, -) +from dstack._internal.core.backends.base.offers import get_catalog_offers from dstack._internal.core.backends.jarvislabs.api_client import JarvisLabsAPIClient from dstack._internal.core.backends.jarvislabs.models import JarvisLabsConfig from dstack._internal.core.errors import ProvisioningError @@ -100,10 +95,10 @@ def get_all_offers_with_availability( for offer in offers ] - def get_offers_modifiers( - self, requirements: Requirements, full_offers: bool - ) -> Iterable[OfferModifier]: - return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)] + def get_offer_disk_size_range( + self, offer: InstanceOfferWithAvailability + ) -> Optional[Range[Memory]]: + return CONFIGURABLE_DISK_SIZE def get_offers_post_filter( self, requirements: Requirements diff --git a/src/dstack/_internal/core/backends/nebius/compute.py b/src/dstack/_internal/core/backends/nebius/compute.py index 4e4e645eab..97cafb6171 100644 --- a/src/dstack/_internal/core/backends/nebius/compute.py +++ b/src/dstack/_internal/core/backends/nebius/compute.py @@ -2,7 +2,6 @@ import random import shlex import time -from collections.abc import Iterable from decimal import ROUND_CEILING, Decimal from functools import cached_property from typing import List, Optional @@ -24,11 +23,7 @@ get_user_data, merge_tags, ) -from dstack._internal.core.backends.base.offers import ( - OfferModifier, - get_catalog_offers, - get_offers_disk_modifier, -) +from dstack._internal.core.backends.base.offers import get_catalog_offers from dstack._internal.core.backends.nebius import resources from dstack._internal.core.backends.nebius.models import ( NebiusConfig, @@ -58,7 +53,7 @@ PlacementStrategy, ) from dstack._internal.core.models.resources import Memory, Range -from dstack._internal.core.models.runs import JobProvisioningData, Requirements +from dstack._internal.core.models.runs import JobProvisioningData from dstack._internal.utils.logging import get_logger logger = get_logger(__name__) @@ -145,10 +140,10 @@ def get_all_offers_with_availability( for offer in offers ] - def get_offers_modifiers( - self, requirements: Requirements, full_offers: bool - ) -> Iterable[OfferModifier]: - return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)] + def get_offer_disk_size_range( + self, offer: InstanceOfferWithAvailability + ) -> Optional[Range[Memory]]: + return CONFIGURABLE_DISK_SIZE def create_instance( self, diff --git a/src/dstack/_internal/core/backends/oci/compute.py b/src/dstack/_internal/core/backends/oci/compute.py index b4705bb0a9..783454f03e 100644 --- a/src/dstack/_internal/core/backends/oci/compute.py +++ b/src/dstack/_internal/core/backends/oci/compute.py @@ -1,4 +1,4 @@ -from collections.abc import Container, Iterable +from collections.abc import Container from concurrent.futures import ThreadPoolExecutor from functools import cached_property, partial from typing import List, Optional @@ -15,11 +15,7 @@ generate_unique_instance_name, get_user_data, ) -from dstack._internal.core.backends.base.offers import ( - OfferModifier, - get_catalog_offers, - get_offers_disk_modifier, -) +from dstack._internal.core.backends.base.offers import get_catalog_offers from dstack._internal.core.backends.oci import resources from dstack._internal.core.backends.oci.models import OCIConfig from dstack._internal.core.backends.oci.region import make_region_clients_map @@ -33,7 +29,7 @@ ) from dstack._internal.core.models.placement import PlacementGroup from dstack._internal.core.models.resources import Memory, Range -from dstack._internal.core.models.runs import JobProvisioningData, Requirements +from dstack._internal.core.models.runs import JobProvisioningData SUPPORTED_SHAPE_FAMILIES = [ "VM.Standard2.", @@ -106,10 +102,10 @@ def get_all_offers_with_availability( return offers_with_availability - def get_offers_modifiers( - self, requirements: Requirements, full_offers: bool - ) -> Iterable[OfferModifier]: - return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)] + def get_offer_disk_size_range( + self, offer: InstanceOfferWithAvailability + ) -> Optional[Range[Memory]]: + return CONFIGURABLE_DISK_SIZE def terminate_instance( self, instance_id: str, region: str, backend_data: Optional[str] = None diff --git a/src/dstack/_internal/core/backends/runpod/compute.py b/src/dstack/_internal/core/backends/runpod/compute.py index 63d677f990..5f248d3753 100644 --- a/src/dstack/_internal/core/backends/runpod/compute.py +++ b/src/dstack/_internal/core/backends/runpod/compute.py @@ -1,6 +1,5 @@ import json import uuid -from collections.abc import Iterable from datetime import timedelta from typing import Callable, List, Optional @@ -17,17 +16,11 @@ get_job_instance_name, ) from dstack._internal.core.backends.base.models import JobConfiguration -from dstack._internal.core.backends.base.offers import ( - OfferModifier, - get_catalog_offers, - get_offers_disk_modifier, -) +from dstack._internal.core.backends.base.offers import get_catalog_offers from dstack._internal.core.backends.runpod.api_client import RunpodApiClient, RunpodApiClientError from dstack._internal.core.backends.runpod.models import RunpodConfig from dstack._internal.core.consts import DSTACK_RUNNER_SSH_PORT -from dstack._internal.core.errors import ( - ComputeError, -) +from dstack._internal.core.errors import ComputeError from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.common import CoreModel, RegistryAuth, validate_extra_ignore from dstack._internal.core.models.compute_groups import ComputeGroup, ComputeGroupProvisioningData @@ -91,27 +84,16 @@ def get_all_offers_with_availability( ] return offers - def get_offers_modifiers( - self, requirements: Requirements, full_offers: bool - ) -> Iterable[OfferModifier]: - gpu_disk_modifier = get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements) - - def disk_modifier( - offer: InstanceOfferWithAvailability, - ) -> Optional[InstanceOfferWithAvailability]: - if len(offer.instance.resources.gpus) > 0: - return gpu_disk_modifier(offer) - - # For Runpod CPU offers, gpuhunt disk is the per-flavor max. - # Choose requested disk within [1GB, max] or filter the offer out. - cpu_max_disk_size_gb = Memory(offer.instance.resources.disk.size_mib / 1024) - cpu_configurable_disk_size = Range[Memory]( - min=Memory.parse("1GB"), - max=cpu_max_disk_size_gb, - ) - return get_offers_disk_modifier(cpu_configurable_disk_size, requirements)(offer) - - return [disk_modifier] + def get_offer_disk_size_range( + self, offer: InstanceOfferWithAvailability + ) -> Optional[Range[Memory]]: + if len(offer.instance.resources.gpus) > 0: + return CONFIGURABLE_DISK_SIZE + # For Runpod CPU offers, gpuhunt disk is the per-flavor max. + return Range[Memory]( + min=Memory.parse("1GB"), + max=Memory(offer.instance.resources.disk.size_mib / 1024), + ) def get_offers_post_filter( self, requirements: Requirements diff --git a/src/dstack/_internal/core/backends/verda/compute.py b/src/dstack/_internal/core/backends/verda/compute.py index c8b71b4058..e5b988e69e 100644 --- a/src/dstack/_internal/core/backends/verda/compute.py +++ b/src/dstack/_internal/core/backends/verda/compute.py @@ -1,4 +1,3 @@ -from collections.abc import Iterable from typing import Dict, List, Optional from verda import VerdaClient @@ -14,11 +13,7 @@ generate_unique_instance_name, get_shim_commands, ) -from dstack._internal.core.backends.base.offers import ( - OfferModifier, - get_catalog_offers, - get_offers_disk_modifier, -) +from dstack._internal.core.backends.base.offers import get_catalog_offers from dstack._internal.core.backends.verda.models import VerdaConfig from dstack._internal.core.errors import ( BackendError, @@ -36,7 +31,7 @@ ) from dstack._internal.core.models.placement import PlacementGroup from dstack._internal.core.models.resources import Memory, Range -from dstack._internal.core.models.runs import JobProvisioningData, Requirements +from dstack._internal.core.models.runs import JobProvisioningData from dstack._internal.utils.logging import get_logger logger = get_logger("verda.compute") @@ -74,10 +69,10 @@ def get_all_offers_with_availability( offers_with_availability = self._get_offers_with_availability(offers) return offers_with_availability - def get_offers_modifiers( - self, requirements: Requirements, full_offers: bool - ) -> Iterable[OfferModifier]: - return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)] + def get_offer_disk_size_range( + self, offer: InstanceOfferWithAvailability + ) -> Optional[Range[Memory]]: + return CONFIGURABLE_DISK_SIZE def _get_offers_with_availability( self, offers: List[InstanceOffer] diff --git a/src/dstack/_internal/server/services/backends/__init__.py b/src/dstack/_internal/server/services/backends/__init__.py index 0c5709f746..8472375b7a 100644 --- a/src/dstack/_internal/server/services/backends/__init__.py +++ b/src/dstack/_internal/server/services/backends/__init__.py @@ -475,7 +475,7 @@ async def get_backend_offers( exclude_not_available: bool = False, ) -> Iterable[Tuple[Backend, InstanceOfferWithAvailability]]: """ - Yields backend offers satisfying `requirements` sorted by price. + Returns a lazy iterable of backend offers satisfying `requirements` sorted by price. """ def get_filtered_offers_with_backends( diff --git a/src/dstack/_internal/server/services/offers.py b/src/dstack/_internal/server/services/offers.py index 7fbe24f5c3..eb826e03ee 100644 --- a/src/dstack/_internal/server/services/offers.py +++ b/src/dstack/_internal/server/services/offers.py @@ -25,6 +25,7 @@ from dstack._internal.core.models.volumes import Volume, VolumeConfigurationWithRegion from dstack._internal.server.models import ProjectModel from dstack._internal.server.services import backends as backends_services +from dstack._internal.utils.common import run_async async def get_offers_by_requirements( @@ -123,6 +124,14 @@ async def get_offers_by_requirements( if max_offers is not None: offers = itertools.islice(offers, max_offers) + # Sorting consumes the lazy offers chain, which is CPU-bound, + # so it must not run on the event loop. + return await run_async(_sort_offers_by_availability, offers) + + +def _sort_offers_by_availability( + offers: Iterable[Tuple[Backend, InstanceOfferWithAvailability]], +) -> List[Tuple[Backend, InstanceOfferWithAvailability]]: # Put NOT_AVAILABLE and NO_QUOTA offers at the end. # We have to do this after taking max_offers to avoid processing all offers # if all/most offers are unavailable. diff --git a/src/tests/_internal/core/backends/base/test_compute.py b/src/tests/_internal/core/backends/base/test_compute.py index 4fac843447..c885bafcc3 100644 --- a/src/tests/_internal/core/backends/base/test_compute.py +++ b/src/tests/_internal/core/backends/base/test_compute.py @@ -5,6 +5,7 @@ import pytest from dstack._internal.core.backends.base.compute import ( + ComputeWithAllOffersCached, GoArchType, generate_unique_backend_name, generate_unique_gateway_instance_name, @@ -12,13 +13,76 @@ generate_unique_volume_name, normalize_arch, ) +from dstack._internal.core.models.instances import Disk, InstanceOfferWithAvailability +from dstack._internal.core.models.resources import Memory, Range, ResourcesSpec +from dstack._internal.core.models.runs import Requirements from dstack._internal.server.testing.common import ( get_gateway_replica_configuration, get_instance_configuration, + get_instance_offer_with_availability, get_volume, ) +class _ComputeWithCachedOffers(ComputeWithAllOffersCached): + def __init__( + self, + offers: list[InstanceOfferWithAvailability], + disk_size_ranges: dict[str, Optional[Range[Memory]]], + ) -> None: + super().__init__() + self.offers = offers + self.disk_size_ranges = disk_size_ranges + self.disk_size_range_calls: list[InstanceOfferWithAvailability] = [] + + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> list[InstanceOfferWithAvailability]: + return self.offers + + def get_offer_disk_size_range( + self, offer: InstanceOfferWithAvailability + ) -> Optional[Range[Memory]]: + self.disk_size_range_calls.append(offer) + return self.disk_size_ranges[offer.instance.name] + + +class TestComputeWithAllOffersCachedGetOffers: + def test_fits_disk_only_for_offers_matching_other_requirements(self): + gpu_offer = get_instance_offer_with_availability( + instance_type="gpu", gpu_count=1, disk_gib=100 + ) + cpu_offer = get_instance_offer_with_availability( + instance_type="cpu", gpu_count=0, disk_gib=100 + ) + configurable = Range[Memory](min=Memory.parse("50GB"), max=None) + compute = _ComputeWithCachedOffers( + [gpu_offer, cpu_offer], {"gpu": configurable, "cpu": configurable} + ) + requirements = Requirements(resources=ResourcesSpec(gpu=1, disk="200GB..")) + + offers = list(compute.get_offers(requirements, False, False)) + + assert [o.instance.resources.disk for o in offers] == [Disk(size_mib=200 * 1024)] + assert compute.disk_size_range_calls == [gpu_offer] + assert gpu_offer.instance.resources.disk == Disk(size_mib=100 * 1024) + + def test_excludes_offers_whose_disk_does_not_fit(self): + configurable_offer = get_instance_offer_with_availability( + instance_type="configurable", disk_gib=100 + ) + fixed_offer = get_instance_offer_with_availability(instance_type="fixed", disk_gib=100) + compute = _ComputeWithCachedOffers( + [configurable_offer, fixed_offer], + {"configurable": Range[Memory](min=Memory.parse("50GB"), max=None), "fixed": None}, + ) + requirements = Requirements(resources=ResourcesSpec(disk="200GB..")) + + offers = list(compute.get_offers(requirements, False, False)) + + assert [o.instance.name for o in offers] == ["configurable"] + + class TestGenerateUniqueInstanceName: def test_generates_name(self): configuration = get_instance_configuration( diff --git a/src/tests/_internal/core/backends/base/test_offers.py b/src/tests/_internal/core/backends/base/test_offers.py index ba3829a40f..cfd928431a 100644 --- a/src/tests/_internal/core/backends/base/test_offers.py +++ b/src/tests/_internal/core/backends/base/test_offers.py @@ -1,7 +1,10 @@ +from typing import Optional + import gpuhunt import pytest from dstack._internal.core.backends.base.offers import ( + OfferDiskFitter, filter_offers_by_requirements, gpu_matches_gpu_spec, ) @@ -13,8 +16,9 @@ InstanceType, Resources, ) -from dstack._internal.core.models.resources import GPUSpec, ResourcesSpec +from dstack._internal.core.models.resources import GPUSpec, Memory, Range, ResourcesSpec from dstack._internal.core.models.runs import Requirements +from dstack._internal.server.testing.common import get_instance_offer_with_availability NVIDIA = gpuhunt.AcceleratorVendor.NVIDIA AMD = gpuhunt.AcceleratorVendor.AMD @@ -152,3 +156,71 @@ def test_known_disk_size_within_range_matches(self): requirements = Requirements(resources=ResourcesSpec(disk="100GB..200GB")) assert list(filter_offers_by_requirements([offer], requirements)) == [offer] + + def test_ignore_disk_skips_disk_requirement(self): + offer = self.make_offer(disk_size_mib=50 * 1024) + requirements = Requirements(resources=ResourcesSpec(disk="100GB..")) + + assert list(filter_offers_by_requirements([offer], requirements, ignore_disk=True)) == [ + offer + ] + + +class TestOfferDiskFitter: + def test_sets_min_disk_size_within_configurable_and_required_ranges(self): + offer = get_instance_offer_with_availability(gpu_count=1, disk_gib=100) + requirements = Requirements(resources=ResourcesSpec(disk="200GB..")) + fitter = OfferDiskFitter(requirements) + + fitted_a = fitter.fit(offer, Range[Memory](min=Memory.parse("10GB"), max=None)) + fitted_b = fitter.fit(offer, Range[Memory](min=Memory.parse("300GB"), max=None)) + fitted_c = fitter.fit(offer, Range[Memory](min=Memory.parse("10GB"), max=None)) + + assert fitted_a is not None and fitted_b is not None and fitted_c is not None + assert fitted_a.instance.resources.disk == Disk(size_mib=200 * 1024) + assert fitted_b.instance.resources.disk == Disk(size_mib=300 * 1024) + assert fitted_c.instance.resources.disk is fitted_a.instance.resources.disk + assert fitted_a.instance.resources.gpus is offer.instance.resources.gpus + assert offer.instance.resources.disk == Disk(size_mib=100 * 1024) + + def test_uses_default_disk_if_requirements_have_no_disk(self): + offer = get_instance_offer_with_availability(disk_gib=1) + requirements = Requirements(resources=ResourcesSpec(disk=None)) + configurable = Range[Memory](min=Memory.parse("10GB"), max=None) + + fitted = OfferDiskFitter(requirements).fit(offer, configurable) + + assert fitted is not None + assert fitted.instance.resources.disk == Disk(size_mib=100 * 1024) + + def test_returns_none_if_configurable_range_does_not_fit(self): + offer = get_instance_offer_with_availability() + requirements = Requirements(resources=ResourcesSpec(disk="..50GB")) + configurable = Range[Memory](min=Memory.parse("100GB"), max=None) + + assert OfferDiskFitter(requirements).fit(offer, configurable) is None + + def test_returns_same_offer_if_disk_size_unchanged(self): + offer = get_instance_offer_with_availability(disk_gib=100) + requirements = Requirements(resources=ResourcesSpec(disk="100GB..")) + configurable = Range[Memory](min=Memory.parse("10GB"), max=None) + + assert OfferDiskFitter(requirements).fit(offer, configurable) is offer + + @pytest.mark.parametrize( + ["disk_gib", "required_disk", "fits"], + [ + pytest.param(150, "100GB..200GB", True, id="within-range"), + pytest.param(50, "100GB..", False, id="below-min"), + pytest.param(300, "..200GB", False, id="above-max"), + pytest.param(0, "100GB..200GB", True, id="unknown-size"), + pytest.param(50, None, True, id="no-disk-requirement"), + ], + ) + def test_checks_fixed_disk(self, disk_gib: int, required_disk: Optional[str], fits: bool): + offer = get_instance_offer_with_availability(disk_gib=disk_gib) + requirements = Requirements(resources=ResourcesSpec(disk=required_disk)) + + fitted = OfferDiskFitter(requirements).fit(offer, None) + + assert fitted is (offer if fits else None)