From 7bab83fd80f4c36105785eb5fc89079f50b7a1e5 Mon Sep 17 00:00:00 2001 From: Dmitry Meyer Date: Wed, 22 Jul 2026 15:19:58 +0000 Subject: [PATCH] Add `full_offers` option to `Compute.get_offers()` When set to `True`, instructs the Compute not to adjust offers by requirements. Ignored by all backends except for Kubernetes and Slurm, which return full allocatable node resources if `True` or a purely synthetic offer reflecting resources that would be actually requested if `False` (the old/current default behavior). Exposed as `full_offers` request body argument in HTTP API and `--full-offers` flag in `dstack apply`/`dstack offer` commands. --- frontend/src/pages/Offers/List/index.tsx | 1 + frontend/src/types/gpu.d.ts | 1 + src/dstack/_internal/cli/commands/offer.py | 7 ++ .../cli/services/configurators/run.py | 7 ++ .../_internal/core/backends/aws/compute.py | 4 +- .../_internal/core/backends/azure/compute.py | 4 +- .../_internal/core/backends/base/compute.py | 60 +++++++++--- .../_internal/core/backends/crusoe/compute.py | 4 +- .../_internal/core/backends/gcp/compute.py | 4 +- .../core/backends/jarvislabs/compute.py | 4 +- .../core/backends/kubernetes/compute.py | 6 +- .../_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 | 6 +- .../core/backends/template/compute.py.jinja | 2 +- .../_internal/core/backends/vastai/compute.py | 2 +- .../_internal/core/backends/verda/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 | 5 +- src/dstack/_internal/server/schemas/runs.py | 3 + .../server/services/backends/__init__.py | 9 +- src/dstack/_internal/server/services/gpus.py | 8 +- .../_internal/server/services/offers.py | 2 + .../server/services/runs/__init__.py | 2 + .../_internal/server/services/runs/plan.py | 19 ++++ src/dstack/api/_public/runs.py | 11 ++- src/dstack/api/server/_gpus.py | 2 + src/dstack/api/server/_runs.py | 8 +- .../core/backends/jarvislabs/test_compute.py | 8 +- .../core/backends/kubernetes/test_compute.py | 90 ++++++++++++++++++ .../core/backends/slurm/test_compute.py | 93 +++++++++++++++++++ .../core/backends/vastai/test_compute.py | 6 +- .../_internal/server/routers/test_gpus.py | 40 ++++++++ .../_internal/server/routers/test_runs.py | 57 +++++++++++- .../server/services/runs/test_plan.py | 3 + 39 files changed, 461 insertions(+), 39 deletions(-) create mode 100644 src/tests/_internal/core/backends/kubernetes/test_compute.py create mode 100644 src/tests/_internal/core/backends/slurm/test_compute.py diff --git a/frontend/src/pages/Offers/List/index.tsx b/frontend/src/pages/Offers/List/index.tsx index d9abef6b02..e248510288 100644 --- a/frontend/src/pages/Offers/List/index.tsx +++ b/frontend/src/pages/Offers/List/index.tsx @@ -67,6 +67,7 @@ const getRequestParams = ({ profile: { name: 'default', default: false }, ssh_key_pub: '(dummy)', }, + full_offers: true, }; }; diff --git a/frontend/src/types/gpu.d.ts b/frontend/src/types/gpu.d.ts index 494f5712e0..507386aea6 100644 --- a/frontend/src/types/gpu.d.ts +++ b/frontend/src/types/gpu.d.ts @@ -88,6 +88,7 @@ declare type TGpusListQueryParams = { profile?: { name: string; default?: boolean }; ssh_key_pub: string; }; + full_offers: boolean; }; declare type TGpusListQueryResponse = { diff --git a/src/dstack/_internal/cli/commands/offer.py b/src/dstack/_internal/cli/commands/offer.py index 0781e80b96..5c67ea76f5 100644 --- a/src/dstack/_internal/cli/commands/offer.py +++ b/src/dstack/_internal/cli/commands/offer.py @@ -59,6 +59,11 @@ def _register(self): type=int, default=50, ) + self._parser.add_argument( + "--full-offers", + action="store_true", + help="Show full offers not adjusted by requirements", + ) resources_group = self._parser.add_argument_group("Resources") register_resources_args(resources_group) # TODO: register only relevant options @@ -79,6 +84,7 @@ def _list_offers(self, args: argparse.Namespace) -> None: project_name=self.api.project, run_spec=run_spec, max_offers=args.max_offers, + full_offers=args.full_offers, ) job_plan = run_plan.job_plans[0] if args.format == "plain": @@ -103,6 +109,7 @@ def _list_gpus(self, args: argparse.Namespace, group_by: list[str]) -> None: project_name=self.api.project, run_spec=run_spec, group_by=[g for g in group_by if g != "gpu"], + full_offers=args.full_offers, ) 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 f71193bd5a..352c5932c6 100644 --- a/src/dstack/_internal/cli/services/configurators/run.py +++ b/src/dstack/_internal/cli/services/configurators/run.py @@ -143,6 +143,8 @@ def get_plan( configuration_path=configuration_path, profile=profile, ssh_identity_file=configurator_args.ssh_identity_file, + max_offers=configurator_args.max_offers, + full_offers=configurator_args.full_offers, ) return run_plan, repo @@ -387,6 +389,11 @@ def register_args(cls, parser: argparse.ArgumentParser): type=int, default=3, ) + configuration_group.add_argument( + "--full-offers", + action="store_true", + help="Show full offers not adjusted by requirements", + ) 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 f2420ecf16..106b353322 100644 --- a/src/dstack/_internal/core/backends/aws/compute.py +++ b/src/dstack/_internal/core/backends/aws/compute.py @@ -181,7 +181,9 @@ def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability ) return availability_offers - def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]: + def get_offers_modifiers( + self, requirements: Requirements, full_offers: bool + ) -> Iterable[OfferModifier]: return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)] def _get_offers_cached_key(self, requirements: Requirements) -> int: diff --git a/src/dstack/_internal/core/backends/azure/compute.py b/src/dstack/_internal/core/backends/azure/compute.py index d2843c7f22..20d62bd5d1 100644 --- a/src/dstack/_internal/core/backends/azure/compute.py +++ b/src/dstack/_internal/core/backends/azure/compute.py @@ -115,7 +115,9 @@ def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability ) return offers_with_availability - def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]: + def get_offers_modifiers( + self, requirements: Requirements, full_offers: bool + ) -> Iterable[OfferModifier]: return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)] def create_instance( diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index 7cb7c84b35..9ead731bb1 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, Dict, List, Optional +from typing import Callable, ClassVar, Dict, List, Optional import git import requests @@ -109,12 +109,23 @@ class Compute(ABC): """ @abstractmethod - def get_offers(self, requirements: Requirements) -> Iterator[InstanceOfferWithAvailability]: + def get_offers( + self, requirements: Requirements, full_offers: bool + ) -> Iterator[InstanceOfferWithAvailability]: """ Returns offers with availability matching `requirements`. If the provider is added to gpuhunt, typically gets offers using `base.offers.get_catalog_offers()` and extends them with availability info. It is called from async code in executor. It can block on call but not between yields. + + if `full_offers` set to `True`, the method should not adjust offer's resources according to + `requirements`. For most backends, this flag has no meaning, as they work with predefined + provider offers (even configurable disk size reflects the actual disk created once the + instance is provisioned), but some backends such as Kubernetes and Slurm allocates flexible + slices of instances (nodes) according to the requested resources; such Computes usually + 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`). """ pass @@ -190,11 +201,15 @@ def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability """ pass - def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]: + def get_offers_modifiers( + self, requirements: Requirements, full_offers: bool + ) -> Iterable[OfferModifier]: """ 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. + + See `Compute.get_offers()` for the `full_offers` argument description. """ return [] @@ -207,12 +222,16 @@ def get_offers_post_filter( """ return None - def get_offers(self, requirements: Requirements) -> Iterator[InstanceOfferWithAvailability]: + def get_offers( + self, requirements: Requirements, full_offers: 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() - offers = self.__apply_modifiers(cached_offers, self.get_offers_modifiers(requirements)) + offers = self.__apply_modifiers( + cached_offers, self.get_offers_modifiers(requirements, full_offers) + ) offers = filter_offers_by_requirements(offers, requirements) post_filter = self.get_offers_post_filter(requirements) if post_filter is not None: @@ -246,6 +265,12 @@ class ComputeWithFilteredOffersCached(ABC): It caches offers using requirements as key. """ + full_offers_argument_has_effect: ClassVar[bool] = False + """ + Set to `True` if `get_offers_by_requirements()` produces different results based on + the `full_offers` value. Doubles the amount of cached data. + """ + def __init__(self) -> None: super().__init__() self._offers_cache_lock = threading.Lock() @@ -253,19 +278,30 @@ def __init__(self) -> None: @abstractmethod def get_offers_by_requirements( - self, requirements: Requirements + self, + requirements: Requirements, + full_offers: bool, ) -> List[InstanceOfferWithAvailability]: """ Returns backend offers with availability matching 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. """ pass - def get_offers(self, requirements: Requirements) -> Iterator[InstanceOfferWithAvailability]: - return iter(self._get_offers_cached(requirements)) + def get_offers( + self, requirements: Requirements, full_offers: bool + ) -> Iterator[InstanceOfferWithAvailability]: + return iter(self._get_offers_cached(requirements, full_offers)) - def _get_offers_cached_key(self, requirements: Requirements) -> int: + def _get_offers_cached_key(self, requirements: Requirements, full_offers: bool) -> int: # Requirements is not hashable, so we use a hack to get arguments hash - return hash(requirements.json()) + hashable_requirements = requirements.json() + if self.full_offers_argument_has_effect: + return hash((hashable_requirements, full_offers)) + return hash(hashable_requirements) @cachedmethod( cache=lambda self: self._offers_cache, @@ -273,9 +309,9 @@ def _get_offers_cached_key(self, requirements: Requirements) -> int: lock=lambda self: self._offers_cache_lock, ) def _get_offers_cached( - self, requirements: Requirements + self, requirements: Requirements, full_offers: bool ) -> List[InstanceOfferWithAvailability]: - return self.get_offers_by_requirements(requirements) + return self.get_offers_by_requirements(requirements, full_offers) class ComputeWithCreateInstanceSupport(ABC): diff --git a/src/dstack/_internal/core/backends/crusoe/compute.py b/src/dstack/_internal/core/backends/crusoe/compute.py index fe1411fe7c..86ab767449 100644 --- a/src/dstack/_internal/core/backends/crusoe/compute.py +++ b/src/dstack/_internal/core/backends/crusoe/compute.py @@ -180,7 +180,9 @@ def _get_quota_map(self) -> dict[str, int]: result[prog_name] = available return result - def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]: + 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) diff --git a/src/dstack/_internal/core/backends/gcp/compute.py b/src/dstack/_internal/core/backends/gcp/compute.py index 3be1255d8e..0974ac9838 100644 --- a/src/dstack/_internal/core/backends/gcp/compute.py +++ b/src/dstack/_internal/core/backends/gcp/compute.py @@ -168,7 +168,9 @@ def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability offer_with_availability.region = region return offers_with_availability - def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]: + def get_offers_modifiers( + self, requirements: Requirements, full_offers: bool + ) -> Iterable[OfferModifier]: modifiers = [] if requirements.reservation: diff --git a/src/dstack/_internal/core/backends/jarvislabs/compute.py b/src/dstack/_internal/core/backends/jarvislabs/compute.py index 24e13f8501..cc995c0239 100644 --- a/src/dstack/_internal/core/backends/jarvislabs/compute.py +++ b/src/dstack/_internal/core/backends/jarvislabs/compute.py @@ -91,7 +91,9 @@ def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability for offer in offers ] - def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]: + def get_offers_modifiers( + self, requirements: Requirements, full_offers: bool + ) -> Iterable[OfferModifier]: return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)] def create_instance( diff --git a/src/dstack/_internal/core/backends/kubernetes/compute.py b/src/dstack/_internal/core/backends/kubernetes/compute.py index b69e546a30..d5f578c474 100644 --- a/src/dstack/_internal/core/backends/kubernetes/compute.py +++ b/src/dstack/_internal/core/backends/kubernetes/compute.py @@ -169,7 +169,11 @@ def get_all_offers_with_availability(self) -> list[InstanceOfferWithAvailability offers.extend(cluster_offers) return offers - def get_offers_modifiers(self, requirements: Requirements) -> list[OfferModifier]: + def get_offers_modifiers( + self, requirements: Requirements, full_offers: bool + ) -> list[OfferModifier]: + if full_offers: + return [] resource_requests = ResourceRequests.from_resources_spec(requirements.resources) return [partial(_offer_modifier, resource_requests)] diff --git a/src/dstack/_internal/core/backends/nebius/compute.py b/src/dstack/_internal/core/backends/nebius/compute.py index a90b4a4c1f..d23d46a854 100644 --- a/src/dstack/_internal/core/backends/nebius/compute.py +++ b/src/dstack/_internal/core/backends/nebius/compute.py @@ -133,7 +133,9 @@ def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability offer.with_availability(availability=InstanceAvailability.UNKNOWN) for offer in offers ] - def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]: + def get_offers_modifiers( + self, requirements: Requirements, full_offers: bool + ) -> Iterable[OfferModifier]: return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)] def create_instance( diff --git a/src/dstack/_internal/core/backends/oci/compute.py b/src/dstack/_internal/core/backends/oci/compute.py index ceaffcf030..f158163cdf 100644 --- a/src/dstack/_internal/core/backends/oci/compute.py +++ b/src/dstack/_internal/core/backends/oci/compute.py @@ -100,7 +100,9 @@ def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability return offers_with_availability - def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]: + def get_offers_modifiers( + self, requirements: Requirements, full_offers: bool + ) -> Iterable[OfferModifier]: return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)] def terminate_instance( diff --git a/src/dstack/_internal/core/backends/runpod/compute.py b/src/dstack/_internal/core/backends/runpod/compute.py index bafdb91b89..d66b78fb54 100644 --- a/src/dstack/_internal/core/backends/runpod/compute.py +++ b/src/dstack/_internal/core/backends/runpod/compute.py @@ -89,7 +89,9 @@ def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability ] return offers - def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]: + 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( diff --git a/src/dstack/_internal/core/backends/slurm/compute.py b/src/dstack/_internal/core/backends/slurm/compute.py index 9d7b40c539..2739b2e75c 100644 --- a/src/dstack/_internal/core/backends/slurm/compute.py +++ b/src/dstack/_internal/core/backends/slurm/compute.py @@ -114,7 +114,11 @@ def get_all_offers_with_availability(self) -> list[InstanceOfferWithAvailability offers.extend(cluster_offers) return offers - def get_offers_modifiers(self, requirements: Requirements) -> list[OfferModifier]: + def get_offers_modifiers( + self, requirements: Requirements, full_offers: bool + ) -> list[OfferModifier]: + if full_offers: + return [] requested_resources = get_requested_resources_from_resources_spec(requirements.resources) return [partial(self._offer_modifier, requested_resources)] diff --git a/src/dstack/_internal/core/backends/template/compute.py.jinja b/src/dstack/_internal/core/backends/template/compute.py.jinja index 046e429ea2..67af917bcf 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 + self, requirements: Requirements, full_offers: 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 6843135ed9..d44241f0c1 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 + self, requirements: Requirements, full_offers: 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 c3ff2a05a9..820b914379 100644 --- a/src/dstack/_internal/core/backends/verda/compute.py +++ b/src/dstack/_internal/core/backends/verda/compute.py @@ -72,7 +72,9 @@ def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability offers_with_availability = self._get_offers_with_availability(offers) return offers_with_availability - def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]: + def get_offers_modifiers( + self, requirements: Requirements, full_offers: bool + ) -> Iterable[OfferModifier]: return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)] def _get_offers_with_availability( diff --git a/src/dstack/_internal/core/compatibility/gpus.py b/src/dstack/_internal/core/compatibility/gpus.py index 3885001273..c7ed783c6a 100644 --- a/src/dstack/_internal/core/compatibility/gpus.py +++ b/src/dstack/_internal/core/compatibility/gpus.py @@ -7,6 +7,8 @@ def get_list_gpus_excludes(request: ListGpusRequest) -> Optional[IncludeExcludeDictType]: list_gpus_excludes: IncludeExcludeDictType = {} + if not request.full_offers: + list_gpus_excludes["full_offers"] = 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 ab07a4d03b..e71005095e 100644 --- a/src/dstack/_internal/core/compatibility/runs.py +++ b/src/dstack/_internal/core/compatibility/runs.py @@ -70,6 +70,8 @@ def get_get_plan_excludes(request: GetRunPlanRequest) -> Optional[IncludeExclude clients backward-compatibility with older servers. """ get_plan_excludes: IncludeExcludeDictType = {} + if not request.full_offers: + get_plan_excludes["full_offers"] = 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 7c244e327e..9cfcc225cc 100644 --- a/src/dstack/_internal/server/routers/gpus.py +++ b/src/dstack/_internal/server/routers/gpus.py @@ -37,6 +37,7 @@ async def list_gpus( project=project, run_spec=body.run_spec, group_by=body.group_by, + full_offers=body.full_offers, ) 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 ab468c16d5..cb11ad120c 100644 --- a/src/dstack/_internal/server/routers/runs.py +++ b/src/dstack/_internal/server/routers/runs.py @@ -139,6 +139,7 @@ async def get_plan( user=user, run_spec=body.run_spec, max_offers=body.max_offers, + full_offers=body.full_offers, 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 16d1191483..af790e3e99 100644 --- a/src/dstack/_internal/server/schemas/gpus.py +++ b/src/dstack/_internal/server/schemas/gpus.py @@ -1,4 +1,4 @@ -from typing import List, Literal, Optional +from typing import Annotated, List, Literal, Optional from pydantic import Field @@ -16,6 +16,9 @@ class ListGpusRequest(CoreModel): description="List of fields to group by. Valid values: 'backend', 'region', 'count'. " "Note: 'region' can only be used together with 'backend'.", ) + full_offers: Annotated[ + bool, Field(description="Don't adjust backend offers by requirements") + ] = False class ListGpusResponse(CoreModel): diff --git a/src/dstack/_internal/server/schemas/runs.py b/src/dstack/_internal/server/schemas/runs.py index 12b0c36bac..84861a6a93 100644 --- a/src/dstack/_internal/server/schemas/runs.py +++ b/src/dstack/_internal/server/schemas/runs.py @@ -45,6 +45,9 @@ class GetRunPlanRequest(CoreModel): max_offers: Optional[int] = Field( description="The maximum number of offers to return", ge=1, le=10000 ) + full_offers: Annotated[ + bool, Field(description="Return full offers not adjusted by requirements") + ] = 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 5dd3d91ae1..8ea35c215c 100644 --- a/src/dstack/_internal/server/services/backends/__init__.py +++ b/src/dstack/_internal/server/services/backends/__init__.py @@ -475,6 +475,7 @@ async def get_project_backend_with_model_by_id_or_error( async def get_backend_offers( backends: List[Backend], requirements: Requirements, + full_offers: bool, exclude_not_available: bool = False, ) -> Iterable[Tuple[Backend, InstanceOfferWithAvailability]]: """ @@ -490,7 +491,9 @@ def get_filtered_offers_with_backends( yield (backend, offer) logger.debug("Requesting instance offers from backends: %s", [b.TYPE.value for b in backends]) - tasks = [run_async(get_offers_tracked, backend, requirements) for backend in backends] + tasks = [ + run_async(get_offers_tracked, backend, requirements, full_offers) 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)): if isinstance(result, BackendError): @@ -521,10 +524,10 @@ def check_backend_type_available(backend_type: BackendType): def get_offers_tracked( - backend: Backend, requirements: Requirements + backend: Backend, requirements: Requirements, full_offers: bool ) -> Iterator[InstanceOfferWithAvailability]: start = time.time() - res = backend.compute().get_offers(requirements) + res = backend.compute().get_offers(requirements, full_offers) 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 3e0ebec51c..2ba42ae209 100644 --- a/src/dstack/_internal/server/services/gpus.py +++ b/src/dstack/_internal/server/services/gpus.py @@ -26,9 +26,12 @@ async def list_gpus_grouped( project: ProjectModel, run_spec: RunSpec, group_by: Optional[List[Literal["backend", "region", "count"]]] = None, + full_offers: 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) + offers = await _get_gpu_offers( + session=session, project=project, run_spec=run_spec, full_offers=full_offers + ) backend_gpus = _process_offers_into_backend_gpus(offers) group_by_set = set(group_by) if group_by else set() if "region" in group_by_set and "backend" not in group_by_set: @@ -58,6 +61,7 @@ async def _get_gpu_offers( session: AsyncSession, project: ProjectModel, run_spec: RunSpec, + full_offers: 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 @@ -85,6 +89,7 @@ async def _get_gpu_offers( run_spec=run_spec, job=job, skip_backend_offers=skip_backend_offers, + full_offers=full_offers, ) else: instance_offers, backend_offers = await get_non_fleet_offers( @@ -93,6 +98,7 @@ async def _get_gpu_offers( run_spec=run_spec, job=job, skip_backend_offers=skip_backend_offers, + full_offers=full_offers, ) 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 6fd739f13e..9d0d2d7005 100644 --- a/src/dstack/_internal/server/services/offers.py +++ b/src/dstack/_internal/server/services/offers.py @@ -40,6 +40,7 @@ async def get_offers_by_requirements( placement_group: Optional[PlacementGroup] = None, blocks: Union[int, Literal["auto"]] = 1, max_offers: Optional[int] = None, + full_offers: bool = False, ) -> List[Tuple[Backend, InstanceOfferWithAvailability]]: backends: List[Backend] = await backends_services.get_project_backends(project=project) @@ -91,6 +92,7 @@ async def get_offers_by_requirements( offers = await backends_services.get_backend_offers( backends=backends, requirements=requirements, + full_offers=full_offers, 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 39d4815d08..6afea7c651 100644 --- a/src/dstack/_internal/server/services/runs/__init__.py +++ b/src/dstack/_internal/server/services/runs/__init__.py @@ -532,6 +532,7 @@ async def get_plan( user: UserModel, run_spec: RunSpec, max_offers: Optional[int], + full_offers: bool, legacy_repo_dir: bool = False, ) -> RunPlan: # Spec must be copied by parsing to calculate merged_profile @@ -570,6 +571,7 @@ async def get_plan( project=project, run_spec=effective_run_spec, max_offers=max_offers, + full_offers=full_offers, ) 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 39275a41a6..e876d05bde 100644 --- a/src/dstack/_internal/server/services/runs/plan.py +++ b/src/dstack/_internal/server/services/runs/plan.py @@ -87,6 +87,7 @@ async def get_job_plans( project: ProjectModel, run_spec: RunSpec, max_offers: Optional[int], + full_offers: bool, ) -> list[JobPlan]: """ Returns job plans for the given run spec. @@ -145,6 +146,7 @@ async def get_job_plans( replica_group_name=replica_group_name, ) if candidate_fleet_models is not None: + # Regular job planning fleet_model, instance_offers, backend_offers = await find_optimal_fleet_with_offers( project=project, fleet_models=candidate_fleet_models, @@ -155,8 +157,10 @@ async def get_job_plans( volumes=volumes, exclude_not_available=False, skip_backend_offers=skip_backend_offers, + full_offers=full_offers, ) elif run_spec.merged_profile.instances is not None: + # Regular job planning or offer collection instance_offers = await get_targeted_instance_offers( session=session, project=project, @@ -166,6 +170,7 @@ async def get_job_plans( ) backend_offers = [] elif run_spec.merged_profile.fleets is not None: + # Offer collection instance_offers, backend_offers = await get_offers_in_run_candidate_fleets( session=session, project=project, @@ -173,8 +178,10 @@ async def get_job_plans( job=jobs[0], volumes=volumes, skip_backend_offers=skip_backend_offers, + full_offers=full_offers, ) else: + # Offer collection instance_offers, backend_offers = await get_non_fleet_offers( session=session, project=project, @@ -182,6 +189,7 @@ async def get_job_plans( job=jobs[0], volumes=volumes, skip_backend_offers=skip_backend_offers, + full_offers=full_offers, ) for job in jobs: @@ -324,6 +332,7 @@ async def find_optimal_fleet_with_offers( exclude_not_available: bool, skip_backend_offers: bool = False, skip_backend_offers_on_pool_capacity: bool = False, + full_offers: bool = False, ) -> tuple[ Optional[FleetModel], list[tuple[InstanceModel, InstanceOfferWithAvailability]], @@ -427,6 +436,7 @@ async def find_optimal_fleet_with_offers( job=job, volumes=volumes, max_offers=_PER_FLEET_MAX_OFFERS, + full_offers=full_offers, ) available_backend_offers = _exclude_non_available_backend_offers(backend_offers) candidates_with_backend_offers.append( @@ -459,6 +469,7 @@ async def find_optimal_fleet_with_offers( job=job, volumes=volumes, max_offers=None, + full_offers=full_offers, ) if exclude_not_available: backend_offers = _exclude_non_available_backend_offers(backend_offers) @@ -715,6 +726,7 @@ async def _get_backend_offers_in_fleet( volumes: Optional[list[list[Volume]]], fleet_spec: Optional[FleetSpec] = None, max_offers: Optional[int] = None, + full_offers: bool = False, ) -> list[tuple[Backend, InstanceOfferWithAvailability]]: if fleet_spec is None: fleet_spec = get_fleet_spec(fleet_model) @@ -746,6 +758,7 @@ async def _get_backend_offers_in_fleet( privileged=job.job_spec.privileged, instance_mounts=check_run_spec_requires_instance_mounts(run_spec), max_offers=max_offers, + full_offers=full_offers, ) return backend_offers @@ -793,6 +806,7 @@ async def get_non_fleet_offers( job: Job, volumes: Optional[list[list[Volume]]] = None, skip_backend_offers: bool = False, + full_offers: bool = False, ) -> tuple[ list[tuple[InstanceModel, InstanceOfferWithAvailability]], list[tuple[Backend, InstanceOfferWithAvailability]], @@ -821,6 +835,7 @@ async def get_non_fleet_offers( volumes=volumes, privileged=job.job_spec.privileged, instance_mounts=check_run_spec_requires_instance_mounts(run_spec), + full_offers=full_offers, ) return instance_offers, backend_offers @@ -832,6 +847,7 @@ async def get_backend_offers_in_run_candidate_fleets( job: Job, volumes: Optional[list[list[Volume]]], max_offers_per_fleet: Optional[int] = None, + full_offers: bool = False, ) -> list[tuple[Backend, InstanceOfferWithAvailability]]: """ Returns backend offers across the run's selected candidate fleets. @@ -860,6 +876,7 @@ async def get_backend_offers_in_run_candidate_fleets( job=job, volumes=volumes, max_offers=max_offers_per_fleet, + full_offers=full_offers, ): offer_identity = _get_backend_offer_identity(offer) if offer_identity not in seen_offer_identities: @@ -876,6 +893,7 @@ async def get_offers_in_run_candidate_fleets( job: Job, volumes: Optional[list[list[Volume]]] = None, skip_backend_offers: bool = False, + full_offers: bool = False, ) -> tuple[ list[tuple[InstanceModel, InstanceOfferWithAvailability]], list[tuple[Backend, InstanceOfferWithAvailability]], @@ -923,6 +941,7 @@ async def get_offers_in_run_candidate_fleets( job=job, volumes=volumes, max_offers_per_fleet=None, + full_offers=full_offers, ) return instance_offers, backend_offers diff --git a/src/dstack/api/_public/runs.py b/src/dstack/api/_public/runs.py index 7e3e29cabc..752eca07a9 100644 --- a/src/dstack/api/_public/runs.py +++ b/src/dstack/api/_public/runs.py @@ -475,6 +475,8 @@ def get_run_plan( configuration_path: Optional[str] = None, repo_dir: Union[Deprecated, str, None] = Deprecated.PLACEHOLDER, ssh_identity_file: Optional[PathLike] = None, + max_offers: Optional[int] = None, + full_offers: bool = False, ) -> RunPlan: """ Get a run plan. @@ -491,6 +493,8 @@ def get_run_plan( (`.pub` file) is read and included in the run plan, allowing SSH access to the instances. If the `.pub` file does not exist, it is generated automatically. 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. Returns: Run plan. @@ -541,7 +545,12 @@ def get_run_plan( ssh_key_pub=ssh_key_pub, ) logger.debug("Getting run plan") - run_plan = self._api_client.runs.get_plan(self._project, run_spec) + run_plan = self._api_client.runs.get_plan( + project_name=self._project, + run_spec=run_spec, + max_offers=max_offers, + full_offers=full_offers, + ) return run_plan def apply_plan( diff --git a/src/dstack/api/server/_gpus.py b/src/dstack/api/server/_gpus.py index 253410604d..d471a8e7bd 100644 --- a/src/dstack/api/server/_gpus.py +++ b/src/dstack/api/server/_gpus.py @@ -15,10 +15,12 @@ def list_gpus( project_name: str, run_spec: RunSpec, group_by: Optional[List[str]] = None, + full_offers: 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, ) 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 0543f3384a..6018a98f38 100644 --- a/src/dstack/api/server/_runs.py +++ b/src/dstack/api/server/_runs.py @@ -72,9 +72,13 @@ def get( return parse_obj_as(Run.__response__, resp.json()) def get_plan( - self, project_name: str, run_spec: RunSpec, max_offers: Optional[int] = None + self, + project_name: str, + run_spec: RunSpec, + max_offers: Optional[int] = None, + full_offers: bool = False, ) -> RunPlan: - body = GetRunPlanRequest(run_spec=run_spec, max_offers=max_offers) + body = GetRunPlanRequest(run_spec=run_spec, max_offers=max_offers, full_offers=full_offers) 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 6ee60dfc25..d938dc113d 100644 --- a/src/tests/_internal/core/backends/jarvislabs/test_compute.py +++ b/src/tests/_internal/core/backends/jarvislabs/test_compute.py @@ -161,8 +161,12 @@ def test_get_offers_reuses_all_offers_cache_and_modifies_disk_size(): return_value=[_cpu_offer(disk_size_mib=100 * 1024)] ) - offers_250gb = list(compute.get_offers(Requirements(resources=ResourcesSpec(disk="250GB")))) - offers_300gb = list(compute.get_offers(Requirements(resources=ResourcesSpec(disk="300GB")))) + offers_250gb = list( + compute.get_offers(Requirements(resources=ResourcesSpec(disk="250GB")), False) + ) + offers_300gb = list( + compute.get_offers(Requirements(resources=ResourcesSpec(disk="300GB")), False) + ) assert len(offers_250gb) == 1 assert offers_250gb[0].instance.resources.disk.size_mib == 250 * 1024 diff --git a/src/tests/_internal/core/backends/kubernetes/test_compute.py b/src/tests/_internal/core/backends/kubernetes/test_compute.py new file mode 100644 index 0000000000..c8c6748e85 --- /dev/null +++ b/src/tests/_internal/core/backends/kubernetes/test_compute.py @@ -0,0 +1,90 @@ +from unittest.mock import MagicMock, patch + +from dstack._internal.core.backends.kubernetes.compute import KubernetesCompute +from dstack._internal.core.backends.kubernetes.models import KubeconfigConfig, KubernetesConfig +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.instances import ( + Disk, + Gpu, + InstanceAvailability, + InstanceOfferWithAvailability, + InstanceType, + Resources, +) +from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.core.models.runs import Requirements + + +def _compute() -> KubernetesCompute: + with patch( + "dstack._internal.core.backends.kubernetes.compute.get_clusters_from_backend_config", + return_value=[], + ): + return KubernetesCompute( + KubernetesConfig( + kubeconfig=KubeconfigConfig(data="mocked", filename="-"), + contexts=["ctx"], + ) + ) + + +def _node_offer() -> InstanceOfferWithAvailability: + return InstanceOfferWithAvailability( + backend=BackendType.KUBERNETES, + instance=InstanceType( + name="ctx-node", + resources=Resources( + cpus=8, + memory_mib=64 * 1024, + gpus=[Gpu(name="A100", memory_mib=80 * 1024) for _ in range(4)], + spot=False, + disk=Disk(size_mib=200 * 1024), + ), + ), + region="ctx", + price=0.0, + availability=InstanceAvailability.AVAILABLE, + ) + + +def test_get_offers_modifiers_are_skipped_with_full_offers(): + compute = _compute() + requirements = Requirements(resources=ResourcesSpec(cpu="2", memory="8GB", gpu="1")) + + assert compute.get_offers_modifiers(requirements, full_offers=True) == [] + assert compute.get_offers_modifiers(requirements, full_offers=False) != [] + + +def test_get_offers_with_full_offers_keeps_full_node_resources(): + compute = _compute() + compute.get_all_offers_with_availability = MagicMock(return_value=[_node_offer()]) + # Open-ended requirements so the full node satisfies them without an upper bound. + requirements = Requirements( + resources=ResourcesSpec(cpu="2..", memory="8GB..", gpu="1..", disk="100GB..") + ) + + full_offers = list(compute.get_offers(requirements, full_offers=True)) + + assert len(full_offers) == 1 + full_resources = full_offers[0].instance.resources + assert full_resources.cpus == 8 + assert full_resources.memory_mib == 64 * 1024 + assert len(full_resources.gpus) == 4 + assert full_resources.disk.size_mib == 200 * 1024 + + +def test_get_offers_without_full_offers_adjusts_to_requested_slice(): + compute = _compute() + compute.get_all_offers_with_availability = MagicMock(return_value=[_node_offer()]) + requirements = Requirements( + resources=ResourcesSpec(cpu="2..", memory="8GB..", gpu="1..", disk="100GB..") + ) + + adjusted_offers = list(compute.get_offers(requirements, full_offers=False)) + + assert len(adjusted_offers) == 1 + adjusted_resources = adjusted_offers[0].instance.resources + assert adjusted_resources.cpus == 2 + assert adjusted_resources.memory_mib == 8 * 1024 + assert len(adjusted_resources.gpus) == 1 + assert adjusted_resources.disk.size_mib == 100 * 1024 diff --git a/src/tests/_internal/core/backends/slurm/test_compute.py b/src/tests/_internal/core/backends/slurm/test_compute.py new file mode 100644 index 0000000000..d7b39c9273 --- /dev/null +++ b/src/tests/_internal/core/backends/slurm/test_compute.py @@ -0,0 +1,93 @@ +from unittest.mock import MagicMock, patch + +from dstack._internal.core.backends.slurm.compute import SlurmCompute +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.instances import ( + Disk, + Gpu, + InstanceAvailability, + InstanceOfferWithAvailability, + InstanceType, + Resources, +) +from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.core.models.runs import Requirements + + +def _compute() -> SlurmCompute: + # Cluster discovery is mocked out, so the config itself is never inspected. + with patch( + "dstack._internal.core.backends.slurm.compute.get_clusters_from_backend_config", + return_value=[], + ): + return SlurmCompute(MagicMock()) + + +def _node_offer() -> InstanceOfferWithAvailability: + return InstanceOfferWithAvailability( + backend=BackendType.SLURM, + instance=InstanceType( + name="slurm-node", + resources=Resources( + cpus=16, + memory_mib=128 * 1024, + gpus=[Gpu(name="H100", memory_mib=80 * 1024) for _ in range(8)], + spot=False, + disk=Disk(size_mib=500 * 1024), + ), + ), + region="cluster1", + price=0.0, + availability=InstanceAvailability.AVAILABLE, + availability_zones=["partition1"], + ) + + +def test_get_offers_modifiers_are_skipped_with_full_offers(): + compute = _compute() + requirements = Requirements(resources=ResourcesSpec(cpu="2", memory="8GB", gpu="1")) + + assert compute.get_offers_modifiers(requirements, full_offers=True) == [] + assert compute.get_offers_modifiers(requirements, full_offers=False) != [] + + +def test_get_offers_with_full_offers_keeps_full_node_resources(): + compute = _compute() + compute.get_all_offers_with_availability = MagicMock(return_value=[_node_offer()]) + # Open-ended requirements so the full node satisfies them without an upper bound. + requirements = Requirements( + resources=ResourcesSpec(cpu="2..", memory="8GB..", gpu="1..", disk="100GB..") + ) + + full_offers = list(compute.get_offers(requirements, full_offers=True)) + + assert len(full_offers) == 1 + full_resources = full_offers[0].instance.resources + assert full_resources.cpus == 16 + assert full_resources.memory_mib == 128 * 1024 + assert len(full_resources.gpus) == 8 + assert full_resources.disk.size_mib == 500 * 1024 + + +def test_get_offers_without_full_offers_adjusts_to_requested_slice(): + compute = _compute() + compute.get_all_offers_with_availability = MagicMock(return_value=[_node_offer()]) + compute._get_cluster = MagicMock() + requirements = Requirements( + resources=ResourcesSpec(cpu="2..", memory="8GB..", gpu="1..", disk="100GB..") + ) + + # Slicing needs cluster/partition lookups, which are otherwise backed by live cluster state. + with patch( + "dstack._internal.core.backends.slurm.compute._get_cluster_partitions", + return_value={"partition1"}, + ): + adjusted_offers = list(compute.get_offers(requirements, full_offers=False)) + + assert len(adjusted_offers) == 1 + adjusted_resources = adjusted_offers[0].instance.resources + assert adjusted_resources.cpus == 2 + assert adjusted_resources.memory_mib == 8 * 1024 + assert len(adjusted_resources.gpus) == 1 + assert adjusted_resources.disk.size_mib == 100 * 1024 + assert adjusted_offers[0].availability_zones == ["partition1"] diff --git a/src/tests/_internal/core/backends/vastai/test_compute.py b/src/tests/_internal/core/backends/vastai/test_compute.py index 48a2230d40..0b0da5401b 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())) + list(compute.get_offers(_requirements(), 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())) + list(compute.get_offers(_requirements(), 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())) + list(compute.get_offers(_requirements(), 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 a1e2afd734..b54b48b0be 100644 --- a/src/tests/_internal/server/routers/test_gpus.py +++ b/src/tests/_internal/server/routers/test_gpus.py @@ -138,11 +138,14 @@ async def call_gpus_api( run_spec: RunSpec, group_by: Optional[List[str]] = None, client_version: Optional[str] = None, + full_offers: Optional[bool] = None, ): """Helper to call the GPUs API with standard parameters.""" json_data = {"run_spec": run_spec.dict()} if group_by is not None: json_data["group_by"] = group_by + if full_offers is not None: + json_data["full_offers"] = full_offers headers = get_auth_headers(user_token) if client_version is not None: headers["X-API-Version"] = client_version @@ -190,6 +193,43 @@ async def test_returns_gpus_without_group_by( assert isinstance(response_data["gpus"], list) assert len(response_data["gpus"]) >= 1 + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + @pytest.mark.parametrize( + ("body_full_offers", "expected_full_offers"), + [ + 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_full_offers_to_compute_get_offers( + self, + test_db, + session: AsyncSession, + client: AsyncClient, + body_full_offers: Optional[bool], + expected_full_offers: 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, full_offers=body_full_offers + ) + + 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) + assert all( + call_args.args[1] is expected_full_offers + 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 e98e0b1502..e2c1750aea 100644 --- a/src/tests/_internal/server/routers/test_runs.py +++ b/src/tests/_internal/server/routers/test_runs.py @@ -1572,6 +1572,61 @@ async def test_returns_run_plan_privileged_false( assert response.status_code == 200, response.json() assert response.json() == run_plan_dict + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + @pytest.mark.parametrize( + ("body_full_offers", "expected_full_offers"), + [ + 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_full_offers_to_compute_get_offers( + self, + test_db, + session: AsyncSession, + client: AsyncClient, + body_full_offers: Optional[bool], + expected_full_offers: 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_full_offers is not None: + body["full_offers"] = body_full_offers + + 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) + assert all( + call_args.args[1] is expected_full_offers + 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( @@ -1745,7 +1800,7 @@ 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): + def offers_by_requirements(requirements: Requirements, full_offers: 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 01ff0c4e3d..21375b2509 100644 --- a/src/tests/_internal/server/services/runs/test_plan.py +++ b/src/tests/_internal/server/services/runs/test_plan.py @@ -134,6 +134,7 @@ async def test_skips_backend_offers_by_creation_policy( project=project, run_spec=run_spec, max_offers=None, + full_offers=False, ) find_optimal_fleet_with_offers_mock.assert_awaited_once() @@ -172,6 +173,7 @@ async def test_excludes_backend_offers_when_instances_specified( project=project, run_spec=run_spec, max_offers=None, + full_offers=False, ) get_targeted_instance_offers_mock.assert_awaited_once() @@ -211,6 +213,7 @@ async def test_empty_dev_environment_with_fleet_does_not_use_targeted_instances( user=user, run_spec=run_spec, max_offers=None, + full_offers=False, ) select_instances_mock.assert_not_awaited()