ODCR iam

Signed-off-by: Hanwen <[email protected]>

Author: hanwen-cluster

cli/src/pcluster/config/cluster_config.py

@@ -2482,3 +2482,13 @@ def image_dict(self):
             self.__image_dict[queue.name] = queue.queue_ami or self.image.custom_ami or self.official_ami
 
         return self.__image_dict
+
+    @property
+    def capacity_reservation_targets(self):
+        """Return a list of capacity reservation targets from all compute resources with the section."""
+        capacity_reservation_targets_list = []
+        for queue in self.scheduling.queues:
+            for compute_resource in queue.compute_resources:
+                if compute_resource.capacity_reservation_target:
+                    capacity_reservation_targets_list.append(compute_resource.capacity_reservation_target)
+        return capacity_reservation_targets_list

cli/src/pcluster/schemas/cluster_schema.py

@@ -951,6 +951,19 @@ def make_resource(self, data, **kwargs):
         """Generate resource."""
         return CapacityReservationTarget(**data)
 
+    @validates_schema
+    def no_coexist_instance_type_flexibility(self, data, **kwargs):
+        """Validate that 'capacity_reservation_id' and 'capacity_reservation_resource_group_arn' do not co-exist."""
+        if self.fields_coexist(
+            data,
+            ["capacity_reservation_id", "capacity_reservation_resource_group_arn"],
+            one_required=True,
+            **kwargs,
+        ):
+            raise ValidationError(
+                "A Capacity Reservation Target needs to specify either Capacity Reservation ID or Capacity Reservation Resource Group ARN."
+            )
+
 
 class ClusterDevSettingsSchema(BaseDevSettingsSchema):
     """Represent the schema of Dev Setting."""

cli/src/pcluster/templates/cdk_builder_utils.py

@@ -654,6 +654,31 @@ def _build_policy(self) -> List[iam.PolicyStatement]:
                             ]
                         )
 
+        if self._config.scheduling.scheduler == "slurm":
+            if self._config.capacity_reservation_targets:
+                for capacity_reservation_target in self._config.capacity_reservation_targets:
+                    if capacity_reservation_target.capacity_reservation_id:
+                        policy.append(
+                            iam.PolicyStatement(
+                                actions=["ec2:RunInstances"],
+                                effect=iam.Effect.ALLOW,
+                                resources=[
+                                    self._format_arn(
+                                        service="ec2",
+                                        resource=f"capacity-reservation/{capacity_reservation_target.capacity_reservation_id}",
+                                    )
+                                ],
+                            )
+                        )
+                    elif capacity_reservation_target.capacity_reservation_resource_group_arn:
+                        policy.append(
+                            iam.PolicyStatement(
+                                actions=["ec2:RunInstances"],
+                                effect=iam.Effect.ALLOW,
+                                resources=[self.capacity_reservation_target.capacity_reservation_resource_group_arn],
+                            )
+                        )
+
         if self._config.directory_service:
             policy.append(
                 iam.PolicyStatement(