Speed up offer lookups and consume them off the event loop

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

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,

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):
     """

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:

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,

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]]:

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