diff --git a/tests/common.py b/tests/common.py index bd4969642..95fd16ccb 100644 --- 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 index 4b2f205e2..1b5df838f 100644 --- a/tests/slurm_plugin/test_fleet_manager.py +++ b/tests/slurm_plugin/test_fleet_manager.py @@ -297,6 +297,25 @@ 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"), [ @@ -304,6 +323,8 @@ class TestEc2CreateFleetManager: (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/tests/slurm_plugin/test_fleet_manager/TestCreateFleetManager/test_evaluate_launch_params/fleet_capacity_block/expected_launch_params.json b/tests/slurm_plugin/test_fleet_manager/TestCreateFleetManager/test_evaluate_launch_params/fleet_capacity_block/expected_launch_params.json new file mode 100644 index 000000000..67b68a1b7 --- /dev/null +++ b/tests/slurm_plugin/test_fleet_manager/TestCreateFleetManager/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" +} \ No newline at end of file