temporary commit

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

Author: hanwen-cluster

cli/src/pcluster/aws/aws_resources.py

@@ -37,6 +37,7 @@ def __init__(self, stack_data: dict):
         self._params = self._stack_data.get("Parameters", [])
         self.tags = self._stack_data.get("Tags", [])
         self.outputs = self._stack_data.get("Outputs", [])
+        self.__resources = None
 
     @property
     def id(self):
@@ -48,6 +49,15 @@ def name(self):
         """Return the name of the stack."""
         return self._stack_data.get("StackName")
 
+    @property
+    def resources(self):
+        """Return the resources of the stack."""
+        if not self.__resources:
+            from pcluster.aws.aws_api import AWSApi
+
+            self.__resources = AWSApi.instance().cfn.describe_stack_resources(self.name)
+        return self.__resources
+
     @property
     def status(self):
         """Return the status of the stack."""
@@ -90,6 +100,10 @@ def _get_param(self, key_name):
         param_value = next((par["ParameterValue"] for par in self._params if par["ParameterKey"] == key_name), None)
         return None if param_value is None else param_value.strip()
 
+    def get_resource(self, resource_logical_id: str):
+        """Return the resource information."""
+        return self.resources.get(resource_logical_id)
+
 
 class InstanceInfo:
     """Object to store Instance information, initialized with a describe_instances call."""

cli/src/pcluster/aws/cfn.py

@@ -147,6 +147,12 @@ def describe_stack_resource(self, stack_name: str, logic_resource_id: str):
                 function_name="describe_stack_resource", message=f"No resource {logic_resource_id} found."
             )
 
+    @AWSExceptionHandler.handle_client_exception
+    def describe_stack_resources(self, stack_name: str):
+        """Get stack resources information."""
+        response = self._client.describe_stack_resources(StackName=stack_name).get("StackResources")
+        return {resource["LogicalResourceId"]: resource for resource in response}  # Build dictionary for better query.
+
     @AWSExceptionHandler.handle_client_exception
     def get_imagebuilder_stacks(self, next_token=None):
         """List existing imagebuilder stacks."""

cli/src/pcluster/aws/ec2.py

@@ -136,7 +136,7 @@ def get_subnet_vpc(self, subnet_id):
     @Cache.cached
     def get_subnet_cidr(self, subnet_id):
         """Return cidr block  of the given subnet."""
-        subnets = self._client.describe_subnets(SubnetIds=[subnet_id]).get("Subnets")
+        subnets = self.describe_subnets([subnet_id])
         if subnets:
             return subnets[0].get("CidrBlock")
         raise AWSClientError(function_name="describe_subnets", message=f"Subnet {subnet_id} not found")

cli/src/pcluster/config/cluster_config.py

@@ -1261,6 +1261,8 @@ def __init__(
         self._official_ami = None
         self.imds = imds or TopLevelImds(implied="v1.0")
         self.deployment_settings = deployment_settings
+        self.managed_head_node_security_group = None
+        self.managed_compute_security_group = None
 
     def _register_validators(self, context: ValidatorContext = None):  # noqa: D102 #pylint: disable=unused-argument
         self._register_validator(RegionValidator, region=self.region)
@@ -1387,15 +1389,15 @@ def _register_storage_validators(self):
                             EfsIdValidator,
                             efs_id=storage.file_system_id,
                             avail_zones_mapping=self.availability_zones_subnets_mapping,
-                            are_all_security_groups_customized=self.are_all_security_groups_customized,
+                            security_groups_by_nodes=self.security_groups_by_nodes,
                         )
                     else:
                         new_storage_count["efs"] += 1
             self._register_validator(
                 ExistingFsxNetworkingValidator,
                 file_system_ids=list(existing_fsx),
-                head_node_subnet_id=self.head_node.networking.subnet_id,
-                are_all_security_groups_customized=self.are_all_security_groups_customized,
+                subnet_ids=[self.head_node.networking.subnet_id] + self.compute_subnet_ids,
+                security_groups_by_nodes=self.security_groups_by_nodes,
             )
 
             self._validate_max_storage_count(ebs_count, existing_storage_count, new_storage_count)
@@ -1624,17 +1626,29 @@ def is_dcv_enabled(self):
         return self.head_node.dcv and self.head_node.dcv.enabled
 
     @property
-    def are_all_security_groups_customized(self):
+    def security_groups_by_nodes(self):
         """Return True if all head node and queues have (additional) security groups specified."""
         head_node_networking = self.head_node.networking
+        security_groups_for_head_node = set()
+        if head_node_networking.security_groups:
+            security_groups_for_head_node.update(set(head_node_networking.security_groups))
+        if head_node_networking.additional_security_groups:
+            security_groups_for_head_node.update(set(head_node_networking.additional_security_groups))
         if not (head_node_networking.security_groups or head_node_networking.additional_security_groups):
-            return False
+            security_groups_for_head_node.add(self.managed_head_node_security_group)
+        security_groups_for_all_nodes = {frozenset(security_groups_for_head_node)}
         for queue in self.scheduling.queues:
             queue_networking = queue.networking
             if isinstance(queue_networking, _QueueNetworking):
-                if not (queue_networking.security_groups or queue_networking.additional_security_groups):
-                    return False
-        return True
+                security_groups_for_compute_node = set()
+                if queue_networking.security_groups:
+                    security_groups_for_compute_node.update(set(queue_networking.security_groups))
+                else:
+                    security_groups_for_compute_node.add(self.managed_compute_security_group)
+                if queue_networking.additional_security_groups:
+                    security_groups_for_compute_node.update(set(queue_networking.additional_security_groups))
+                security_groups_for_all_nodes.add(frozenset(security_groups_for_compute_node))
+        return security_groups_for_all_nodes
 
     @property
     def extra_chef_attributes(self):

cli/src/pcluster/models/cluster.py

@@ -440,6 +440,9 @@ def _validate_and_parse_config(
             Cluster._load_additional_instance_type_data(cluster_config_dict)
             config = self._load_config(cluster_config_dict)
             config.official_ami = self.__official_ami
+            if context.during_update:
+                config.managed_head_node_security_group = self.stack.get_resource("HeadNodeSecurityGroup")
+                config.managed_compute_security_group = self.stack.get_resource("ComputeSecurityGroup")
 
             validation_failures = config.validate(validator_suppressors, context)
             if any(f.level.value >= FailureLevel(validation_failure_level).value for f in validation_failures):
@@ -851,7 +854,7 @@ def validate_update_request(
             validator_suppressors=validator_suppressors,
             validation_failure_level=validation_failure_level,
             config_text=target_source_config,
-            context=ValidatorContext(head_node_instance_id=self.head_node_instance.id),
+            context=ValidatorContext(head_node_instance_id=self.head_node_instance.id, during_update=True),
         )
         changes = self._validate_patch(force, target_config)