[Test] Install Nvidia driver in createami tests

Reason: This increases test coverage of dependencies files in S3 buckets

How
1. Use GPU instances and use availability zones with best GPU capacity
2. Change build-image config file

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

Author: hanwen-cluster

tests/integration-tests/configs/develop.yaml

@@ -149,16 +149,16 @@ test-suites:
           oss: [{{ OS_X86_3 }}]
     test_createami.py::test_build_image:
       dimensions:
-        - regions: ["eu-west-3"]
-          instances: {{ common.INSTANCES_DEFAULT_X86 }}
+        - regions: ["euw3-az1"]
+          instances: ["g4dn.2xlarge"]
           schedulers: [ "slurm" ]
           oss: {{ common.OSS_COMMERCIAL_X86 }}
-        - regions: ["cn-north-1"]
-          instances: {{ common.INSTANCES_DEFAULT_X86 }}
+        - regions: ["cnn1-az1"]
+          instances: ["g4dn.2xlarge"]
           schedulers: ["slurm"]
           oss: [{{ OS_X86_5 }}]
         - regions: ["us-gov-west-1"]
-          instances: {{ common.INSTANCES_DEFAULT_X86 }}
+          instances: ["g4dn.2xlarge"]
           schedulers: ["slurm"]
           oss: [{{ OS_X86_7 }}]
     test_createami.py::test_build_image_custom_components:

tests/integration-tests/tests/createami/test_createami.py

@@ -27,7 +27,7 @@
 from retrying import retry
 from time_utils import minutes, seconds
 from troposphere import Template, iam
-from utils import generate_stack_name, get_arn_partition
+from utils import generate_stack_name, get_arn_partition, get_gpu_count
 
 from tests.common.assertions import (
     assert_head_node_is_running,
@@ -94,10 +94,11 @@ def test_invalid_config(
     assert_that(suppressed.message).contains("Request would have succeeded")
 
 
-@pytest.mark.usefixtures("instance", "scheduler")
+@pytest.mark.usefixtures("scheduler")
 def test_build_image(
     region,
     os,
+    instance,
     pcluster_config_reader,
     architecture,
     s3_bucket_factory,
@@ -138,9 +139,12 @@ def test_build_image(
     else:
         # Test vanilla AMIs.
         base_ami = retrieve_latest_ami(region, os, ami_type="official", architecture=architecture)
-
     image_config = pcluster_config_reader(
-        config_file="image.config.yaml", parent_image=base_ami, instance_role=instance_role, bucket_name=bucket_name
+        config_file="image.config.yaml",
+        parent_image=base_ami,
+        instance_role=instance_role,
+        bucket_name=bucket_name,
+        gpu_count=get_gpu_count(instance),
     )
 
     image = images_factory(image_id, image_config, region)

tests/integration-tests/tests/createami/test_createami/test_build_image/image.config.yaml

@@ -19,12 +19,12 @@ Build:
     UpdateOsPackages:
         Enabled: true
 {% endif %}
-{% if os in ["ubuntu2204", "rhel9", "rocky9"] %}
-# Disable Lustre installation because these newer operating systems release new kernels more often. Lustre usually does not support the latest kernels
     Installation:
         LustreClient:
-            Enabled: false
-{% endif %}
+            # Disable Lustre installation because these newer operating systems release new kernels more often. Lustre usually does not support the latest kernels
+            Enabled: {% if os in ["ubuntu2204", "rhel9", "rocky9"] %} false {% else %} true {% endif %}
+        NvidiaSoftware:
+            Enabled: {% if gpu_count > 0 %} true {% else %} false {% endif %}
 
 CustomS3Bucket: {{ bucket_name }}
 

tests/integration-tests/utils.py

@@ -612,6 +612,18 @@ def get_network_interfaces_count(instance_type, region_name=None):
     return get_instance_info(instance_type, region_name).get("NetworkInfo").get("MaximumNetworkCards", 1)
 
 
+def get_gpu_count(instance_type, region_name=None):
+    """Return the number of GPU for the provided instance type."""
+    gpu_info = get_instance_info(instance_type, region_name).get("GpuInfo", None)
+    gpu_count = 0
+    if gpu_info:
+        for gpu in gpu_info.get("Gpus", []):
+            manufacturer = gpu.get("Manufacturer", "")
+            if manufacturer.upper() == "NVIDIA":
+                gpu_count += gpu.get("Count", 0)
+    return gpu_count
+
+
 def get_root_volume_id(instance_id, region, os):
     """Return the root EBS volume's ID for the given EC2 instance."""
     logging.info("Getting root volume for instance %s", instance_id)