[Test] Skip Nvidia driver installation if Deep Learning AMI is used

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

Author: hanwen-cluster

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

@@ -124,10 +124,12 @@ def test_build_image(
     bucket_name = s3_bucket_factory()
     _set_s3_bucket_policy(bucket_name, get_arn_partition(region), region)
 
+    enable_nvidia = True
     # Get base AMI
     if os in ["alinux2", "ubuntu2004"]:
         # Test Deep Learning AMIs
         base_ami = retrieve_latest_ami(region, os, ami_type="remarkable", architecture=architecture)
+        enable_nvidia = False  # Deep learning AMIs have Nvidia pre-installed
     elif "rhel" in os or "rocky" in os or "ubuntu" in os:
         # Test AMIs from first stage build. Because RHEL/Rocky and Ubuntu have specific requirement of kernel versions.
         try:
@@ -144,7 +146,7 @@ def test_build_image(
         parent_image=base_ami,
         instance_role=instance_role,
         bucket_name=bucket_name,
-        gpu_count=get_gpu_count(instance),
+        enable_nvidia=str(enable_nvidia and get_gpu_count(instance) > 0).lower(),
     )
 
     image = images_factory(image_id, image_config, region)

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

@@ -24,7 +24,7 @@ Build:
             # 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 %}
+            Enabled: {{ enable_nvidia }}
 
 CustomS3Bucket: {{ bucket_name }}