AWS Capacity Reservation support

docs/docs/reference/server/config.yml.md

@@ -69,12 +69,15 @@ There are two ways to configure AWS: using an access key or using the default cr
                     "ec2:AttachVolume",
                     "ec2:AuthorizeSecurityGroupEgress",
                     "ec2:AuthorizeSecurityGroupIngress",
+                    "ec2:CreatePlacementGroup",
                     "ec2:CancelSpotInstanceRequests",
                     "ec2:CreateSecurityGroup",
                     "ec2:CreateTags",
                     "ec2:CreateVolume",
+                    "ec2:DeletePlacementGroup",
                     "ec2:DeleteVolume",
                     "ec2:DescribeAvailabilityZones",
+                    "ec2:DescribeCapacityReservations"
                     "ec2:DescribeImages",
                     "ec2:DescribeInstances",
                     "ec2:DescribeInstanceAttribute",

src/dstack/_internal/cli/services/configurators/fleet.py

@@ -270,6 +270,7 @@ def th(s: str) -> str:
     fleet_type = "cloud"
     nodes = plan.spec.configuration.nodes or "-"
     placement = plan.spec.configuration.placement or InstanceGroupPlacement.ANY
+    reservation = plan.spec.configuration.reservation
     backends = None
     if plan.spec.configuration.backends is not None:
         backends = ", ".join(b.value for b in plan.spec.configuration.backends)
@@ -297,6 +298,8 @@ def th(s: str) -> str:
         configuration_table.add_row(th("Resources"), resources)
     if spot_policy is not None:
         configuration_table.add_row(th("Spot policy"), spot_policy)
+    if reservation is not None:
+        configuration_table.add_row(th("Reservation"), reservation)
 
     offers_table = Table(box=None)
     offers_table.add_column("#")

src/dstack/_internal/cli/utils/fleet.py

@@ -19,6 +19,8 @@ def get_fleets_table(
 ) -> Table:
     table = Table(box=None)
     table.add_column("FLEET", no_wrap=True)
+    if verbose:
+        table.add_column("RESERVATION")
     table.add_column("INSTANCE")
     table.add_column("BACKEND")
     table.add_column("RESOURCES")
@@ -61,6 +63,12 @@ def get_fleets_table(
                 format_date(instance.created),
             ]
 
+            if verbose and i == 0:
+                row.insert(
+                    1,
+                    fleet.spec.configuration.reservation if i == 0 else "",
+                )
+
             if verbose:
                 error = ""
                 if instance.status == InstanceStatus.TERMINATED and instance.termination_reason:

src/dstack/_internal/cli/utils/run.py

@@ -62,6 +62,7 @@ def th(s: str) -> str:
     props.add_row(th("Creation policy"), creation_policy)
     props.add_row(th("Termination policy"), termination_policy)
     props.add_row(th("Termination idle time"), termination_idle_time)
+    props.add_row(th("Reservation"), run_plan.run_spec.configuration.reservation)
 
     offers = Table(box=None)
     offers.add_column("#")
@@ -121,6 +122,8 @@ def get_runs_table(
     if verbose:
         table.add_column("INSTANCE", no_wrap=True)
     table.add_column("RESOURCES")
+    if verbose:
+        table.add_column("RESERVATION", no_wrap=True)
     table.add_column("PRICE", no_wrap=True)
     table.add_column("STATUS", no_wrap=True)
     table.add_column("SUBMITTED", style="grey58", no_wrap=True)
@@ -154,6 +157,7 @@ def get_runs_table(
                         "BACKEND": f"{jpd.backend.value.replace('remote', 'ssh')} ({jpd.region})",
                         "INSTANCE": jpd.instance_type.name,
                         "RESOURCES": jpd.instance_type.resources.pretty_format(include_spot=True),
+                        "RESERVATION": jpd.reservation,
                         "PRICE": f"${jpd.price:.4}",
                     }
                 )

src/dstack/_internal/core/backends/__init__.py

@@ -21,6 +21,10 @@
 BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT = [
     BackendType.AWS,
 ]
