From 55543583170b91ef45698424ff4062955f8c12b3 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Fri, 24 Jul 2026 13:14:05 +0000 Subject: [PATCH] Add `unallocated_resources` option to `Compute.get_offers()` Adds a new `unallocated_resources` argument alongside `full_offers`. When set, backends that pack multiple jobs onto a node (Kubernetes; not yet Slurm) subtract already-allocated resources so offers reflect only unallocated (available) capacity. Exposed as the `--unallocated` CLI flag (`dstack offer` and `dstack apply`) and the `unallocated_resources` API parameter of `get_plan`/`list_gpus`. It is a discovery-only knob: the scheduling path always passes `False`. Kubernetes previously subtracted allocated resources unconditionally. That is now off: the cached allocation view can lag reality (e.g. a node whose previous job already terminated still looks allocated), which blocked provisioning of nodes that were actually free. Provisioning is now optimistic and lets the scheduler decide; the short `JOB_POD_SCHEDULING_TIMEOUT` makes a retry cheap. --- src/dstack/_internal/cli/commands/offer.py | 7 ++ .../cli/services/configurators/run.py | 6 ++ .../_internal/core/backends/aws/compute.py | 16 ++-- .../_internal/core/backends/azure/compute.py | 4 +- .../_internal/core/backends/base/compute.py | 76 +++++++++++++++---- .../core/backends/cloudrift/compute.py | 4 +- .../_internal/core/backends/crusoe/compute.py | 4 +- .../backends/digitalocean_base/compute.py | 4 +- .../_internal/core/backends/gcp/compute.py | 4 +- .../core/backends/hotaisle/compute.py | 4 +- .../core/backends/jarvislabs/compute.py | 4 +- .../core/backends/kubernetes/compute.py | 8 +- .../core/backends/kubernetes/resources.py | 14 +++- .../core/backends/lambdalabs/compute.py | 4 +- .../_internal/core/backends/nebius/compute.py | 4 +- .../_internal/core/backends/oci/compute.py | 4 +- .../_internal/core/backends/runpod/compute.py | 4 +- .../_internal/core/backends/slurm/compute.py | 7 +- .../core/backends/template/compute.py.jinja | 2 +- .../_internal/core/backends/vastai/compute.py | 2 +- .../_internal/core/backends/verda/compute.py | 4 +- .../_internal/core/backends/vultr/compute.py | 4 +- .../_internal/core/compatibility/gpus.py | 2 + .../_internal/core/compatibility/runs.py | 2 + src/dstack/_internal/server/routers/gpus.py | 1 + src/dstack/_internal/server/routers/runs.py | 1 + src/dstack/_internal/server/schemas/gpus.py | 4 + src/dstack/_internal/server/schemas/runs.py | 4 + .../server/services/backends/__init__.py | 8 +- src/dstack/_internal/server/services/gpus.py | 10 ++- .../_internal/server/services/offers.py | 2 + .../server/services/runs/__init__.py | 2 + .../_internal/server/services/runs/plan.py | 15 ++++ src/dstack/api/_public/runs.py | 4 + src/dstack/api/server/_gpus.py | 2 + src/dstack/api/server/_runs.py | 8 +- .../core/backends/jarvislabs/test_compute.py | 14 +++- .../core/backends/kubernetes/test_compute.py | 8 +- .../core/backends/slurm/test_compute.py | 8 +- .../core/backends/vastai/test_compute.py | 6 +- .../_internal/server/routers/test_gpus.py | 44 +++++++++++ .../_internal/server/routers/test_runs.py | 59 +++++++++++++- .../server/services/runs/test_plan.py | 3 + 43 files changed, 339 insertions(+), 58 deletions(-) diff --git a/src/dstack/_internal/cli/commands/offer.py b/src/dstack/_internal/cli/commands/offer.py index 5c67ea76f5..56c7cd2b56 100644 --- a/src/dstack/_internal/cli/commands/offer.py +++ b/src/dstack/_internal/cli/commands/offer.py @@ -64,6 +64,11 @@ def _register(self): action="store_true", help="Show full offers not adjusted by requirements", ) + self._parser.add_argument( + "--unallocated", + action="store_true", + help="Subtract allocated resources to show only unallocated resources", + ) resources_group = self._parser.add_argument_group("Resources") register_resources_args(resources_group) # TODO: register only relevant options @@ -85,6 +90,7 @@ def _list_offers(self, args: argparse.Namespace) -> None: run_spec=run_spec, max_offers=args.max_offers, full_offers=args.full_offers, + unallocated_resources=args.unallocated, ) job_plan = run_plan.job_plans[0] if args.format == "plain": @@ -110,6 +116,7 @@ def _list_gpus(self, args: argparse.Namespace, group_by: list[str]) -> None: run_spec=run_spec, group_by=[g for g in group_by if g != "gpu"], full_offers=args.full_offers, + unallocated_resources=args.unallocated, ) if args.format == "plain": print_gpu_table(gpus, run_spec, group_by, self.api.project) diff --git a/src/dstack/_internal/cli/services/configurators/run.py b/src/dstack/_internal/cli/services/configurators/run.py index 352c5932c6..41e01fb9db 100644 --- a/src/dstack/_internal/cli/services/configurators/run.py +++ b/src/dstack/_internal/cli/services/configurators/run.py @@ -145,6 +145,7 @@ def get_plan( ssh_identity_file=configurator_args.ssh_identity_file, max_offers=configurator_args.max_offers, full_offers=configurator_args.full_offers, + unallocated_resources=configurator_args.unallocated, ) return run_plan, repo @@ -394,6 +395,11 @@ def register_args(cls, parser: argparse.ArgumentParser): action="store_true", help="Show full offers not adjusted by requirements", ) + configuration_group.add_argument( + "--unallocated", + action="store_true", + help="Subtract allocated resources to show only unallocated resources", + ) cls.register_env_args(configuration_group) register_resources_args(configuration_group) register_profile_args(parser) diff --git a/src/dstack/_internal/core/backends/aws/compute.py b/src/dstack/_internal/core/backends/aws/compute.py index 7157c63bdb..61620165c0 100644 --- a/src/dstack/_internal/core/backends/aws/compute.py +++ b/src/dstack/_internal/core/backends/aws/compute.py @@ -157,7 +157,9 @@ def __init__( self._security_group_cache = ComputeTTLCache(cache=TTLCache(maxsize=100, ttl=600)) self._image_id_and_username_cache = ComputeTTLCache(cache=TTLCache(maxsize=100, ttl=600)) - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: offers = get_catalog_offers( backend=BackendType.AWS, locations=self.config.regions, @@ -186,17 +188,21 @@ def get_offers_modifiers( ) -> Iterable[OfferModifier]: return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)] - def _get_offers_cached_key(self, requirements: Requirements) -> int: + def get_offers_post_filter( + self, requirements: Requirements + ) -> Optional[Callable[[InstanceOfferWithAvailability], bool]]: + return self._get_offers_post_filter_cached(requirements) + + def _get_offers_post_filter_cached_key(self, requirements: Requirements) -> int: # Requirements is not hashable, so we use a hack to get arguments hash return hash(requirements.json()) - # For `pyright: ignore` directive, see: https://github.com/tkem/cachetools/issues/394 @cachedmethod( cache=lambda self: self._offers_post_filter_cache.cache, - key=_get_offers_cached_key, + key=_get_offers_post_filter_cached_key, lock=lambda self: self._offers_post_filter_cache.lock, ) - def get_offers_post_filter( # pyright: ignore[reportIncompatibleMethodOverride] + def _get_offers_post_filter_cached( self, requirements: Requirements ) -> Optional[Callable[[InstanceOfferWithAvailability], bool]]: if requirements.reservation: diff --git a/src/dstack/_internal/core/backends/azure/compute.py b/src/dstack/_internal/core/backends/azure/compute.py index 20d62bd5d1..e8ed2d74b5 100644 --- a/src/dstack/_internal/core/backends/azure/compute.py +++ b/src/dstack/_internal/core/backends/azure/compute.py @@ -102,7 +102,9 @@ def __init__(self, config: AzureConfig, credential: TokenCredential): credential=credential, subscription_id=config.subscription_id ) - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: offers = get_catalog_offers( backend=BackendType.AZURE, locations=self.config.regions, diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index 7545e3651d..14455f6017 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/src/dstack/_internal/core/backends/base/compute.py @@ -10,7 +10,7 @@ from enum import Enum from functools import lru_cache from pathlib import Path -from typing import Callable, ClassVar, Dict, List, Optional +from typing import Callable, ClassVar, Dict, List, Optional, Union import git import requests @@ -110,7 +110,7 @@ class Compute(ABC): @abstractmethod def get_offers( - self, requirements: Requirements, full_offers: bool + self, requirements: Requirements, full_offers: bool, unallocated_resources: bool ) -> Iterator[InstanceOfferWithAvailability]: """ Returns offers with availability matching `requirements`. @@ -126,6 +126,13 @@ def get_offers( generate synthetic offers from discovered nodes on the fly; these synthetic offers should reflect either resources that would be allocated based on `requirements` (`full_offers=False`) or full allocatable node resources (`full_offers=True`). + + if `unallocated_resources` set to `True`, the method should exclude (subtract) already + allocated resources from offer's resources. As with `full_offers`, this flag has no meaning + for most backends and is intended for backends such as Kubernetes and Slurm where many jobs + can coexist on a single node. With such backends, the method should return full allocatable + node resources if `unallocated_resources=False` and only unallocated (available) resources + if `unallocated_resources=True`. """ pass @@ -188,16 +195,29 @@ class ComputeWithAllOffersCached(ABC): It caches all offers with availability and post-filters by requirements. """ + unallocated_resources_argument_has_effect: ClassVar[bool] = False + """ + Set to `True` if `get_all_offers_with_availability()` produces different results based on + the `unallocated_resources` value. Doubles the amount of cached data. + """ + def __init__(self) -> None: super().__init__() self._offers_cache_lock = threading.Lock() self._offers_cache_execution_lock = threading.Lock() - self._offers_cache = TTLCache(maxsize=1, ttl=180) + self._offers_cache = TTLCache( + maxsize=(2 if self.unallocated_resources_argument_has_effect else 1), + ttl=180, + ) @abstractmethod - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: """ Returns all backend offers with availability. + + See `Compute.get_offers()` for the `unallocated_resources` argument description. """ pass @@ -223,12 +243,12 @@ def get_offers_post_filter( return None def get_offers( - self, requirements: Requirements, full_offers: bool + self, requirements: Requirements, full_offers: bool, unallocated_resources: bool ) -> Iterator[InstanceOfferWithAvailability]: with self._offers_cache_execution_lock: # Cache lock does not prevent concurrent execution. # We use a separate lock to avoid requesting offers in parallel, re-doing the work and hitting rate limits. - cached_offers = self._get_all_offers_with_availability_cached() + cached_offers = self._get_all_offers_with_availability_cached(unallocated_resources) offers = self.__apply_modifiers( cached_offers, self.get_offers_modifiers(requirements, full_offers) ) @@ -238,12 +258,20 @@ def get_offers( offers = (o for o in offers if post_filter(o)) return offers + def _get_all_offers_with_availability_cached_key(self, unallocated_resources: bool) -> int: + if self.unallocated_resources_argument_has_effect: + return hash(unallocated_resources) + return hash(None) + @cachedmethod( cache=lambda self: self._offers_cache, + key=_get_all_offers_with_availability_cached_key, lock=lambda self: self._offers_cache_lock, ) - def _get_all_offers_with_availability_cached(self) -> List[InstanceOfferWithAvailability]: - return self.get_all_offers_with_availability() + def _get_all_offers_with_availability_cached( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: + return self.get_all_offers_with_availability(unallocated_resources) @staticmethod def __apply_modifiers( @@ -271,6 +299,12 @@ class ComputeWithFilteredOffersCached(ABC): the `full_offers` value. Doubles the amount of cached data. """ + unallocated_resources_argument_has_effect: ClassVar[bool] = False + """ + Set to `True` if `get_offers_by_requirements()` produces different results based on + the `unallocated_resources` value. Doubles the amount of cached data. + """ + def __init__(self) -> None: super().__init__() self._offers_cache_lock = threading.Lock() @@ -281,6 +315,7 @@ def get_offers_by_requirements( self, requirements: Requirements, full_offers: bool, + unallocated_resources: bool, ) -> List[InstanceOfferWithAvailability]: """ Returns backend offers with availability matching requirements. @@ -288,20 +323,29 @@ def get_offers_by_requirements( See `Compute.get_offers()` for the `full_offers` argument description. Set the class variable `full_offers_argument_has_effect` to `True` if the `full_offers` value has an effect on the offers produced by this method. + + See `Compute.get_offers()` for the `unallocated_resources` argument description. + Set the class variable `unallocated_resources_argument_has_effect` to `True` if + the `unallocated_resources` value has an effect on the offers produced by this method. """ pass def get_offers( - self, requirements: Requirements, full_offers: bool + self, requirements: Requirements, full_offers: bool, unallocated_resources: bool ) -> Iterator[InstanceOfferWithAvailability]: - return iter(self._get_offers_cached(requirements, full_offers)) + return iter(self._get_offers_cached(requirements, full_offers, unallocated_resources)) - def _get_offers_cached_key(self, requirements: Requirements, full_offers: bool) -> int: + def _get_offers_cached_key( + self, requirements: Requirements, full_offers: bool, unallocated_resources: bool + ) -> int: + hash_items: list[Union[str, bool]] = [] # Requirements is not hashable, so we use a hack to get arguments hash - hashable_requirements = requirements.json() + hash_items.append(requirements.json()) if self.full_offers_argument_has_effect: - return hash((hashable_requirements, full_offers)) - return hash(hashable_requirements) + hash_items.append(full_offers) + if self.unallocated_resources_argument_has_effect: + hash_items.append(unallocated_resources) + return hash(tuple(hash_items)) @cachedmethod( cache=lambda self: self._offers_cache, @@ -309,9 +353,9 @@ def _get_offers_cached_key(self, requirements: Requirements, full_offers: bool) lock=lambda self: self._offers_cache_lock, ) def _get_offers_cached( - self, requirements: Requirements, full_offers: bool + self, requirements: Requirements, full_offers: bool, unallocated_resources: bool ) -> List[InstanceOfferWithAvailability]: - return self.get_offers_by_requirements(requirements, full_offers) + return self.get_offers_by_requirements(requirements, full_offers, unallocated_resources) class ComputeWithCreateInstanceSupport(ABC): diff --git a/src/dstack/_internal/core/backends/cloudrift/compute.py b/src/dstack/_internal/core/backends/cloudrift/compute.py index b968c340eb..6c2f4437bf 100644 --- a/src/dstack/_internal/core/backends/cloudrift/compute.py +++ b/src/dstack/_internal/core/backends/cloudrift/compute.py @@ -40,7 +40,9 @@ def __init__(self, config: CloudRiftConfig): self.config = config self.client = RiftClient(self.config.creds.api_key) - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: offers = get_catalog_offers( backend=BackendType.CLOUDRIFT, locations=self.config.regions or None, diff --git a/src/dstack/_internal/core/backends/crusoe/compute.py b/src/dstack/_internal/core/backends/crusoe/compute.py index 86ab767449..97b3b85a31 100644 --- a/src/dstack/_internal/core/backends/crusoe/compute.py +++ b/src/dstack/_internal/core/backends/crusoe/compute.py @@ -143,7 +143,9 @@ def __init__(self, config: CrusoeConfig): ) ) - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: offers = get_catalog_offers( backend=BackendType.CRUSOE, locations=self.config.regions or None, diff --git a/src/dstack/_internal/core/backends/digitalocean_base/compute.py b/src/dstack/_internal/core/backends/digitalocean_base/compute.py index 1128cc5306..8f672480bc 100644 --- a/src/dstack/_internal/core/backends/digitalocean_base/compute.py +++ b/src/dstack/_internal/core/backends/digitalocean_base/compute.py @@ -56,7 +56,9 @@ def __init__(self, config: BaseDigitalOceanConfig, api_url: str, type: BackendTy DigitalOceanProvider(api_key=config.creds.api_key, api_url=api_url) ) - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: offers = get_catalog_offers( backend=self.BACKEND_TYPE, locations=self.config.regions, diff --git a/src/dstack/_internal/core/backends/gcp/compute.py b/src/dstack/_internal/core/backends/gcp/compute.py index 0974ac9838..54aa640bfb 100644 --- a/src/dstack/_internal/core/backends/gcp/compute.py +++ b/src/dstack/_internal/core/backends/gcp/compute.py @@ -134,7 +134,9 @@ def __init__(self, config: GCPConfig): # Smaller TTL since we check the reservation's in_use_count, which can change often self._reservation_cache = ComputeTTLCache(cache=TTLCache(maxsize=8, ttl=20)) - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: regions = get_or_error(self.config.regions) zones_by_key: Dict[Tuple, List[str]] = {} catalog_item_filter = _make_catalog_item_filter(regions, zones_by_key) diff --git a/src/dstack/_internal/core/backends/hotaisle/compute.py b/src/dstack/_internal/core/backends/hotaisle/compute.py index b206a71d47..c23ab717ce 100644 --- a/src/dstack/_internal/core/backends/hotaisle/compute.py +++ b/src/dstack/_internal/core/backends/hotaisle/compute.py @@ -52,7 +52,9 @@ def __init__(self, config: HotAisleConfig): HotAisleProvider(api_key=config.creds.api_key, team_handle=config.team_handle) ) - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: offers = get_catalog_offers( backend=BackendType.HOTAISLE, locations=self.config.regions or None, diff --git a/src/dstack/_internal/core/backends/jarvislabs/compute.py b/src/dstack/_internal/core/backends/jarvislabs/compute.py index cc995c0239..f6468ab89f 100644 --- a/src/dstack/_internal/core/backends/jarvislabs/compute.py +++ b/src/dstack/_internal/core/backends/jarvislabs/compute.py @@ -79,7 +79,9 @@ def __init__(self, config: JarvisLabsConfig): self._catalog = gpuhunt.Catalog(balance_resources=False, auto_reload=False) self._catalog.add_provider(JarvisLabsProvider(api_key=self.config.creds.api_key)) - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: offers = get_catalog_offers( backend=BackendType.JARVISLABS, locations=self.config.regions or None, diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index d5f578c474..7b3990714b 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/src/dstack/_internal/core/backends/kubernetes/compute.py @@ -140,12 +140,16 @@ class KubernetesCompute( ComputeWithMultinodeSupport, Compute, ): + unallocated_resources_argument_has_effect = True + def __init__(self, config: KubernetesConfig): super().__init__() self.region_cluster_map = {c.region: c for c in get_clusters_from_backend_config(config)} self.skip_offer_cache = RegionalSkipOfferCache(ttl=60) - def get_all_offers_with_availability(self) -> list[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> list[InstanceOfferWithAvailability]: offers: list[InstanceOfferWithAvailability] = [] with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor: future_cluster_map: dict[ @@ -153,7 +157,7 @@ def get_all_offers_with_availability(self) -> list[InstanceOfferWithAvailability ] = {} for region, cluster in self.region_cluster_map.items(): api = client.CoreV1Api(cluster.api_client) - future = executor.submit(get_instance_offers, api, region) + future = executor.submit(get_instance_offers, api, region, unallocated_resources) future_cluster_map[future] = cluster for future in concurrent.futures.as_completed(future_cluster_map): try: diff --git a/src/dstack/_internal/core/backends/kubernetes/resources.py b/src/dstack/_internal/core/backends/kubernetes/resources.py index 4dcfedc997..cff54cbdda 100644 --- a/src/dstack/_internal/core/backends/kubernetes/resources.py +++ b/src/dstack/_internal/core/backends/kubernetes/resources.py @@ -385,8 +385,12 @@ def is_taint_tolerated(taint: V1Taint) -> bool: return taint.key in (NVIDIA_GPU_NODE_TAINT, AMD_GPU_NODE_TAINT) -def get_instance_offers(api: CoreV1Api, region: str) -> list[InstanceOfferWithAvailability]: - nodes_allocated_resources = _get_nodes_allocated_resources(api) +def get_instance_offers( + api: CoreV1Api, region: str, unallocated_resources: bool +) -> list[InstanceOfferWithAvailability]: + nodes_allocated_resources: Optional[dict[str, KubernetesResources]] = None + if unallocated_resources: + nodes_allocated_resources = _get_nodes_allocated_resources(api) offers: list[InstanceOfferWithAvailability] = [] for node in api.list_node().items: if (node_name := get_node_name(node)) is None: @@ -394,7 +398,11 @@ def get_instance_offers(api: CoreV1Api, region: str) -> list[InstanceOfferWithAv offer = _get_instance_offer_from_node( node=node, node_name=node_name, - node_allocated_resources=nodes_allocated_resources.get(node_name), + node_allocated_resources=( + None + if nodes_allocated_resources is None + else nodes_allocated_resources.get(node_name) + ), region=region, ) if offer is not None: diff --git a/src/dstack/_internal/core/backends/lambdalabs/compute.py b/src/dstack/_internal/core/backends/lambdalabs/compute.py index 1007d33bea..f749e284c6 100644 --- a/src/dstack/_internal/core/backends/lambdalabs/compute.py +++ b/src/dstack/_internal/core/backends/lambdalabs/compute.py @@ -42,7 +42,9 @@ def __init__(self, config: LambdaConfig): self.config = config self.api_client = LambdaAPIClient(config.creds.api_key) - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: offers = get_catalog_offers( backend=BackendType.LAMBDA, locations=self.config.regions or None, diff --git a/src/dstack/_internal/core/backends/nebius/compute.py b/src/dstack/_internal/core/backends/nebius/compute.py index d23d46a854..530d08e48e 100644 --- a/src/dstack/_internal/core/backends/nebius/compute.py +++ b/src/dstack/_internal/core/backends/nebius/compute.py @@ -123,7 +123,9 @@ def _get_subnet_id(self, region: str) -> str: ).metadata.id return self._subnet_id_cache[region] - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: offers = get_catalog_offers( backend=BackendType.NEBIUS, locations=list(self._region_to_project_id), diff --git a/src/dstack/_internal/core/backends/oci/compute.py b/src/dstack/_internal/core/backends/oci/compute.py index f158163cdf..ab9be4908b 100644 --- a/src/dstack/_internal/core/backends/oci/compute.py +++ b/src/dstack/_internal/core/backends/oci/compute.py @@ -69,7 +69,9 @@ def __init__(self, config: OCIConfig): def shapes_quota(self) -> resources.ShapesQuota: return resources.ShapesQuota.load(self.regions, self.config.compartment_id) - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: offers = get_catalog_offers( backend=BackendType.OCI, locations=self.config.regions, diff --git a/src/dstack/_internal/core/backends/runpod/compute.py b/src/dstack/_internal/core/backends/runpod/compute.py index d66b78fb54..b918139e1b 100644 --- a/src/dstack/_internal/core/backends/runpod/compute.py +++ b/src/dstack/_internal/core/backends/runpod/compute.py @@ -76,7 +76,9 @@ def __init__(self, config: RunpodConfig): self.config = config self.api_client = RunpodApiClient(config.creds.api_key) - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: offers = get_catalog_offers( backend=BackendType.RUNPOD, locations=self.config.regions or None, diff --git a/src/dstack/_internal/core/backends/slurm/compute.py b/src/dstack/_internal/core/backends/slurm/compute.py index 2739b2e75c..bef16aad65 100644 --- a/src/dstack/_internal/core/backends/slurm/compute.py +++ b/src/dstack/_internal/core/backends/slurm/compute.py @@ -80,6 +80,9 @@ class SlurmCompute( ComputeWithGroupProvisioningSupport, Compute, ): + # TODO: support allocated resource accounting and change to True + unallocated_resources_argument_has_effect = False + def __init__(self, config: SlurmConfig): super().__init__() self._region_to_cluster_map = { @@ -91,7 +94,9 @@ def __init__(self, config: SlurmConfig): # configuration -> the same zones) self._skip_offer_cache = RegionalSkipOfferCache(ttl=60) - def get_all_offers_with_availability(self) -> list[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> list[InstanceOfferWithAvailability]: offers: list[InstanceOfferWithAvailability] = [] with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor: future_to_cluster_map: dict[ diff --git a/src/dstack/_internal/core/backends/template/compute.py.jinja b/src/dstack/_internal/core/backends/template/compute.py.jinja index 67af917bcf..99b0336bb0 100644 --- a/src/dstack/_internal/core/backends/template/compute.py.jinja +++ b/src/dstack/_internal/core/backends/template/compute.py.jinja @@ -49,7 +49,7 @@ class {{ backend_name }}Compute( self.config = config def get_offers( - self, requirements: Requirements, full_offers: bool + self, requirements: Requirements, full_offers: bool, unallocated_resources: bool ) -> Iterator[InstanceOfferWithAvailability]: # If the provider is added to gpuhunt, you'd typically get offers # using `get_catalog_offers()` and extend them with availability info. diff --git a/src/dstack/_internal/core/backends/vastai/compute.py b/src/dstack/_internal/core/backends/vastai/compute.py index d44241f0c1..d51e40d297 100644 --- a/src/dstack/_internal/core/backends/vastai/compute.py +++ b/src/dstack/_internal/core/backends/vastai/compute.py @@ -90,7 +90,7 @@ def _make_catalog(self, options: VastAIProfileOptions) -> gpuhunt.Catalog: return catalog def get_offers_by_requirements( - self, requirements: Requirements, full_offers: bool + self, requirements: Requirements, full_offers: bool, unallocated_resources: bool ) -> List[InstanceOfferWithAvailability]: vastai_options = ( get_backend_profile_options(requirements.backend_options, VastAIProfileOptions) diff --git a/src/dstack/_internal/core/backends/verda/compute.py b/src/dstack/_internal/core/backends/verda/compute.py index 820b914379..f6cb4c19ac 100644 --- a/src/dstack/_internal/core/backends/verda/compute.py +++ b/src/dstack/_internal/core/backends/verda/compute.py @@ -64,7 +64,9 @@ def __init__(self, config: VerdaConfig, backend_type: BackendType): ) self.backend_type = backend_type - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: offers = get_catalog_offers( backend=self.backend_type, locations=self.config.regions, diff --git a/src/dstack/_internal/core/backends/vultr/compute.py b/src/dstack/_internal/core/backends/vultr/compute.py index 2ae3210621..8afcb79369 100644 --- a/src/dstack/_internal/core/backends/vultr/compute.py +++ b/src/dstack/_internal/core/backends/vultr/compute.py @@ -47,7 +47,9 @@ def __init__(self, config: VultrConfig): self.config = config self.api_client = VultrApiClient(config.creds.api_key) - def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]: + def get_all_offers_with_availability( + self, unallocated_resources: bool + ) -> List[InstanceOfferWithAvailability]: offers = get_catalog_offers( backend=BackendType.VULTR, requirements=None, diff --git a/src/dstack/_internal/core/compatibility/gpus.py b/src/dstack/_internal/core/compatibility/gpus.py index c7ed783c6a..ba33a46687 100644 --- a/src/dstack/_internal/core/compatibility/gpus.py +++ b/src/dstack/_internal/core/compatibility/gpus.py @@ -9,6 +9,8 @@ def get_list_gpus_excludes(request: ListGpusRequest) -> Optional[IncludeExcludeD list_gpus_excludes: IncludeExcludeDictType = {} if not request.full_offers: list_gpus_excludes["full_offers"] = True + if not request.unallocated_resources: + list_gpus_excludes["unallocated_resources"] = True run_spec_excludes = get_run_spec_excludes(request.run_spec) if run_spec_excludes is not None: list_gpus_excludes["run_spec"] = run_spec_excludes diff --git a/src/dstack/_internal/core/compatibility/runs.py b/src/dstack/_internal/core/compatibility/runs.py index e71005095e..847e7be303 100644 --- a/src/dstack/_internal/core/compatibility/runs.py +++ b/src/dstack/_internal/core/compatibility/runs.py @@ -72,6 +72,8 @@ def get_get_plan_excludes(request: GetRunPlanRequest) -> Optional[IncludeExclude get_plan_excludes: IncludeExcludeDictType = {} if not request.full_offers: get_plan_excludes["full_offers"] = True + if not request.unallocated_resources: + get_plan_excludes["unallocated_resources"] = True run_spec_excludes = get_run_spec_excludes(request.run_spec) if run_spec_excludes is not None: get_plan_excludes["run_spec"] = run_spec_excludes diff --git a/src/dstack/_internal/server/routers/gpus.py b/src/dstack/_internal/server/routers/gpus.py index 9cfcc225cc..0866beb097 100644 --- a/src/dstack/_internal/server/routers/gpus.py +++ b/src/dstack/_internal/server/routers/gpus.py @@ -38,6 +38,7 @@ async def list_gpus( run_spec=body.run_spec, group_by=body.group_by, full_offers=body.full_offers, + unallocated_resources=body.unallocated_resources, ) patch_list_gpus_response(resp, client_version) return resp diff --git a/src/dstack/_internal/server/routers/runs.py b/src/dstack/_internal/server/routers/runs.py index cb11ad120c..994b91fa29 100644 --- a/src/dstack/_internal/server/routers/runs.py +++ b/src/dstack/_internal/server/routers/runs.py @@ -140,6 +140,7 @@ async def get_plan( run_spec=body.run_spec, max_offers=body.max_offers, full_offers=body.full_offers, + unallocated_resources=body.unallocated_resources, legacy_repo_dir=legacy_repo_dir, ) patch_run_plan(run_plan, client_version) diff --git a/src/dstack/_internal/server/schemas/gpus.py b/src/dstack/_internal/server/schemas/gpus.py index af790e3e99..28c0cf18a8 100644 --- a/src/dstack/_internal/server/schemas/gpus.py +++ b/src/dstack/_internal/server/schemas/gpus.py @@ -19,6 +19,10 @@ class ListGpusRequest(CoreModel): full_offers: Annotated[ bool, Field(description="Don't adjust backend offers by requirements") ] = False + unallocated_resources: Annotated[ + bool, + Field(description="Subtract allocated resources to return only unallocated resources"), + ] = False class ListGpusResponse(CoreModel): diff --git a/src/dstack/_internal/server/schemas/runs.py b/src/dstack/_internal/server/schemas/runs.py index 84861a6a93..2163320fd8 100644 --- a/src/dstack/_internal/server/schemas/runs.py +++ b/src/dstack/_internal/server/schemas/runs.py @@ -48,6 +48,10 @@ class GetRunPlanRequest(CoreModel): full_offers: Annotated[ bool, Field(description="Return full offers not adjusted by requirements") ] = False + unallocated_resources: Annotated[ + bool, + Field(description="Subtract allocated resources to return only unallocated resources"), + ] = False class SubmitRunRequest(CoreModel): diff --git a/src/dstack/_internal/server/services/backends/__init__.py b/src/dstack/_internal/server/services/backends/__init__.py index 8ea35c215c..98dcb6b2e8 100644 --- a/src/dstack/_internal/server/services/backends/__init__.py +++ b/src/dstack/_internal/server/services/backends/__init__.py @@ -476,6 +476,7 @@ async def get_backend_offers( backends: List[Backend], requirements: Requirements, full_offers: bool, + unallocated_resources: bool, exclude_not_available: bool = False, ) -> Iterable[Tuple[Backend, InstanceOfferWithAvailability]]: """ @@ -492,7 +493,8 @@ def get_filtered_offers_with_backends( logger.debug("Requesting instance offers from backends: %s", [b.TYPE.value for b in backends]) tasks = [ - run_async(get_offers_tracked, backend, requirements, full_offers) for backend in backends + run_async(get_offers_tracked, backend, requirements, full_offers, unallocated_resources) + for backend in backends ] offers_by_backend: list[Iterable[tuple[Backend, InstanceOfferWithAvailability]]] = [] for backend, result in zip(backends, await asyncio.gather(*tasks, return_exceptions=True)): @@ -524,10 +526,10 @@ def check_backend_type_available(backend_type: BackendType): def get_offers_tracked( - backend: Backend, requirements: Requirements, full_offers: bool + backend: Backend, requirements: Requirements, full_offers: bool, unallocated_resources: bool ) -> Iterator[InstanceOfferWithAvailability]: start = time.time() - res = backend.compute().get_offers(requirements, full_offers) + res = backend.compute().get_offers(requirements, full_offers, unallocated_resources) duration = time.time() - start logger.debug("Got offers from %s in %.6fs", backend.TYPE.value, duration) return res diff --git a/src/dstack/_internal/server/services/gpus.py b/src/dstack/_internal/server/services/gpus.py index 2ba42ae209..f139423f69 100644 --- a/src/dstack/_internal/server/services/gpus.py +++ b/src/dstack/_internal/server/services/gpus.py @@ -27,10 +27,15 @@ async def list_gpus_grouped( run_spec: RunSpec, group_by: Optional[List[Literal["backend", "region", "count"]]] = None, full_offers: bool = False, + unallocated_resources: bool = False, ) -> ListGpusResponse: """Retrieves available GPU specifications based on a run spec, with optional grouping.""" offers = await _get_gpu_offers( - session=session, project=project, run_spec=run_spec, full_offers=full_offers + session=session, + project=project, + run_spec=run_spec, + full_offers=full_offers, + unallocated_resources=unallocated_resources, ) backend_gpus = _process_offers_into_backend_gpus(offers) group_by_set = set(group_by) if group_by else set() @@ -62,6 +67,7 @@ async def _get_gpu_offers( project: ProjectModel, run_spec: RunSpec, full_offers: bool, + unallocated_resources: bool, ) -> list[InstanceOfferWithAvailability]: """Fetches all available instance offers that match the run spec's GPU requirements.""" # NOTE: Basically, this is a simplified version of get_job_plans(); keep them in sync @@ -90,6 +96,7 @@ async def _get_gpu_offers( job=job, skip_backend_offers=skip_backend_offers, full_offers=full_offers, + unallocated_resources=unallocated_resources, ) else: instance_offers, backend_offers = await get_non_fleet_offers( @@ -99,6 +106,7 @@ async def _get_gpu_offers( job=job, skip_backend_offers=skip_backend_offers, full_offers=full_offers, + unallocated_resources=unallocated_resources, ) return [offer for _, offer in instance_offers] + [offer for _, offer in backend_offers] diff --git a/src/dstack/_internal/server/services/offers.py b/src/dstack/_internal/server/services/offers.py index 9d0d2d7005..a35f1bef93 100644 --- a/src/dstack/_internal/server/services/offers.py +++ b/src/dstack/_internal/server/services/offers.py @@ -41,6 +41,7 @@ async def get_offers_by_requirements( blocks: Union[int, Literal["auto"]] = 1, max_offers: Optional[int] = None, full_offers: bool = False, + unallocated_resources: bool = False, ) -> List[Tuple[Backend, InstanceOfferWithAvailability]]: backends: List[Backend] = await backends_services.get_project_backends(project=project) @@ -93,6 +94,7 @@ async def get_offers_by_requirements( backends=backends, requirements=requirements, full_offers=full_offers, + unallocated_resources=unallocated_resources, exclude_not_available=exclude_not_available, ) diff --git a/src/dstack/_internal/server/services/runs/__init__.py b/src/dstack/_internal/server/services/runs/__init__.py index 5438ef8461..b7668f7d20 100644 --- a/src/dstack/_internal/server/services/runs/__init__.py +++ b/src/dstack/_internal/server/services/runs/__init__.py @@ -533,6 +533,7 @@ async def get_plan( run_spec: RunSpec, max_offers: Optional[int], full_offers: bool, + unallocated_resources: bool, legacy_repo_dir: bool = False, ) -> RunPlan: # Spec must be copied by parsing to calculate merged_profile @@ -572,6 +573,7 @@ async def get_plan( run_spec=effective_run_spec, max_offers=max_offers, full_offers=full_offers, + unallocated_resources=unallocated_resources, ) run_plan = RunPlan( project_name=project.name, diff --git a/src/dstack/_internal/server/services/runs/plan.py b/src/dstack/_internal/server/services/runs/plan.py index e876d05bde..97d3185088 100644 --- a/src/dstack/_internal/server/services/runs/plan.py +++ b/src/dstack/_internal/server/services/runs/plan.py @@ -88,6 +88,7 @@ async def get_job_plans( run_spec: RunSpec, max_offers: Optional[int], full_offers: bool, + unallocated_resources: bool, ) -> list[JobPlan]: """ Returns job plans for the given run spec. @@ -158,6 +159,7 @@ async def get_job_plans( exclude_not_available=False, skip_backend_offers=skip_backend_offers, full_offers=full_offers, + unallocated_resources=unallocated_resources, ) elif run_spec.merged_profile.instances is not None: # Regular job planning or offer collection @@ -179,6 +181,7 @@ async def get_job_plans( volumes=volumes, skip_backend_offers=skip_backend_offers, full_offers=full_offers, + unallocated_resources=unallocated_resources, ) else: # Offer collection @@ -190,6 +193,7 @@ async def get_job_plans( volumes=volumes, skip_backend_offers=skip_backend_offers, full_offers=full_offers, + unallocated_resources=unallocated_resources, ) for job in jobs: @@ -333,6 +337,7 @@ async def find_optimal_fleet_with_offers( skip_backend_offers: bool = False, skip_backend_offers_on_pool_capacity: bool = False, full_offers: bool = False, + unallocated_resources: bool = False, ) -> tuple[ Optional[FleetModel], list[tuple[InstanceModel, InstanceOfferWithAvailability]], @@ -437,6 +442,7 @@ async def find_optimal_fleet_with_offers( volumes=volumes, max_offers=_PER_FLEET_MAX_OFFERS, full_offers=full_offers, + unallocated_resources=unallocated_resources, ) available_backend_offers = _exclude_non_available_backend_offers(backend_offers) candidates_with_backend_offers.append( @@ -470,6 +476,7 @@ async def find_optimal_fleet_with_offers( volumes=volumes, max_offers=None, full_offers=full_offers, + unallocated_resources=unallocated_resources, ) if exclude_not_available: backend_offers = _exclude_non_available_backend_offers(backend_offers) @@ -727,6 +734,7 @@ async def _get_backend_offers_in_fleet( fleet_spec: Optional[FleetSpec] = None, max_offers: Optional[int] = None, full_offers: bool = False, + unallocated_resources: bool = False, ) -> list[tuple[Backend, InstanceOfferWithAvailability]]: if fleet_spec is None: fleet_spec = get_fleet_spec(fleet_model) @@ -759,6 +767,7 @@ async def _get_backend_offers_in_fleet( instance_mounts=check_run_spec_requires_instance_mounts(run_spec), max_offers=max_offers, full_offers=full_offers, + unallocated_resources=unallocated_resources, ) return backend_offers @@ -807,6 +816,7 @@ async def get_non_fleet_offers( volumes: Optional[list[list[Volume]]] = None, skip_backend_offers: bool = False, full_offers: bool = False, + unallocated_resources: bool = False, ) -> tuple[ list[tuple[InstanceModel, InstanceOfferWithAvailability]], list[tuple[Backend, InstanceOfferWithAvailability]], @@ -836,6 +846,7 @@ async def get_non_fleet_offers( privileged=job.job_spec.privileged, instance_mounts=check_run_spec_requires_instance_mounts(run_spec), full_offers=full_offers, + unallocated_resources=unallocated_resources, ) return instance_offers, backend_offers @@ -848,6 +859,7 @@ async def get_backend_offers_in_run_candidate_fleets( volumes: Optional[list[list[Volume]]], max_offers_per_fleet: Optional[int] = None, full_offers: bool = False, + unallocated_resources: bool = False, ) -> list[tuple[Backend, InstanceOfferWithAvailability]]: """ Returns backend offers across the run's selected candidate fleets. @@ -877,6 +889,7 @@ async def get_backend_offers_in_run_candidate_fleets( volumes=volumes, max_offers=max_offers_per_fleet, full_offers=full_offers, + unallocated_resources=unallocated_resources, ): offer_identity = _get_backend_offer_identity(offer) if offer_identity not in seen_offer_identities: @@ -894,6 +907,7 @@ async def get_offers_in_run_candidate_fleets( volumes: Optional[list[list[Volume]]] = None, skip_backend_offers: bool = False, full_offers: bool = False, + unallocated_resources: bool = False, ) -> tuple[ list[tuple[InstanceModel, InstanceOfferWithAvailability]], list[tuple[Backend, InstanceOfferWithAvailability]], @@ -942,6 +956,7 @@ async def get_offers_in_run_candidate_fleets( volumes=volumes, max_offers_per_fleet=None, full_offers=full_offers, + unallocated_resources=unallocated_resources, ) return instance_offers, backend_offers diff --git a/src/dstack/api/_public/runs.py b/src/dstack/api/_public/runs.py index 752eca07a9..09c24f028e 100644 --- a/src/dstack/api/_public/runs.py +++ b/src/dstack/api/_public/runs.py @@ -477,6 +477,7 @@ def get_run_plan( ssh_identity_file: Optional[PathLike] = None, max_offers: Optional[int] = None, full_offers: bool = False, + unallocated_resources: bool = False, ) -> RunPlan: """ Get a run plan. @@ -495,6 +496,8 @@ def get_run_plan( If ssh_identity_file is not specified, the user key is used. max_offers: Maximum number of offers returned in the run plan. full_offers: Return full offers not adjusted by requirements. + unallocated_resources: Subtract allocated resources to return only unallocated + resources. Returns: Run plan. @@ -550,6 +553,7 @@ def get_run_plan( run_spec=run_spec, max_offers=max_offers, full_offers=full_offers, + unallocated_resources=unallocated_resources, ) return run_plan diff --git a/src/dstack/api/server/_gpus.py b/src/dstack/api/server/_gpus.py index d471a8e7bd..068538c94e 100644 --- a/src/dstack/api/server/_gpus.py +++ b/src/dstack/api/server/_gpus.py @@ -16,11 +16,13 @@ def list_gpus( run_spec: RunSpec, group_by: Optional[List[str]] = None, full_offers: bool = False, + unallocated_resources: bool = False, ) -> List[GpuGroup]: body = ListGpusRequest( run_spec=run_spec, group_by=cast(Optional[List[Literal["backend", "region", "count"]]], group_by), full_offers=full_offers, + unallocated_resources=unallocated_resources, ) resp = self._request( f"/api/project/{project_name}/gpus/list", diff --git a/src/dstack/api/server/_runs.py b/src/dstack/api/server/_runs.py index 6018a98f38..e0e0bdd48a 100644 --- a/src/dstack/api/server/_runs.py +++ b/src/dstack/api/server/_runs.py @@ -77,8 +77,14 @@ def get_plan( run_spec: RunSpec, max_offers: Optional[int] = None, full_offers: bool = False, + unallocated_resources: bool = False, ) -> RunPlan: - body = GetRunPlanRequest(run_spec=run_spec, max_offers=max_offers, full_offers=full_offers) + body = GetRunPlanRequest( + run_spec=run_spec, + max_offers=max_offers, + full_offers=full_offers, + unallocated_resources=unallocated_resources, + ) body = copy.deepcopy(body) patch_run_spec(body.run_spec) resp = self._request( diff --git a/src/tests/_internal/core/backends/jarvislabs/test_compute.py b/src/tests/_internal/core/backends/jarvislabs/test_compute.py index d938dc113d..e9b055cd21 100644 --- a/src/tests/_internal/core/backends/jarvislabs/test_compute.py +++ b/src/tests/_internal/core/backends/jarvislabs/test_compute.py @@ -143,7 +143,7 @@ def test_get_all_offers_uses_configurable_disk_size(): "dstack._internal.core.backends.jarvislabs.compute.get_catalog_offers", return_value=[_cpu_catalog_offer()], ) as m: - offers = compute.get_all_offers_with_availability() + offers = compute.get_all_offers_with_availability(False) assert len(offers) == 1 assert offers[0].availability == InstanceAvailability.AVAILABLE @@ -162,10 +162,18 @@ def test_get_offers_reuses_all_offers_cache_and_modifies_disk_size(): ) offers_250gb = list( - compute.get_offers(Requirements(resources=ResourcesSpec(disk="250GB")), False) + compute.get_offers( + Requirements(resources=ResourcesSpec(disk="250GB")), + full_offers=False, + unallocated_resources=False, + ) ) offers_300gb = list( - compute.get_offers(Requirements(resources=ResourcesSpec(disk="300GB")), False) + compute.get_offers( + Requirements(resources=ResourcesSpec(disk="300GB")), + full_offers=False, + unallocated_resources=False, + ) ) assert len(offers_250gb) == 1 diff --git a/src/tests/_internal/core/backends/kubernetes/test_compute.py b/src/tests/_internal/core/backends/kubernetes/test_compute.py index c8c6748e85..39ef170850 100644 --- a/src/tests/_internal/core/backends/kubernetes/test_compute.py +++ b/src/tests/_internal/core/backends/kubernetes/test_compute.py @@ -63,7 +63,9 @@ def test_get_offers_with_full_offers_keeps_full_node_resources(): resources=ResourcesSpec(cpu="2..", memory="8GB..", gpu="1..", disk="100GB..") ) - full_offers = list(compute.get_offers(requirements, full_offers=True)) + full_offers = list( + compute.get_offers(requirements, full_offers=True, unallocated_resources=False) + ) assert len(full_offers) == 1 full_resources = full_offers[0].instance.resources @@ -80,7 +82,9 @@ def test_get_offers_without_full_offers_adjusts_to_requested_slice(): resources=ResourcesSpec(cpu="2..", memory="8GB..", gpu="1..", disk="100GB..") ) - adjusted_offers = list(compute.get_offers(requirements, full_offers=False)) + adjusted_offers = list( + compute.get_offers(requirements, full_offers=False, unallocated_resources=False) + ) assert len(adjusted_offers) == 1 adjusted_resources = adjusted_offers[0].instance.resources diff --git a/src/tests/_internal/core/backends/slurm/test_compute.py b/src/tests/_internal/core/backends/slurm/test_compute.py index d7b39c9273..58beddabcf 100644 --- a/src/tests/_internal/core/backends/slurm/test_compute.py +++ b/src/tests/_internal/core/backends/slurm/test_compute.py @@ -59,7 +59,9 @@ def test_get_offers_with_full_offers_keeps_full_node_resources(): resources=ResourcesSpec(cpu="2..", memory="8GB..", gpu="1..", disk="100GB..") ) - full_offers = list(compute.get_offers(requirements, full_offers=True)) + full_offers = list( + compute.get_offers(requirements, full_offers=True, unallocated_resources=False) + ) assert len(full_offers) == 1 full_resources = full_offers[0].instance.resources @@ -82,7 +84,9 @@ def test_get_offers_without_full_offers_adjusts_to_requested_slice(): "dstack._internal.core.backends.slurm.compute._get_cluster_partitions", return_value={"partition1"}, ): - adjusted_offers = list(compute.get_offers(requirements, full_offers=False)) + adjusted_offers = list( + compute.get_offers(requirements, full_offers=False, unallocated_resources=False) + ) assert len(adjusted_offers) == 1 adjusted_resources = adjusted_offers[0].instance.resources diff --git a/src/tests/_internal/core/backends/vastai/test_compute.py b/src/tests/_internal/core/backends/vastai/test_compute.py index 0b0da5401b..ee5b823d87 100644 --- a/src/tests/_internal/core/backends/vastai/test_compute.py +++ b/src/tests/_internal/core/backends/vastai/test_compute.py @@ -83,7 +83,7 @@ def test_vastai_compute_enables_community_cloud_by_default(): ): catalog_instance = catalog_cls.return_value compute = VastAICompute(_config()) - list(compute.get_offers(_requirements(), False)) + list(compute.get_offers(_requirements(), full_offers=False, unallocated_resources=False)) vast_provider_cls.assert_called_once() assert vast_provider_cls.call_args.kwargs["community_cloud"] is True catalog_instance.add_provider.assert_called_once() @@ -97,7 +97,7 @@ def test_vastai_compute_can_enable_community_cloud(): ): catalog_instance = catalog_cls.return_value compute = VastAICompute(_config(community_cloud=True)) - list(compute.get_offers(_requirements(), False)) + list(compute.get_offers(_requirements(), full_offers=False, unallocated_resources=False)) vast_provider_cls.assert_called_once() assert vast_provider_cls.call_args.kwargs["community_cloud"] is True catalog_instance.add_provider.assert_called_once() @@ -111,7 +111,7 @@ def test_vastai_compute_can_disable_community_cloud(): ): catalog_instance = catalog_cls.return_value compute = VastAICompute(_config(community_cloud=False)) - list(compute.get_offers(_requirements(), False)) + list(compute.get_offers(_requirements(), full_offers=False, unallocated_resources=False)) vast_provider_cls.assert_called_once() assert vast_provider_cls.call_args.kwargs["community_cloud"] is False catalog_instance.add_provider.assert_called_once() diff --git a/src/tests/_internal/server/routers/test_gpus.py b/src/tests/_internal/server/routers/test_gpus.py index b54b48b0be..18cfee03db 100644 --- a/src/tests/_internal/server/routers/test_gpus.py +++ b/src/tests/_internal/server/routers/test_gpus.py @@ -139,6 +139,7 @@ async def call_gpus_api( group_by: Optional[List[str]] = None, client_version: Optional[str] = None, full_offers: Optional[bool] = None, + unallocated_resources: Optional[bool] = None, ): """Helper to call the GPUs API with standard parameters.""" json_data = {"run_spec": run_spec.dict()} @@ -146,6 +147,8 @@ async def call_gpus_api( json_data["group_by"] = group_by if full_offers is not None: json_data["full_offers"] = full_offers + if unallocated_resources is not None: + json_data["unallocated_resources"] = unallocated_resources headers = get_auth_headers(user_token) if client_version is not None: headers["X-API-Version"] = client_version @@ -230,6 +233,47 @@ async def test_forwards_full_offers_to_compute_get_offers( for call_args in get_offers_mock.call_args_list ) + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + @pytest.mark.parametrize( + ("body_unallocated_resources", "expected_unallocated_resources"), + [ + pytest.param(None, False, id="omitted-defaults-to-false"), + pytest.param(True, True, id="true"), + pytest.param(False, False, id="false"), + ], + ) + async def test_forwards_unallocated_resources_to_compute_get_offers( + self, + test_db, + session: AsyncSession, + client: AsyncClient, + body_unallocated_resources: Optional[bool], + expected_unallocated_resources: bool, + ): + user, project, repo, run_spec = await gpu_test_setup(session) + offer = create_gpu_offer(BackendType.AWS, "T4", 16384, 0.50) + mocked_backends = create_mock_backends_with_offers({BackendType.AWS: [offer]}) + + with patch("dstack._internal.server.services.backends.get_project_backends") as m: + m.return_value = mocked_backends + response = await call_gpus_api( + client, + project.name, + user.token, + run_spec, + unallocated_resources=body_unallocated_resources, + ) + + assert response.status_code == 200, response.json() + get_offers_mock = mocked_backends[0].compute.return_value.get_offers + get_offers_mock.assert_called() + # get_offers is called as get_offers(requirements, full_offers, unallocated_resources) + assert all( + call_args.args[2] is expected_unallocated_resources + for call_args in get_offers_mock.call_args_list + ) + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) async def test_filters_gpus_by_multiple_specified_fleets( diff --git a/src/tests/_internal/server/routers/test_runs.py b/src/tests/_internal/server/routers/test_runs.py index e2c1750aea..0eeb4a11ac 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -1627,6 +1627,61 @@ async def test_forwards_full_offers_to_compute_get_offers( for call_args in get_offers_mock.call_args_list ) + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + @pytest.mark.parametrize( + ("body_unallocated_resources", "expected_unallocated_resources"), + [ + pytest.param(None, False, id="omitted-defaults-to-false"), + pytest.param(True, True, id="true"), + pytest.param(False, False, id="false"), + ], + ) + async def test_forwards_unallocated_resources_to_compute_get_offers( + self, + test_db, + session: AsyncSession, + client: AsyncClient, + body_unallocated_resources: Optional[bool], + expected_unallocated_resources: bool, + ): + user = await create_user(session=session, global_role=GlobalRole.USER) + project = await create_project(session=session, owner=user) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + fleet_spec = get_fleet_spec() + fleet_spec.configuration.nodes = FleetNodesSpec(min=0, target=0, max=None) + await create_fleet(session=session, project=project, spec=fleet_spec) + repo = await create_repo(session=session, project_id=project.id) + run_spec = get_run_spec( + repo_id=repo.name, + configuration=DevEnvironmentConfiguration(ide="vscode"), + ) + body: dict = {"run_spec": json.loads(run_spec.json())} + if body_unallocated_resources is not None: + body["unallocated_resources"] = body_unallocated_resources + + with patch("dstack._internal.server.services.backends.get_project_backends") as m: + backend_mock = Mock() + backend_mock.TYPE = BackendType.AWS + get_offers_mock = backend_mock.compute.return_value.get_offers + get_offers_mock.return_value = [] + m.return_value = [backend_mock] + response = await client.post( + f"/api/project/{project.name}/runs/get_plan", + headers=get_auth_headers(user.token), + json=body, + ) + + assert response.status_code == 200, response.json() + get_offers_mock.assert_called() + # get_offers is called as get_offers(requirements, full_offers, unallocated_resources) + assert all( + call_args.args[2] is expected_unallocated_resources + for call_args in get_offers_mock.call_args_list + ) + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) async def test_returns_run_plan_privileged_true( @@ -1800,7 +1855,9 @@ async def test_service_with_two_replica_groups_returns_two_job_plans( ) body = {"run_spec": json.loads(run_spec.json())} - def offers_by_requirements(requirements: Requirements, full_offers: bool): + def offers_by_requirements( + requirements: Requirements, full_offers: bool, unallocated_resources: bool + ): if ( requirements.resources.gpu is not None and requirements.resources.gpu.count.min is not None diff --git a/src/tests/_internal/server/services/runs/test_plan.py b/src/tests/_internal/server/services/runs/test_plan.py index 21375b2509..ac84913852 100644 --- a/src/tests/_internal/server/services/runs/test_plan.py +++ b/src/tests/_internal/server/services/runs/test_plan.py @@ -135,6 +135,7 @@ async def test_skips_backend_offers_by_creation_policy( run_spec=run_spec, max_offers=None, full_offers=False, + unallocated_resources=False, ) find_optimal_fleet_with_offers_mock.assert_awaited_once() @@ -174,6 +175,7 @@ async def test_excludes_backend_offers_when_instances_specified( run_spec=run_spec, max_offers=None, full_offers=False, + unallocated_resources=False, ) get_targeted_instance_offers_mock.assert_awaited_once() @@ -214,6 +216,7 @@ async def test_empty_dev_environment_with_fleet_does_not_use_targeted_instances( run_spec=run_spec, max_offers=None, full_offers=False, + unallocated_resources=False, ) select_instances_mock.assert_not_awaited()