[integ-tests] Find capacity blocks by OS

Author: hanwen-cluster

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

@@ -554,15 +554,22 @@ def terminate_nodes_manually(instance_ids, region):
     logging.info("Terminated nodes: {}".format(instance_ids))
 
 
-def get_capacity_reservation_id(instance_type, region, count):
+def get_capacity_reservation_id(instance_type, region, count, os):
+    os_platform = "Linux/UNIX"
+    if "rhel" in os.lower():
+        os_platform = "Red Hat Enterprise Linux"
     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:
+            if (
+                instance_type == reservation.get("InstanceType")
+                and os_platform == reservation.get("InstancePlatform")
+                and reservation.get("AvailableInstanceCount") >= count
+            ):
                 reservations_ids.append(
                     {
                         "CapacityReservationId": reservation["CapacityReservationId"],

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

@@ -58,7 +58,7 @@ def test_efa(
     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)
+        capacity_reservations_ids = get_capacity_reservation_id(instance, region, max_queue_size, os)
         if capacity_reservations_ids:
             p6_b200_capacity_reservation_id = capacity_reservations_ids[0].get("CapacityReservationId")
         else:

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

@@ -346,7 +346,6 @@ def assert_topology_plugin_completely_disabled(cluster: Cluster):
 
 
 @pytest.mark.usefixtures("serial_execution_by_instance")
-@pytest.mark.usefixtures("os")
 def test_gb200(
     pcluster_config_reader,
     file_reader,
@@ -356,6 +355,7 @@ def test_gb200(
     region,
     instance,
     scheduler,
+    os,
     scheduler_commands_factory,
 ):
     """
@@ -391,7 +391,7 @@ def test_gb200(
     capacity_reservation_id = None
     max_queue_size = 2
     if instance == "p6e-gb200.36xlarge":
-        ultraserver_reservations_ids = get_capacity_reservation_id(instance, region, max_queue_size)
+        ultraserver_reservations_ids = get_capacity_reservation_id(instance, region, max_queue_size, os)
         if ultraserver_reservations_ids:
             capacity_reservation_id = ultraserver_reservations_ids[0].get("CapacityReservationId")
         else: