dstack 0.19.32__py3-none-any.whl → 0.19.34__py3-none-any.whl

dstack/_internal/cli/commands/offer.py

@@ -104,8 +104,8 @@ class OfferCommand(APIBaseCommand):
 
         run_spec = RunSpec(
             configuration=conf,
-            ssh_key_pub="(dummy)",
             profile=profile,
+            ssh_key_pub="(dummy)",  # TODO: Remove since 0.19.40
         )
 
         if args.group_by:

dstack/_internal/cli/services/configurators/run.py

@@ -62,7 +62,6 @@ from dstack._internal.utils.interpolator import InterpolatorError, VariablesInte
 from dstack._internal.utils.logging import get_logger
 from dstack._internal.utils.nested_list import NestedList, NestedListItem
 from dstack._internal.utils.path import is_absolute_posix_path
-from dstack.api._public.repos import get_ssh_keypair
 from dstack.api._public.runs import Run
 from dstack.api.server import APIClient
 from dstack.api.utils import load_profile
@@ -135,10 +134,6 @@ class BaseRunConfigurator(
 
         config_manager = ConfigManager()
         repo = self.get_repo(conf, configuration_path, configurator_args, config_manager)
-        self.api.ssh_identity_file = get_ssh_keypair(
-            configurator_args.ssh_identity_file,
-            config_manager.dstack_key_path,
-        )
         profile = load_profile(Path.cwd(), configurator_args.profile)
         with console.status("Getting apply plan..."):
             run_plan = self.api.runs.get_run_plan(
@@ -146,6 +141,7 @@ class BaseRunConfigurator(
                 repo=repo,
                 configuration_path=configuration_path,
                 profile=profile,
+                ssh_identity_file=configurator_args.ssh_identity_file,
             )
 
         print_run_plan(run_plan, max_offers=configurator_args.max_offers)

dstack/_internal/core/backends/aws/compute.py

@@ -1,4 +1,5 @@
 import threading
+from collections.abc import Iterable
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
@@ -34,7 +35,11 @@ from dstack._internal.core.backends.base.compute import (
     get_user_data,
     merge_tags,
 )
-from dstack._internal.core.backends.base.offers import get_catalog_offers, get_offers_disk_modifier
+from dstack._internal.core.backends.base.offers import (
+    OfferModifier,
+    get_catalog_offers,
+    get_offers_disk_modifier,
+)
 from dstack._internal.core.errors import (
     ComputeError,
     NoCapacityError,
@@ -159,10 +164,8 @@ class AWSCompute(
             )
         return availability_offers
 
-    def get_offers_modifier(
-        self, requirements: Requirements
-    ) -> Callable[[InstanceOfferWithAvailability], Optional[InstanceOfferWithAvailability]]:
-        return get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)
+    def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]:
+        return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)]
 
     def _get_offers_cached_key(self, requirements: Requirements) -> int:
         # Requirements is not hashable, so we use a hack to get arguments hash

dstack/_internal/core/backends/azure/compute.py

@@ -1,8 +1,9 @@
 import base64
 import enum
 import re
+from collections.abc import Iterable
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from typing import Callable, Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 from azure.core.credentials import TokenCredential
 from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError
@@ -51,7 +52,11 @@ from dstack._internal.core.backends.base.compute import (
     merge_tags,
     requires_nvidia_proprietary_kernel_modules,
 )
-from dstack._internal.core.backends.base.offers import get_catalog_offers, get_offers_disk_modifier
+from dstack._internal.core.backends.base.offers import (
+    OfferModifier,
+    get_catalog_offers,
+    get_offers_disk_modifier,
+)
 from dstack._internal.core.consts import DSTACK_OS_IMAGE_WITH_PROPRIETARY_NVIDIA_KERNEL_MODULES
 from dstack._internal.core.errors import ComputeError, NoCapacityError
 from dstack._internal.core.models.backends.base import BackendType
@@ -108,10 +113,8 @@ class AzureCompute(
         )
         return offers_with_availability
 
