Retry a throttled CreateFleet only when its other errors are pool-scoped

src/slurm_plugin/fleet_manager.py

@@ -21,6 +21,7 @@
 from common.utils import setup_logging_filter
 from retrying import retry
 from slurm_plugin.common import print_with_count
+from slurm_plugin.slurm_resources import SlurmNode
 
 logger = logging.getLogger(__name__)
 
@@ -41,6 +42,21 @@
     "Failed to fulfill capacity. Please review errors in the response.",
 )
 
+# Errors confined to one instance type and subnet pool. Reported alongside throttling, they leave the throttled
+# pools able to serve the batch once the rate limit refills, so the throttling is still retried. Any other error
+# alongside throttling would fail the retry on every pool alike and is reported instead.
+POOL_LEVEL_ERROR_CODES = frozenset(
+    SlurmNode.EC2_ICE_ERROR_CODES
+    | {
+        "InsufficientFreeAddressesInSubnet",
+        "InvalidSubnetID.NotFound",
+        "InvalidSubnet",
+        "InsufficientVolumeCapacity",
+        "VolumeTypeNotAvailableInZone",
+        "ServiceUnavailable",
+    }
+)
+
 
 class EC2Instance:
     def __init__(self, id, private_ip, hostname, all_private_ips, launch_time):
@@ -450,14 +466,27 @@ def _launch_instances(self, launch_params):
                 ]
                 if real_errors:
                     err_list = real_errors
-                # A single cause is normally left. Should there be several, prefer throttling as a safety net: it is
-                # the only cause that resolves on its own, and reporting it as insufficient capacity would instead
-                # fail the compute resource over.
+                # A single cause is normally left. Should there be several, throttling is retried only when every
+                # other cause is confined to a pool, since the throttled pools can then still serve the batch once the
+                # rate limit refills. Any other cause would fail the retry on every pool alike and is reported instead.
+                # For example:
+                # - [RequestLimitExceeded, InsufficientInstanceCapacity] -> retry, other pools can still serve the batch
+                # - [RequestLimitExceeded, VcpuLimitExceeded] -> report VcpuLimitExceeded, a retry hits the same limit
                 throttling = next(
                     (err for err in err_list if err.get("ErrorCode") == LAUNCH_THROTTLING_ERROR_CODE), None
                 )
                 if throttling:
-                    raise LaunchInstancesError(throttling.get("ErrorCode"), throttling.get("ErrorMessage"))
+                    blocking = next(
+                        (
+                            err
+                            for err in err_list
+                            if err.get("ErrorCode") != LAUNCH_THROTTLING_ERROR_CODE
+                            and err.get("ErrorCode") not in POOL_LEVEL_ERROR_CODES
+                        ),
+                        None,
+                    )
+                    chosen = blocking or throttling
+                    raise LaunchInstancesError(chosen.get("ErrorCode"), chosen.get("ErrorMessage"))
             # Normally a single cause is left. Reporting the first one of several is a second safety net: the
             # caller otherwise records a hardcoded InsufficientInstanceCapacity, and any code EC2 actually
             # returned is more useful than an invented one, whichever of them the response happens to list first.

tests/slurm_plugin/test_fleet_manager.py

@@ -19,6 +19,7 @@
 from slurm_plugin.fleet_manager import (
     INSTANCE_INFO_RETRIEVAL_MAX_BACKOFF,
     INSTANCE_INFO_RETRIEVAL_TIMEOUT_DEFAULT,
+    POOL_LEVEL_ERROR_CODES,
     Ec2CreateFleetManager,
     EC2Instance,
     Ec2RunInstancesManager,
@@ -37,13 +38,33 @@
     "ErrorCode": "UnfulfillableCapacity",
     "ErrorMessage": "Unable to fulfill request due to MinTargetCapacity constraints. Please adjust your request.",
 }
