Set the default root volume size according to AMIs being used

1. Set the default root volume sizes of all nodes to the size of snapshot in the AMI used for each node respectively. The old behavior was to have a constant 35GB default
2. Users can overwrite the defaults by the `Size` parameter under `RootVolume`
3. Add a validator to make sure the `Size` specified by the user is equal or greater than the size of the snapshot in the AMI.

For example:
```
Region: us-east-1
Image:
  Os: alinux2
HeadNode:
  ...
Scheduling:
  Scheduler: slurm
  SlurmQueues:
  - Name: queue1
    ...
  - Name: queue2
    Image:
      CustomAmi: ami-00f25fb4d0adb70a5
    ...
```
where `ami-00f25fb4d0adb70a5` is a deep learning AMI. The head node and compute nodes in queue1 will have 35GB of root volume, the size of official ParallelCluster AMI. The compute nodes in queue2 will have 167GB of root volume, the size of deep learning AMI

Signed-off-by: Hanwen <hanwenli@amazon.com>

Author: hanwen-cluster

CHANGELOG.md

@@ -30,6 +30,7 @@ x.x.x
 - Add `parallelcluster:cluster-name` tag to all resources created by ParallelCluster.
 - Do not allow setting `PlacementGroup/Id` when `PlacementGroup/Enabled` is explicitly set to `false`.
 - Restrict IPv6 access to IMDS to root and cluster admin users only.
+- Change the default root volume size from 35 GiB to the size of AMIs. The default can be overwritten in cluster configuration file.
 
 **BUG FIXES**
 - Fix default for disable validate and test components when building custom AMI. The default was to disable those components, but it wasn't effective.

cli/src/pcluster/config/cluster_config.py

@@ -90,6 +90,7 @@
     NumberOfStorageValidator,
     OverlappingMountDirValidator,
     RegionValidator,
+    RootVolumeSizeValidator,
     SchedulableMemoryValidator,
     SchedulerOsValidator,
     SharedStorageMountDirValidator,
@@ -159,15 +160,13 @@ class Ebs(Resource):
 
     def __init__(
         self,
-        size: int = None,
         encrypted: bool = None,
         volume_type: str = None,
         iops: int = None,
         throughput: int = None,
         **kwargs,
     ):
         super().__init__(**kwargs)
-        self.size = Resource.init_param(size, default=EBS_VOLUME_SIZE_DEFAULT)
         self.encrypted = Resource.init_param(encrypted, default=True)
         self.volume_type = Resource.init_param(volume_type, default=EBS_VOLUME_TYPE_DEFAULT)
         self.iops = Resource.init_param(iops, default=EBS_VOLUME_TYPE_IOPS_DEFAULT.get(self.volume_type))
@@ -192,8 +191,9 @@ def _register_validators(self):
 class RootVolume(Ebs):
     """Represent the root volume configuration."""
 
-    def __init__(self, delete_on_termination: bool = None, **kwargs):
+    def __init__(self, size: int = None, delete_on_termination: bool = None, **kwargs):
         super().__init__(**kwargs)
+        self.size = Resource.init_param(size)
         # The default delete_on_termination takes effect both on head and compute nodes.
         # If the default of the head node is to be changed, please separate this class for different defaults.
         self.delete_on_termination = Resource.init_param(delete_on_termination, default=True)
