Add unit tests for fleet manager when using capacity block

Signed-off-by: Enrico Usai <[email protected]>
aws · Nov 3, 2023 · 22b3a9b · 22b3a9b
1 parent bda941f
commit 22b3a9b
Show file tree

Hide file tree

Showing 3 changed files with 141 additions and 1 deletion.
diff --git a/tests/common.py b/tests/common.py
@@ -94,6 +94,22 @@ def client_error(error_code):
             "Networking": MULTIPLE_SUBNETS,
         },
     },
+    "queue-cb": {
+        "run-instances-capacity-block": {
+            "Api": "run-instances",
+            "Instances": [{"InstanceType": "c5.xlarge"}],
+            "CapacityType": "capacity-block",
+            "Networking": SINGLE_SUBNET,
+            "CapacityReservationId": "cr-123456",
+        },
+        "fleet-capacity-block": {
+            "Api": "create-fleet",
+            "Instances": [{"InstanceType": "t2.medium"}, {"InstanceType": "t2.large"}],
+            "CapacityType": "capacity-block",
+            "Networking": SINGLE_SUBNET,
+            "CapacityReservationId": "cr-234567",
+        },
+    },
 }
 
 LAUNCH_OVERRIDES = {}
diff --git a/tests/slurm_plugin/test_fleet_manager.py b/tests/slurm_plugin/test_fleet_manager.py
@@ -297,13 +297,34 @@ class TestEc2CreateFleetManager:
         "Type": "instant",
     }
 
+    test_capacity_block_params = {
+        "LaunchTemplateConfigs": [
+            {
+                "LaunchTemplateSpecification": {
+                    "LaunchTemplateName": "queue-cb-fleet-capacity-block",
+                    "Version": "$Latest",
+                },
+            }
+        ],
+        "OnDemandOptions": {
+            "SingleInstanceType": False,
+            "SingleAvailabilityZone": True,
+            "MinTargetCapacity": 1,
+            "CapacityReservationOptions": {"UsageStrategy": "use-capacity-reservations-first"},
+        },
+        "TargetCapacitySpecification": {"TotalTargetCapacity": 5, "DefaultTargetCapacityType": "capacity-block"},
+        "Type": "instant",
+    }
+
     @pytest.mark.parametrize(
         ("batch_size", "queue", "compute_resource", "all_or_nothing", "launch_overrides", "log_assertions"),
         [
             # normal - spot
             (5, "queue1", "fleet-spot", False, {}, None),
             # normal - on-demand
             (5, "queue2", "fleet-ondemand", False, {}, None),
+            # normal - capacity-block
+            (5, "queue-cb", "fleet-capacity-block", False, {}, None),
             # all or nothing
             (5, "queue1", "fleet-spot", True, {}, None),
             # launch_overrides
@@ -342,6 +363,7 @@ class TestEc2CreateFleetManager:
         ids=[
             "fleet_spot",
             "fleet_ondemand",
+            "fleet_capacity_block",
             "all_or_nothing",
             "launch_overrides",
             "fleet-single-az-multi-it-all_or_nothing",
@@ -554,6 +576,68 @@ def test_evaluate_launch_params(
                     }
                 ],
             ),
+            # normal - capacity-block
+            (
+                test_capacity_block_params,
+                [
+                    MockedBoto3Request(
+                        method="create_fleet",
+                        response={
+                            "Instances": [{"InstanceIds": ["i-12345"]}],
+                            "Errors": [
+                                {"ErrorCode": "InsufficientInstanceCapacity", "ErrorMessage": "Insufficient capacity."}
+                            ],
+                            "ResponseMetadata": {"RequestId": "1234-abcde"},
+                        },
+                        expected_params=test_capacity_block_params,
+                    ),
+                    MockedBoto3Request(
+                        method="describe_instances",
+                        response={
+                            "Reservations": [
+                                {
+                                    "Instances": [
+                                        {
+                                            "InstanceId": "i-12345",
+                                            "PrivateIpAddress": "ip-2",
+                                            "PrivateDnsName": "hostname",
+                                            "LaunchTime": datetime(2020, 1, 1, tzinfo=timezone.utc),
+                                            "NetworkInterfaces": [
+                                                {
+                                                    "Attachment": {
+                                                        "DeviceIndex": 0,
+                                                        "NetworkCardIndex": 0,
+                                                    },
+                                                    "PrivateIpAddress": "ip-2",
+                                                },
+                                            ],
+                                        },
+                                    ]
+                                }
+                            ]
+                        },
+                        expected_params={"InstanceIds": ["i-12345"]},
+                        generate_error=False,
+                    ),
+                ],
+                [
+                    {
+                        "InstanceId": "i-12345",
+                        "PrivateIpAddress": "ip-2",
+                        "PrivateDnsName": "hostname",
+                        "LaunchTime": datetime(2020, 1, 1, tzinfo=timezone.utc),
+                        "NetworkInterfaces": [
+                            {
+                                "Attachment": {
+                                    "DeviceIndex": 0,
+                                    "NetworkCardIndex": 0,
+                                },
+                                "PrivateIpAddress": "ip-2",
+                            },
+                        ],
+                    }
+                ],
+            ),
             # create-fleet - throttling
             (
                 test_on_demand_params,
@@ -592,7 +676,14 @@ def test_evaluate_launch_params(
                 [],
             ),
         ],
-        ids=["fleet_spot", "fleet_exception", "fleet_ondemand", "fleet_throttling", "fleet_multiple_errors"],
+        ids=[
+            "fleet_spot",
+            "fleet_exception",
+            "fleet_ondemand",
+            "fleet_capacity_block",
+            "fleet_throttling",
+            "fleet_multiple_errors",
+        ],
     )
     def test_launch_instances(
         self,

diff --git a/...FleetManager/test_evaluate_launch_params/fleet_capacity_block/expected_launch_params.json b/...FleetManager/test_evaluate_launch_params/fleet_capacity_block/expected_launch_params.json
@@ -0,0 +1,33 @@
+{
+  "LaunchTemplateConfigs":[
+    {
+      "LaunchTemplateSpecification":{
+        "LaunchTemplateName":"hit-queue-cb-fleet-capacity-block",
+        "Version":"$Latest"
+      },
+      "Overrides":[
+        {
+          "InstanceType":"t2.medium",
+          "SubnetId":"1234567"
+        },
+        {
+          "InstanceType":"t2.large",
+          "SubnetId":"1234567"
+        }
+      ]
+    }
+  ],
+  "OnDemandOptions":{
+    "SingleInstanceType":false,
+    "SingleAvailabilityZone":true,
+    "MinTargetCapacity":1,
+    "CapacityReservationOptions":{
+      "UsageStrategy":"use-capacity-reservations-first"
+    }
+  },
+  "TargetCapacitySpecification":{
+    "TotalTargetCapacity":5,
+    "DefaultTargetCapacityType":"capacity-block"
+  },
+  "Type":"instant"
+}