diff --git a/src/dstack/_internal/core/backends/verda/compute.py b/src/dstack/_internal/core/backends/verda/compute.py index 511d00914d..1bcd6ae44e 100644 --- a/src/dstack/_internal/core/backends/verda/compute.py +++ b/src/dstack/_internal/core/backends/verda/compute.py @@ -78,19 +78,21 @@ def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModi def _get_offers_with_availability( self, offers: List[InstanceOffer] ) -> List[InstanceOfferWithAvailability]: - raw_availabilities: List[Dict] = self.client.instances.get_availabilities() - region_availabilities = {} - for location in raw_availabilities: - location_code = location["location_code"] - availabilities = location["availabilities"] - for name in availabilities: - key = (name, location_code) - region_availabilities[key] = InstanceAvailability.AVAILABLE + for is_spot in (False, True): + raw_availabilities: List[Dict] = self.client.instances.get_availabilities( + is_spot=is_spot + ) + for location in raw_availabilities: + location_code = location["location_code"] + availabilities = location["availabilities"] + for name in availabilities: + key = (name, location_code, is_spot) + region_availabilities[key] = InstanceAvailability.AVAILABLE availability_offers = [] for offer in offers: - key = (offer.instance.name, offer.region) + key = (offer.instance.name, offer.region, offer.instance.resources.spot) availability = region_availabilities.get(key, InstanceAvailability.NOT_AVAILABLE) availability_offers.append(offer.with_availability(availability=availability)) diff --git a/src/tests/_internal/core/backends/verda/test_compute.py b/src/tests/_internal/core/backends/verda/test_compute.py index e94ee7d12e..b8ae8f494d 100644 --- a/src/tests/_internal/core/backends/verda/test_compute.py +++ b/src/tests/_internal/core/backends/verda/test_compute.py @@ -11,6 +11,25 @@ _create_startup_script, ) from dstack._internal.core.errors import BackendError, NoCapacityError +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.instances import ( + InstanceAvailability, + InstanceOffer, + InstanceType, + Resources, +) + + +def _offer(spot: bool, name: str = "SOME.INSTANCE", region: str = "FIN-01") -> InstanceOffer: + return InstanceOffer( + backend=BackendType.VERDA, + instance=InstanceType( + name=name, + resources=Resources(cpus=8, memory_mib=16384, gpus=[], spot=spot), + ), + region=region, + price=1.0, + ) def _assert_terminate_call(action_mock: MagicMock): @@ -286,6 +305,38 @@ def test_stores_ssh_key_ids_in_backend_data(self): assert backend_data.ssh_key_ids == ["ssh-key-id-1", "ssh-key-id-2"] +class TestGetOffersWithAvailability: + @pytest.mark.parametrize("available_as_spot", [True, False]) + def test_availability_resolved_against_matching_inventory(self, available_as_spot): + compute = VerdaCompute.__new__(VerdaCompute) + compute.client = MagicMock() + + def get_availabilities(is_spot): + names = ["SOME.INSTANCE"] if is_spot == available_as_spot else [] + return [{"location_code": "FIN-01", "availabilities": names}] + + compute.client.instances.get_availabilities.side_effect = get_availabilities + + offers = compute._get_offers_with_availability([_offer(spot=False), _offer(spot=True)]) + availability_by_spot = {o.instance.resources.spot: o.availability for o in offers} + + assert availability_by_spot[available_as_spot] == InstanceAvailability.AVAILABLE + assert availability_by_spot[not available_as_spot] == InstanceAvailability.NOT_AVAILABLE + + def test_queries_both_spot_and_on_demand_availability(self): + compute = VerdaCompute.__new__(VerdaCompute) + compute.client = MagicMock() + compute.client.instances.get_availabilities.return_value = [] + + compute._get_offers_with_availability([_offer(spot=True)]) + + requested_is_spot = { + call.kwargs.get("is_spot") + for call in compute.client.instances.get_availabilities.call_args_list + } + assert requested_is_spot == {True, False} + + class TestTerminateInstance: def test_terminate_instance_without_backend_data(self): compute = VerdaCompute.__new__(VerdaCompute)