[Test] In test_patching_cluster, reduce the risk of ICEs and facilitate troubleshooting

cloudformation/patching/ami-patching.yaml

@@ -21,12 +21,16 @@ Parameters:
   PatchFlavour:
     Description: >-
       Patching flavour passed to the patch script. 'minimal' applies only security
-      errata (smallest set); 'full' applies all available package updates.
+      errata (smallest set); 'full' applies all available package updates. The
+      '-capped' variants additionally cap the kernel to the newest version the FSx
+      Lustre client supports.
     Type: String
     Default: minimal
     AllowedValues:
       - minimal
       - full
+      - minimal-capped
+      - full-capped
 
 Resources:
 

tests/integration-tests/configs/develop.yaml

@@ -176,16 +176,21 @@ test-suites:
   patching:
     test_patching.py::test_patching_cluster:
       dimensions:
-        - regions: [{{ g4dn_8xlarge_CAPACITY_RESERVATION_3_INSTANCES_2_HOURS_NOPG_rhel9 }}]
+        - regions: [{{ g4dn_8xlarge_CAPACITY_RESERVATION_2_INSTANCES_2_HOURS_NOPG_ubuntu2404 }}]
           instances: ["g4dn.8xlarge"]
-          oss: {{ RHEL_OS_X86 }}
+          oss: [{{ UBUNTU_OS_X86_0 }}]
           schedulers: ["slurm"]
-          flags: ["patching:minimal"]
-        - regions: [{{ g4dn_8xlarge_CAPACITY_RESERVATION_3_INSTANCES_2_HOURS_NOPG_ubuntu2404 }}]
+          flags: ["patching:full-capped"]
+        - regions: [{{ g4dn_8xlarge_CAPACITY_RESERVATION_2_INSTANCES_2_HOURS_NOPG_rhel9 }}]
           instances: ["g4dn.8xlarge"]
-          oss: {{ NO_RHEL_OS_X86 }}
-          flags: ["patching:minimal"]
+          oss: [{{ RHEL_OS_X86_0 }}]
+          schedulers: ["slurm"]
+          flags: ["patching:full-capped"]
+        - regions: [{{ g4dn_8xlarge_CAPACITY_RESERVATION_2_INSTANCES_2_HOURS_NOPG_alinux2023 }}]
+          instances: ["g4dn.8xlarge"]
+          oss: ["alinux2023"]
           schedulers: ["slurm"]
+          flags: ["patching:full"]
   custom_resource:
     test_cluster_custom_resource.py::test_cluster_create:
       dimensions:

tests/integration-tests/configs/released.yaml

@@ -118,16 +118,21 @@ test-suites:
   # patching:
   #   test_patching.py::test_patching_cluster:
   #     dimensions:
-  #       - regions: [{{ g4dn_8xlarge_CAPACITY_RESERVATION_3_INSTANCES_2_HOURS_NOPG_rhel9 }}]
+  #       - regions: [{{ g4dn_8xlarge_CAPACITY_RESERVATION_2_INSTANCES_2_HOURS_NOPG_ubuntu2404 }}]
   #         instances: ["g4dn.8xlarge"]
-  #         oss: {{ RHEL_OS_X86 }}
+  #         oss: [{{ UBUNTU_OS_X86_0 }}]
   #         schedulers: ["slurm"]
-  #         flags: ["patching:minimal"]
-  #       - regions: [{{ g4dn_8xlarge_CAPACITY_RESERVATION_3_INSTANCES_2_HOURS_NOPG_ubuntu2404 }}]
+  #         flags: ["patching:full-capped"]
+  #       - regions: [{{ g4dn_8xlarge_CAPACITY_RESERVATION_2_INSTANCES_2_HOURS_NOPG_rhel9 }}]
   #         instances: ["g4dn.8xlarge"]
-  #         oss: {{ NO_RHEL_OS_X86 }}
+  #         oss: [{{ RHEL_OS_X86_0 }}]
   #         schedulers: ["slurm"]
