[Test] Only disable Lustre installation if non-first-stage RHEL9/Rock…

…y9/Ubuntu22 AMIs are used first stage AMIs are used in release.yaml non-first-stage AMIs are used in develop.yaml Signed-off-by: Hanwen <[email protected]>
aws · Dec 19, 2024 · 62c3b2d · 62c3b2d
1 parent 3bd0d7e
commit 62c3b2d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/tests/integration-tests/tests/createami/test_createami.py b/tests/integration-tests/tests/createami/test_createami.py
@@ -126,6 +126,7 @@ def test_build_image(
 
     enable_nvidia = True
     update_os_packages = False
+    enable_lustre_client = True
     # Get base AMI
     if os in ["alinux2", "ubuntu2004"]:
         # Test Deep Learning AMIs
@@ -140,6 +141,8 @@ def test_build_image(
             logging.info("First stage AMI not available, using official AMI instead.")
             base_ami = retrieve_latest_ami(region, os, ami_type="official", architecture=architecture)
             update_os_packages = True
+            if os in ["ubuntu2204", "rhel9", "rocky9"]:
+                enable_lustre_client = False
     else:
         # Test vanilla AMIs.
         base_ami = retrieve_latest_ami(region, os, ami_type="official", architecture=architecture)
@@ -152,6 +155,7 @@ def test_build_image(
         bucket_name=bucket_name,
         enable_nvidia=str(enable_nvidia and get_gpu_count(instance) > 0).lower(),
         update_os_packages=str(update_os_packages).lower(),
+        enable_lustre_client=str(enable_lustre_client).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
@@ -19,7 +19,7 @@ Build:
     Installation:
         LustreClient:
             # 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 %}
+            Enabled: {{ enable_lustre_client }}
         NvidiaSoftware:
             Enabled: {{ enable_nvidia }}