From c293ff0698eb9cd2befe9b045f02a8f3194fcaea Mon Sep 17 00:00:00 2001 From: Hanwen Date: Tue, 10 Dec 2024 06:24:34 -0800 Subject: [PATCH] [Test] Skip Nvidia driver installation if Deep Learning AMI is used Signed-off-by: Hanwen --- tests/integration-tests/tests/createami/test_createami.py | 4 +++- .../test_createami/test_build_image/image.config.yaml | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/integration-tests/tests/createami/test_createami.py b/tests/integration-tests/tests/createami/test_createami.py index 5f05353447..75c465f9ba 100644 --- a/tests/integration-tests/tests/createami/test_createami.py +++ b/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) diff --git a/tests/integration-tests/tests/createami/test_createami/test_build_image/image.config.yaml b/tests/integration-tests/tests/createami/test_createami/test_build_image/image.config.yaml index 621834288d..4a41bc9db2 100644 --- a/tests/integration-tests/tests/createami/test_createami/test_build_image/image.config.yaml +++ b/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 }}