-    def get_offers_modifier(
-        self, requirements: Requirements
-    ) -> Callable[[InstanceOfferWithAvailability], Optional[InstanceOfferWithAvailability]]:
-        return get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)
+    def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]:
+        return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)]
 
     def create_instance(
         self,

dstack/_internal/core/backends/base/compute.py

@@ -17,12 +17,13 @@ from cachetools import TTLCache, cachedmethod
 from gpuhunt import CPUArchitecture
 
 from dstack._internal import settings
-from dstack._internal.core.backends.base.offers import filter_offers_by_requirements
+from dstack._internal.core.backends.base.offers import OfferModifier, filter_offers_by_requirements
 from dstack._internal.core.consts import (
     DSTACK_RUNNER_HTTP_PORT,
     DSTACK_RUNNER_SSH_PORT,
     DSTACK_SHIM_HTTP_PORT,
 )
+from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.configurations import LEGACY_REPO_DIR
 from dstack._internal.core.models.gateways import (
     GatewayComputeConfiguration,
@@ -168,17 +169,13 @@ class ComputeWithAllOffersCached(ABC):
         """
         pass
 
-    def get_offers_modifier(
-        self, requirements: Requirements
-    ) -> Optional[
-        Callable[[InstanceOfferWithAvailability], Optional[InstanceOfferWithAvailability]]
-    ]:
+    def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]:
         """
-        Returns a modifier function that modifies offers before they are filtered by requirements.
-        Can return `None` to exclude the offer.
+        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.
         """
-        return None
+        return []
 
     def get_offers_post_filter(
         self, requirements: Requirements
@@ -191,14 +188,7 @@ class ComputeWithAllOffersCached(ABC):
 
     def get_offers(self, requirements: Requirements) -> List[InstanceOfferWithAvailability]:
         offers = self._get_all_offers_with_availability_cached()
-        modifier = self.get_offers_modifier(requirements)
-        if modifier is not None:
-            modified_offers = []
-            for o in offers:
-                modified_offer = modifier(o)
-                if modified_offer is not None:
-                    modified_offers.append(modified_offer)
-            offers = modified_offers
+        offers = self.__apply_modifiers(offers, self.get_offers_modifiers(requirements))
         offers = filter_offers_by_requirements(offers, requirements)
         post_filter = self.get_offers_post_filter(requirements)
         if post_filter is not None:
@@ -212,6 +202,20 @@ class ComputeWithAllOffersCached(ABC):
     def _get_all_offers_with_availability_cached(self) -> List[InstanceOfferWithAvailability]:
         return self.get_all_offers_with_availability()
 
+    @staticmethod
+    def __apply_modifiers(
+        offers: Iterable[InstanceOfferWithAvailability], modifiers: Iterable[OfferModifier]
+    ) -> list[InstanceOfferWithAvailability]:
+        modified_offers = []
+        for offer in offers:
+            for modifier in modifiers:
+                offer = modifier(offer)
+                if offer is None:
+                    break
+            else:
+                modified_offers.append(offer)
+        return modified_offers
+
 
 class ComputeWithFilteredOffersCached(ABC):
     """
@@ -341,6 +345,15 @@ class ComputeWithMultinodeSupport:
 class ComputeWithReservationSupport:
     """
     Must be subclassed to support provisioning from reservations.
+
+    The following is expected from a backend that supports reservations:
+
+    - `get_offers` respects `Requirements.reservation` if set, and only returns
+      offers that can be provisioned in the configured reservation. It can
+      adjust some offer properties such as `availability` and
+      `availability_zones` if necessary.
+    - `create_instance` respects `InstanceConfig.reservation` if set, and
+      provisions the instance in the configured reservation.
     """
 
     pass
@@ -391,6 +404,16 @@ class ComputeWithPlacementGroupSupport(ABC):
         """
         pass
 
+    def are_placement_groups_compatible_with_reservations(self, backend_type: BackendType) -> bool:
+        """
+        Whether placement groups can be used for instances provisioned in reservations.
+
+        Arguments:
+            backend_type: matches the backend type of this compute, unless this compute is a proxy
+                for other backends (dstack Sky)
+        """
+        return True
+
 
 class ComputeWithGatewaySupport(ABC):
     """

dstack/_internal/core/backends/base/offers.py

@@ -23,6 +23,8 @@ SUPPORTED_GPUHUNT_FLAGS = [
     "oci-spot",
     "lambda-arm",
     "gcp-a4",
+    "gcp-g4",
+    "gcp-dws-calendar-mode",
 ]
 
 
@@ -93,6 +95,7 @@ def catalog_item_to_offer(
         ),
         region=item.location,
         price=item.price,
+        backend_data=item.provider_data,
     )
 
 
@@ -199,9 +202,12 @@ def choose_disk_size_mib(
     return round(disk_size_gib * 1024)
 
 
+OfferModifier = Callable[[InstanceOfferWithAvailability], Optional[InstanceOfferWithAvailability]]
+
+
 def get_offers_disk_modifier(
     configurable_disk_size: Range[Memory], requirements: Requirements
-) -> Callable[[InstanceOfferWithAvailability], Optional[InstanceOfferWithAvailability]]:
+) -> OfferModifier:
     """
     Returns a func that modifies offers disk by setting min value that satisfies both
     `configurable_disk_size` and `requirements`.

dstack/_internal/core/backends/datacrunch/compute.py

@@ -1,4 +1,5 @@
-from typing import Callable, Dict, List, Optional
+from collections.abc import Iterable
+from typing import Dict, List, Optional
 
 from datacrunch import DataCrunchClient
 from datacrunch.exceptions import APIException
@@ -12,7 +13,11 @@ from dstack._internal.core.backends.base.compute import (
     generate_unique_instance_name,
     get_shim_commands,
 )
-from dstack._internal.core.backends.base.offers import get_catalog_offers, get_offers_disk_modifier
+from dstack._internal.core.backends.base.offers import (
+    OfferModifier,
+    get_catalog_offers,
+    get_offers_disk_modifier,
+)
 from dstack._internal.core.backends.datacrunch.models import DataCrunchConfig
 from dstack._internal.core.errors import NoCapacityError
 from dstack._internal.core.models.backends.base import BackendType
@@ -59,10 +64,8 @@ class DataCrunchCompute(
         offers_with_availability = self._get_offers_with_availability(offers)
         return offers_with_availability
 
-    def get_offers_modifier(
-        self, requirements: Requirements
-    ) -> Callable[[InstanceOfferWithAvailability], Optional[InstanceOfferWithAvailability]]:
-        return get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)
+    def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]:
+        return [get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)]
 
     def _get_offers_with_availability(
         self, offers: List[InstanceOffer]

dstack/_internal/core/backends/gcp/compute.py

@@ -1,7 +1,9 @@
 import concurrent.futures
 import json
+import re
 import threading
 from collections import defaultdict
+from collections.abc import Iterable
 from dataclasses import dataclass
 from typing import Callable, Dict, List, Literal, Optional, Tuple
 
@@ -24,6 +26,7 @@ from dstack._internal.core.backends.base.compute import (
     ComputeWithPlacementGroupSupport,
     ComputeWithPrivateGatewaySupport,
     ComputeWithPrivilegedSupport,
+    ComputeWithReservationSupport,
     ComputeWithVolumeSupport,
     generate_unique_gateway_instance_name,
     generate_unique_instance_name,
@@ -35,6 +38,7 @@ from dstack._internal.core.backends.base.compute import (
     requires_nvidia_proprietary_kernel_modules,
 )
 from dstack._internal.core.backends.base.offers import (
+    OfferModifier,
     get_catalog_offers,
     get_offers_disk_modifier,
 )
@@ -78,9 +82,16 @@ logger = get_logger(__name__)
 # pd-balanced disks can be 10GB-64TB, but dstack images are 20GB and cannot grow larger
 # than 32TB because of filesystem settings
 CONFIGURABLE_DISK_SIZE = Range[Memory](min=Memory.parse("20GB"), max=Memory.parse("32TB"))
+# Pattern from https://cloud.google.com/compute/docs/instances/reservations-consume#consuming_instances_from_a_specific_reservation
+RESERVATION_PATTERN = re.compile(
+    r"projects/(?P<project_id>[a-z0-9-]+)/reservations/(?P<reservation_name>[a-z0-9-]+)"
+)
+RESOURCE_NAME_PATTERN = re.compile(r"[a-z0-9-]+")
+TPU_VERSIONS = [tpu.name for tpu in KNOWN_TPUS]
 
 
-TPU_VERSIONS = [tpu.name for tpu in KNOWN_TPUS]
+class GCPOfferBackendData(CoreModel):
+    is_dws_calendar_mode: bool = False
 
 
 class GCPVolumeDiskBackendData(CoreModel):
@@ -93,6 +104,7 @@ class GCPCompute(
     ComputeWithCreateInstanceSupport,
     ComputeWithPrivilegedSupport,
     ComputeWithMultinodeSupport,
+    ComputeWithReservationSupport,
     ComputeWithPlacementGroupSupport,
     ComputeWithGatewaySupport,
     ComputeWithPrivateGatewaySupport,
@@ -113,8 +125,12 @@ class GCPCompute(
         self.resource_policies_client = compute_v1.ResourcePoliciesClient(
             credentials=self.credentials
         )
+        self.reservations_client = compute_v1.ReservationsClient(credentials=self.credentials)
         self._usable_subnets_cache_lock = threading.Lock()
         self._usable_subnets_cache = TTLCache(maxsize=1, ttl=120)
+        self._find_reservation_cache_lock = threading.Lock()
+        # smaller TTL, since we check the reservation's in_use_count, which can change often
+        self._find_reservation_cache = TTLCache(maxsize=8, ttl=20)
 
     def get_all_offers_with_availability(self) -> List[InstanceOfferWithAvailability]:
         regions = get_or_error(self.config.regions)
@@ -149,10 +165,57 @@ class GCPCompute(
             offers_with_availability[-1].region = region
         return offers_with_availability
 
-    def get_offers_modifier(
+    def get_offers_modifiers(self, requirements: Requirements) -> Iterable[OfferModifier]:
+        modifiers = []
+
+        if requirements.reservation:
+            zone_to_reservation = self._find_reservation(requirements.reservation)
+
+            def reservation_modifier(
+                offer: InstanceOfferWithAvailability,
+            ) -> Optional[InstanceOfferWithAvailability]:
+                if offer.instance.resources.spot:
+                    return None
+                assert offer.availability_zones is not None
+                matching_zones = []
+                zones_with_capacity = []
+                for zone in offer.availability_zones:
+                    reservation = zone_to_reservation.get(zone)
+                    if reservation is not None and _offer_matches_reservation(offer, reservation):
+                        matching_zones.append(zone)
+                        if _reservation_has_capacity(reservation):
+                            zones_with_capacity.append(zone)
+                if not matching_zones:
+                    return None
+                offer = offer.copy(deep=True)
+                if zones_with_capacity:
+                    offer.availability_zones = zones_with_capacity
+                else:
+                    offer.availability_zones = matching_zones
+                    offer.availability = InstanceAvailability.NOT_AVAILABLE
+                return offer
+
+            modifiers.append(reservation_modifier)
+
+        modifiers.append(get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements))
+        return modifiers
+
+    def get_offers_post_filter(
         self, requirements: Requirements
-    ) -> Callable[[InstanceOfferWithAvailability], Optional[InstanceOfferWithAvailability]]:
-        return get_offers_disk_modifier(CONFIGURABLE_DISK_SIZE, requirements)
+    ) -> Optional[Callable[[InstanceOfferWithAvailability], bool]]:
+        if requirements.reservation is None:
+
+            def reserved_offers_filter(offer: InstanceOfferWithAvailability) -> bool:
+                """Remove reserved-only offers"""
+                if GCPOfferBackendData.__response__.parse_obj(
+                    offer.backend_data
+                ).is_dws_calendar_mode:
+                    return False
+                return True
+
+            return reserved_offers_filter
+
+        return None
 
     def terminate_instance(
         self, instance_id: str, region: str, backend_data: Optional[str] = None
@@ -305,6 +368,16 @@ class GCPCompute(
         )
 
         for zone in zones:
+            reservation = None
+            if instance_config.reservation:
+                reservation = self._find_reservation(instance_config.reservation).get(zone)
+                if reservation is None:
+                    logger.warning(
+                        "Reservation %s no longer exists in zone %s",
+                        instance_config.reservation,
+                        zone,
+                    )
+                    continue
             request = compute_v1.InsertInstanceRequest()
             request.zone = zone
             request.project = self.config.project_id
@@ -335,6 +408,7 @@ class GCPCompute(
                 roce_subnetworks=roce_subnets,
                 allocate_public_ip=allocate_public_ip,
                 placement_policy=placement_policy,
+                reservation=reservation,
             )
             try:
                 # GCP needs some time to return an error in case of no capacity (< 30s).
@@ -475,6 +549,11 @@ class GCPCompute(
     ) -> bool:
         return placement_group.configuration.region == instance_offer.region
 
+    def are_placement_groups_compatible_with_reservations(self, backend_type: BackendType) -> bool:
+        # Cannot use our own placement policies when provisioning in a reservation.
+        # Instead, we use the placement policy defined in reservation settings.
+        return False
+
     def create_gateway(
         self,
         configuration: GatewayComputeConfiguration,
@@ -880,6 +959,26 @@ class GCPCompute(
             usable_subnets=self._list_usable_subnets(),
         )
 
+    @cachedmethod(
+        cache=lambda self: self._find_reservation_cache,
+        lock=lambda self: self._find_reservation_cache_lock,
+    )
+    def _find_reservation(self, configured_name: str) -> dict[str, compute_v1.Reservation]:
+        if match := RESERVATION_PATTERN.fullmatch(configured_name):
+            project_id = match.group("project_id")
+            name = match.group("reservation_name")
+        elif RESOURCE_NAME_PATTERN.fullmatch(configured_name):
+            project_id = self.config.project_id
+            name = configured_name
+        else:
+            # misconfigured or non-GCP
+            return {}
+        return gcp_resources.find_reservation(
+            reservations_client=self.reservations_client,
+            project_id=project_id,
+            name=name,
+        )
+
 
 def _supported_instances_and_zones(
     regions: List[str],
@@ -922,8 +1021,8 @@ def _has_gpu_quota(quotas: Dict[str, float], resources: Resources) -> bool:
     gpu = resources.gpus[0]
     if _is_tpu(gpu.name):
         return True
-    if gpu.name in ["B200", "H100"]:
-        # B200, H100 and H100_MEGA quotas are not returned by `regions_client.list`
+    if gpu.name in ["B200", "H100", "RTXPRO6000"]:
+        # B200, H100, H100_MEGA, and RTXPRO6000 quotas are not returned by `regions_client.list`
         return True
     quota_name = f"NVIDIA_{gpu.name}_GPUS"
     if gpu.name == "A100" and gpu.memory_mib == 80 * 1024:
@@ -933,6 +1032,52 @@ def _has_gpu_quota(quotas: Dict[str, float], resources: Resources) -> bool:
     return len(resources.gpus) <= quotas.get(quota_name, 0)
 
 
+def _offer_matches_reservation(
+    offer: InstanceOfferWithAvailability, reservation: compute_v1.Reservation
+) -> bool:
+    if (
+        reservation.specific_reservation is None
+        or reservation.specific_reservation.instance_properties is None
+    ):
+        return False
+    properties = reservation.specific_reservation.instance_properties
+    if properties.machine_type != offer.instance.name:
+        return False
+    accelerators = properties.guest_accelerators or []
+    if not accelerators and offer.instance.resources.gpus:
+        return False
+    if len(accelerators) > 1:
+        logger.warning(
+            "Expected 0 or 1 accelerator types per instance,"
+            f" but {properties.machine_type} has {len(accelerators)}."
+            f" Ignoring reservation {reservation.self_link}"
+        )
+        return False
+    if accelerators:
+        if accelerators[0].accelerator_count != len(offer.instance.resources.gpus):
+            return False
+        if (
+            offer.instance.resources.gpus
+            and gcp_resources.find_accelerator_name(
+                offer.instance.resources.gpus[0].name,
+                offer.instance.resources.gpus[0].memory_mib,
+            )
+            != accelerators[0].accelerator_type
+        ):
+            return False
+    return True
+
+
+def _reservation_has_capacity(reservation: compute_v1.Reservation) -> bool:
+    return (
+        reservation.specific_reservation is not None
+        and reservation.specific_reservation.in_use_count is not None
+        and reservation.specific_reservation.assured_count is not None
+        and reservation.specific_reservation.in_use_count
+        < reservation.specific_reservation.assured_count
+    )
+
+
 def _unique_instance_name(instance: InstanceType) -> str:
     if instance.resources.spot:
         name = f"{instance.name}-spot"

dstack/_internal/core/backends/gcp/models.py

@@ -89,6 +89,16 @@ class GCPBackendConfig(CoreModel):
             description="The tags (labels) that will be assigned to resources created by `dstack`"
         ),
     ] = None
+    preview_features: Annotated[
+        Optional[List[Literal["g4"]]],
+        Field(
+            description=(
+                "The list of preview GCP features to enable."
+                " There are currently no preview features"
+            ),
+            max_items=1,
+        ),
+    ] = None
 
 
 class GCPBackendConfigWithCreds(GCPBackendConfig):