gen_workers.py

#!/usr/bin/env python3

import boto3
import datetime
import statistics
import itertools
import jinja2
import os

list_of_templates = [
    'ci-configuration/worker-pools.yml.jinja2',
    'mozilla-central/taskcluster/ci/test/test-platforms.yml.jinja2',
]

IO1_GB_PRICE = 0.125
GP2_GB_PRICE = 0.10
ST1_GB_PRICE = 0.045
IOPS_PRICE = 0.065

TEMPLATE_DIR = '/home/vagrant/work'

MAX_WORKER_PRICE = 0.15

def get_spot_prices():
    ec2 = boto3.client('ec2')
    instance_types = ['m5.metal', 'm5d.metal', 'c5.metal', 'c5d.metal', 'r5.metal', 'r5d.metal']

    resp = ec2.describe_spot_price_history(
        StartTime = datetime.datetime.now() - datetime.timedelta(days=30),
        EndTime = datetime.datetime.now(),
        InstanceTypes = instance_types,
    )

    spot_prices = []
    spot_prices.append((x['InstanceType'], float(x['SpotPrice'])) for x in resp['SpotPriceHistory'])

    while resp['NextToken']:
        resp = ec2.describe_spot_price_history(NextToken = resp['NextToken'])
        spot_prices.append((x['InstanceType'], float(x['SpotPrice'])) for x in resp['SpotPriceHistory'])

    spot_prices = list(itertools.chain(*spot_prices))
    
    return {
        instance_type: statistics.mean(price for it, price in spot_prices if it == instance_type)
            for instance_type in instance_types
    }

class InstanceConfig:
    spot_prices = {
        'm5.metal': 2.1333382333553064,
        'm5d.metal': 2.2497,
        'c5.metal': 2.7161,
        'c5d.metal': 2.7699908235294117,
        'r5.metal': 2.128372814910026,
        'r5d.metal': 2.5997717747216766,
    } # get_spot_prices()

    def __init__(self, instance_type, number_of_tasks, disk_type=None, iops=None):
        if instance_type.endswith('5d.metal'):
            disk_type = 'ssd'
            self.spot = False
        else:
            self.spot = True

        if disk_type == 'io1' and iops is None:
            raise ValueError('Disk type is io1 but Iops was not given')

        if disk_type is None:
            raise ValueError('Must provide a disk_type')

        self.instance_type = instance_type
        self.number_of_tasks = number_of_tasks
        self.disk_type = disk_type
        self.iops = iops

    def disk_size(self):
        return self.number_of_tasks * 40

    def price(self):
        if self.disk_type == 'gp2':
            disk_price = self.disk_size() * GP2_GB_PRICE
        elif self.disk_type == 'io1':
            disk_price = self.disk_size() * IO1_GB_PRICE + self.iops * IOPS_PRICE
        elif self.disk_type == 'st1':
            disk_price = self.disk_size() * ST1_GB_PRICE
        elif self.disk_type == 'ssd':
            disk_price = 0
        else:
            raise ValueError(f'Unrecognized disk type {self.disk_type}')
        
        return (self.spot_prices[self.instance_type] + disk_price / 730) / self.number_of_tasks

    def worker_name(self):
        # remove the '.metal' suffix
        instance_name = self.instance_type[:self.instance_type.index('.')]

        wn = f'metal-{instance_name}-{self.number_of_tasks}-{self.disk_type}-{self.disk_size()}'
        if self.disk_type == 'io1':
            wn += f'-iops-{self.iops}'
        return wn

    def __str__(self):
        return self.worker_name()

    def __repr__(self):
        return self.worker_name()

def process_templates():
    tasks = range(16, 29, 4)
    iops = range(16000, 33000, 4000)

    print('Instances prices:')
    print(InstanceConfig.spot_prices)

    spot_instances = InstanceConfig.spot_prices.keys()

    configs = [
        InstanceConfig(i, t, 'io1', iops) for i, t, iops in itertools.product(
            filter(lambda i: i.endswith('5.metal'), spot_instances),
            tasks,
            iops,
        )
    ]

    configs.extend(InstanceConfig(i, t) for i, t in itertools.product(
        filter(lambda i: i.endswith('5d.metal'), spot_instances),
        tasks,
    ))

    configs.extend(InstanceConfig(i, 15, 'gp2') for i in
        filter(lambda i: i.endswith('5.metal'), spot_instances)
    )

    configs = tuple(filter(lambda c: c.price() <= MAX_WORKER_PRICE, configs))

    print('Configs:')
    for c in configs:
        print(f'{c}: $ {c.price()}')

    env = jinja2.Environment(loader=jinja2.FileSystemLoader(TEMPLATE_DIR))
    for tpl in list_of_templates:
        template = env.get_template(tpl)
        result = template.render(configs=configs, variants = ['opt', 'debug'])
        output_name = os.path.join(TEMPLATE_DIR, os.path.splitext(tpl)[0])
        with open(output_name, 'w') as f:
            f.write(result)

if __name__ == '__main__':
    process_templates()