"""Kickoff task runs.""" import argparse import time import uuid import boto3 from botocore.exceptions import ClientError DMS_MAP = { 1: 'Apple', 187: 'Amazon', 399: 'Tidal', 838: 'Kakao M', 1279: 'Genie Music', 1454: 'All Saints' } def main(): """Entrypoint.""" args = get_args() env = args.env dms_ids = args.dms_ids stereo_upc = args.stereo_upc spatial_upc = args.spatial_upc # constants service_name = 'dolby-atmos-packager-alpha' cluster_name = f'{env}-{service_name}' aws_region = 'us-east-1' num_tasks_to_run = 1 task_started_by = 'jenkins' task_poll_timeout = 1800 poll_interval = 5 # boto clients ec2_client = boto3.client('ec2', region_name=aws_region) ecs_client = boto3.client('ecs', region_name=aws_region) # necessary data to run task vpc_id = get_vpc_id(ec2_client, env) security_group_id = get_security_group(ec2_client, env, service_name, vpc_id) # noqa:E501 subnet_ids = get_subnets(ec2_client, vpc_id) # start multiple tasks tasks = dict() for dms_id in dms_ids: job_id = str(uuid.uuid4()) overrides = { 'STEREO_UPC': stereo_upc, 'SPATIAL_UPC': spatial_upc, 'DMS_ID': dms_id, 'JOB_ID': job_id } override_list = [ {'name': k, 'value': str(v)} for k, v in overrides.items() ] task_arn = run_task( ecs_client, env, cluster_name, service_name, task_started_by, num_tasks_to_run, security_group_id, subnet_ids, override_list ) tasks[task_arn] = { 'dms_id': dms_id, 'job_id': job_id, 'stopped': False } print(f'Started: {_format_dms_details(dms_id)} | {job_id}') task_error = False poll_end_time = time.time() + task_poll_timeout while True: # waited too long, exit if time.time() >= poll_end_time: raise SystemExit('Task(s) have timed out') # find tasks still running task_arns = [k for k, v in tasks.items() if not v['stopped']] # no more tasks running, end if not task_arns: break # pause time.sleep(poll_interval) # get current statuses statuses = get_tasks_status(ecs_client, cluster_name, task_arns) # look for stopped tasks for task_arn, task in tasks.items(): status = statuses.get(task_arn, {}) if not status: continue # task is stopped mark is as such stop_code = status.get('stopCode') if stop_code: # report on task results if stop_code != 'EssentialContainerExited': task_error = True message = f'unexpected task stop code {stop_code}' else: container = next( container for container in status['containers'] if container['name'] == service_name ) container_exit_code = container['exitCode'] if container_exit_code != 0: task_error = True message = f'container exit code {container_exit_code}' print(f"Stopped: {_format_dms_details(task['dms_id'])} | {task['job_id']} => {message}") # noqa:E501 # stop querying for status task['stopped'] = True # fail script if any tasks failed exit_code = 1 if task_error else 0 exit(exit_code) def _format_dms_details(dms_id): max_id_len = max([len(str(k)) for k, _ in DMS_MAP.items()]) max_name_len = max([len(v) for _, v in DMS_MAP.items()]) padded_dms_id = str(dms_id).ljust(max_id_len, ' ') padded_dms_name = DMS_MAP[dms_id].ljust(max_name_len, ' ') return f'{padded_dms_id} | {padded_dms_name}' def get_args(): """Parse CLI args.""" parser = argparse.ArgumentParser() parser.add_argument( 'env', type=str, choices=['qa', 'prod'], help='environment to run task in' ) parser.add_argument( 'stereo_upc', type=str, help='stereo upc to package' ) parser.add_argument( 'spatial_upc', type=str, help='spatial upc to package' ) parser.add_argument( 'dms_ids', type=int, nargs='+', choices=[k for k, _ in DMS_MAP.items()], help='integer dms ids to package for' ) return parser.parse_args() def get_tasks_status(ecs_client, cluster_name, task_arns): """Query for status of tasks by arns.""" return { x['taskArn']: x for x in ecs_client.describe_tasks( cluster=cluster_name, tasks=task_arns )['tasks'] } """ Functions inspired by python-deployment-utils/fargate/run_fargate_task.py """ def run_task( ecs_client, env, cluster_name, service_name, task_started_by, num_tasks_to_run, security_group_id, subnet_ids, environment_overrides): """Run task and returns info about it, once started.""" if environment_overrides: container_overrides = [ { 'name': service_name, 'environment': environment_overrides, } ] else: container_overrides = [] try: response = ecs_client.run_task( cluster=cluster_name, count=num_tasks_to_run, launchType='FARGATE', networkConfiguration={ 'awsvpcConfiguration': { 'subnets': subnet_ids, 'securityGroups': [ security_group_id, ], 'assignPublicIp': 'DISABLED' } }, overrides={ 'containerOverrides': container_overrides, }, startedBy=task_started_by, tags=[ { 'key': 'environment', 'value': env }, { 'key': 'service_name', 'value': service_name }, ], # Always run the latest task definition taskDefinition=cluster_name ) # If starting more than one task, the first should be representative if response['tasks']: task_arn = response['tasks'][0]['taskArn'] return task_arn elif response['failures']: raise SystemExit('Failure error running task') except ClientError: raise SystemExit('Client error running task') def get_vpc_id(ec2_client, env): """Get VPC corresponding to environment.""" response = ec2_client.describe_vpcs( Filters=[ { 'Name': 'tag:Name', 'Values': ['prod'] } ] ) if response['Vpcs']: vpc_id = response['Vpcs'][0]['VpcId'] return vpc_id else: raise SystemExit('Could not retrieve VPC ID') def get_security_group(ec2_client, env, service_name, vpc_id): """Get security group corresponding to service.""" security_group_name = f'{env}-{service_name}-task-security-group' response = ec2_client.describe_security_groups( Filters=[ { 'Name': 'group-name', 'Values': [ security_group_name, ] }, { 'Name': 'vpc-id', 'Values': [ vpc_id, ] }, ], ) if response['SecurityGroups']: """ Infrastructure automation enforces uniqueness, but just in case, pick the first returned group """ security_group_id = response['SecurityGroups'][0]['GroupId'] return security_group_id else: raise SystemExit('Could not retrieve security group') def get_subnets(ec2_client, vpc_id): """Get subnet IDs with most available IP addresses.""" response = ec2_client.describe_subnets( Filters=[ { 'Name': 'tag:Name', 'Values': [ '*private*', ] }, { 'Name': 'tag:tier', 'Values': [ 'private', ] }, { 'Name': 'vpc-id', 'Values': [ vpc_id, ] }, ], ) if response['Subnets']: """ Find and return up to two subnets with the most available IP addresses """ sorted_subnets = sorted( response['Subnets'], key=lambda k: k['AvailableIpAddressCount'], reverse=True ) subnet_ids = [subnet['SubnetId'] for subnet in sorted_subnets] return subnet_ids[0:2] else: raise SystemExit('Could not retrieve subnet IDs') if __name__ == '__main__': main()