Add validation to make sure the ODCR has the right instance type and availability zone

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

Author: hanwen-cluster

cli/src/pcluster/aws/aws_api.py

@@ -20,6 +20,7 @@
 from pcluster.aws.imagebuilder import ImageBuilderClient
 from pcluster.aws.kms import KmsClient
 from pcluster.aws.logs import LogsClient
+from pcluster.aws.resource_groups import ResourceGroupsClient
 from pcluster.aws.route53 import Route53Client
 from pcluster.aws.s3 import S3Client
 from pcluster.aws.s3_resource import S3Resource
@@ -57,6 +58,7 @@ def __init__(self):
         self._logs = None
         self._route53 = None
         self._secretsmanager = None
+        self._resource_groups = None
 
     @property
     def cfn(self):
@@ -163,6 +165,13 @@ def secretsmanager(self):
             self._secretsmanager = SecretsManagerClient()
         return self._secretsmanager
 
+    @property
+    def resource_groups(self):
+        """Resource Groups client."""
+        if not self._resource_groups:
+            self._resource_groups = ResourceGroupsClient()
+        return self._resource_groups
+
     @staticmethod
     def instance():
         """Return the singleton AWSApi instance."""

cli/src/pcluster/aws/ec2.py

@@ -35,6 +35,7 @@ def __init__(self):
         self.additional_instance_types_data = {}
         self.security_groups_cache = {}
         self.subnets_cache = {}
+        self.capacity_reservations_cache = {}
 
     @AWSExceptionHandler.handle_client_exception
     @Cache.cached
@@ -83,6 +84,30 @@ def describe_subnets(self, subnet_ids):
                 result.append(subnet)
         return result
 
+    @AWSExceptionHandler.handle_client_exception
+    def describe_capacity_reservations(self, capacity_reservation_ids):
+        """Return a list of Capacity Reservations."""
+        result = []
+        missed_capacity_reservations = []
+        for capacity_reservation_id in capacity_reservation_ids:
+            cached_data = self.capacity_reservations_cache.get(capacity_reservation_id)
+            if cached_data:
+                result.append(cached_data)
+            else:
+                missed_capacity_reservations.append(capacity_reservation_id)
+        if missed_capacity_reservations:
+            response = list(
+                self._paginate_results(
+                    self._client.describe_capacity_reservations, CapacityReservationIds=missed_capacity_reservations
+                )
+            )
+            for capacity_reservation in response:
+                self.capacity_reservations_cache[
+                    capacity_reservation.get("CapacityReservationId")
+                ] = capacity_reservation
+                result.append(capacity_reservation)
+        return result
+
     @AWSExceptionHandler.handle_client_exception
     @Cache.cached
     def get_subnet_avail_zone(self, subnet_id):

cli/src/pcluster/aws/resource_groups.py

@@ -0,0 +1,31 @@
+# Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions and
+# limitations under the License.
+import re
+
+from pcluster.aws.common import AWSExceptionHandler, Boto3Client, Cache
+
+
+class ResourceGroupsClient(Boto3Client):
+    """Implement Resource Groups Boto3 client."""
+
+    def __init__(self):
+        super().__init__("resource-groups")
+
+    @AWSExceptionHandler.handle_client_exception
+    @Cache.cached
+    def get_capacity_reservation_ids_from_group_resources(self, group):
+        """Return a list of capacity reservation ids."""
+        capacity_reservation_ids = []
+        resources = self._client.list_group_resources(Group=group)["Resources"]
+        for resource in resources:
+            if resource["Identifier"]["ResourceType"] == "AWS::EC2::CapacityReservation":
+                capacity_reservation_ids.append(re.match("(.*)(cr-.*)", resource["Identifier"]["ResourceArn"]).group(2))
+        return capacity_reservation_ids

cli/src/pcluster/config/cluster_config.py

@@ -114,6 +114,8 @@
 )
 from pcluster.validators.ec2_validators import (
     AmiOsCompatibleValidator,
+    CapacityReservationResourceGroupValidator,
+    CapacityReservationValidator,
     CapacityTypeValidator,
     InstanceTypeBaseAMICompatibleValidator,
     InstanceTypeMemoryInfoValidator,
@@ -2439,6 +2441,8 @@ def __init__(self, cluster_name: str, scheduling: SlurmScheduling, **kwargs):
         super().__init__(cluster_name, **kwargs)
         self.scheduling = scheduling
         self.__image_dict = None
+        # Cache capacity reservations information together to reduce number of boto3 calls
+        AWSApi.instance().ec2.describe_capacity_reservations(self.all_relevant_capacity_reservation_ids)
 
     def get_instance_types_data(self):
         """Get instance type infos for all instance types used in the configuration file."""
@@ -2477,6 +2481,27 @@ def _register_validators(self):
                             instance_type=instance_type,
                             instance_type_data=instance_types_data[compute_resource.instance_type],
                         )
