@pcrespov review: improve message and exception handling

Author: sanderegg

packages/aws-library/src/aws_library/ec2/_client.py

@@ -36,6 +36,7 @@
     check_max_number_of_instances_not_exceeded,
     compose_user_data,
     ec2_instance_data_from_aws_instance,
+    get_subnet_azs,
     get_subnet_capacity,
 )
 
@@ -169,7 +170,7 @@ async def launch_instances(
             # and also allows to circumvent a moto bug that does not raise
             # InsufficientInstanceCapacity when a subnet is full
             subnet_id_to_available_ips = await get_subnet_capacity(
-                self, subnet_ids=instance_config.subnet_ids
+                self.client, subnet_ids=instance_config.subnet_ids
             )
 
             total_available_ips = sum(subnet_id_to_available_ips.values())
@@ -244,14 +245,11 @@ async def launch_instances(
                     raise
 
             else:
-                # All subnets failed with capacity errors
-                _logger.error(
-                    "All subnets failed with insufficient capacity for instance type %s",
-                    instance_config.type.name,
+                subnet_zones = await get_subnet_azs(
+                    self.client, subnet_ids=subnet_ids_with_capacity
                 )
-
                 raise EC2InsufficientCapacityError(
-                    subnet_ids=instance_config.subnet_ids,
+                    availability_zones=subnet_zones,
                     instance_type=instance_config.type.name,
                 )
             instance_ids = [

packages/aws-library/src/aws_library/ec2/_error_handler.py

@@ -68,7 +68,7 @@ def _map_botocore_client_exception(
                 availability_zone = match.group("failed_az")
 
             raise EC2InsufficientCapacityError(
-                availability_zone=availability_zone, instance_type=instance_type
+                availability_zones=availability_zone, instance_type=instance_type
             )
         case _:
             return EC2AccessError(

packages/aws-library/src/aws_library/ec2/_errors.py

@@ -40,7 +40,7 @@ class EC2TooManyInstancesError(EC2AccessError):
 
 class EC2InsufficientCapacityError(EC2AccessError):
     msg_template: str = (
-        "Insufficient capacity in {availability_zone} for {instance_type}"
+        "Insufficient capacity in {availability_zones} for {instance_type}"
     )
 
 

packages/aws-library/src/aws_library/ec2/_utils.py

@@ -1,6 +1,7 @@
 from textwrap import dedent
 from typing import TYPE_CHECKING, cast
 
+from types_aiobotocore_ec2 import EC2Client
 from types_aiobotocore_ec2.type_defs import (
     InstanceTypeDef,
     SubnetTypeDef,
@@ -67,10 +68,9 @@ async def check_max_number_of_instances_not_exceeded(
 
 
 async def get_subnet_capacity(
-    ec2_client: "SimcoreEC2API", *, subnet_ids: list[str]
+    aioboto_ec2_client: EC2Client, *, subnet_ids: list[str]
 ) -> dict[str, int]:
-
-    subnets = await ec2_client.client.describe_subnets(SubnetIds=subnet_ids)
+    subnets = await aioboto_ec2_client.describe_subnets(SubnetIds=subnet_ids)
     assert "Subnets" in subnets  # nosec
     subnet_id_to_subnet_map: dict[str, SubnetTypeDef] = {
         subnet["SubnetId"]: subnet  # pyright: ignore[reportTypedDictNotRequiredAccess]
@@ -85,3 +85,23 @@ async def get_subnet_capacity(
         for subnet_id in subnet_ids
     }
     return subnet_id_to_available_ips
+
+
+async def get_subnet_azs(
+    aioboto_ec2_client: EC2Client, *, subnet_ids: list[str]
+) -> list[str]:
+    subnets = await aioboto_ec2_client.describe_subnets(SubnetIds=subnet_ids)
+    assert "Subnets" in subnets  # nosec
+    subnet_id_to_subnet_map: dict[str, SubnetTypeDef] = {
+        subnet["SubnetId"]: subnet  # pyright: ignore[reportTypedDictNotRequiredAccess]
+        for subnet in subnets["Subnets"]
+    }
+    # preserve the order of instance_config.subnet_ids
+
+    subnet_azs: list[str] = [
+        subnet_id_to_subnet_map[subnet_id][
+            "AvailabilityZone"
+        ]  # pyright: ignore[reportTypedDictNotRequiredAccess]
+        for subnet_id in subnet_ids
+    ]
+    return subnet_azs