[Subnet Prioritization] Support capacity-optimized-prioritized and prioritized Allocation Strategy

CHANGELOG.md

@@ -13,7 +13,7 @@ This file is used to list changes made in each version of the aws-parallelcluste
 ------
 
 **CHANGES**
-- There were no changes for this version.
+- Support prioritized|capacity-optimized-prioritized Allocation Strategy and EnableSingleAvailabilityZone
 
 3.13.0
 ------

src/slurm_plugin/fleet_manager.py

@@ -296,11 +296,21 @@ def _evaluate_template_overrides(self) -> list:
             if self._compute_resource_config.get("MaxPrice"):
                 overrides.update({"MaxPrice": str(self._compute_resource_config["MaxPrice"])})
 
+        priority = 0.0
         for instance_type in self._compute_resource_config["Instances"]:
             subnet_ids = self._compute_resource_config["Networking"]["SubnetIds"]
             for subnet_id in subnet_ids:
-                overrides.update({"InstanceType": instance_type["InstanceType"], "SubnetId": subnet_id})
+                if (
+                    self._compute_resource_config.get("AllocationStrategy") == "prioritized"
+                    or self._compute_resource_config.get("AllocationStrategy") == "capacity-optimized-prioritized"
+                ):
+                    overrides.update(
+                        {"InstanceType": instance_type["InstanceType"], "SubnetId": subnet_id, "Priority": priority}
+                    )
+                else:
+                    overrides.update({"InstanceType": instance_type["InstanceType"], "SubnetId": subnet_id})
                 template_overrides.append(copy.deepcopy(overrides))
+                priority += 1.0
         return template_overrides
 
     def _uses_single_instance_type(self):
@@ -317,8 +327,11 @@ def _evaluate_launch_params(self, count):
         try:
             common_launch_options = {
                 "SingleInstanceType": self._uses_single_instance_type(),
-                "SingleAvailabilityZone": self._uses_single_az(),  # If using Multi-AZ (by specifying multiple subnets),
-                # set SingleAvailabilityZone to False
+                "SingleAvailabilityZone": (
+                    self._compute_resource_config["Networking"]["SingleAvailabilityZone"]
+                    if self._compute_resource_config["Networking"]["SingleAvailabilityZone"] is not None
+                    else self._uses_single_az()
+                ),
             }
             allocation_strategy = self._compute_resource_config.get("AllocationStrategy")
             if allocation_strategy:

tests/common.py

@@ -31,8 +31,9 @@ def client_error(error_code):
     return ClientError({"Error": {"Code": error_code}}, "failed_operation")
 
 
-SINGLE_SUBNET = {"SubnetIds": ["1234567"]}
-MULTIPLE_SUBNETS = {"SubnetIds": ["1234567", "7654321"]}
+SINGLE_SUBNET = {"SubnetIds": ["1234567"], "SingleAvailabilityZone": None}
+MULTIPLE_SUBNETS = {"SubnetIds": ["1234567", "7654321"], "SingleAvailabilityZone": None}
+MULTIPLE_SUBNET_ENABLE_SINGLE_AVAILABILITY_ZONE = {"SubnetIds": ["1234567", "7654321"], "SingleAvailabilityZone": True}
 
 FLEET_CONFIG = {
     "queue": {"c5xlarge": {"Api": "run-instances", "Instances": [{"InstanceType": "c5.xlarge"}]}},
@@ -110,6 +111,36 @@ def client_error(error_code):
             "CapacityReservationId": "cr-234567",
         },
     },
+    "queue-single-az": {
+        "c5xlarge": {"Api": "run-instances", "Instances": [{"InstanceType": "c5.xlarge"}]},
+        "fleet1": {
+            "Api": "create-fleet",
+            "Instances": [{"InstanceType": "t2.medium"}, {"InstanceType": "t2.large"}],
+            "AllocationStrategy": "lowest-price",
+            "CapacityType": "on-demand",
+            "Networking": MULTIPLE_SUBNET_ENABLE_SINGLE_AVAILABILITY_ZONE,
+        },
+    },
+    "queue-prioritized": {
+        "c5xlarge": {"Api": "run-instances", "Instances": [{"InstanceType": "c5.xlarge"}]},
+        "fleet1": {
+            "Api": "create-fleet",
+            "Instances": [{"InstanceType": "t2.medium"}, {"InstanceType": "t2.large"}],
+            "AllocationStrategy": "prioritized",
+            "CapacityType": "on-demand",
+            "Networking": MULTIPLE_SUBNET_ENABLE_SINGLE_AVAILABILITY_ZONE,
+        },
+    },
+    "queue-capacity-optimized-prioritized": {
+        "c5xlarge": {"Api": "run-instances", "Instances": [{"InstanceType": "c5.xlarge"}]},
+        "fleet1": {
+            "Api": "create-fleet",
+            "Instances": [{"InstanceType": "t2.medium"}, {"InstanceType": "t2.large"}],
+            "AllocationStrategy": "capacity-optimized-prioritized",
+            "CapacityType": "on-demand",
+            "Networking": MULTIPLE_SUBNET_ENABLE_SINGLE_AVAILABILITY_ZONE,
+        },
+    },
 }
 
 LAUNCH_OVERRIDES = {}

tests/slurm_plugin/test_fleet_manager.py