+                    self._register_validator(
+                        InstanceTypeMemoryInfoValidator,
+                        instance_type=compute_resource.instance_type,
+                        instance_type_data=instance_types_data[compute_resource.instance_type],
+                    )
+                # The validation below has to be in cluster config class instead of queue class
+                # to make sure the subnet APIs are cached by previous validations.
+                if compute_resource.capacity_reservation_target:
+                    cr_target = compute_resource.capacity_reservation_target
+                    self._register_validator(
+                        CapacityReservationValidator,
+                        capacity_reservation_id=cr_target.capacity_reservation_id,
+                        instance_type=compute_resource.instance_type,
+                        subnet=queue.networking.subnet_ids[0],
+                    )
+                    self._register_validator(
+                        CapacityReservationResourceGroupValidator,
+                        capacity_reservation_resource_group_arn=cr_target.capacity_reservation_resource_group_arn,
+                        instance_type=compute_resource.instance_type,
+                        subnet=queue.networking.subnet_ids[0],
+                    )
 
     @property
     def image_dict(self):
@@ -2519,3 +2544,15 @@ def capacity_reservation_resource_group_arns(self):
             if capacity_reservation_target.capacity_reservation_resource_group_arn:
                 result.add(capacity_reservation_target.capacity_reservation_resource_group_arn)
         return list(result)
+
+    @property
+    def all_relevant_capacity_reservation_ids(self):
+        """Return a list of capacity reservation ids specified in the config or used by resource groups."""
+        capacity_reservation_ids = set(self.capacity_reservation_ids)
+        for capacity_reservation_resource_group_arn in self.capacity_reservation_resource_group_arns:
+            capacity_reservation_ids.update(
+                AWSApi.instance().resource_groups.get_capacity_reservation_ids_from_group_resources(
+                    capacity_reservation_resource_group_arn
+                )
+            )
+        return list(capacity_reservation_ids)

cli/src/pcluster/validators/ec2_validators.py

@@ -165,3 +165,49 @@ def _validate(self, os: str, image_id: str):
                 f"they are compatible before cluster creation and update operations.",
                 FailureLevel.WARNING,
             )
+
+
+class CapacityReservationValidator(Validator):
+    """Validate capacity reservation can be used with the instance type and subnet."""
+
+    def _validate(self, capacity_reservation_id: str, instance_type: str, subnet: str):
+        if capacity_reservation_id:
+            capacity_reservation = AWSApi.instance().ec2.describe_capacity_reservations([capacity_reservation_id])[0]
+            if capacity_reservation["InstanceType"] != instance_type:
+                self._add_failure(
+                    f"Capacity reservation {capacity_reservation_id} must has the same instance type "
+                    f"as {instance_type}.",
+                    FailureLevel.ERROR,
+                )
+            if capacity_reservation["AvailabilityZone"] != AWSApi.instance().ec2.get_subnet_avail_zone(subnet):
+                self._add_failure(
+                    f"Capacity reservation {capacity_reservation_id} must use the same availability zone "
+                    f"as subnet {subnet}.",
+                    FailureLevel.ERROR,
+                )
+
+
+class CapacityReservationResourceGroupValidator(Validator):
+    """Validate at least one capacity reservation in the resource group can be used with the instance and subnet."""
+
+    def _validate(self, capacity_reservation_resource_group_arn: str, instance_type: str, subnet: str):
+        if capacity_reservation_resource_group_arn:
+            capacity_reservation_ids = (
+                AWSApi.instance().resource_groups.get_capacity_reservation_ids_from_group_resources(
+                    capacity_reservation_resource_group_arn
+                )
+            )
+            capacity_reservations = AWSApi.instance().ec2.describe_capacity_reservations(capacity_reservation_ids)
+            found_qualified_capacity_reservation = False
+            for capacity_reservation in capacity_reservations:
+                if capacity_reservation["InstanceType"] == instance_type and capacity_reservation[
+                    "AvailabilityZone"
+                ] == AWSApi.instance().ec2.get_subnet_avail_zone(subnet):
+                    found_qualified_capacity_reservation = True
+                    break
+            if not found_qualified_capacity_reservation:
+                self._add_failure(
+                    f"Capacity reservation resource group {capacity_reservation_resource_group_arn} must have at least "
+                    f"one capacity reservation for {instance_type} in the same availability zone as subnet {subnet}.",
+                    FailureLevel.ERROR,
+                )

cli/tests/pcluster/aws/dummy_aws_api.py

@@ -20,6 +20,7 @@
 from pcluster.aws.imagebuilder import ImageBuilderClient
 from pcluster.aws.kms import KmsClient
 from pcluster.aws.logs import LogsClient
+from pcluster.aws.resource_groups import ResourceGroupsClient
 from pcluster.aws.route53 import Route53Client
 from pcluster.aws.s3 import S3Client
 from pcluster.aws.s3_resource import S3Resource