+THROTTLING_ERROR = {"ErrorCode": "RequestLimitExceeded", "ErrorMessage": "Request limit exceeded."}
+VCPU_LIMIT_ERROR = {"ErrorCode": "VcpuLimitExceeded", "ErrorMessage": "vCPU limit"}
+SUBNET_FULL_ERROR = {"ErrorCode": "InsufficientFreeAddressesInSubnet", "ErrorMessage": "Not enough free addresses."}
 
 
 def _raises_launch_error(err_list):
     """Mirror when _launch_instances turns a CreateFleet response with no instances into an exception."""
     return bool(err_list)
 
 
+def _expected_launch_error_code(err_list):
+    """Mirror which error _launch_instances reports: a non pool-level cause beats throttling, which beats the rest."""
+    real_errors = [err for err in err_list if err != UNFULFILLED_OVERRIDE] or err_list
+    throttling = next((err for err in real_errors if err["ErrorCode"] == "RequestLimitExceeded"), None)
+    if throttling:
+        blocking = next(
+            (
+                err
+                for err in real_errors
+                if err["ErrorCode"] != "RequestLimitExceeded" and err["ErrorCode"] not in POOL_LEVEL_ERROR_CODES
+            ),
+            None,
+        )
+        return (blocking or throttling)["ErrorCode"]
+    return real_errors[0]["ErrorCode"]
+
+
 def _expected_describe_attempts(timeout):
     """Compute DescribeInstances attempts for a never-converging instance, mirroring _get_instances_info."""
     attempts = 0
@@ -809,7 +830,9 @@ def test_launch_instances(
         elif not expected_assigned_nodes and _raises_launch_error(mocked_boto3_request[0].response.get("Errors", [])):
             with pytest.raises(LaunchInstancesError) as e:
                 fleet_manager._launch_instances(launch_params)
-            assert_that(e.value.code).is_equal_to(mocked_boto3_request[0].response.get("Errors")[0].get("ErrorCode"))
+            assert_that(e.value.code).is_equal_to(
+                _expected_launch_error_code(mocked_boto3_request[0].response.get("Errors"))
+            )
         else:
             assigned_nodes = fleet_manager._launch_instances(launch_params)
             assert_that(assigned_nodes.get("Instances", [])).is_equal_to(expected_assigned_nodes)
@@ -1411,6 +1434,29 @@ def test_launch_ec2_instances_retries_on_throttling(self, mocker):
         assert_that(create_fleet.call_count).is_equal_to(2)
         assert_that(launched).is_length(1)
 
+    def test_launch_ec2_instances_does_not_retry_throttling_with_a_blocking_error(self, mocker):
+        """A cause that fails every pool alike is reported instead of waiting for the rate limit to refill."""
+        mocker.patch("time.sleep")
+        fleet_manager = FleetManagerFactory.get_manager(
+            "hit", "region", "boto3_config", FLEET_CONFIG, "queue2", "fleet-ondemand", True, {}, {}
+        )
+        mocker.patch.object(fleet_manager, "_evaluate_launch_params", return_value={})
+        mocker.patch.object(fleet_manager, "_get_instances_info", return_value=([], []))
+        create_fleet = mocker.patch(
+            "slurm_plugin.fleet_manager.create_fleet",
+            return_value={
+                "Instances": [],
+                "Errors": [THROTTLING_ERROR, VCPU_LIMIT_ERROR] + [UNFULFILLED_OVERRIDE] * 34,
+                "ResponseMetadata": {"RequestId": "1234-abcde"},
+            },
+        )
+
+        with pytest.raises(LaunchInstancesError) as e:
+            fleet_manager.launch_ec2_instances(1)
+
+        assert_that(e.value.code).is_equal_to("VcpuLimitExceeded")
+        assert_that(create_fleet.call_count).is_equal_to(1)
+
     @pytest.mark.parametrize(
         ("err_list", "expected_error_code"),
         [
@@ -1426,6 +1472,18 @@ def test_launch_ec2_instances_retries_on_throttling(self, mocker):
             # Nothing to prefer: report the first entry rather than let a hardcoded code be recorded.
             ([UNFULFILLED_OVERRIDE] * 36, "UnfulfillableCapacity"),
             ([UNSUPPORTED_ERROR, {"ErrorCode": "VcpuLimitExceeded", "ErrorMessage": "vCPU limit"}], "Unsupported"),
+            # Throttling is retried when every other cause is confined to a pool: other pools can serve the batch.
+            ([THROTTLING_ERROR, UNSUPPORTED_ERROR] + [UNFULFILLED_OVERRIDE] * 34, "RequestLimitExceeded"),
+            ([THROTTLING_ERROR, SUBNET_FULL_ERROR], "RequestLimitExceeded"),
+            (
+                [THROTTLING_ERROR, {"ErrorCode": "VolumeTypeNotAvailableInZone", "ErrorMessage": "io2"}],
+                "RequestLimitExceeded",
+            ),
+            # Any other cause would fail the retry on every pool alike, so it is reported instead of the throttling.
+            ([THROTTLING_ERROR, VCPU_LIMIT_ERROR] + [UNFULFILLED_OVERRIDE] * 34, "VcpuLimitExceeded"),
+            ([UNFULFILLED_OVERRIDE] * 34 + [VCPU_LIMIT_ERROR, THROTTLING_ERROR], "VcpuLimitExceeded"),
+            ([THROTTLING_ERROR, UNSUPPORTED_ERROR, VCPU_LIMIT_ERROR], "VcpuLimitExceeded"),
+            ([THROTTLING_ERROR, {"ErrorCode": "SomeFutureCode", "ErrorMessage": "?"}], "SomeFutureCode"),
             # An empty error list is the only case left to the caller, which records insufficient capacity.
             ([], None),
         ],
@@ -1437,6 +1495,13 @@ def test_launch_ec2_instances_retries_on_throttling(self, mocker):
             "single_override_min_target_capacity",
             "only_unfulfilled_overrides",
             "two_real_causes",
+            "throttling_retried_with_capacity_error",
+            "throttling_retried_with_subnet_error",
+            "throttling_retried_with_az_volume_type_error",
+            "quota_error_blocks_retry",
+            "quota_error_blocks_retry_any_order",
+            "first_non_pool_error_reported",
+            "unknown_code_blocks_retry",
             "no_errors_reported",
         ],
     )

