From 03201985483ccc1aefc50a4e41cb149c03d16b29 Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Wed, 7 Oct 2026 11:00:33 +0500 Subject: [PATCH 1/3] Skip placement groups for AWS Capacity Blocks AWS rejects launches into Capacity Blocks that specify a placement group. Fixes #4344. --- .../_internal/core/backends/aws/compute.py | 34 +++++++++++ .../_internal/core/backends/aws/resources.py | 5 +- .../_internal/core/backends/base/compute.py | 16 +++-- .../_internal/core/backends/gcp/compute.py | 6 +- .../instances/cloud_provisioning.py | 8 ++- .../pipeline_tasks/jobs_submitted.py | 8 ++- .../_internal/server/services/placement.py | 14 +++++ .../test_instances/test_cloud_provisioning.py | 61 +++++++++++++++++++ .../pipeline_tasks/test_submitted_jobs.py | 60 ++++++++++++++++++ 9 files changed, 199 insertions(+), 13 deletions(-) diff --git a/src/dstack/_internal/core/backends/aws/compute.py b/src/dstack/_internal/core/backends/aws/compute.py index 9579b31ba6..96dd28af52 100644 --- a/src/dstack/_internal/core/backends/aws/compute.py +++ b/src/dstack/_internal/core/backends/aws/compute.py @@ -169,6 +169,7 @@ def __init__( self._subnets_availability_zones_cache = ComputeCache(cache=Cache(maxsize=100)) self._security_group_cache = ComputeTTLCache(cache=TTLCache(maxsize=100, ttl=600)) self._image_id_and_username_cache = ComputeTTLCache(cache=TTLCache(maxsize=100, ttl=600)) + self._is_capacity_block_cache = ComputeCache(cache=Cache(maxsize=100)) def get_all_offers_with_availability( self, unallocated_resources: bool @@ -528,6 +529,22 @@ def is_suitable_placement_group( return False return placement_group.configuration.region == instance_offer.region + def are_placement_groups_compatible_with_reservation( + self, + instance_offer: InstanceOffer, + reservation: str, + ) -> bool: + # AWS rejects launches into Capacity Blocks that specify a placement group. + # Capacity Block instances are already placed close together in EC2 UltraClusters. + try: + is_capacity_block = self._is_capacity_block( + region=instance_offer.region, reservation_id=reservation + ) + except botocore.exceptions.ClientError as e: + logger.warning("Failed to get reservation %s: %s", reservation, e) + return True + return not is_capacity_block + def create_gateway_replica( self, configuration: GatewayReplicaConfiguration, @@ -1231,6 +1248,23 @@ def _get_gateway_lb_subnets_ids( az_to_subnet_id_map = {az: subnet_id for subnet_id, az in subnet_id_to_az_map.items()} return list(az_to_subnet_id_map.values()) + def _is_capacity_block_cache_key(self, region: str, reservation_id: str) -> tuple: + return hashkey(region, reservation_id) + + @cachedmethod( + cache=lambda self: self._is_capacity_block_cache.cache, + key=_is_capacity_block_cache_key, + lock=lambda self: self._is_capacity_block_cache.lock, + ) + def _is_capacity_block(self, region: str, reservation_id: str) -> bool: + reservation = aws_resources.get_reservation( + ec2_client=self.session.client("ec2", region_name=region), + reservation_id=reservation_id, + is_capacity_block=True, + active_only=False, + ) + return reservation is not None + def get_vpc_id_subnets_ids_or_error( ec2_client: botocore.client.BaseClient, diff --git a/src/dstack/_internal/core/backends/aws/resources.py b/src/dstack/_internal/core/backends/aws/resources.py index 5ee3f63191..c9aff710b5 100644 --- a/src/dstack/_internal/core/backends/aws/resources.py +++ b/src/dstack/_internal/core/backends/aws/resources.py @@ -678,8 +678,11 @@ def get_reservation( instance_count: int = 0, instance_types: Optional[List[str]] = None, is_capacity_block: bool = False, + active_only: bool = True, ) -> Optional[Dict[str, Any]]: - filters = [{"Name": "state", "Values": ["active"]}] + filters = [] + if active_only: + filters.append({"Name": "state", "Values": ["active"]}) if instance_types: filters.append({"Name": "instance-type", "Values": instance_types}) try: diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index 9bf202d3f9..8dee44b2bd 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/src/dstack/_internal/core/backends/base/compute.py @@ -26,7 +26,6 @@ DSTACK_RUNNER_SSH_PORT, DSTACK_SHIM_HTTP_PORT, ) -from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.compute_groups import ComputeGroup, ComputeGroupProvisioningData from dstack._internal.core.models.gateways import ( GatewayLoadBalancerConfiguration, @@ -559,13 +558,20 @@ def is_suitable_placement_group( """ pass - def are_placement_groups_compatible_with_reservations(self, backend_type: BackendType) -> bool: + def are_placement_groups_compatible_with_reservation( + self, + instance_offer: InstanceOffer, + reservation: str, + ) -> bool: """ - Whether placement groups can be used for instances provisioned in reservations. + Whether a placement group can be used for an instance provisioned in the reservation. + + May perform API calls. Arguments: - backend_type: matches the backend type of this compute, unless this compute is a proxy - for other backends (dstack Sky) + instance_offer: the offer to provision. Its backend matches the backend type of this + compute, unless this compute is a proxy for other backends (dstack Sky) + reservation: the reservation to provision the instance in """ return True diff --git a/src/dstack/_internal/core/backends/gcp/compute.py b/src/dstack/_internal/core/backends/gcp/compute.py index 8ab25584f6..662a5455c3 100644 --- a/src/dstack/_internal/core/backends/gcp/compute.py +++ b/src/dstack/_internal/core/backends/gcp/compute.py @@ -594,7 +594,11 @@ def is_suitable_placement_group( ) -> bool: return placement_group.configuration.region == instance_offer.region - def are_placement_groups_compatible_with_reservations(self, backend_type: BackendType) -> bool: + def are_placement_groups_compatible_with_reservation( + self, + instance_offer: InstanceOffer, + reservation: str, + ) -> bool: # Cannot use our own placement policies when provisioning in a reservation. # Instead, we use the placement policy defined in reservation settings. return False diff --git a/src/dstack/_internal/server/background/pipeline_tasks/instances/cloud_provisioning.py b/src/dstack/_internal/server/background/pipeline_tasks/instances/cloud_provisioning.py index cb87d0294a..bd869b8325 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/instances/cloud_provisioning.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/instances/cloud_provisioning.py @@ -47,6 +47,7 @@ from dstack._internal.server.services.logging import fmt from dstack._internal.server.services.offers import get_instance_offer_with_restricted_az from dstack._internal.server.services.placement import ( + can_use_placement_groups, get_fleet_placement_group_models, placement_group_model_to_placement_group, placement_group_model_to_placement_group_optional, @@ -140,9 +141,10 @@ async def create_cloud_instance(instance_model: InstanceModel) -> ProcessResult: and cluster_context.is_current_instance_master and instance_offer.backend in BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT and isinstance(compute, ComputeWithPlacementGroupSupport) - and ( - compute.are_placement_groups_compatible_with_reservations(instance_offer.backend) - or instance_configuration.reservation is None + and await can_use_placement_groups( + compute=compute, + instance_offer=instance_offer, + reservation=instance_configuration.reservation, ) ): ( diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py index b0f592b0e6..23edeab065 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py @@ -122,6 +122,7 @@ ) from dstack._internal.server.services.pipelines import PipelineHinterProtocol from dstack._internal.server.services.placement import ( + can_use_placement_groups, find_or_create_suitable_placement_group, get_placement_group_model_for_job, placement_group_model_to_placement_group_optional, @@ -2457,9 +2458,10 @@ async def _provision_new_capacity( and is_cloud_cluster(fleet_model) and offer.backend in BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT and isinstance(compute, ComputeWithPlacementGroupSupport) - and ( - compute.are_placement_groups_compatible_with_reservations(offer.backend) - or job.job_spec.requirements.reservation is None + and await can_use_placement_groups( + compute=compute, + instance_offer=offer, + reservation=job.job_spec.requirements.reservation, ) ): placement_group_model = await find_or_create_suitable_placement_group( diff --git a/src/dstack/_internal/server/services/placement.py b/src/dstack/_internal/server/services/placement.py index 6544442fff..9d169c3299 100644 --- a/src/dstack/_internal/server/services/placement.py +++ b/src/dstack/_internal/server/services/placement.py @@ -121,6 +121,20 @@ def get_placement_group_model_for_job( return placement_group_model +async def can_use_placement_groups( + compute: ComputeWithPlacementGroupSupport, + instance_offer: InstanceOffer, + reservation: Optional[str], +) -> bool: + if reservation is None: + return True + return await run_async( + compute.are_placement_groups_compatible_with_reservation, + instance_offer, + reservation, + ) + + async def find_or_create_suitable_placement_group( fleet_model: FleetModel, placement_groups: list[PlacementGroupModel], diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_cloud_provisioning.py b/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_cloud_provisioning.py index c559931ac8..45f29c8976 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_cloud_provisioning.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_instances/test_cloud_provisioning.py @@ -31,6 +31,7 @@ create_project, get_fleet_configuration, get_fleet_spec, + get_instance_configuration, get_instance_offer_with_availability, get_job_provisioning_data, get_placement_group_provisioning_data, @@ -749,6 +750,66 @@ async def test_create_placement_group_if_placement_cluster( assert backend_mock.compute.return_value.create_placement_group.call_count == 0 assert len(placement_groups) == 0 + @pytest.mark.parametrize("compatible", [True, False]) + async def test_creates_placement_group_in_reservation_only_if_compatible( + self, + test_db, + session: AsyncSession, + worker: InstanceWorker, + compatible: bool, + ) -> None: + project = await create_project(session=session) + fleet = await create_fleet( + session, + project, + spec=get_fleet_spec( + conf=get_fleet_configuration( + placement=InstanceGroupPlacement.CLUSTER, + nodes=FleetNodesSpec(min=1, target=1, max=1), + ) + ), + ) + instance_configuration = get_instance_configuration() + instance_configuration.reservation = "test-reservation" + instance = await create_instance( + session=session, + project=project, + fleet=fleet, + status=InstanceStatus.PENDING, + offer=None, + job_provisioning_data=None, + instance_configuration=instance_configuration, + ) + await _set_current_master_instance(session, fleet, instance) + offer = get_instance_offer_with_availability() + backend_mock = Mock() + backend_mock.TYPE = BackendType.AWS + compute_mock = Mock(spec=ComputeMockSpec) + backend_mock.compute.return_value = compute_mock + compute_mock.get_offers.return_value = [offer] + compute_mock.are_placement_groups_compatible_with_reservation.return_value = compatible + compute_mock.create_instance.return_value = get_job_provisioning_data() + compute_mock.create_placement_group.return_value = get_placement_group_provisioning_data() + with patch("dstack._internal.server.services.backends.get_project_backends") as m: + m.return_value = [backend_mock] + await process_instance(session, worker, instance) + + await session.refresh(instance) + assert instance.status == InstanceStatus.PROVISIONING + compute_mock.are_placement_groups_compatible_with_reservation.assert_called_once_with( + offer, "test-reservation" + ) + placement_groups = (await session.execute(select(PlacementGroupModel))).scalars().all() + created_placement_group = compute_mock.create_instance.call_args[0][2] + if compatible: + assert compute_mock.create_placement_group.call_count == 1 + assert len(placement_groups) == 1 + assert isinstance(created_placement_group, PlacementGroup) + else: + assert compute_mock.create_placement_group.call_count == 0 + assert len(placement_groups) == 0 + assert created_placement_group is None + @pytest.mark.parametrize("can_reuse", [True, False]) async def test_reuses_placement_group_between_offers_if_the_group_is_suitable( self, diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py index e511241d7f..536df52cc1 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_submitted_jobs.py @@ -803,6 +803,66 @@ async def test_creates_placement_group_for_cluster_fleet( placement_group = (await session.execute(select(PlacementGroupModel))).scalar() assert placement_group is not None + @pytest.mark.parametrize("compatible", [True, False]) + async def test_creates_placement_group_in_reservation_only_if_compatible( + self, test_db, session: AsyncSession, worker: JobSubmittedWorker, compatible: bool + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + fleet_spec = get_fleet_spec() + fleet_spec.configuration.placement = InstanceGroupPlacement.CLUSTER + fleet_spec.configuration.nodes = FleetNodesSpec(min=0, target=0, max=None) + fleet = await create_fleet(session=session, project=project, spec=fleet_spec) + run_spec = get_run_spec(run_name="test-run", repo_id=repo.name) + run_spec.configuration.reservation = "test-reservation" + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + fleet=fleet, + run_name="test-run", + run_spec=run_spec, + ) + job = await create_job(session=session, run=run, instance_assigned=True) + offer = get_instance_offer_with_availability(backend=BackendType.AWS) + + with patch("dstack._internal.server.services.backends.get_project_backends") as m: + backend_mock = Mock() + compute_mock = Mock(spec=ComputeMockSpec) + backend_mock.TYPE = BackendType.AWS + backend_mock.compute.return_value = compute_mock + m.return_value = [backend_mock] + compute_mock.get_offers.return_value = [offer] + compute_mock.are_placement_groups_compatible_with_reservation.return_value = compatible + compute_mock.run_job.return_value = get_job_provisioning_data( + backend=BackendType.AWS, + ) + compute_mock.create_placement_group.return_value = ( + get_placement_group_provisioning_data() + ) + + await _process_job(session=session, worker=worker, job_model=job) + + await session.refresh(job) + assert job.status == JobStatus.PROVISIONING + compute_mock.are_placement_groups_compatible_with_reservation.assert_called_once() + assert ( + compute_mock.are_placement_groups_compatible_with_reservation.call_args[0][1] + == "test-reservation" + ) + compute_mock.run_job.assert_called_once() + placement_group = (await session.execute(select(PlacementGroupModel))).scalar() + if compatible: + compute_mock.create_placement_group.assert_called_once() + assert isinstance(compute_mock.run_job.call_args[0][6], PlacementGroup) + assert placement_group is not None + else: + compute_mock.create_placement_group.assert_not_called() + assert compute_mock.run_job.call_args[0][6] is None + assert placement_group is None + async def test_marks_unused_existing_placement_groups_for_cleanup( self, test_db, session: AsyncSession, worker: JobSubmittedWorker ): From a0475efee3c41abbce28048ef78bbe43cce5696a Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Wed, 7 Oct 2026 11:48:23 +0500 Subject: [PATCH 2/3] Drop AWS Capacity Block lookup cache and handle missing ReservationType AWS omits ReservationType for regular capacity reservations. --- .../_internal/core/backends/aws/compute.py | 27 +++++-------------- .../_internal/core/backends/aws/resources.py | 2 +- 2 files changed, 7 insertions(+), 22 deletions(-) diff --git a/src/dstack/_internal/core/backends/aws/compute.py b/src/dstack/_internal/core/backends/aws/compute.py index 96dd28af52..b5e9af0bdc 100644 --- a/src/dstack/_internal/core/backends/aws/compute.py +++ b/src/dstack/_internal/core/backends/aws/compute.py @@ -169,7 +169,6 @@ def __init__( self._subnets_availability_zones_cache = ComputeCache(cache=Cache(maxsize=100)) self._security_group_cache = ComputeTTLCache(cache=TTLCache(maxsize=100, ttl=600)) self._image_id_and_username_cache = ComputeTTLCache(cache=TTLCache(maxsize=100, ttl=600)) - self._is_capacity_block_cache = ComputeCache(cache=Cache(maxsize=100)) def get_all_offers_with_availability( self, unallocated_resources: bool @@ -537,13 +536,16 @@ def are_placement_groups_compatible_with_reservation( # AWS rejects launches into Capacity Blocks that specify a placement group. # Capacity Block instances are already placed close together in EC2 UltraClusters. try: - is_capacity_block = self._is_capacity_block( - region=instance_offer.region, reservation_id=reservation + capacity_block = aws_resources.get_reservation( + ec2_client=self.session.client("ec2", region_name=instance_offer.region), + reservation_id=reservation, + is_capacity_block=True, + active_only=False, ) except botocore.exceptions.ClientError as e: logger.warning("Failed to get reservation %s: %s", reservation, e) return True - return not is_capacity_block + return capacity_block is None def create_gateway_replica( self, @@ -1248,23 +1250,6 @@ def _get_gateway_lb_subnets_ids( az_to_subnet_id_map = {az: subnet_id for subnet_id, az in subnet_id_to_az_map.items()} return list(az_to_subnet_id_map.values()) - def _is_capacity_block_cache_key(self, region: str, reservation_id: str) -> tuple: - return hashkey(region, reservation_id) - - @cachedmethod( - cache=lambda self: self._is_capacity_block_cache.cache, - key=_is_capacity_block_cache_key, - lock=lambda self: self._is_capacity_block_cache.lock, - ) - def _is_capacity_block(self, region: str, reservation_id: str) -> bool: - reservation = aws_resources.get_reservation( - ec2_client=self.session.client("ec2", region_name=region), - reservation_id=reservation_id, - is_capacity_block=True, - active_only=False, - ) - return reservation is not None - def get_vpc_id_subnets_ids_or_error( ec2_client: botocore.client.BaseClient, diff --git a/src/dstack/_internal/core/backends/aws/resources.py b/src/dstack/_internal/core/backends/aws/resources.py index c9aff710b5..aa0648fc46 100644 --- a/src/dstack/_internal/core/backends/aws/resources.py +++ b/src/dstack/_internal/core/backends/aws/resources.py @@ -710,7 +710,7 @@ def get_reservation( if instance_count > 0 and reservation["AvailableInstanceCount"] < instance_count: return None - if is_capacity_block and reservation["ReservationType"] != "capacity-block": + if is_capacity_block and reservation.get("ReservationType") != "capacity-block": return None return reservation From dd668b4a03dd3d68d6ab019b6c0ddaa22850762e Mon Sep 17 00:00:00 2001 From: Victor Skvortsov Date: Wed, 7 Oct 2026 12:11:39 +0500 Subject: [PATCH 3/3] Use fleet reservation when provisioning runs --- src/dstack/_internal/core/backends/base/compute.py | 2 +- .../server/background/pipeline_tasks/jobs_submitted.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/dstack/_internal/core/backends/base/compute.py b/src/dstack/_internal/core/backends/base/compute.py index 8dee44b2bd..2a952aff7e 100644 --- a/src/dstack/_internal/core/backends/base/compute.py +++ b/src/dstack/_internal/core/backends/base/compute.py @@ -414,7 +414,7 @@ def run_job( user=run.user, ssh_keys=[SSHKey(public=project_ssh_public_key.strip())], volumes=volumes, - reservation=job.job_spec.requirements.reservation, + reservation=requirements.reservation, tags=run.run_spec.merged_profile.tags, ) instance_offer = instance_offer.model_copy() diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py index 23edeab065..1ab7875d04 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_submitted.py @@ -2461,7 +2461,7 @@ async def _provision_new_capacity( and await can_use_placement_groups( compute=compute, instance_offer=offer, - reservation=job.job_spec.requirements.reservation, + reservation=requirements.reservation, ) ): placement_group_model = await find_or_create_suitable_placement_group(