[integ-tests] Add p6-b200 to test_efa

p6-b200 requires a capacity block. This commit makes the test to automatically find the capacity block by reusing the same logic as test_gb200. This requires a little bit of refactor

Using capacity block requires using RunInstances instead of CreateFleet. Therefore, this commit changes the cluster configuration file of test_efa to always use RunInstances

Author: hanwen-cluster

tests/integration-tests/configs/develop.yaml

@@ -264,6 +264,10 @@ test-suites:
           instances: ["p4d.24xlarge"]
           oss: [{{ NO_RHEL_OS_X86_1 }}]  # The capacity reservation cannot use RHEL operating system
           schedulers: ["slurm"]
+        - regions: [ "use1-az6" ]  # do not move, unless capacity reservation is moved as well
+          instances: [ "p6-b200.48xlarge" ]
+          oss: ["alinux2023", "ubuntu2204", "ubuntu2404", "rocky8", "rocky9"]
+          schedulers: [ "slurm" ]
         - regions: [{{ c6gn_16xlarge_CAPACITY_RESERVATION_2_INSTANCES_2_HOURS_YESPG_OS_ARM_0 }}]
           instances: ["c6gn.16xlarge"]
           oss: [{{ OS_ARM_0 }}]

tests/integration-tests/tests/common/utils.py

@@ -552,3 +552,22 @@ def terminate_nodes_manually(instance_ids, region):
         assert_that(instance_states.get("InstanceId")).is_equal_to(instance_id)
         assert_that(instance_states.get("CurrentState").get("Name")).is_in("shutting-down", "terminated")
     logging.info("Terminated nodes: {}".format(instance_ids))
+
+
+def get_capacity_reservation_id(instance_type, region, count):
+    ec2_client = boto3.client("ec2", region_name=region)
+    paginator = ec2_client.get_paginator("describe_capacity_reservations")
+    # List to store matching reservation IDs
+    reservations_ids = []
+    # Paginate through the results
+    for page in paginator.paginate():
+        for reservation in page.get("CapacityReservations", []):
+            if instance_type == reservation.get("InstanceType") and reservation.get("AvailableInstanceCount") >= count:
+                reservations_ids.append(
+                    {
+                        "CapacityReservationId": reservation["CapacityReservationId"],
+                        "TotalInstanceCount": reservation["TotalInstanceCount"],
+                        "AvailableInstanceCount": reservation["AvailableInstanceCount"],
+                    }
+                )
+    return reservations_ids

tests/integration-tests/tests/efa/test_efa.py

@@ -11,6 +11,7 @@
 # See the License for the specific language governing permissions and limitations under the License.
 import logging
 
+import pytest
 import xmltodict
 from assertpy import assert_that, soft_assertions
 from remote_command_executor import RemoteCommandExecutor
@@ -19,13 +20,20 @@
 from tests.common.assertions import assert_no_errors_in_logs
 from tests.common.mpi_common import _test_mpi
 from tests.common.nccl_common import install_and_run_nccl_benchmarks
-from tests.common.utils import fetch_instance_slots, read_remote_file, run_system_analyzer, wait_process_completion
+from tests.common.utils import (
+    fetch_instance_slots,
+    get_capacity_reservation_id,
+    read_remote_file,
+    run_system_analyzer,
+    wait_process_completion,
+)
 
 FABTESTS_BASIC_TESTS = ["rdm_tagged_bw", "rdm_tagged_pingpong"]
 
 FABTESTS_GDRCOPY_TESTS = ["runt"]
 
 
