Skip to content

Commit

Permalink
[Test] Skip Nvidia driver installation if Deep Learning AMI is used
Browse files Browse the repository at this point in the history
Signed-off-by: Hanwen <hanwenli@amazon.com>
  • Loading branch information
hanwen-cluster committed Dec 10, 2024
1 parent ada534f commit c293ff0
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 2 deletions.
4 changes: 3 additions & 1 deletion tests/integration-tests/tests/createami/test_createami.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 }}

Expand Down

0 comments on commit c293ff0

Please sign in to comment.