Add missing fields to launch template validators

This patch adds the missing fields to the launch template validators.
For the head node, these are the `BlockDeviceMappings` and `KeyName`,
while for the compute resources these are the `BlockDeviceMappings`,
`InstanceMetadataOptions` and `CapacityReservationSpecification`.

Author: ermanno

cli/src/pcluster/config/cluster_config.py

@@ -1112,6 +1112,10 @@ def __init__(self, name: str, capacity_type: str = None):
         _capacity_type = CapacityType[capacity_type.upper()] if capacity_type else None
         self.capacity_type = Resource.init_param(_capacity_type, default=CapacityType.ONDEMAND)
 
+    def is_spot(self):
+        """Return True if the queue has SPOT capacity."""
+        return self.capacity_type == CapacityType.SPOT
+
     def _register_validators(self, context: ValidatorContext = None):  # noqa: D102 #pylint: disable=unused-argument
         self._register_validator(NameValidator, name=self.name)
 
@@ -1193,6 +1197,7 @@ def _register_validators(self, context: ValidatorContext = None):  # noqa: D102
         self._register_validator(
             HeadNodeLaunchTemplateValidator,
             head_node=self.head_node,
+            os=self.image.os,
             ami_id=self.head_node_ami,
             tags=self.get_cluster_tags(),
         )
@@ -2470,6 +2475,7 @@ def _register_validators(self, context: ValidatorContext = None):
                 ComputeResourceLaunchTemplateValidator,
                 queue=queue,
                 ami_id=queue_image,
+                os=self.image.os,
                 tags=self.get_cluster_tags(),
             )
             ami_volume_size = AWSApi.instance().ec2.describe_image(queue_image).volume_size

cli/src/pcluster/launch_template_utils.py

@@ -0,0 +1,56 @@
+from abc import ABC, abstractmethod
+
+from pcluster.constants import OS_MAPPING
+
+
+class _LaunchTemplateBuilder(ABC):
+    """Abstract class with methods with the common logic for launch template builders."""
+
+    def get_block_device_mappings(self, root_volume, image_os):
+        """Return a list of block device mappings."""
+        block_device_mappings = []
+        for _, (device_name_index, virtual_name_index) in enumerate(zip(list(map(chr, range(97, 121))), range(0, 24))):
+            device_name = "/dev/xvdb{0}".format(device_name_index)
+            virtual_name = "ephemeral{0}".format(virtual_name_index)
+            block_device_mappings.append(self._block_device_mapping_for_virt(device_name, virtual_name))
+
+        block_device_mappings.append(
+            self._block_device_mapping_for_ebs(OS_MAPPING[image_os]["root-device"], root_volume)
+        )
+        return block_device_mappings
+
+    def get_instance_market_options(self, queue, compute_resource):
+        """Return the instance market options for spot instances."""
+        instance_market_options = None
+        if queue.is_spot():
+            instance_market_options = self._instance_market_option(
+                market_type="spot",
+                spot_instance_type="one-time",
+                instance_interruption_behavior="terminate",
+                max_price=None if compute_resource.spot_price is None else str(compute_resource.spot_price),
+            )
+        return instance_market_options
+
+    def get_capacity_reservation(self, queue, compute_resource):
+        """Return the capacity reservation if a target is defined at the queue or compute resource level."""
+        capacity_reservation = None
+        cr_target = compute_resource.capacity_reservation_target or queue.capacity_reservation_target
+        if cr_target:
+            capacity_reservation = self._capacity_reservation(cr_target)
+        return capacity_reservation
+
+    @abstractmethod
+    def _block_device_mapping_for_virt(self, device_name, virtual_name):
+        pass
+
+    @abstractmethod
+    def _block_device_mapping_for_ebs(self, device_name, volume):
+        pass
+
+    @abstractmethod
+    def _instance_market_option(self, market_type, spot_instance_type, instance_interruption_behavior, max_price):
+        pass
+
+    @abstractmethod
+    def _capacity_reservation(self, cr_target):
+        pass

cli/src/pcluster/templates/cdk_builder_utils.py

@@ -26,7 +26,6 @@
     BaseComputeResource,
     BaseQueue,
     HeadNode,
-    LocalStorage,
     SharedStorageType,
     SlurmClusterConfig,
     SlurmQueue,
@@ -35,11 +34,11 @@
     COOKBOOK_PACKAGES_VERSIONS,
     CW_LOGS_RETENTION_DAYS_DEFAULT,
     IAM_ROLE_PATH,
-    OS_MAPPING,
     PCLUSTER_CLUSTER_NAME_TAG,
     PCLUSTER_DYNAMODB_PREFIX,
     PCLUSTER_NODE_TYPE_TAG,
 )
