added start instances

Author: sanderegg

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

@@ -8,7 +8,7 @@
 import botocore.exceptions
 from aiobotocore.session import ClientCreatorContext
 from aiocache import cached  # type: ignore[import-untyped]
-from pydantic import ByteSize, PositiveInt, parse_obj_as
+from pydantic import ByteSize, PositiveInt
 from servicelib.logging_utils import log_context
 from settings_library.ec2 import EC2Settings
 from types_aiobotocore_ec2 import EC2Client
@@ -29,7 +29,7 @@
     EC2Tags,
     Resources,
 )
-from .utils import compose_user_data
+from .utils import compose_user_data, ec2_instance_data_from_aws_instance
 
 _logger = logging.getLogger(__name__)
 
@@ -193,27 +193,17 @@ async def launch_instances(
             await waiter.wait(InstanceIds=instance_ids)
             _logger.info("instances %s exists now.", instance_ids)
 
-            # get the private IPs
+            # NOTE: waiting for pending ensure we get all the IPs back
             described_instances = await self.client.describe_instances(
                 InstanceIds=instance_ids
             )
+            assert "Instances" in described_instances["Reservations"][0]  # nosec
             instance_datas = [
-                EC2InstanceData(
-                    launch_time=instance["LaunchTime"],
-                    id=instance["InstanceId"],
-                    aws_private_dns=instance["PrivateDnsName"],
-                    aws_public_ip=instance.get("PublicIpAddress", None),
-                    type=instance["InstanceType"],
-                    state=instance["State"]["Name"],
-                    tags=parse_obj_as(
-                        EC2Tags, {tag["Key"]: tag["Value"] for tag in instance["Tags"]}
-                    ),
-                    resources=instance_config.type.resources,
-                )
-                for instance in described_instances["Reservations"][0]["Instances"]
+                await ec2_instance_data_from_aws_instance(self, i)
+                for i in described_instances["Reservations"][0]["Instances"]
             ]
             _logger.info(
-                "%s is available, happy computing!!",
+                "%s is pending now, happy computing!!",
                 f"{instance_datas=}",
             )
             return instance_datas
@@ -245,38 +235,51 @@ async def get_instances(
         all_instances = []
         for reservation in instances["Reservations"]:
             assert "Instances" in reservation  # nosec
-            for instance in reservation["Instances"]:
-                assert "LaunchTime" in instance  # nosec
-                assert "InstanceId" in instance  # nosec
-                assert "PrivateDnsName" in instance  # nosec
-                assert "InstanceType" in instance  # nosec
-                assert "State" in instance  # nosec
-                assert "Name" in instance["State"]  # nosec
-                ec2_instance_types = await self.get_ec2_instance_capabilities(
-                    {instance["InstanceType"]}
-                )
-                assert len(ec2_instance_types) == 1  # nosec
-                assert "Tags" in instance  # nosec
-                all_instances.append(
-                    EC2InstanceData(
-                        launch_time=instance["LaunchTime"],
-                        id=instance["InstanceId"],
-                        aws_private_dns=instance["PrivateDnsName"],
-                        aws_public_ip=instance.get("PublicIpAddress", None),
-                        type=instance["InstanceType"],
-                        state=instance["State"]["Name"],
-                        resources=ec2_instance_types[0].resources,
-                        tags=parse_obj_as(
-                            EC2Tags,
-                            {tag["Key"]: tag["Value"] for tag in instance["Tags"]},
-                        ),
-                    )
-                )
+            all_instances.extend(
+                [
+                    await ec2_instance_data_from_aws_instance(self, i)
+                    for i in reservation["Instances"]
+                ]
+            )
         _logger.debug(
             "received: %s instances with %s", f"{len(all_instances)}", f"{state_names=}"
         )
         return all_instances
 
+    async def start_instances(
+        self, instance_datas: Iterable[EC2InstanceData]
+    ) -> list[EC2InstanceData]:
+        try:
+            instance_ids = [i.id for i in instance_datas]
+            with log_context(
+                _logger,
+                logging.INFO,
+                msg=f"starting instances {instance_ids}",
+            ):
+                await self.client.start_instances(InstanceIds=instance_ids)
+                # wait for the instance to be in a pending state
+                # NOTE: reference to EC2 states https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/ec2-instance-lifecycle.html
+                waiter = self.client.get_waiter("instance_exists")
+                await waiter.wait(InstanceIds=instance_ids)
+                _logger.info("instances %s exists now.", instance_ids)
+                # NOTE: waiting for pending ensure we get all the IPs back
+                aws_instances = await self.client.describe_instances(
+                    InstanceIds=instance_ids
+                )
+                assert len(aws_instances["Reservations"]) == 1  # nosec
+                assert "Instances" in aws_instances["Reservations"][0]  # nosec
+                return [
+                    await ec2_instance_data_from_aws_instance(self, i)
+                    for i in aws_instances["Reservations"][0]["Instances"]
+                ]
+        except botocore.exceptions.ClientError as exc:
+            if (
+                exc.response.get("Error", {}).get("Code", "")
+                == "InvalidInstanceID.NotFound"
+            ):
+                raise EC2InstanceNotFoundError from exc
+            raise  # pragma: no cover
+
     async def stop_instances(self, instance_datas: Iterable[EC2InstanceData]) -> None:
         try:
             with log_context(

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

@@ -1,4 +1,12 @@
 from textwrap import dedent
+from typing import TYPE_CHECKING, cast
+
+from types_aiobotocore_ec2.type_defs import InstanceTypeDef
+
+from .models import EC2InstanceData, EC2Tags
+
+if TYPE_CHECKING:
+    from .client import SimcoreEC2API
 
 
 def compose_user_data(docker_join_bash_command: str) -> str:
@@ -8,3 +16,30 @@ def compose_user_data(docker_join_bash_command: str) -> str:
 {docker_join_bash_command}
 """
     )
+
+
+async def ec2_instance_data_from_aws_instance(
+    ec2_client: "SimcoreEC2API",
+    instance: InstanceTypeDef,
+) -> EC2InstanceData:
+    assert "LaunchTime" in instance  # nosec
+    assert "InstanceId" in instance  # nosec
+    assert "PrivateDnsName" in instance  # nosec
+    assert "InstanceType" in instance  # nosec
+    assert "State" in instance  # nosec
+    assert "Name" in instance["State"]  # nosec
+    ec2_instance_types = await ec2_client.get_ec2_instance_capabilities(
+        {instance["InstanceType"]}
+    )
+    assert len(ec2_instance_types) == 1  # nosec
+    assert "Tags" in instance  # nosec
+    return EC2InstanceData(
+        launch_time=instance["LaunchTime"],
+        id=instance["InstanceId"],
+        aws_private_dns=instance["PrivateDnsName"],
+        aws_public_ip=instance.get("PublicIpAddress", None),
+        type=instance["InstanceType"],
+        state=instance["State"]["Name"],
+        resources=ec2_instance_types[0].resources,
+        tags=cast(EC2Tags, {tag["Key"]: tag["Value"] for tag in instance["Tags"]}),
+    )

packages/aws-library/tests/test_ec2_client.py

@@ -5,6 +5,7 @@
 
 import random
 from collections.abc import AsyncIterator, Callable
+from dataclasses import fields
 from typing import cast, get_args
 
 import botocore.exceptions
@@ -343,7 +344,7 @@ async def test_get_instances(
             assert not instance_received
 
 
-async def test_stop_instances(
+async def test_stop_start_instances(
     simcore_ec2_api: SimcoreEC2API,
     ec2_client: EC2Client,
     faker: Faker,
@@ -389,6 +390,25 @@ async def test_stop_instances(
         expected_state="stopped",
     )
 
+    # start the instances now
+    started_instances = await simcore_ec2_api.start_instances(created_instances)
+    await _assert_instances_in_ec2(
+        ec2_client,
+        expected_num_reservations=1,
+        expected_num_instances=num_instances,
+        expected_instance_type=ec2_instance_config.type,
+        expected_tags=ec2_instance_config.tags,
+        expected_state="running",
+    )
+    # the public IPs change when the instances are stopped and started
+    for s, c in zip(started_instances, created_instances, strict=True):
+        # the rest shall be the same
+        for f in fields(EC2InstanceData):
+            if f.name == "aws_public_ip":
+                assert getattr(s, f.name) != getattr(c, f.name)
+            else:
+                assert getattr(s, f.name) == getattr(c, f.name)
+
 
 async def test_terminate_instance(
     simcore_ec2_api: SimcoreEC2API,