@@ -100,6 +101,7 @@ def __init__(self):
         self._logs = _DummyLogsClient()
         self._ddb_resource = _DummyDynamoResource()
         self._route53 = _DummyRoute53Client()
+        self._resource_groups = _DummyResourceGroupsClient()
 
 
 class _DummyCfnClient(CfnClient):
@@ -111,7 +113,12 @@ def __init__(self):
 class _DummyEc2Client(Ec2Client):
     def __init__(self):
         """Override Parent constructor. No real boto3 client is created."""
-        pass
+        self.capacity_reservations_cache = {
+            "cr-54321fcdbd5971234": {"InstanceType": "t2.micro", "AvailabilityZone": "string"},
+            "cr-321456cdbd597f551": {"InstanceType": "t2.micro", "AvailabilityZone": "string"},
+            "cr-123": {"InstanceType": "t2.micro", "AvailabilityZone": "string"},
+            "cr-234": {"InstanceType": "t2.micro", "AvailabilityZone": "string"},
+        }
 
     def get_official_image_id(self, os, architecture, filters=None):
         return "dummy-ami-id"
@@ -263,6 +270,16 @@ def __init__(self):
         pass
 
 
+class _DummyResourceGroupsClient(ResourceGroupsClient):
+    def __init__(self):
+        """Override Parent constructor. No real boto3 client is created."""
+        pass
+
+    def get_capacity_reservation_ids_from_group_resources(self, group):
+        """Return a list of capacity reservation ids."""
+        return ["cr-123", "cr-234"]
+
+
 def mock_aws_api(mocker, mock_instance_type_info=True):
     """Mock AWS Api."""
     mocker.patch("pcluster.aws.aws_api.AWSApi.instance", return_value=_DummyAWSApi())

cli/tests/pcluster/aws/test_ec2.py

@@ -368,6 +368,53 @@ def test_describe_subnets_cache(boto3_stubber):
     assert_that(AWSApi.instance().ec2.describe_subnets([subnet])[0]["State"]).is_equal_to("available")
 
 
+def get_describe_capacity_reservation_mocked_request(capacity_reservations, state):
+    return MockedBoto3Request(
+        method="describe_capacity_reservations",
+        response={
+            "CapacityReservations": [
+                {"CapacityReservationId": capacity_reservation, "State": state}
+                for capacity_reservation in capacity_reservations
+            ]
+        },
+        expected_params={"CapacityReservationIds": capacity_reservations},
+    )
+
+
+def test_describe_capacity_reservations_cache(boto3_stubber):
+    # First boto3 call. Nothing has been cached
+    capacity_reservation = "cr-123"
+    additional_capacity_reservation = "cr-234"
+    # The first mocked request and the third are about the same cr. However, the state of the cr changes
+    # from pending to available. The second mocked request is about another cr
+    mocked_requests = [
+        get_describe_capacity_reservation_mocked_request([capacity_reservation], "pending"),
+        get_describe_capacity_reservation_mocked_request([additional_capacity_reservation], "pending"),
+        get_describe_capacity_reservation_mocked_request([capacity_reservation], "active"),
+    ]
+    boto3_stubber("ec2", mocked_requests)
+    assert_that(AWSApi.instance().ec2.describe_capacity_reservations([capacity_reservation])[0]["State"]).is_equal_to(
+        "pending"
+    )
+
+    # Second boto3 call with more subnets. The cr already cached should not be included in the boto3 call.
+    response = AWSApi.instance().ec2.describe_capacity_reservations(
+        [capacity_reservation, additional_capacity_reservation]
+    )
+    assert_that(response).is_length(2)
+
+    # Third boto3 call. The result should be from cache even if the state of the cr is different
+    assert_that(AWSApi.instance().ec2.describe_capacity_reservations([capacity_reservation])[0]["State"]).is_equal_to(
+        "pending"
+    )
+
+    # Fourth boto3 call after resetting the AWSApi instance. The latest cr state should be retrieved from boto3
+    AWSApi.reset()
+    assert_that(AWSApi.instance().ec2.describe_capacity_reservations([capacity_reservation])[0]["State"]).is_equal_to(
+        "active"
+    )
+
+
 def get_describe_security_groups_mocked_request(security_groups, ip_permissions):
     return MockedBoto3Request(
         method="describe_security_groups",

cli/tests/pcluster/validators/test_all_validators/test_slurm_all_validators_are_called/slurm_2.yaml

@@ -20,6 +20,10 @@ Scheduling:
       ComputeResources:
         - Name: compute_resource1
           InstanceType: c5.xlarge
+        - Name: compute_resource2
+          InstanceType: t2.micro
+          CapacityReservationTarget:
+            CapacityReservationId: cr-54321fcdbd5971234
 SharedStorage:
   - MountDir: /my/mount/point2
     Name: name1