+@pytest.mark.usefixtures("serial_execution_by_instance")
 def test_efa(
     os,
     region,
@@ -47,9 +55,21 @@ def test_efa(
         head_node_instance = "c5.18xlarge"
     else:
         head_node_instance = "c6g.16xlarge"
+    max_queue_size = 2
+    p6_b200_capacity_reservation_id = None
+    if instance == "p6-b200.48xlarge":
+        capacity_reservations_ids = get_capacity_reservation_id(instance, region, max_queue_size)
+        if capacity_reservations_ids:
+            p6_b200_capacity_reservation_id = capacity_reservations_ids[0].get("CapacityReservationId")
+        else:
+            pytest.skip(f"Skipping the test No Capacity Block for {instance} was found in {region}")
 
     slots_per_instance = fetch_instance_slots(region, instance, multithreading_disabled=True)
-    cluster_config = pcluster_config_reader(head_node_instance=head_node_instance)
+    cluster_config = pcluster_config_reader(
+        head_node_instance=head_node_instance,
+        max_queue_size=max_queue_size,
+        p6_b200_capacity_reservation_id=p6_b200_capacity_reservation_id,
+    )
     cluster = clusters_factory(cluster_config)
     remote_command_executor = RemoteCommandExecutor(cluster)
     scheduler_commands = scheduler_commands_factory(remote_command_executor)

tests/integration-tests/tests/efa/test_efa/test_efa/pcluster.config.yaml

@@ -13,22 +13,22 @@ Scheduling:
   Scheduler: {{ scheduler }}
   SlurmQueues:
     - Name: efa-enabled
+      {% if p6_b200_capacity_reservation_id %}
+      CapacityType: CAPACITY_BLOCK
+      CapacityReservationTarget:
+        CapacityReservationId: {{ p6_b200_capacity_reservation_id }}
+      {% endif %}
       Networking:
         PlacementGroup:
-          Enabled: {% if instance not in ["p4d.24xlarge"] %}true{% else %}false{% endif %}
+          Enabled: {% if instance not in ["p4d.24xlarge", "p6-b200.48xlarge"] %}true{% else %}false{% endif %}
           {% if instance in ["c5n.18xlarge", "c6gn.16xlarge"] %}Name: {{ capacity_reservation_framework_placement_group }}{% endif %}
         SubnetIds:
           - {{ private_subnet_id }}
       ComputeResources:
         - Name: efa-enabled-i1
-          {% if "us-iso" in region %}
           InstanceType: {{ instance }}
-          {% else %}
-          Instances:
-            - InstanceType: {{ instance }}
-          {% endif %}
-          MaxCount: 2
-          MinCount: 2
+          MaxCount: {{ max_queue_size }}
+          MinCount: {{ max_queue_size }}
           DisableSimultaneousMultithreading: true
           Efa:
             Enabled: true

tests/integration-tests/tests/ultraserver/test_gb200.py

@@ -25,7 +25,13 @@
 from tests.common.mpi_common import _test_mpi
 from tests.common.nccl_common import install_and_run_nccl_benchmarks
 from tests.common.schedulers_common import SlurmCommands
-from tests.common.utils import fetch_instance_slots, is_existing_remote_file, read_remote_file, terminate_nodes_manually
+from tests.common.utils import (
+    fetch_instance_slots,
+    get_capacity_reservation_id,
+    is_existing_remote_file,
+    read_remote_file,
+    terminate_nodes_manually,
+)
 
 # We use placeholder IPs just to get IMEX started.
 # These values are hardwired in the cookbook.
@@ -339,25 +345,6 @@ def assert_topology_plugin_completely_disabled(cluster: Cluster):
     logging.info("TopologyPlugin correctly completely disabled")
 
 
-def get_capacity_reservation_id(instance_type, region, count):
-    ec2_client = boto3.client("ec2", region_name=region)
-    paginator = ec2_client.get_paginator("describe_capacity_reservations")
-    # List to store matching reservation IDs
-    reservations_ids = []
-    # Paginate through the results
-    for page in paginator.paginate():
-        for reservation in page.get("CapacityReservations", []):
-            if instance_type == reservation.get("InstanceType") and reservation.get("AvailableInstanceCount") >= count:
-                reservations_ids.append(
-                    {
-                        "CapacityReservationId": reservation["CapacityReservationId"],
-                        "TotalInstanceCount": reservation["TotalInstanceCount"],
-                        "AvailableInstanceCount": reservation["AvailableInstanceCount"],
-                    }
-                )
-    return reservations_ids
-
-
 @pytest.mark.usefixtures("serial_execution_by_instance")
 @pytest.mark.usefixtures("os")
 def test_gb200(