@@ -241,6 +241,7 @@ def __init__(
         self,
         mount_dir: str,
         name: str,
+        size: int = None,
         kms_key_id: str = None,
         snapshot_id: str = None,
         volume_id: str = None,
@@ -249,6 +250,7 @@ def __init__(
         **kwargs,
     ):
         super().__init__(**kwargs)
+        self.size = Resource.init_param(size, default=EBS_VOLUME_SIZE_DEFAULT)
         self.kms_key_id = Resource.init_param(kms_key_id)
         self.mount_dir = Resource.init_param(mount_dir)
         self.name = Resource.init_param(name)
@@ -1178,6 +1180,11 @@ def _register_validators(self):
         self._register_validator(
             HeadNodeImdsValidator, imds_secured=self.head_node.imds.secured, scheduler=self.scheduling.scheduler
         )
+        self._register_validator(
+            RootVolumeSizeValidator,
+            root_volume_size=self.head_node.local_storage.root_volume.size,
+            ami_id=self.head_node_ami,
+        )
 
     def _register_storage_validators(self):
         if self.shared_storage:

cli/src/pcluster/schemas/cluster_schema.py

@@ -108,7 +108,6 @@
 from pcluster.constants import (
     DELETION_POLICIES,
     DELETION_POLICIES_WITH_SNAPSHOT,
-    EBS_VOLUME_SIZE_DEFAULT,
     FSX_LUSTRE,
     FSX_ONTAP,
     FSX_OPENZFS,
@@ -169,14 +168,6 @@ def make_resource(self, data, **kwargs):
         """Generate resource."""
         return RootVolume(**data)
 
-    @validates("size")
-    def validate_size(self, value):
-        """Validate the size of root volume."""
-        if value < EBS_VOLUME_SIZE_DEFAULT:
-            raise ValidationError(
-                f"Root volume size {value} is invalid. It must be at least {EBS_VOLUME_SIZE_DEFAULT}."
-            )
-
 
 class QueueRootVolumeSchema(BaseSchema):
     """Represent the RootVolume schema for the queue."""
@@ -192,14 +183,6 @@ def make_resource(self, data, **kwargs):
         """Generate resource."""
         return RootVolume(**data)
 
-    @validates("size")
-    def validate_size(self, value):
-        """Validate the size of root volume."""
-        if value < EBS_VOLUME_SIZE_DEFAULT:
-            raise ValidationError(
-                f"Root volume size {value} is invalid. It must be at least {EBS_VOLUME_SIZE_DEFAULT}."
-            )
-
 
 class RaidSchema(BaseSchema):
     """Represent the schema of the parameters specific to Raid. It is a child of EBS schema."""

cli/src/pcluster/validators/cluster_validators.py

@@ -1078,6 +1078,20 @@ def _test_compute_resource(
         )
 
 
+class RootVolumeSizeValidator(Validator):
+    """Verify the root volume size is equal or greater to the size of the snapshot of the AMI."""
+
+    def _validate(self, root_volume_size, ami_id):
+        if root_volume_size:
+            ami_volume_size = AWSApi.instance().ec2.describe_image(ami_id).volume_size
+            if root_volume_size < ami_volume_size:
+                self._add_failure(
+                    f"Root volume size {root_volume_size} GiB must be equal or greater than the volume size of "
+                    f"the AMI {ami_id}: {ami_volume_size} GiB.",
+                    FailureLevel.ERROR,
+                )
+
+
 class HostedZoneValidator(Validator):
     """Validate custom private domain in the same VPC as head node."""
 

cli/src/pcluster/validators/ebs_validators.py

@@ -45,7 +45,7 @@ class EbsVolumeTypeSizeValidator(Validator):
     """
 
     def _validate(self, volume_type: str, volume_size: int):
-        if volume_type in EBS_VOLUME_TYPE_TO_VOLUME_SIZE_BOUNDS:
+        if volume_size is not None and volume_type in EBS_VOLUME_TYPE_TO_VOLUME_SIZE_BOUNDS:
             min_size, max_size = EBS_VOLUME_TYPE_TO_VOLUME_SIZE_BOUNDS.get(volume_type)
             if volume_size > max_size:
                 self._add_failure(

cli/tests/pcluster/schemas/test_schema_validators.py

@@ -71,7 +71,6 @@ def test_mount_dir_validator(mount_dir, expected_message):
         ("", "Not a valid integer"),
         ("NONE", "Not a valid integer"),
         ("wrong_value", "Not a valid integer"),
-        (19, "must be at least 35"),
         (36, None),
     ],
 )

cli/tests/pcluster/validators/test_cluster_validators.py

@@ -12,7 +12,7 @@
 from assertpy import assert_that
 from munch import DefaultMunch
 
-from pcluster.aws.aws_resources import InstanceTypeInfo
+from pcluster.aws.aws_resources import ImageInfo, InstanceTypeInfo
 from pcluster.config.cluster_config import PlacementGroup, Tag
 from pcluster.constants import PCLUSTER_NAME_MAX_LENGTH
 from pcluster.validators.cluster_validators import (
@@ -41,6 +41,7 @@
     NumberOfStorageValidator,
     OverlappingMountDirValidator,
     RegionValidator,
+    RootVolumeSizeValidator,
     SchedulableMemoryValidator,
     SchedulerOsValidator,
     SharedStorageMountDirValidator,
@@ -1137,3 +1138,24 @@ def test_mixed_security_group_overwrite_validator(head_node_security_groups, que
     )
     expected_message = "make sure.*cluster nodes are reachable" if expect_warning else None
     assert_failure_messages(actual_failures, expected_message)
+
+
+@pytest.mark.parametrize(
+    "root_volume_size, ami_size, expected_message",
+    [
+        (65, 50, None),
+        (
+            25,
+            50,
+            "Root volume size 25 GiB must be equal or greater than .* 50 GiB.",
+        ),
+    ],
+)
+def test_root_volume_size_validator(mocker, root_volume_size, ami_size, expected_message):
+    mock_aws_api(mocker)
+    mocker.patch(
+        "pcluster.aws.ec2.Ec2Client.describe_image",
+        return_value=ImageInfo({"BlockDeviceMappings": [{"Ebs": {"VolumeSize": ami_size}}]}),
+    )
+    actual_failures = RootVolumeSizeValidator().execute(root_volume_size, "ami-123456789a8a37250")
+    assert_failure_messages(actual_failures, expected_message)

cli/tests/pcluster/validators/test_ebs_validators.py

@@ -94,24 +94,30 @@ def test_ebs_volume_iops_validators(volume_type, volume_size, volume_iops, expec
         ("standard", 15, None),
         ("standard", 0, "The size of standard volumes must be at least 1 GiB"),
         ("standard", 1025, "The size of standard volumes can not exceed 1024 GiB"),
+        ("standard", None, None),
         ("io1", 15, None),
         ("io1", 3, "The size of io1 volumes must be at least 4 GiB"),
         ("io1", 16385, "The size of io1 volumes can not exceed 16384 GiB"),
+        ("io1", None, None),
         ("io2", 15, None),
         ("io2", 3, "The size of io2 volumes must be at least 4 GiB"),
         ("io2", 65537, "The size of io2 volumes can not exceed 65536 GiB"),
+        ("io2", None, None),
         ("gp2", 15, None),
         ("gp2", 0, "The size of gp2 volumes must be at least 1 GiB"),
         ("gp2", 16385, "The size of gp2 volumes can not exceed 16384 GiB"),
+        ("gp2", None, None),
         ("gp3", 15, None),
         ("gp3", 0, "The size of gp3 volumes must be at least 1 GiB"),
         ("gp3", 16385, "The size of gp3 volumes can not exceed 16384 GiB"),
+        ("gp3", None, None),
         ("st1", 500, None),
         ("st1", 20, "The size of st1 volumes must be at least 500 GiB"),
         ("st1", 16385, "The size of st1 volumes can not exceed 16384 GiB"),
         ("sc1", 500, None),
         ("sc1", 20, "The size of sc1 volumes must be at least 500 GiB"),
         ("sc1", 16385, "The size of sc1 volumes can not exceed 16384 GiB"),
+        ("sc1", None, None),
     ],
 )
 def test_ebs_volume_type_size_validator(volume_type, volume_size, expected_message):