[Test] In test_essential_feature use flexible GPU-powered instance types to reduce the risk of ICEs.

Flexible instance type are cached to reduce the number of EC2 requests.

Author: gmarciani

tests/integration-tests/tests/basic/test_essential_features.py

@@ -15,7 +15,7 @@
 from assertpy import assert_that, soft_assertions
 from constants import UNSUPPORTED_OSES_FOR_DCV
 from remote_command_executor import RemoteCommandExecutor
-from utils import check_status, is_dcv_supported, test_cluster_health_metric
+from utils import check_status, get_flexible_gpu_instance_types, is_dcv_supported, test_cluster_health_metric
 
 from tests.basic.disable_hyperthreading_utils import _test_disable_hyperthreading_settings
 from tests.basic.log_rotation_utils import _test_compute_log_rotation, _test_headnode_log_rotation
@@ -65,6 +65,7 @@ def test_essential_features(
         dcv_enabled=dcv_enabled,
         max_queue_size=max_queue_size,
         scaledown_idletime=scaledown_idletime,
+        flexible_gpu_instance_types=get_flexible_gpu_instance_types(instance, region),
     )
     cluster = clusters_factory(cluster_config)
 

tests/integration-tests/tests/basic/test_essential_features/test_essential_features/pcluster.config.yaml

@@ -128,7 +128,9 @@ Scheduling:
       ComputeResources:
         - Name: cr1
           Instances:
-            - InstanceType: g4dn.2xlarge
+            {% for instance_type in flexible_gpu_instance_types %}
+            - InstanceType: {{ instance_type }}
+            {% endfor %}
           MinCount: 0
           MaxCount: 1
   SlurmSettings:

tests/integration-tests/utils.py

@@ -19,6 +19,7 @@
 import string
 import subprocess
 from datetime import datetime, timedelta
+from functools import cache
 from hashlib import sha1
 
 import boto3
@@ -1073,6 +1074,14 @@ def get_similar_instance_types(instance_type: str, region: str = None, max_items
     return similar_instances
 
 
+@cache
+def get_flexible_gpu_instance_types(instance, region):
+    """Return a list of NVIDIA GPU instance types compatible with ``instance``'s architecture."""
+    architecture = get_architecture_supported_by_instance_type(instance, region)
+    gpu_instance_type = "g4dn.2xlarge" if architecture == "x86_64" else "g5g.2xlarge"
+    return list({gpu_instance_type, *get_similar_instance_types(gpu_instance_type, region, 5)})
+
+
 def verify_cluster_node_config_version_in_ddb(region, cluster_name, instance_id, expected_version):
     """
     Verify that a cluster node has the correct config version in DynamoDB.