dstack 0.19.32__py3-none-any.whl → 0.19.33__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,7 @@ SUPPORTED_GPUHUNT_FLAGS = [
     "oci-spot",
     "lambda-arm",
     "gcp-a4",
+    "gcp-g4-preview",
 ]
 
 
@@ -199,9 +200,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,8 +82,11 @@ 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]
 
 
@@ -93,6 +100,7 @@ class GCPCompute(
     ComputeWithCreateInstanceSupport,
     ComputeWithPrivilegedSupport,
     ComputeWithMultinodeSupport,
+    ComputeWithReservationSupport,
     ComputeWithPlacementGroupSupport,
     ComputeWithGatewaySupport,
     ComputeWithPrivateGatewaySupport,
@@ -113,8 +121,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)
@@ -130,13 +142,19 @@ class GCPCompute(
         offer_keys_to_offers = {}
         offers_with_availability = []
         for offer in offers:
+            preview = False
+            if offer.instance.name.startswith("g4-standard-"):
+                if self.config.preview_features and "g4" in self.config.preview_features:
+                    preview = True
+                else:
+                    continue
             region = offer.region[:-2]  # strip zone
             key = (_unique_instance_name(offer.instance), region)
             if key in offer_keys_to_offers:
                 offer_keys_to_offers[key].availability_zones.append(offer.region)
                 continue
             availability = InstanceAvailability.NO_QUOTA
-            if _has_gpu_quota(quotas[region], offer.instance.resources):
+            if preview or _has_gpu_quota(quotas[region], offer.instance.resources):
                 availability = InstanceAvailability.UNKNOWN
             # todo quotas: cpu, memory, global gpu, tpu
             offer_with_availability = InstanceOfferWithAvailability(
@@ -149,10 +167,40 @@ class GCPCompute(
             offers_with_availability[-1].region = region
         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]:
+        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 terminate_instance(
         self, instance_id: str, region: str, backend_data: Optional[str] = None
@@ -305,6 +353,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 +393,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 +534,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 +944,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],
@@ -933,6 +1017,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,13 @@ 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. Supported values: `g4`"),
+            max_items=1,
+        ),
+    ] = None
 
 
 class GCPBackendConfigWithCreds(GCPBackendConfig):

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

@@ -26,9 +26,35 @@ supported_accelerators = [
     {"accelerator_name": "nvidia-tesla-t4", "gpu_name": "T4", "memory_mb": 1024 * 16},
     {"accelerator_name": "nvidia-tesla-v100", "gpu_name": "V100", "memory_mb": 1024 * 16},
     {"accelerator_name": "nvidia-tesla-p100", "gpu_name": "P100", "memory_mb": 1024 * 16},
+    {"accelerator_name": "nvidia-rtx-pro-6000", "gpu_name": "RTXPRO6000", "memory_mb": 1024 * 96},
 ]
 
 
