[Subnet Prioritization] Update unit tests to test that priority is correctly set with prioritized|capacity-optimized-prioritized AllocationStrategy and all_or_nothing ScalingStrategy

Signed-off-by: Hanxuan Zhang <[email protected]>

Author: Allenz5

tests/common.py

@@ -140,6 +140,26 @@ def client_error(error_code):
             "Networking": MULTIPLE_SUBNETS,
         },
     },
+    "queue-prioritized-all-or-nothing": {
+        "c5xlarge": {"Api": "run-instances", "Instances": [{"InstanceType": "c5.xlarge"}]},
+        "fleet1": {
+            "Api": "create-fleet",
+            "Instances": [{"InstanceType": "t2.medium"}],
+            "AllocationStrategy": "prioritized",
+            "CapacityType": "on-demand",
+            "Networking": MULTIPLE_SUBNETS,
+        },
+    },
+    "queue-capacity-optimized-prioritized-all-or-nothing": {
+        "c5xlarge": {"Api": "run-instances", "Instances": [{"InstanceType": "c5.xlarge"}]},
+        "fleet1": {
+            "Api": "create-fleet",
+            "Instances": [{"InstanceType": "t2.medium"}],
+            "AllocationStrategy": "capacity-optimized-prioritized",
+            "CapacityType": "spot",
+            "Networking": MULTIPLE_SUBNETS,
+        },
+    },
 }
 
 LAUNCH_OVERRIDES = {}

tests/slurm_plugin/test_fleet_manager.py

@@ -396,10 +396,10 @@ class TestEc2CreateFleetManager:
             # Use "capacity-optimized-prioritized" Allocation Strategy AND Launch Override with Priority
             (5, "queue-capacity-optimized-prioritized", "fleet1", False, {}, None),
             # Use "prioritized" Allocation Strategy AND Launch Override with Priority AND all_or_nothing is True
-            (5, "queue-prioritized", "fleet1", True, {}, None),
+            (5, "queue-prioritized-all-or-nothing", "fleet1", True, {}, None),
             # Use "capacity-optimized-prioritized" Allocation Strategy
             # AND Launch Override with Priority AND all_or_nothing is True
-            (5, "queue-capacity-optimized-prioritized", "fleet1", True, {}, None),
+            (5, "queue-capacity-optimized-prioritized-all-or-nothing", "fleet1", True, {}, None),
         ],
         ids=[
             "fleet_spot",

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

@@ -2,7 +2,7 @@
   "LaunchTemplateConfigs": [
     {
       "LaunchTemplateSpecification": {
-        "LaunchTemplateName": "hit-queue-capacity-optimized-prioritized-fleet1",
+        "LaunchTemplateName": "hit-queue-capacity-optimized-prioritized-all-or-nothing-fleet1",
         "Version": "$Latest"
       },
       "Overrides": [
@@ -15,24 +15,15 @@
           "InstanceType": "t2.medium",
           "SubnetId": "7654321",
           "Priority": 1.0
-        },
-        {
-          "InstanceType": "t2.large",
-          "SubnetId": "1234567",
-          "Priority": 2.0
-        },
-        {
-          "InstanceType": "t2.large",
-          "SubnetId": "7654321",
-          "Priority": 3.0
         }
       ]
     }
   ],
   "SpotOptions": {
-    "SingleInstanceType": false,
+    "SingleInstanceType": true,
     "SingleAvailabilityZone": false,
-    "AllocationStrategy": "capacity-optimized-prioritized"
+    "AllocationStrategy": "capacity-optimized-prioritized",
+    "MinTargetCapacity": 5
   },
   "TargetCapacitySpecification": {
     "TotalTargetCapacity": 5,

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

@@ -2,7 +2,7 @@
   "LaunchTemplateConfigs": [
     {
       "LaunchTemplateSpecification": {
-        "LaunchTemplateName": "hit-queue-prioritized-fleet1",
+        "LaunchTemplateName": "hit-queue-prioritized-all-or-nothing-fleet1",
         "Version": "$Latest"
       },
       "Overrides": [
@@ -15,27 +15,18 @@
           "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,
+    "SingleInstanceType": true,
     "SingleAvailabilityZone": false,
     "CapacityReservationOptions": {
       "UsageStrategy": "use-capacity-reservations-first"
-    }
+    },
+    "MinTargetCapacity": 5
   },
   "TargetCapacitySpecification": {
     "TotalTargetCapacity": 5,