Update ECS instance types using spot info

update_ecs_instance_types.py

"""Update ECS Instance Types.

# Check if spotinfo command is installed
if ! command -v spotinfo &> /dev/null
then
  echo "spotinfo could not be found, please install it"
  # Brew install spotinfo
  brew tap alexei-led/spotinfo
  brew install spotinfo
fi
python infrastructure/scripts/update_ecs_instance_types.py
"""
import csv
import math
import os
import re
import shlex
import subprocess
from distutils.util import strtobool

import yaml

DEFAULT_MEMORY = 1024 * 4
DEFAULT_CPU = 1024
DEFAULT_PRICE = 0.20
AMOUNT_OF_INSTANCE_TYPES = 10
DEFAULT_USE_ARM_INSTANCES = True

NON_GRAVITON_MATCHER = re.compile(r"[a-z]\d[^g]?\..+")


def run() -> None:
    print("Updating ECS Instance Types...")
    for dir_path, _, file_names in os.walk("./infrastructure"):
        for file_name in file_names:
            if file_name.endswith(".yml"):
                file_path = os.path.join(dir_path, file_name)
                with open(file_path) as f:
                    try:
                        yaml_data = yaml.safe_load(f)
                    except yaml.constructor.ConstructorError:
                        continue

                if isinstance(yaml_data, list):
                    continue

                if yaml_data.get("parameters") and yaml_data["parameters"].get(
                    "EcsInstanceTypes"
                ):
                    memory = yaml_data["parameters"].get(
                        "MemoryLimit",
                        yaml_data["parameters"].get(
                            "Memory",
                            yaml_data["parameters"].get("MemoryReservation"),
                        ),
                    )
                    cpu_reservation = yaml_data["parameters"].get("CPUReservation")

                    arm_defined = DEFAULT_USE_ARM_INSTANCES or strtobool(
                        yaml_data["parameters"].get("UseArmInstances", "False")
                    )
                    memory_in_gb = int(
                        math.ceil((int(memory or DEFAULT_MEMORY) + 512) / 1024)
                    )
                    cpu_in_count = int(
                        math.ceil(
                            (int(cpu_reservation or (DEFAULT_CPU / 2)) * 2.0) / 1024
                        )
                    )

                    # Call the spotinfo command using subprocess
                    if arm_defined:
                        output = subprocess.run(
                            [
                                "/usr/local/bin/spotinfo",
                                "--type",
                                r"^[[:alnum:]]{2}g\.\S*",
                                "--memory",
                                shlex.quote(str(int(memory_in_gb))),
                                "--cpu",
                                shlex.quote(str(int(cpu_in_count))),
                                "--price",
                                shlex.quote(str(float(DEFAULT_PRICE))),
                                "--region",
                                "eu-west-1",
                                "--output",
                                "csv",
                                "--sort",
                                "price",
                                "--order",
                                "asc",
                            ],
                            capture_output=True,
                            shell=False,
                        )
                    else:
                        output = subprocess.run(
                            [
                                "/usr/local/bin/spotinfo",
                                "--memory",
                                shlex.quote(str(int(memory_in_gb))),
                                "--cpu",
                                shlex.quote(str(int(cpu_in_count))),
                                "--price",
                                shlex.quote(str(float(DEFAULT_PRICE))),
                                "--region",
                                "eu-west-1",
                                "--output",
                                "csv",
                                "--sort",
                                "price",
                                "--order",
                                "asc",
                            ],
                            capture_output=True,
                            shell=False,
                        )

                    output_bytes = output.stdout.decode("utf-8")
                    csv_reader = csv.reader(output_bytes.splitlines(), delimiter=",")
                    data = list(csv_reader)[2:]

                    # If not USE_GRAVITON, then filter out graviton instances
                    if not arm_defined:
                        # Filter out graviton instances
                        data = [
                            item for item in data if NON_GRAVITON_MATCHER.match(item[0])
                        ]

                    # else:
                    #     data = [item for item in data if "7" not in item[0]]

                    # Filter to not include .metal instances
                    data = [item for item in data if ".metal" not in item[0]]

                    data = data[:AMOUNT_OF_INSTANCE_TYPES]

                    instance_names = [row[0] for row in data]
                    average_price = sum([float(row[5]) for row in data]) / len(data)
                    instance_declaration = ",".join(instance_names)
                    print(
                        file_name,
                        f"{memory_in_gb}gb",
                        "x86-x64" if not arm_defined else "arm",
                        cpu_in_count,
                        instance_declaration,
                        f"{round(average_price * 100) / 100}$ hour",
                        f"{round(average_price * 24 * 365 / 12 * 100) / 100}$ month",
                        sep="\t\t",
                    )

                    if (
                        not instance_names
                        or len(instance_names) < AMOUNT_OF_INSTANCE_TYPES
                    ):
                        print("Not enough instance types found, skipping")
                        continue

                    # Write new ECS Instance Types to yaml file
                    yaml_data["parameters"]["EcsInstanceTypes"] = instance_declaration
                    with open(file_path, "w") as f:
                        # Write the new yaml data to the file, nicely formatted
                        yaml.dump(yaml_data, f, sort_keys=False, indent=4)
    print("Done updating ECS Instance Types.")


if __name__ == "__main__":
    run()