+def find_accelerator_name(gpu_name: str, memory_mib: int) -> Optional[str]:
+    for acc in supported_accelerators:
+        if gpu_name == acc["gpu_name"] and memory_mib == acc["memory_mb"]:
+            return acc["accelerator_name"]
+    return None
+
+
+def sanitize_filter_value(value: str) -> str:
+    """
+    Escape characters that could break the Compute Engine API filter string.
+    """
+    return value.replace("\\", "\\\\").replace('"', '\\"')
+
+
+def get_resource_project(resource_url: str) -> str:
+    """
+    Extract the project ID from a URL like
+    https://www.googleapis.com/compute/v1/projects/proj-id/zones/us-central1-a/instances/vm-name
+    """
+    matches = re.findall(r"/projects/(?P<project_id>[a-z0-9-]+)/", resource_url)
+    if not matches:
+        raise BackendError(f"Invalid resource URL {resource_url}")
+    return matches[0]
+
+
 def get_availability_zones(
     regions_client: compute_v1.RegionsClient,
     project_id: str,
@@ -123,6 +149,7 @@ def create_instance_struct(
     roce_subnetworks: Optional[List[Tuple[str, str]]] = None,
     allocate_public_ip: bool = True,
     placement_policy: Optional[str] = None,
+    reservation: Optional[compute_v1.Reservation] = None,
 ) -> compute_v1.Instance:
     instance = compute_v1.Instance()
     instance.name = instance_name
@@ -147,6 +174,25 @@ def create_instance_struct(
         initialize_params.disk_type = f"zones/{zone}/diskTypes/hyperdisk-balanced"
     disk.initialize_params = initialize_params
     instance.disks = [disk]
+    if (
+        reservation is not None
+        and reservation.specific_reservation is not None
+        and reservation.specific_reservation.instance_properties is not None
+        and reservation.specific_reservation.instance_properties.local_ssds is not None
+    ):
+        for local_ssd in reservation.specific_reservation.instance_properties.local_ssds:
+            instance.disks.append(
+                compute_v1.AttachedDisk(
+                    auto_delete=True,
+                    boot=False,
+                    type_="SCRATCH",
+                    initialize_params=compute_v1.AttachedDiskInitializeParams(
+                        disk_type=f"zones/{zone}/diskTypes/local-ssd",
+                        disk_size_gb=local_ssd.disk_size_gb,
+                    ),
+                    interface=local_ssd.interface,
+                )
+            )
 
     if accelerators:
         instance.guest_accelerators = accelerators
@@ -162,6 +208,8 @@ def create_instance_struct(
 
     if placement_policy is not None:
         instance.resource_policies = [placement_policy]
+    elif reservation is not None and "placement" in reservation.resource_policies:
+        instance.resource_policies = [reservation.resource_policies["placement"]]
 
     if spot:
         instance.scheduling = compute_v1.Scheduling()
@@ -187,6 +235,17 @@ def create_instance_struct(
             )
         ]
 
+    if reservation is not None:
+        reservation_project = get_resource_project(reservation.self_link)
+        instance.reservation_affinity = compute_v1.ReservationAffinity()
+        instance.reservation_affinity.consume_reservation_type = (
+            compute_v1.ReservationAffinity.ConsumeReservationType.SPECIFIC_RESERVATION.name
+        )
+        instance.reservation_affinity.key = "compute.googleapis.com/reservation-name"
+        instance.reservation_affinity.values = [
+            f"projects/{reservation_project}/reservations/{reservation.name}"
+        ]
+
     return instance
 
 
@@ -350,11 +409,8 @@ def get_accelerators(
         return []
     accelerator_config = compute_v1.AcceleratorConfig()
     accelerator_config.accelerator_count = len(gpus)
-    for acc in supported_accelerators:
-        if gpus[0].name == acc["gpu_name"] and gpus[0].memory_mib == acc["memory_mb"]:
-            accelerator_name = acc["accelerator_name"]
-            break
-    else:
+    accelerator_name = find_accelerator_name(gpus[0].name, gpus[0].memory_mib)
+    if accelerator_name is None:
         raise ValueError(f"Unsupported GPU: {gpus[0].name} {gpus[0].memory_mib} MiB")
     accelerator_config.accelerator_type = (
         f"projects/{project_id}/zones/{zone}/acceleratorTypes/{accelerator_name}"
@@ -362,6 +418,31 @@ def get_accelerators(
     return [accelerator_config]
 
 
+def find_reservation(
+    reservations_client: compute_v1.ReservationsClient,
+    project_id: str,
+    name: str,
+) -> dict[str, compute_v1.Reservation]:
+    request = compute_v1.AggregatedListReservationsRequest(
+        project=project_id,
+        filter=(
+            f'(name = "{sanitize_filter_value(name)}")'
+            ' AND (status = "READY")'
+            " AND (specificReservationRequired = true)"
+        ),
+    )
+    try:
+        aggregated_reservations = reservations_client.aggregated_list(request=request)
+    except (google.api_core.exceptions.NotFound, google.api_core.exceptions.Forbidden) as e:
+        logger.warning("Could not find reservation: %s", e)
+        return {}
+    zone_to_reservation = {}
+    for zone, zone_reservations in aggregated_reservations:
+        if zone_reservations.reservations:
+            zone_to_reservation[zone.split("/")[-1]] = zone_reservations.reservations[0]
+    return zone_to_reservation
+
+
 def filter_invalid_labels(labels: Dict[str, str]) -> Dict[str, str]:
     filtered_labels = {}
     for k, v in labels.items():
@@ -499,5 +580,6 @@ def instance_type_supports_persistent_disk(instance_type_name: str) -> bool:
             "h3-",
             "v6e",
             "a4-",
+            "g4-",
         ]
     )