+from pcluster.launch_template_utils import _LaunchTemplateBuilder
 from pcluster.models.s3_bucket import S3Bucket, parse_bucket_url
 from pcluster.utils import (
     get_attr,
@@ -52,34 +51,6 @@
 PCLUSTER_LAMBDA_PREFIX = "pcluster-"
 
 
-def get_block_device_mappings(local_storage: LocalStorage, os: str):
-    """Return block device mapping."""
-    block_device_mappings = []
-    for _, (device_name_index, virtual_name_index) in enumerate(zip(list(map(chr, range(97, 121))), range(0, 24))):
-        device_name = "/dev/xvdb{0}".format(device_name_index)
-        virtual_name = "ephemeral{0}".format(virtual_name_index)
-        block_device_mappings.append(
-            ec2.CfnLaunchTemplate.BlockDeviceMappingProperty(device_name=device_name, virtual_name=virtual_name)
-        )
-
-    root_volume = local_storage.root_volume
-
-    block_device_mappings.append(
-        ec2.CfnLaunchTemplate.BlockDeviceMappingProperty(
-            device_name=OS_MAPPING[os]["root-device"],
-            ebs=ec2.CfnLaunchTemplate.EbsProperty(
-                volume_size=root_volume.size,
-                encrypted=root_volume.encrypted,
-                volume_type=root_volume.volume_type,
-                iops=root_volume.iops,
-                throughput=root_volume.throughput,
-                delete_on_termination=root_volume.delete_on_termination,
-            ),
-        )
-    )
-    return block_device_mappings
-
-
 def create_hash_suffix(string_to_hash: str):
     """Create 16digit hash string."""
     return (
@@ -814,3 +785,41 @@ def _stack_unique_id(self):
 
     def _format_arn(self, **kwargs):
         return Stack.of(self).format_arn(**kwargs)
+
+
+class CdkLaunchTemplateBuilder(_LaunchTemplateBuilder):
+    """Concrete class for building a CDK launch template."""
+
+    def _block_device_mapping_for_ebs(self, device_name, volume):
+        return ec2.CfnLaunchTemplate.BlockDeviceMappingProperty(
+            device_name=device_name,
+            ebs=ec2.CfnLaunchTemplate.EbsProperty(
+                volume_size=volume.size,
+                encrypted=volume.encrypted,
+                volume_type=volume.volume_type,
+                iops=volume.iops,
+                throughput=volume.throughput,
+                delete_on_termination=volume.delete_on_termination,
+            ),
+        )
+
+    def _block_device_mapping_for_virt(self, device_name, virtual_name):
+        return ec2.CfnLaunchTemplate.BlockDeviceMappingProperty(device_name=device_name, virtual_name=virtual_name)
+
+    def _instance_market_option(self, market_type, spot_instance_type, instance_interruption_behavior, max_price):
+        return ec2.CfnLaunchTemplate.InstanceMarketOptionsProperty(
+            market_type=market_type,
+            spot_options=ec2.CfnLaunchTemplate.SpotOptionsProperty(
+                spot_instance_type=spot_instance_type,
+                instance_interruption_behavior=instance_interruption_behavior,
+                max_price=max_price,
+            ),
+        )
+
+    def _capacity_reservation(self, cr_target):
+        return ec2.CfnLaunchTemplate.CapacityReservationSpecificationProperty(
+            capacity_reservation_target=ec2.CfnLaunchTemplate.CapacityReservationTargetProperty(
+                capacity_reservation_id=cr_target.capacity_reservation_id,
+                capacity_reservation_resource_group_arn=cr_target.capacity_reservation_resource_group_arn,
+            )
+        )

cli/src/pcluster/templates/cluster_stack.py

@@ -45,7 +45,6 @@
 from pcluster.config.cluster_config import (
     AwsBatchClusterConfig,
     BaseSharedFsx,
-    CapacityType,
     ExistingFsxOntap,
     ExistingFsxOpenZfs,
     SchedulerPluginQueue,
@@ -73,6 +72,7 @@
 from pcluster.models.s3_bucket import S3Bucket
 from pcluster.templates.awsbatch_builder import AwsBatchConstruct
 from pcluster.templates.cdk_builder_utils import (
+    CdkLaunchTemplateBuilder,
     ComputeNodeIamResources,
     HeadNodeIamResources,
     PclusterLambdaConstruct,
@@ -81,7 +81,6 @@
     convert_deletion_policy,
     create_hash_suffix,
     generate_launch_template_version_cfn_parameter_hash,
-    get_block_device_mappings,
     get_cloud_watch_logs_policy_statement,
     get_cloud_watch_logs_retention_days,
     get_common_user_data_env,
@@ -119,6 +118,7 @@ def __init__(
     ) -> None:
         super().__init__(scope, construct_id, **kwargs)
         self._stack_name = stack_name
+        self._launch_template_builder = CdkLaunchTemplateBuilder()
         self.config = cluster_config
         self.bucket = bucket
         self.timestamp = datetime.utcnow().strftime("%Y%m%d%H%M%S")
@@ -874,7 +874,9 @@ def _add_head_node(self):
             "HeadNodeLaunchTemplate",
             launch_template_data=ec2.CfnLaunchTemplate.LaunchTemplateDataProperty(
                 instance_type=head_node.instance_type,
-                block_device_mappings=get_block_device_mappings(head_node.local_storage, self.config.image.os),
+                block_device_mappings=self._launch_template_builder.get_block_device_mappings(
+                    head_node.local_storage.root_volume, self.config.image.os
+                ),
                 key_name=head_node.ssh.key_name,
                 network_interfaces=head_lt_nw_interfaces,
                 image_id=self.config.head_node_ami,
@@ -1317,6 +1319,7 @@ def __init__(
         self._dynamodb_table = dynamodb_table
         self._compute_node_instance_profiles = compute_node_instance_profiles
         self._head_eni = head_eni
+        self._launch_template_builder = CdkLaunchTemplateBuilder()
         self._add_resources()
 
     # -- Utility methods --------------------------------------------------------------------------------------------- #
@@ -1443,41 +1446,19 @@ def _add_compute_resource_launch_template(
                 )
             )
 
-        instance_market_options = None
-        if queue.capacity_type == CapacityType.SPOT:
-            instance_market_options = ec2.CfnLaunchTemplate.InstanceMarketOptionsProperty(
-                market_type="spot",
-                spot_options=ec2.CfnLaunchTemplate.SpotOptionsProperty(
-                    spot_instance_type="one-time",
-                    instance_interruption_behavior="terminate",
-                    max_price=None if compute_resource.spot_price is None else str(compute_resource.spot_price),
-                ),
-            )
-
         conditional_template_properties = {}
-
         if compute_resource.is_ebs_optimized:
             conditional_template_properties.update({"ebs_optimized": True})
         if isinstance(compute_resource, SlurmComputeResource):
             conditional_template_properties.update({"instance_type": compute_resource.instance_type})
 
-        capacity_reservation_specification = None
-        cr_target = compute_resource.capacity_reservation_target or queue.capacity_reservation_target
-        if cr_target:
-            capacity_reservation_specification = ec2.CfnLaunchTemplate.CapacityReservationSpecificationProperty(
-                capacity_reservation_target=ec2.CfnLaunchTemplate.CapacityReservationTargetProperty(
-                    capacity_reservation_id=cr_target.capacity_reservation_id,
-                    capacity_reservation_resource_group_arn=cr_target.capacity_reservation_resource_group_arn,
-                )
-            )
-
         return ec2.CfnLaunchTemplate(
             self,
             f"LaunchTemplate{create_hash_suffix(queue.name + compute_resource.name)}",
             launch_template_name=f"{self.stack_name}-{queue.name}-{compute_resource.name}",
             launch_template_data=ec2.CfnLaunchTemplate.LaunchTemplateDataProperty(
-                block_device_mappings=get_block_device_mappings(
-                    queue.compute_settings.local_storage, self._config.image.os
+                block_device_mappings=self._launch_template_builder.get_block_device_mappings(
+                    queue.compute_settings.local_storage.root_volume, self._config.image.os
                 ),
                 # key_name=,
                 network_interfaces=compute_lt_nw_interfaces,
@@ -1486,9 +1467,14 @@ def _add_compute_resource_launch_template(
                 iam_instance_profile=ec2.CfnLaunchTemplate.IamInstanceProfileProperty(
                     name=instance_profiles[queue.name]
                 ),
-                instance_market_options=instance_market_options,
+                instance_market_options=self._launch_template_builder.get_instance_market_options(
+                    queue, compute_resource
+                ),
                 instance_initiated_shutdown_behavior="terminate",
-                capacity_reservation_specification=capacity_reservation_specification,
+                capacity_reservation_specification=self._launch_template_builder.get_capacity_reservation(
+                    queue,
+                    compute_resource,
+                ),
                 metadata_options=ec2.CfnLaunchTemplate.MetadataOptionsProperty(
                     http_tokens=get_http_tokens_setting(self._config.imds.imds_support)
                 ),

cli/src/pcluster/utils.py

@@ -361,3 +361,8 @@ def yaml_no_duplicates_constructor(loader, node, deep=False):
 def get_http_tokens_setting(imds_support):
     """Get http tokens settings for supported IMDS version."""
     return "required" if imds_support == "v2.0" else "optional"
+
+
+def remove_none_values(original_dictionary):
+    """Return a dictionary without entries with None value."""
+    return {key: value for key, value in original_dictionary.items() if value is not None}