tasks_per_container.py

#!/usr/bin/env python3
import datetime
import boto3

def main():
    # parse args...
    ecs_client = boto3.client("ecs", "eu-west-1")
    tasks_list = ecs_client.list_tasks(cluster='my_cluster')
    tasks_per_instance = _tasks_per_instance(
        ecs_client, 'my_cluster', tasks_list['taskArns'])

    # Nothing to do if we only have one instance
    if len(tasks_per_instance) == 1:
        return

    if _is_unbalanced(tasks_per_instance.values()):
        for task in tasks_list['taskArns']:
            ecs_client.stop_task(
                cluster='my_cluster',
                task=task,
                reason='Unbalanced instances' # Shown in the ECS UI and logs
            )

def _get_instances(ecs_client) -> dict:
    running_instances = ecs_client.list_container_instances(
        cluster='baker-omfg')['containerInstanceArns']
    return {instance: 0 for instance in running_instances}

def _tasks_per_instance(ecs_client, cluster, tasks_list: list) -> dict:
    instances = _get_instances(ecs_client)
    tasks_desc = ecs_client.describe_tasks(cluster=cluster, tasks=tasks_list)
    for t in tasks_desc['tasks']:
        instances[t['containerInstanceArn']] += 1
    response = ecs_client.describe_container_instances(
        cluster=cluster, containerInstances=list(instances.keys()))
    instances_as_ids = {
        i['ec2InstanceId']: instances[i['containerInstanceArn']]
        for i in response['containerInstances'] if i['status'] == 'ACTIVE'
    }
    return instances_as_ids

def _is_unbalanced(values: list) -> bool:
    if max(values) < 10: # too few instances
        return False
    avg = sum(values) / float(len(values))
    for v in values:
        if v == 0: # Quick check to see if an instance is totally missing tasks
            return True
        x = abs((v-avg) / avg)
        if x > 0.30:
            return True
    return False

if __name__ == '__main__':
    main()