@@ -391,6 +391,12 @@ class TestEc2CreateFleetManager:
                 {},
                 "All-or-Nothing is only available with single instance type compute resources or single subnet queues",
             ),
+            # Enable Single Availability Zone
+            (5, "queue-single-az", "fleet1", False, {}, None),
+            # Use "prioritized" Allocation Strategy AND Launch Override with Priority
+            (5, "queue-prioritized", "fleet1", False, {}, None),
+            # Use "capacity-optimized-prioritized" Allocation Strategy AND Launch Override with Priority
+            (5, "queue-capacity-optimized-prioritized", "fleet1", False, {}, None),
         ],
         ids=[
             "fleet_spot",
@@ -402,6 +408,9 @@ class TestEc2CreateFleetManager:
             "fleet-multi-az-single-it-all_or_nothing",
             "fleet-multi-az-multi-it",
             "fleet-multi-az-multi-it-all_or_nothing",
+            "single_az",
+            "prioritized",
+            "capacity_optimized_prioritized",
         ],
     )
     def test_evaluate_launch_params(

tests/slurm_plugin/test_fleet_manager/TestEc2CreateFleetManager/test_evaluate_launch_params/capacity_optimized_prioritized/expected_launch_params.json

@@ -0,0 +1,45 @@
+{
+  "LaunchTemplateConfigs": [
+    {
+      "LaunchTemplateSpecification": {
+        "LaunchTemplateName": "hit-queue-capacity-optimized-prioritized-fleet1",
+        "Version": "$Latest"
+      },
+      "Overrides": [
+        {
+          "InstanceType": "t2.medium",
+          "SubnetId": "1234567",
+          "Priority": 0.0
+        },
+        {
+          "InstanceType": "t2.medium",
+          "SubnetId": "7654321",
+          "Priority": 1.0
+        },
+        {
+          "InstanceType": "t2.large",
+          "SubnetId": "1234567",
+          "Priority": 2.0
+        },
+        {
+          "InstanceType": "t2.large",
+          "SubnetId": "7654321",
+          "Priority": 3.0
+        }
+      ]
+    }
+  ],
+  "OnDemandOptions": {
+    "AllocationStrategy": "capacity-optimized-prioritized",
+    "SingleInstanceType": false,
+    "SingleAvailabilityZone": true,
+    "CapacityReservationOptions": {
+      "UsageStrategy": "use-capacity-reservations-first"
+    }
+  },
+  "TargetCapacitySpecification": {
+    "TotalTargetCapacity": 5,
+    "DefaultTargetCapacityType": "on-demand"
+  },
+  "Type": "instant"
+}

tests/slurm_plugin/test_fleet_manager/TestEc2CreateFleetManager/test_evaluate_launch_params/prioritized/expected_launch_params.json

@@ -0,0 +1,45 @@
+{
+  "LaunchTemplateConfigs": [
+    {
+      "LaunchTemplateSpecification": {
+        "LaunchTemplateName": "hit-queue-prioritized-fleet1",
+        "Version": "$Latest"
+      },
+      "Overrides": [
+        {
+          "InstanceType": "t2.medium",
+          "SubnetId": "1234567",
+          "Priority": 0.0
+        },
+        {
+          "InstanceType": "t2.medium",
+          "SubnetId": "7654321",
+          "Priority": 1.0
+        },
+        {
+          "InstanceType": "t2.large",
+          "SubnetId": "1234567",
+          "Priority": 2.0
+        },
+        {
+          "InstanceType": "t2.large",
+          "SubnetId": "7654321",
+          "Priority": 3.0
+        }
+      ]
+    }
+  ],
+  "OnDemandOptions": {
+    "AllocationStrategy": "prioritized",
+    "SingleInstanceType": false,
+    "SingleAvailabilityZone": true,
+    "CapacityReservationOptions": {
+      "UsageStrategy": "use-capacity-reservations-first"
+    }
+  },
+  "TargetCapacitySpecification": {
+    "TotalTargetCapacity": 5,
+    "DefaultTargetCapacityType": "on-demand"
+  },
+  "Type": "instant"
+}

tests/slurm_plugin/test_fleet_manager/TestEc2CreateFleetManager/test_evaluate_launch_params/single_az/expected_launch_params.json

@@ -0,0 +1,41 @@
+{
+  "LaunchTemplateConfigs": [
+    {
+      "LaunchTemplateSpecification": {
+        "LaunchTemplateName": "hit-queue-single-az-fleet1",
+        "Version": "$Latest"
+      },
+      "Overrides": [
+        {
+          "InstanceType": "t2.medium",
+          "SubnetId": "1234567"
+        },
+        {
+          "InstanceType": "t2.medium",
+          "SubnetId": "7654321"
+        },
+        {
+          "InstanceType": "t2.large",
+          "SubnetId": "1234567"
+        },
+        {
+          "InstanceType": "t2.large",
+          "SubnetId": "7654321"
+        }
+      ]
+    }
+  ],
+  "OnDemandOptions": {
+    "AllocationStrategy": "lowest-price",
+    "SingleInstanceType": false,
+    "SingleAvailabilityZone": true,
+    "CapacityReservationOptions": {
+      "UsageStrategy": "use-capacity-reservations-first"
+    }
+  },
+  "TargetCapacitySpecification": {
+    "TotalTargetCapacity": 5,
+    "DefaultTargetCapacityType": "on-demand"
+  },
+  "Type": "instant"
+}