albatross/aws.py at master · vishalkuo/albatross · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
import constants
from typing import Optional
import time

# TODO(vishalkuo): refactor to remove custom filters
devserver_filter = {"Name": "tag:application", "Values": [constants.DEVSERVER]}
state_filter = {
    "Name": "instance-state-name",
    "Values": ["pending", "running", "stopping", "stopped"],
}
image_filter = {
    "Name": f"tag:{constants.ALBATROSS_STATUS}",
    "Values": [constants.MARKED_FOR_DELETION],
}


def create_image(client, instance_id: str):
    return client.create_image(
        InstanceId=instance_id, Name=f"albatross-{int(time.time())}"
    )


def get_images(client, include_deleted: bool = True):
    filters = [devserver_filter]
    if include_deleted:
        filters.append(image_filter)
    return client.describe_images(Filters=filters)


def find_devserver(client, non_terminated=False) -> Optional[any]:
    reservations = get_ec2_instances(client, non_terminated=non_terminated)[
        "Reservations"
    ]
    if (
        not reservations
        or not reservations[0]["Instances"]
        or not reservations[0]["Instances"][0]
    ):
        return None

    return reservations[0]["Instances"][0]


def get_ec2_instances(client, non_terminated=False):
    filters = [devserver_filter]
    if non_terminated:
        filters.append(state_filter)
    return client.describe_instances(Filters=filters)


def spawn_devserver(resource, image_id):
    return resource.create_instances(
        MaxCount=1,
        MinCount=1,
        ImageId=image_id,
        SecurityGroups=constants.SECURITY_GROUPS,
        TagSpecifications=[
            {
                "ResourceType": "instance",
                "Tags": [{"Key": "application", "Value": constants.DEVSERVER}],
            }
        ],
        InstanceType=constants.INSTANCE_TYPE,
        KeyName=constants.KEY_NAME,
    )