-  #         flags: ["patching:minimal"]
+  #         flags: ["patching:full-capped"]
+  #       - regions: [{{ g4dn_8xlarge_CAPACITY_RESERVATION_2_INSTANCES_2_HOURS_NOPG_alinux2023 }}]
+  #         instances: ["g4dn.8xlarge"]
+  #         oss: ["alinux2023"]
+  #         schedulers: ["slurm"]
+  #         flags: ["patching:full"]
   custom_resource:
     test_cluster_custom_resource.py::test_cluster_1_click:
       dimensions:

tests/integration-tests/conftest.py

@@ -1911,49 +1911,25 @@ def _copy_image(image_id, test_name):
             logging.error("Delete copied AMI snapshot failed due to %s", e)
 
 
-def _log_ami_patching_build_output(region, stack_name):  # noqa: C901
-    """Fetch and log the EC2 Image Builder execution logs for the patched-AMI build.
-
-    Image Builder streams the output of every recipe step (including the patch
-    script stdout from the ApplyPatches ExecuteBash step) to the CloudWatch log
-    group /aws/imagebuilder/pcluster-ami-patching-recipe-<stack_name>. Surfacing it
-    here mirrors how the head node patching output is logged by the test, so the AMI
-    patching output is available in the test logs too. Any failure to retrieve the
-    logs is swallowed so it never fails the test.
+def _log_ami_patching_build_output(region, stack_name):
+    """Log the EC2 Image Builder execution logs for the patched-AMI build.
+
+    Image Builder streams the output of every recipe step (including the patch script
+    output) to the CloudWatch log group
+    /aws/imagebuilder/pcluster-ami-patching-recipe-<stack_name>. Logging it surfaces the
+    patch output (and, on failure, the root cause) in the test logs. Any failure to
+    retrieve the logs is swallowed so it never fails the test.
     """
     log_group_name = f"/aws/imagebuilder/pcluster-ami-patching-recipe-{stack_name}"
     logs_client = boto3.client("logs", region_name=region)
     try:
-        streams = []
-        next_token = None
-        while True:
-            kwargs = {"logGroupName": log_group_name, "orderBy": "LogStreamName"}
-            if next_token:
-                kwargs["nextToken"] = next_token
-            response = logs_client.describe_log_streams(**kwargs)
-            streams.extend(response.get("logStreams", []))
-            next_token = response.get("nextToken")
-            if not next_token:
-                break
-        if not streams:
-            logging.info("No AMI patching Image Builder log streams found in log group %s", log_group_name)
-            return
-        for stream in streams:
-            stream_name = stream["logStreamName"]
-            messages = []
-            prev_token = None
-            event_token = None
-            while True:
-                event_kwargs = {"logGroupName": log_group_name, "logStreamName": stream_name, "startFromHead": True}
-                if event_token:
-                    event_kwargs["nextToken"] = event_token
-                event_response = logs_client.get_log_events(**event_kwargs)
-                messages.extend(event["message"] for event in event_response.get("events", []))
-                event_token = event_response.get("nextForwardToken")
-                if event_token == prev_token:
-                    break
-                prev_token = event_token
-            logging.info("AMI patching Image Builder log (stream %s):\n%s", stream_name, "\n".join(messages))
+        messages = []
+        for page in logs_client.get_paginator("filter_log_events").paginate(logGroupName=log_group_name):
+            messages.extend(event["message"] for event in page.get("events", []))
+        if messages:
+            logging.info("AMI patching Image Builder logs:\n%s", "\n".join(messages))
+        else:
+            logging.info("No AMI patching Image Builder logs found in log group %s", log_group_name)
     except Exception as e:  # noqa: BLE001
         logging.warning("Could not retrieve AMI patching Image Builder logs from log group %s: %s", log_group_name, e)
 
@@ -2019,7 +1995,11 @@ def _build(base_ami, builder_instance, flavour="minimal"):
         )
         # create_stack blocks until CREATE_COMPLETE, i.e. until Image Builder has
         # finished building the patched AMI.
-        cfn_stacks_factory.create_stack(stack)
+        try:
+            cfn_stacks_factory.create_stack(stack)
+        except Exception:  # noqa: BLE001
+            _log_ami_patching_build_output(region, stack_name)
+            raise
         ami_id = stack.cfn_outputs["AmiId"]
         built.append((ami_id, stack_name))
         logging.info("Patched AMI %s is available", ami_id)

tests/integration-tests/framework/tests_configuration/config_renderer.py

@@ -93,6 +93,9 @@ def _get_os_parameters(config=None, args=None):
 
     rhel_oss = [os for os in SUPPORTED_OSES if "rhel" in os]
     _propagate_os_jinja_variables("RHEL_", result, rotation_seed, rhel_oss)
+
+    ubuntu_oss = [os for os in SUPPORTED_OSES if "ubuntu" in os]
+    _propagate_os_jinja_variables("UBUNTU_", result, rotation_seed, ubuntu_oss)
     return result
 
 

tests/integration-tests/tests/patching/test_patching.py

@@ -22,6 +22,7 @@
 from tests.common.utils import (
     COMPUTE_NODE,
     GPU_JOB_SCRIPT,
+    HEAD_NODE,
     LOGIN_NODE,
     NODE_TYPES,
     reboot_head_node,
@@ -30,7 +31,7 @@
 )
 
 # Time budget (seconds) for the patching to complete on the head node.
-PATCHING_TIMEOUT = 1800
+PATCHING_TIMEOUT = 3600
 
 # Patching flavour passed to patch_node.sh: "minimal" applies only security patches,
 # "full" applies all available OS package updates. It is selected via the test's
@@ -42,6 +43,54 @@
 # so it is defined in a single place.
 FSX_LUSTRE_MOUNT_DIR = "/shared-fsxlustre"
 
+# The head node is patched in place, which recompiles EFA against the new kernel. Use at least
+# the memory of a c5.4xlarge: a known issue in the EFA compilation can cause an OOM error on
+# smaller instance types.
+HEAD_NODE_INSTANCE = "c5.4xlarge"
+LOGIN_NODE_INSTANCE = "c5.xlarge"
+
+# Kernel modules that are loaded lazily (e.g. by ss, systemd-networkd, or other daemons that
+# may or may not have run by the time we snapshot). We force-load these after patching so that
+# the before/after comparison does not flag them as missing.
+# Maintained per OS and per node type. Keep module names alphabetically sorted.
+LAZY_KERNEL_MODULES = {
+    "alinux2023": {
+        HEAD_NODE: ["tls"],
+        COMPUTE_NODE: ["tls"],
+        LOGIN_NODE: ["tls"],
+    },
+    "rhel8": {
+        HEAD_NODE: ["af_packet_diag", "inet_diag", "tcp_diag", "tls", "udp_diag"],
+        COMPUTE_NODE: ["tls"],
+        LOGIN_NODE: ["tls"],
+    },
+    "rhel9": {
+        HEAD_NODE: ["tls"],
+        COMPUTE_NODE: ["tls"],
+        LOGIN_NODE: ["tls"],
+    },
+    "rocky8": {
+        HEAD_NODE: ["af_packet_diag", "inet_diag", "tcp_diag", "tls", "udp_diag"],
+        COMPUTE_NODE: ["tls"],
+        LOGIN_NODE: ["tls"],
+    },
+    "rocky9": {
+        HEAD_NODE: ["tls"],
+        COMPUTE_NODE: ["tls"],
+        LOGIN_NODE: ["tls"],
+    },
+    "ubuntu2204": {
+        HEAD_NODE: ["tls"],
+        COMPUTE_NODE: ["tls"],
+        LOGIN_NODE: ["tls"],
+    },
+    "ubuntu2404": {
+        HEAD_NODE: ["tls"],
+        COMPUTE_NODE: ["tls"],
+        LOGIN_NODE: ["tls"],
+    },
+}
+
 
 def test_patching_cluster(
     region,
@@ -83,7 +132,11 @@ def test_patching_cluster(
     # Start the cluster creation but do not block on it: the AMI patching below
     # runs concurrently while the cluster comes up.
     create_config = pcluster_config_reader(
-        output_file="pcluster.config.create.yaml", login_nodes_count=1, fsx_lustre_mount_dir=FSX_LUSTRE_MOUNT_DIR
+        output_file="pcluster.config.create.yaml",
+        login_nodes_count=1,
+        head_node_instance=HEAD_NODE_INSTANCE,
+        login_node_instance=LOGIN_NODE_INSTANCE,
+        fsx_lustre_mount_dir=FSX_LUSTRE_MOUNT_DIR,
     )
     cluster = clusters_factory(create_config, wait=False)
 
@@ -123,6 +176,8 @@ def test_patching_cluster(
         output_file="pcluster.config.stop-login.yaml",
         login_nodes_count=0,
         base_ami=base_ami_pin,
+        head_node_instance=HEAD_NODE_INSTANCE,
+        login_node_instance=LOGIN_NODE_INSTANCE,
         fsx_lustre_mount_dir=FSX_LUSTRE_MOUNT_DIR,
     )
     cluster.update(str(stop_login_config))
@@ -137,6 +192,8 @@ def test_patching_cluster(
         login_nodes_count=1,
         base_ami=base_ami_pin,
         patched_ami=patched_ami,
+        head_node_instance=HEAD_NODE_INSTANCE,
+        login_node_instance=LOGIN_NODE_INSTANCE,
         fsx_lustre_mount_dir=FSX_LUSTRE_MOUNT_DIR,
     )
     cluster.update(str(update_config))
@@ -173,7 +230,7 @@ def test_patching_cluster(
     _run_gpu_workload(cluster, scheduler_commands_factory, use_login_node=True)
 
     # Trigger lazily-loaded modules so that we can compare pre v/s post reboot kernel modules.
-    _trigger_lazy_kernel_modules(cluster, scheduler_commands_factory)
+    _trigger_lazy_kernel_modules(cluster, scheduler_commands_factory, os)
     kernel_modules_after = _collect_loaded_kernel_modules(cluster, scheduler_commands_factory)
     logging.info("Kernel modules loaded after patching: %s", kernel_modules_after)
     with soft_assertions():
@@ -244,7 +301,7 @@ def _run_gpu_workload(cluster, scheduler_commands_factory, use_login_node):
     logging.info("GPU validation job %s submitted from the %s succeeded", job_id, source)
 
 
-def _trigger_lazy_kernel_modules(cluster, scheduler_commands_factory):
+def _trigger_lazy_kernel_modules(cluster, scheduler_commands_factory, os):
     """Ensure on-demand kernel modules are loaded before the post-reboot snapshot.
 
     Some modules load lazily and are absent right after the reboot until something
@@ -255,8 +312,9 @@ def _trigger_lazy_kernel_modules(cluster, scheduler_commands_factory):
         executor = _node_executor(cluster, scheduler_commands_factory, node_type)
         # Read the FSx for Lustre mountpoint to trigger the loading of lustre kernel module.
         executor.run_remote_command(f"ls {FSX_LUSTRE_MOUNT_DIR}")
-        # Load the kernel TLS module directly
-        executor.run_remote_command("sudo modprobe tls")
+        # Force-load known lazily-loaded modules for this OS and node type.
+        for module in LAZY_KERNEL_MODULES.get(os, {}).get(node_type, []):
+            executor.run_remote_command(f"sudo modprobe {module}")
 
 
 def _collect_loaded_kernel_modules(cluster, scheduler_commands_factory):