tests/slurm_plugin/test_instance_manager.py

@@ -4065,6 +4065,37 @@ def test_launch_instances_reports_throttling_not_insufficient_capacity(self, moc
 
         assert_that(instance_manager.failed_nodes).is_equal_to({"RequestLimitExceeded": {"queue2-dy-fleet-ondemand-1"}})
 
+    def test_launch_instances_reports_blocking_error_instead_of_retrying_throttling(self, mocker, instance_manager):
+        """A cause that fails every pool alike is recorded instead of waiting for the rate limit to refill."""
+        mocker.patch("time.sleep")
+        mocker.patch(
+            "slurm_plugin.fleet_manager.create_fleet",
+            return_value={
+                "Instances": [],
+                "Errors": [
+                    {"ErrorCode": "RequestLimitExceeded", "ErrorMessage": "Request limit exceeded."},
+                    {"ErrorCode": "VcpuLimitExceeded", "ErrorMessage": "vCPU limit"},
+                ]
+                + [
+                    {
+                        "ErrorCode": "UnfulfillableCapacity",
+                        "ErrorMessage": "Failed to fulfill capacity. Please review errors in the response.",
+                    }
+                ]
+                * 28,
+                "ResponseMetadata": {"RequestId": "1234-abcde"},
+            },
+        )
+
+        instance_manager._launch_instances(
+            job=None,
+            nodes_to_launch={"queue2": {"fleet-ondemand": ["queue2-dy-fleet-ondemand-1"]}},
+            launch_batch_size=1,
+            scaling_strategy=ScalingStrategy.BEST_EFFORT,
+        )
+
+        assert_that(instance_manager.failed_nodes).is_equal_to({"VcpuLimitExceeded": {"queue2-dy-fleet-ondemand-1"}})
+
     @pytest.mark.parametrize(
         "job_list, launch_batch_size, assign_node_batch_size, update_node_address, "
         "expected_single_nodes_no_oversubscribe, scaling_strategy",