+BACKENDS_WITH_RESERVATION_SUPPORT = [
+    BackendType.AWS,
+]
+
 BACKENDS_WITH_GATEWAY_SUPPORT = [
     BackendType.AWS,
     BackendType.AZURE,

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

@@ -72,12 +72,40 @@ def __init__(self, config: AWSConfig):
     def get_offers(
         self, requirements: Optional[Requirements] = None
     ) -> List[InstanceOfferWithAvailability]:
+        filter = _supported_instances
+        if requirements and requirements.reservation:
+            region_to_reservation = {}
+            for region in self.config.regions:
+                reservation = aws_resources.get_reservation(
+                    ec2_client=self.session.client("ec2", region_name=region),
+                    reservation_id=requirements.reservation,
+                    instance_count=1,
+                )
+                if reservation is not None:
+                    region_to_reservation[region] = reservation
+
+            def _supported_instances_with_reservation(offer: InstanceOffer) -> bool:
+                # Filter: only instance types supported by dstack
+                if not _supported_instances(offer):
+                    return False
+                # Filter: Spot instances can't be used with reservations
+                if offer.instance.resources.spot:
+                    return False
+                region = offer.region
+                reservation = region_to_reservation.get(region)
+                # Filter: only instance types matching the capacity reservation
+                if not bool(reservation and offer.instance.name == reservation["InstanceType"]):
+                    return False
+                return True
+
+            filter = _supported_instances_with_reservation
+
         offers = get_catalog_offers(
             backend=BackendType.AWS,
             locations=self.config.regions,
             requirements=requirements,
             configurable_disk_size=CONFIGURABLE_DISK_SIZE,
-            extra_filter=_supported_instances,
+            extra_filter=filter,
         )
         regions = set(i.region for i in offers)
 
@@ -148,6 +176,7 @@ def create_instance(
             ec2_client=ec2_client, instance_type=instance_offer.instance.name
         )
         enable_efa = max_efa_interfaces > 0
+        is_capacity_block = False
         try:
             vpc_id, subnet_ids = get_vpc_id_subnet_id_or_error(
                 ec2_client=ec2_client,
@@ -160,6 +189,22 @@ def create_instance(
                 ec2_client=ec2_client,
                 subnet_ids=subnet_ids,
             )
+            if instance_config.reservation:
+                reservation = aws_resources.get_reservation(
+                    ec2_client=ec2_client,
+                    reservation_id=instance_config.reservation,
+                    instance_count=1,
+                )
+                if reservation is not None:
+                    # Filter out az different from capacity reservation
+                    subnet_id_to_az_map = {
+                        k: v
+                        for k, v in subnet_id_to_az_map.items()
+                        if v == reservation["AvailabilityZone"]
+                    }
+                    if reservation.get("ReservationType") == "capacity-block":
+                        is_capacity_block = True
+
         except botocore.exceptions.ClientError as e:
             logger.warning("Got botocore.exceptions.ClientError: %s", e)
             raise NoCapacityError()
@@ -193,6 +238,8 @@ def create_instance(
                         allocate_public_ip=allocate_public_ip,
                         placement_group_name=instance_config.placement_group_name,
                         enable_efa=enable_efa,
+                        reservation_id=instance_config.reservation,
+                        is_capacity_block=is_capacity_block,
                     )
                 )
                 instance = response[0]
@@ -212,6 +259,7 @@ def create_instance(
                     internal_ip=instance.private_ip_address,
                     region=instance_offer.region,
                     availability_zone=az,
+                    reservation=instance.capacity_reservation_id,
                     price=instance_offer.price,
                     username=username,
                     ssh_port=22,
@@ -240,6 +288,7 @@ def run_job(
                 SSHKey(public=project_ssh_public_key.strip()),
             ],
             user=run.user,
+            reservation=run.run_spec.configuration.reservation,
         )
         if len(volumes) > 0:
             volume = volumes[0]
@@ -706,6 +755,7 @@ def _supported_instances(offer: InstanceOffer) -> bool:
         "g4dn.",
         "g5.",
         "g6.",
+        "g6e.",
         "gr6.",
         "p3.",
         "p4d.",

src/dstack/_internal/core/backends/aws/resources.py

@@ -140,6 +140,8 @@ def create_instances_struct(
     allocate_public_ip: bool = True,
     placement_group_name: Optional[str] = None,
     enable_efa: bool = False,
+    reservation_id: Optional[str] = None,
+    is_capacity_block: bool = False,
 ) -> Dict[str, Any]:
     struct: Dict[str, Any] = dict(
         BlockDeviceMappings=[
@@ -173,6 +175,9 @@ def create_instances_struct(
                 "InstanceInterruptionBehavior": "terminate",
             },
         }
+
+    if is_capacity_block:
+        struct["InstanceMarketOptions"] = {"MarketType": "capacity-block"}
     if enable_efa and not subnet_id:
         raise ComputeError("EFA requires subnet")
     # AWS allows specifying either NetworkInterfaces for specific subnet_id
@@ -205,6 +210,11 @@ def create_instances_struct(
             "GroupName": placement_group_name,
         }
 
+    if reservation_id is not None:
+        struct["CapacityReservationSpecification"] = {
+            "CapacityReservationTarget": {"CapacityReservationId": reservation_id}
+        }
+
     return struct
 
 
@@ -596,3 +606,44 @@ def _is_private_subnet_with_internet_egress(
                     return True
 
     return False
+
+
+def get_reservation(
+    ec2_client: botocore.client.BaseClient,
+    reservation_id: str,
+    instance_count: int = 0,
+    instance_types: Optional[List[str]] = None,
+    is_capacity_block: bool = False,
+) -> Optional[Dict[str, Any]]:
+    filters = [{"Name": "state", "Values": ["active"]}]
+    if instance_types:
+        filters.append({"Name": "instance-type", "Values": instance_types})
+    try:
+        response = ec2_client.describe_capacity_reservations(
+            CapacityReservationIds=[reservation_id], Filters=filters
+        )
+    except botocore.exceptions.ParamValidationError as e:
+        logger.debug(
+            "Skipping reservation %s. Parameter validation error: %s", reservation_id, repr(e)
+        )
+        return None
+    except botocore.exceptions.ClientError as e:
+        error_code = e.response.get("Error", {}).get("Code")
+        if error_code == "InvalidCapacityReservationId.Malformed":
+            logger.debug("Skipping reservation %s. Malformed ID.", reservation_id)
+            return None
+        if error_code == "InvalidCapacityReservationId.NotFound":
+            logger.debug(
+                "Skipping reservation %s. Capacity Reservation not found.", reservation_id
+            )
+            return None
+        raise
+    reservation = response["CapacityReservations"][0]
+
+    if instance_count > 0 and reservation["AvailableInstanceCount"] < instance_count:
+        return None
+
+    if is_capacity_block and reservation["ReservationType"] != "capacity-block":
+        return None
+
+    return reservation

src/dstack/_internal/core/models/fleets.py

@@ -128,6 +128,10 @@ class InstanceGroupParams(CoreModel):
         Optional[InstanceGroupPlacement],
         Field(description="The placement of instances: `any` or `cluster`"),
     ] = None
+    reservation: Annotated[
+        Optional[str],
+        Field(description="The existing reservation for the instances"),
+    ] = None
     resources: Annotated[
         Optional[ResourcesSpec],
         Field(description="The resources requirements"),

src/dstack/_internal/core/models/instances.py

@@ -99,6 +99,7 @@ class InstanceConfiguration(CoreModel):
     instance_id: Optional[str] = None
     availability_zone: Optional[str] = None
     placement_group_name: Optional[str] = None
+    reservation: Optional[str] = None
 
     def get_public_keys(self) -> List[str]:
         return [ssh_key.public.strip() for ssh_key in self.ssh_keys]

src/dstack/_internal/core/models/profiles.py

@@ -109,6 +109,10 @@ class ProfileParams(CoreModel):
             description="The cloud-specific instance types to consider for provisioning (e.g., `[p3.8xlarge, n1-standard-4]`)"
         ),
     ]
+    reservation: Annotated[
+        Optional[str],
+        Field(description="The existing reservation for the instances"),
+    ]
     spot_policy: Annotated[
         Optional[SpotPolicy],
         Field(

src/dstack/_internal/core/models/runs.py

@@ -149,6 +149,7 @@ class Requirements(CoreModel):
     resources: ResourcesSpec
     max_price: Optional[float]
     spot: Optional[bool]
+    reservation: Optional[str]
 
     def pretty_format(self, resources_only: bool = False):
         res = self.resources.pretty_format()
@@ -211,6 +212,7 @@ class JobProvisioningData(CoreModel):
     instance_network: Optional[str] = None
     region: str
     availability_zone: Optional[str] = None
+    reservation: Optional[str] = None
     price: float
     username: str
     # ssh_port be different from 22 for some backends.

src/dstack/_internal/server/background/tasks/process_submitted_jobs.py

@@ -451,6 +451,7 @@ def _get_or_create_fleet_model_for_job(
         configuration=FleetConfiguration(
             name=run.run_spec.run_name,
             placement=placement,
+            reservation=run.run_spec.configuration.reservation,
         ),
         profile=run.run_spec.merged_profile,
         autocreated=True,

src/dstack/_internal/server/services/fleets.py

@@ -350,6 +350,7 @@ async def create_fleet(
                     pool=pool,
                     spec=spec,
                     placement_group_name=placement_group_name,
+                    reservation=spec.configuration.reservation,
                     instance_num=i,
                 )
                 fleet_model.instances.append(instance_model)
@@ -364,6 +365,7 @@ async def create_fleet_instance_model(
     pool: PoolModel,
     spec: FleetSpec,
     placement_group_name: Optional[str],
+    reservation: Optional[str],
     instance_num: int,
 ) -> InstanceModel:
     profile = spec.merged_profile
@@ -378,6 +380,7 @@ async def create_fleet_instance_model(
         instance_name=f"{spec.configuration.name}-{instance_num}",
         instance_num=instance_num,
         placement_group_name=placement_group_name,
+        reservation=reservation,
     )
     return instance_model
 
@@ -738,6 +741,7 @@ def _get_fleet_requirements(fleet_spec: FleetSpec) -> Requirements:
         resources=fleet_spec.configuration.resources or ResourcesSpec(),
         max_price=profile.max_price,
         spot=get_policy_map(profile.spot_policy, default=SpotPolicy.ONDEMAND),
+        reservation=fleet_spec.configuration.reservation,
     )
     return requirements
 

src/dstack/_internal/server/services/jobs/configurators/base.py

@@ -187,6 +187,7 @@ def _requirements(self) -> Requirements:
             resources=self.run_spec.configuration.resources,
             max_price=self.run_spec.merged_profile.max_price,
             spot=None if spot_policy == SpotPolicy.AUTO else (spot_policy == SpotPolicy.SPOT),
+            reservation=self.run_spec.merged_profile.reservation,
         )
 
     def _retry(self) -> Optional[Retry]:

src/dstack/_internal/server/services/offers.py

@@ -3,6 +3,7 @@
 from dstack._internal.core.backends import (
     BACKENDS_WITH_CREATE_INSTANCE_SUPPORT,
     BACKENDS_WITH_MULTINODE_SUPPORT,
+    BACKENDS_WITH_RESERVATION_SUPPORT,
 )
 from dstack._internal.core.backends.base import Backend
 from dstack._internal.core.models.backends.base import BackendType
@@ -53,6 +54,11 @@ async def get_offers_by_requirements(
             backend_types = BACKENDS_WITH_CREATE_INSTANCE_SUPPORT
         backend_types = [b for b in backend_types if b in BACKENDS_WITH_CREATE_INSTANCE_SUPPORT]
 
+    if profile.reservation is not None:
+        if not backend_types:
+            backend_types = BACKENDS_WITH_RESERVATION_SUPPORT
+        backend_types = [b for b in backend_types if b in BACKENDS_WITH_RESERVATION_SUPPORT]
+
     # For multi-node, restrict backend and region.
     # The default behavior is to provision all nodes in the same backend and region.
     if master_job_provisioning_data is not None: