"""Bulk Asset Scheduler.""" from datetime import datetime import pytz import logging import math import os import random from collections import OrderedDict from sys import exit as sysexit from time import sleep from typing import List, Dict, Tuple from constants import db from constants.fields import ( ACTIVE_STATUSES, SFN_ORPHANED_STATUSES, MAX_REQUEST_TIMEOUT_SECS, WAKEUP_KEY ) from connectors.snowflake import close_connection from utils import ( sfn_utils, snowflake_utils ) import config if not config.SFN_NAME: sysexit('You must export SFN_NAME in your environment with the SFN name.') SLEEP_DELAY = 10 # create logger with 'spam_application' logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) # create console handler with a higher log level ch = logging.StreamHandler() ch.setLevel(logging.INFO) # create formatter and add it to the handlers formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') # noqa ch.setFormatter(formatter) # add the handlers to the logger logger.addHandler(ch) # Get "today" at runtime TODAY = datetime.now(pytz.utc).date().isoformat() def main(): """Main method.""" # TODO: Check the number of instances of SFN_NAME # TODO: Ensure an execution is not running with the same input data # Get the SFN_ARN from the name sfn_arn = \ sfn_utils.get_sfn_arn_by_name(config.SFN_NAME)['stateMachineArn'] execution_count = sfn_utils.get_full_execution_count_by_arn(sfn_arn) # Create queue tables snowflake_utils.create_table_queue() snowflake_utils.create_sfn_queue() # Add a new table to the queue if one is specified in the environment if config.INPUT_TABLE_NAME: # Check if table is in table queue if not snowflake_utils.is_table_in_table_queue( table_name=config.INPUT_TABLE_NAME): # Check if table is in sfn queue logger.info(f'Table name \'{config.INPUT_TABLE_NAME}\' passed to ' 'job; checking for active executions.') # Check if a wake-up table has any executions in the sfn queue table_name_prefix = \ config.ENVIRONMENT+'_'+config.INPUT_TABLE_NAME if not snowflake_utils.get_rows_in_sfn_queue_by_fuzzy_table( table_name_prefix=table_name_prefix, status=ACTIVE_STATUSES): logger.info( f'Adding {config.INPUT_TABLE_NAME} to the queue.') snowflake_utils.add_table_queue_item( table_name=config.INPUT_TABLE_NAME ) else: logger.info( f'{config.INPUT_TABLE_NAME} already has associated ' 'executions in the sfn queue. Not added to table queue') else: logger.info(f'{config.INPUT_TABLE_NAME} already in the queue.') # Look for new source tables, make sub-tables, and queue them logger.info('Looking for source ingestion tables to split.') parent_table_list = snowflake_utils.get_all_tables_to_ingest() for parent_table in parent_table_list: parent_table_name = parent_table[db.SOURCE_TABLE] logger.info(f'Table found: {parent_table_name}') release_list = \ snowflake_utils.get_release_sizes(table_name=parent_table_name) logger.info( f'{parent_table_name} contains {len(release_list)} releases.') if config.DEBUG_WAKEUP: logger.info('DEBUG_WAKEUP is set. Processing only two rows.') table_dict = process_two_rows(release_list, parent_table_name) wakeup_table_prefix = list(table_dict.keys())[0] # Check if table is in sfn queue logger.info( f'Checking for active executions for {wakeup_table_prefix}.') if snowflake_utils.get_rows_in_sfn_queue_by_table( table_name=wakeup_table_prefix, status=ACTIVE_STATUSES): logger.info( f'{wakeup_table_prefix} already has associated ' 'executions in the sfn queue. Not processing') else: # Process all rows table_dict = process_release_list(release_list, parent_table_name) # Remove asset info from table_dict table_dict = OrderedDict( [(k, [upc for upc, _ in v]) for k, v in table_dict.items()]) logger.info(f'{len(table_dict)} tables prepared for creation.') # Create sub-tables logger.info(f'Creating {len(table_dict)} new tables.') for table_name, upc_list in table_dict.items(): snowflake_utils.create_asset_sub_table( upc_list=upc_list, source_table=parent_table_name, sub_table=table_name, wakeup=config.DEBUG_WAKEUP) logger.info(f'\'{table_name}\' created.') # Queue each table in the SFN queue logger.info( f'Queueing new SFN Jobs from {parent_table_name} sub tables.') for table_name in table_dict.keys(): snowflake_utils.add_sfn_queue_item( sfn_name=config.SFN_NAME, source_table=table_name, execution_name=None, status='QUEUED' ) # if not config.DEBUG_WAKEUP: # Removed for now - causing problems # Delete table from table queue logger.info(f'Removing {parent_table_name} from table queue.') snowflake_utils.delete_table_queue_item( table_name=parent_table_name) # Find and log excutions without execution names logger.info('Finding orphaned executions.') orphaned_executions = snowflake_utils.get_orphaned_ingestions() orphan_count = len(orphaned_executions) if orphan_count: logger.info( f'{orphan_count} orphaned executions need execution names.') for execution in orphaned_executions: orphan_found = False # Fuse table_name = execution[db.SOURCE_TABLE] logger.info(f'Finding execution name for {table_name}.') # Get the list of execution names execution_name_list = \ sfn_utils.get_full_execution_list_by_arn( sfn_arn, status=SFN_ORPHANED_STATUSES) # Check for matching executions. for execution in execution_name_list: exec_name = execution['name'] exec_status = execution['status'] exec_start_time = execution['start'] exec_end_time = execution['end'] if table_name in exec_name: logger.info( f'Found execution name {exec_name} for orphaned_table ' f'{table_name}.') snowflake_utils.update_queue_field_by_table( field_name=db.EXECUTION_NAME, value=exec_name, # noqa table_name=table_name ) snowflake_utils.update_queue_field_by_exec( field_name=db.STATUS, value=exec_status, # noqa execution_name=exec_name ) snowflake_utils.update_queue_field_by_exec( field_name=db.START_TIME, value=exec_start_time, # noqa execution_name=exec_name ) if exec_end_time: snowflake_utils.update_queue_field_by_exec( field_name=db.END_TIME, value=exec_end_time, # noqa execution_name=exec_name ) orphan_found = True # Break fuse break if not orphan_found: # If it's been MAX_REQUEST_TIMEOUT_SECS since the start time, # mark it as 'NEVER_STARTED' start_time = execution[db.START_TIME] # Ensure start_time is aware if start_time.tzinfo is None: start_time = start_time.replace(tzinfo=pytz.utc) # Ensure datetime.now() is also aware current_time = datetime.now(pytz.utc) latency = (current_time - start_time).seconds if latency > MAX_REQUEST_TIMEOUT_SECS: logger.info(f'Marking {table_name} as NEVER_STARTED.') snowflake_utils.update_queue_field_by_table( field_name=db.STATUS, value='NEVER_STARTED', # noqa table_name=table_name ) snowflake_utils.update_queue_field_by_table( field_name=db.END_TIME, value=current_time, # noqa table_name=table_name ) # Update in-progress SFN execution statuses in the queue logger.info('Checking in-progress ingestions.') in_progress_queue = snowflake_utils.get_in_progress_ingestions() if len(in_progress_queue): logger.info(f'{len(in_progress_queue)} items in progress.') # Iterate through statuses done_status = ['SUCCEEDED', 'FAILED', 'TIMED_OUT', 'ABORTED'] for status in done_status: finished_exec_list = [] # Find finished SFN executions using AWS API. finished_exec_list.extend( sfn_utils.get_full_execution_list_by_arn( sfn_arn, status) or []) # Create a dictionary for quick lookup of finished executions finished_exec_dict = { exec['name']: exec for exec in finished_exec_list} # find any executions that are in the in_progress_queue # in the finished_exec_list. Use the start and end times # from the finished_exec_list in the # snowflake_utils.update_queue_field_by_exec calls for in_progress in in_progress_queue: exec_name = in_progress[db.EXEC_NAME] if exec_name in finished_exec_dict.keys(): finished_exec = finished_exec_dict[exec_name] logger.info( f'Marking execution {exec_name} completed.') if finished_exec['end']: snowflake_utils.update_queue_field_by_exec( field_name=db.END_TIME, value=finished_exec['end'], # noqa execution_name=exec_name ) if finished_exec['start']: snowflake_utils.update_queue_field_by_exec( field_name=db.START_TIME, value=finished_exec['start'], # noqa execution_name=exec_name ) snowflake_utils.update_queue_field_by_exec( field_name=db.STATUS, value=status, # noqa execution_name=exec_name ) # Execute SFN's for each table in the SFN queue - up to MAX_CONCURRENT_SFN logger.info('Finding TODO items in the queue.') todo_tables = snowflake_utils.get_todo_ingestions() logger.info(f'{len(todo_tables)} items waiting in the queue.') # todo_table_names = [t[db.SOURCE_TABLE] for t in todo_tables] # Drop a JSON for each table for table in todo_tables: table_name = table[db.SOURCE_TABLE] if execution_count < config.MAX_CONCURRENT_SFN: # Example ABOVEBOARD_ASSET_ALL_BATCH02_001 file_name = f'{config.SNOWFLAKE_SCHEMA}_{table_name}' bucket = config.JSON_DROP_BUCKET prefix = config.JSON_DROP_KEY_PREFIX # Execute state machine logger.info(f'Executing state machine for: {table_name}') snowflake_utils.update_queue_field_by_table( field_name=db.TRIGGER_JSON_KEY, value=f's3://{os.path.join(bucket, prefix, file_name)}', table_name=table_name ) snowflake_utils.update_queue_field_by_table( field_name=db.START_TIME, value=datetime.now(pytz.utc), # noqa table_name=table_name ) snowflake_utils.update_queue_field_by_table( field_name=db.STATUS, value='REQUESTED', # noqa table_name=table_name ) snowflake_utils.drop_json_trigger( bucket=bucket, prefix=prefix, file_name=file_name, table_name=table_name ) execution_count += 1 else: break logger.info(f'Sleeping {SLEEP_DELAY} seconds.') # Sleep while executions start sleep(SLEEP_DELAY) # TODO: refactor to method with backoff countdown = 0 exec_name_found = False while not exec_name_found: logger.info(f'Finding execution name for {table_name}.') # Get the list of execution names execution_name_list = \ sfn_utils.get_full_execution_name_list_by_arn(sfn_arn) # Check for matching executions. for exec_name in execution_name_list: if table_name in exec_name: logger.info( f'Found execution name {exec_name} for {table_name}.') snowflake_utils.update_queue_field_by_table( field_name=db.EXECUTION_NAME, value=exec_name, table_name=table_name ) snowflake_utils.update_queue_field_by_exec( field_name=db.STATUS, value='RUNNING', execution_name=exec_name ) exec_name_found = True # Break fuse # If we didn't find the execution name, sleep and try again 5 times if not exec_name_found and countdown <= 5 * SLEEP_DELAY: countdown += SLEEP_DELAY logger.info('Execution name not found. Rechecking in ' f'{SLEEP_DELAY} seconds.') sleep(SLEEP_DELAY) elif not exec_name_found: logger.info('Execution name not found after 5 checks. ' 'Exiting.') exec_name_found = True # Break fuse break def get_assets_per_table(table_name): """Calculate the number of assets per table. Args: table_name (str): Name of the table. Returns: int: Number of assets per table. """ if config.ASSETS_PER_TABLE: # Use the config value logger.info( f'Using config value for max assets per table: {config.ASSETS_PER_TABLE}') # noqa return config.ASSETS_PER_TABLE else: # Calculate based on MAX_CONCURRENT_SFN logger.info( f'Calculating assets per table based on MAX_CONCURRENT_SFN={config.MAX_CONCURRENT_SFN}') # noqa row_count = snowflake_utils.get_table_row_count(table_name=table_name) logger.info(f'Total rows in table: {row_count}') # We remove one because we want to ensure that non-uniform # distributions don't result in small row sets. By adding one, we # ensure that the distribution is eased. This will either result in the # MAX_CONCURRENT_SFN being deployed, or a better distribution across # extra tables as needed. The 500 row limit is a hard limit, so we can # be more aggressive in our calculations. assets_per_table = int(math.ceil(row_count / (config.MAX_CONCURRENT_SFN-1))) # noqa if assets_per_table > 500: assets_per_table = 500 logger.info(f'{assets_per_table} is too high. Max assets per ' 'table set to 500.') else: logger.info(f'Max assets per table: {assets_per_table}') return assets_per_table def process_release_list( release_list: List[Dict], table_name: str ) -> OrderedDict[str, List[Tuple[str, int]]]: """Convert release list into optimally chunked tables using binary search. Uses high/low asset count partitioning for optimal packing. Args: release_list: List of release dicts with UPC and asset counts table_name: Base name for generated table chunks Returns: OrderedDict mapping table names to lists of (UPC, asset_count) tuples """ # Initialize tracking table_dict = OrderedDict() table_num = 1 asset_count = 0 total_assets_seen = 0 release_count = 0 upc_asset_list = [] # Calculate total assets for progress tracking total_asset_count = sum(r[db.ASSET_COUNT] for r in release_list) # Get chunk size limits assets_per_table = get_assets_per_table(table_name) # Process releases while any remain remaining = release_list.copy() while remaining: left_to_go = len(remaining) logger.info(f'{left_to_go} releases left to process.') assets_remaining = sum(r[db.ASSET_COUNT] for r in remaining) logger.info(f'{assets_remaining} assets remaining.') midpoint = left_to_go // 2 # Split and sort by asset count sorted_releases = sorted( remaining, key=lambda x: x[db.ASSET_COUNT], reverse=True ) high_assets = sorted_releases[:midpoint] # Reverse so as to process low assets largest-first low_assets = list(reversed(sorted_releases[midpoint:])) used_releases = set() # Process high asset releases for release in high_assets: if asset_count + release[db.ASSET_COUNT] <= assets_per_table: asset_count += release[db.ASSET_COUNT] total_assets_seen += release[db.ASSET_COUNT] release_count += 1 upc_asset_list.append(( release[db.UPC], release[db.ASSET_COUNT] )) used_releases.add(release[db.UPC]) else: # Fill remaining space with largest of small asset counts for low_release in low_assets: if (low_release[db.UPC] not in used_releases and asset_count + low_release[db.ASSET_COUNT] <= assets_per_table): asset_count += low_release[db.ASSET_COUNT] total_assets_seen += low_release[db.ASSET_COUNT] release_count += 1 upc_asset_list.append(( low_release[db.UPC], low_release[db.ASSET_COUNT] )) used_releases.add(low_release[db.UPC]) # Create table if we added anything if upc_asset_list: tmp_table_name = \ f"{config.ENVIRONMENT}_{table_name}_{str(table_num).zfill(3)}" # noqa logger.info( f'Adding table {tmp_table_name} with ' f'{len(upc_asset_list)} UPCs and {asset_count} assets.' ) table_dict[tmp_table_name] = upc_asset_list table_num += 1 upc_asset_list = [] asset_count = 0 # If we failed to add any releases this iteration, it means all # remaining releases are larger than the assets_per_table limit. # To avoid an infinite loop, assign each oversized release to its # own table (allowed overflow). if not used_releases and not upc_asset_list: oversized = [ r for r in remaining if r[db.ASSET_COUNT] > assets_per_table ] if oversized: logger.warning( 'All remaining releases exceed assets_per_table limit. ' 'Assigning each oversized release to its own table.' ) for release in oversized: tmp_table_name = \ f"{config.ENVIRONMENT}_{table_name}_{str(table_num).zfill(3)}" # noqa logger.info( f'Adding OVERSIZE table {tmp_table_name} with 1 UPC ' f"and {release[db.ASSET_COUNT]} assets (exceeds limit " f"{assets_per_table})." ) table_dict[tmp_table_name] = [( release[db.UPC], release[db.ASSET_COUNT] )] total_assets_seen += release[db.ASSET_COUNT] release_count += 1 used_releases.add(release[db.UPC]) table_num += 1 # Ensure counters are reset for next iteration asset_count = 0 upc_asset_list = [] # Report progress logger.info(f'{total_assets_seen} assets chunked.') logger.info( f'{total_asset_count - total_assets_seen} assets remaining.' ) logger.info(f'{release_count} releases chunked.') logger.info( f'{left_to_go - len(used_releases)} releases remaining.' ) # Remove processed releases remaining = [ r for r in remaining if r[db.UPC] not in used_releases ] # If there are any low assets left, add them to the last table if len(remaining) == 1: logger.info('Last release found. It has ' f'{remaining[0][db.ASSET_COUNT]} assets.') # Add the last release to the list upc_asset_list.append(( remaining[0][db.UPC], remaining[0][db.ASSET_COUNT] )) total_assets_seen += remaining[0][db.ASSET_COUNT] asset_count += remaining[0][db.ASSET_COUNT] release_count += 1 remaining.pop() # Create final table if needed if upc_asset_list: tmp_table_name = \ f"{config.ENVIRONMENT}_{table_name}_{str(table_num).zfill(3)}" logger.info( f'Adding table {tmp_table_name} with ' f'{len(upc_asset_list)} UPCs and {asset_count} assets.' ) table_dict[tmp_table_name] = upc_asset_list table_num += 1 # Final stats logger.info(f'{total_assets_seen} total assets chunked.') logger.info(f'{release_count} total releases chunked.') logger.info(f'{table_num-1} tables prepared.') # Log final table stats for name, upc_list in table_dict.items(): logger.info( f'{name}: {len(upc_list)} upcs, ' f'{sum(count for _, count in upc_list)} assets' ) return table_dict def process_two_rows(release_list, table_name): """Find the release with the least assets, >= 2, assume it has an audio and a cover, and then return that as the tabled_dict. Args: release_list (list): List of releases. table_name (str): Name of the table. """ table_dict = {} upc_list = [] single_asset_list = [] # Find the release_list item with the fewest assets # and ensure it has at least 2 assets release_list = sorted(release_list, key=lambda x: x[db.ASSET_COUNT]) while not upc_list: if not len(release_list) > 0: if single_asset_list: logger.info('No releases with at least 2 assets found. ' 'Using releases with 1 asset.') upc_list = single_asset_list break # Choose a random UPC from the list random_index = random.randint(0, len(release_list)-1) release = release_list.pop(random_index) if release[db.ASSET_COUNT] >= 2: upc_list.append((release[db.UPC], release[db.ASSET_COUNT])) if release[db.ASSET_COUNT] == 1 and len(single_asset_list) < 2: single_asset_list.append((release[db.UPC], release[db.ASSET_COUNT])) # noqa if not len(upc_list): raise ValueError('No releases with at least 2 assets found.') # Add environment to table names to prevent collisions tmp_table_name = config.ENVIRONMENT+'_'+table_name+'_'+WAKEUP_KEY table_dict[tmp_table_name] = upc_list upc_list = [] logger.info( f'{release[db.ASSET_COUNT]} assets chunked for WAKEUP execution.') return table_dict if __name__ == '__main__': main() close_connection()