#!/usr/bin/env python3 # Copyright 2020 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ The STS Job Manager. This tool creates STS Jobs and records each job's state. """ import argparse import json import logging import os import time from datetime import datetime from typing import Dict, List, Optional from google.cloud import bigquery, monitoring_v3 from constants import schemas from constants.status import (KNOWN_STATUSES, STATUS, sts_operation_status_to_table_status) from lib.options import STSJobManagerOptions from lib.services import Services from lib.table_util import get_table_identifier, get_table_ref logging.basicConfig() logger = logging.getLogger(__name__) logger.setLevel(os.environ.get("LOGLEVEL", "INFO").upper()) class Job: def __init__(self, data): self.prefix: str = data.prefix self.status: str = data.status self.job_name: str = data.job_name self.last_updated: datetime = data.last_updated self.operation_results: Optional[dict] = getattr( data, 'operation_results', None) def run_query(query: str, params: Optional[List], services: Services, options: STSJobManagerOptions): """ Runs a given query with optional params. """ job_config = bigquery.QueryJobConfig() job_config.query_parameters = params if params else [] return services.bigquery.query( query, location=options.bigquery_options.dataset_location, job_config=job_config ) def get_jobs_by_prefix(services: Services, options: STSJobManagerOptions) \ -> Dict[str, Job]: """ Retrieves jobs from the database and returns them in a key-value format where the `key` is the prefix and the value is a `Job` object. """ table = get_table_identifier( services, options.bigquery_options, options.bigquery_options.table_name['job']) # API does not support table names for preparameterized queries # https://cloud.google.com/bigquery/docs/parameterized-queries query = f""" SELECT * FROM `{table}` """ # nosec results = run_query(query, None, services, options) prefixToStatus: Dict[str, Job] = {} for row in results: prefixToStatus[row.prefix] = Job(row) return prefixToStatus def set_prefixes_to_status(prefixes: List[str], status: str, services: Services, options: STSJobManagerOptions): """ Sets a list of prefixes to a given status in the database. """ logger.info(f'Updating {len(prefixes)} prefixes to `{status}` status') table = get_table_identifier( services, options.bigquery_options, options.bigquery_options.table_name['job']) # API does not support table names for preparameterized queries # https://cloud.google.com/bigquery/docs/parameterized-queries query = f""" UPDATE `{table}` SET status = @status, last_updated = CURRENT_TIMESTAMP() WHERE prefix IN UNNEST(@prefixes) """ params = [ bigquery.ScalarQueryParameter("status", "STRING", status), bigquery.ArrayQueryParameter("prefixes", "STRING", prefixes) ] run_query(query, params, services, options).result() def set_job_name(prefix: str, job_name: str, services: Services, options: STSJobManagerOptions): """ Set's a prefix's transfer operation job name in the database. """ logger.info( f'Updating the prefix `{prefix}` with job name `{job_name}`...') table = get_table_identifier( services, options.bigquery_options, options.bigquery_options.table_name['job']) # API does not support table names for preparameterized queries # https://cloud.google.com/bigquery/docs/parameterized-queries query = f""" UPDATE `{table}` SET job_name = @job_name, last_updated = CURRENT_TIMESTAMP() WHERE prefix = @prefix """ params = [ bigquery.ScalarQueryParameter("prefix", "STRING", prefix), bigquery.ScalarQueryParameter("job_name", "STRING", job_name) ] run_query(query, params, services, options).result() logger.info( f'...updated the prefix `{prefix}` with job name `{job_name}`.') def insert_history(rows: List[object], services: Services, options: STSJobManagerOptions): """ Inserts a list of rows into the job history table. Each object provided in the list matches the `JOB_HISTORY` schema """ logger.info(f'Inserting {len(rows)} row(s) into the history table') table_ref = get_table_ref( services.bigquery, options.bigquery_options, options.bigquery_options.table_name['job_history']) errors = services.bigquery.insert_rows( table_ref, rows, selected_fields=schemas.JOB_HISTORY) if errors: logger.error('errors were found:') for row in errors: logger.error(row) raise Exception('Error inserting one or more rows') def get_latest_operation_by_prefix(services: Services, options: STSJobManagerOptions): """ Gets the latest transfer operation cooresponding to a prefix. Returns a key-value object where the key is a prefix and the value is a [TransferOperation](https://cloud.google.com/storage-transfer/docs/reference/rest/v1/transferOperations#resource-transferoperation). """ job_filter = json.dumps({"project_id": services.bigquery.project}) request = services.sts.transferOperations().list( name='transferOperations', filter=job_filter, pageSize=256) latest_operation_by_prefix: Dict[str, dict] = {} operation_to_prefix: Dict[str, str] = {} while request is not None: response = request.execute() if not response: break for operation in response['operations']: transfer_spec = operation['metadata']['transferSpec'] if 'objectConditions' not in transfer_spec: continue object_conditions = transfer_spec['objectConditions'] if 'includePrefixes' not in object_conditions: continue if 'gcsDataSource' not in operation['metadata']['transferSpec']: continue if 'gcsDataSink' not in operation['metadata']['transferSpec']: continue if options.source_bucket != operation['metadata']['transferSpec'][ 'gcsDataSource']['bucketName']: continue if options.destination_bucket != operation['metadata'][ 'transferSpec']['gcsDataSink']['bucketName']: continue for prefix in object_conditions['includePrefixes']: operation_to_set_for_prefix = None if prefix not in latest_operation_by_prefix: # The prefix does not have an operation, let's use this one operation_to_set_for_prefix = operation elif 'endTime' not in operation['metadata'] or \ 'endTime' not in latest_operation_by_prefix[prefix][ 'metadata']: # if end time is not available, use the start time if operation['metadata']['startTime'] > \ latest_operation_by_prefix[prefix]['metadata'][ 'startTime']: latest_operation_by_prefix[prefix] = operation elif operation['metadata']['endTime'] > \ latest_operation_by_prefix[prefix]['metadata'][ 'endTime']: # This operation is newer than the assigned operation operation_to_set_for_prefix = operation # Set the operation for the prefix if operation_to_set_for_prefix: # unreference existing operation to prefix, if exists operation_to_prefix.pop(operation['name'], None) latest_operation_by_prefix[prefix] = operation operation_to_prefix[operation['name']] = prefix request = services.sts.transferOperations().list_next( previous_request=request, previous_response=response) # If the latest transferOperation is from a deleted job, we should not # consider the operation for state management deleted_job_request = services.sts.transferJobs().list( filter=json.dumps({ "project_id": services.bigquery.project, "jobStatuses": ["DELETED"] }), pageSize=256) while deleted_job_request is not None: deleted_job_response = deleted_job_request.execute() if not deleted_job_response: break for transferJob in deleted_job_response['transferJobs']: if 'latestOperationName' not in transferJob: continue operation_to_remove = transferJob['latestOperationName'] prefix = operation_to_prefix.pop(operation_to_remove, None) if prefix: latest_operation_by_prefix.pop(prefix, None) deleted_job_request = services.sts.transferJobs().list_next( previous_request=deleted_job_request, previous_response=deleted_job_response) return latest_operation_by_prefix def manage_state(services: Services, options: STSJobManagerOptions): """ Gathers all prefix information from both STS and the database, then updates the corresponding rows where necessary. """ logger.info('Checking state...') # jobs from the database jobs = get_jobs_by_prefix(services, options) # transfer operations from STS latest_operation_by_prefix = get_latest_operation_by_prefix( services, options) history_rows: List[object] = [] job_status_to_update: Dict[str, List[str]] = { STATUS.DONE: [], STATUS.ERROR: [], STATUS.PAUSED: [], STATUS.RUNNING: [], STATUS.WAITING: [] } def append_history(job: Job, operation_results: object): history_rows.append({ 'prefix': job.prefix, 'status': job.status, 'job_name': job.job_name, 'operation_results': json.dumps(operation_results), 'timestamp': datetime.now() }) for prefix in jobs: if prefix in latest_operation_by_prefix: operation_status = \ latest_operation_by_prefix[prefix]['metadata']['status'] expected_status = jobs[prefix].status actual_status = sts_operation_status_to_table_status( operation_status) actual_job_name = latest_operation_by_prefix[prefix]['name'] if actual_status != expected_status: # Capture the history for running jobs logger.info( f'Status for prefix `{prefix}` has changed from \ `{expected_status}` to `{actual_status}`') jobs[prefix].status = actual_status job_status_to_update[actual_status].append(prefix) append_history( jobs[prefix], latest_operation_by_prefix[prefix]) elif actual_status == STATUS.RUNNING: # Capture the history for running jobs append_history( jobs[prefix], latest_operation_by_prefix[prefix]) if actual_job_name != jobs[prefix].job_name: set_job_name(prefix, actual_job_name, services, options) # sleep to avoid rate limiting # https://cloud.google.com/bigquery/quotas#standard_tables time.sleep(2) # Assign the latest `operation_results` jobs[prefix].operation_results = latest_operation_by_prefix[prefix] if history_rows: insert_history(history_rows, services, options) for status in job_status_to_update: if job_status_to_update[status]: set_prefixes_to_status( job_status_to_update[status], status, services, options) # sleep to avoid rate limiting # https://cloud.google.com/bigquery/quotas#standard_tables time.sleep(2) logger.info('...state is up to date.') return jobs def run_jobs(count: int, services: Services, options: STSJobManagerOptions): """ Pulls pending prefixes from the database and either create a new transfer operation or resume an existing one. The `manage_state` function will handle the updates in the job statuses; this keeps DML usage to a minimum """ table = get_table_identifier( services, options.bigquery_options, options.bigquery_options.table_name['job']) # API does not support table names for preparameterized queries # https://cloud.google.com/bigquery/docs/parameterized-queries query = f""" SELECT * FROM `{table}` WHERE status IN UNNEST(@statuses) LIMIT @count """ # nosec pending_statuses = [STATUS.WAITING, STATUS.PAUSED] tryable_statuses = [STATUS.WAITING, STATUS.PAUSED, STATUS.ERROR] statuses = pending_statuses if options.no_retry_on_job_error \ else tryable_statuses params = [ bigquery.ArrayQueryParameter("statuses", "STRING", statuses), bigquery.ScalarQueryParameter("count", "INT64", count), ] results = run_query(query, params, services, options) for row in results: job = Job(row) if job.status == STATUS.PAUSED: operation_request = services.sts.transferOperations().resume( name=job.job_name, body={}) operation_request.execute() logger.info(f'Resumed `{job.prefix}` (job name: {job.job_name}).') else: utc_now = datetime.utcnow() if job.status == STATUS.ERROR: logger.error( f'Retrying errored prefix `{job.prefix}`. \ Previous failed job: {job.job_name}') transfer_job_body = { 'description': f'Created via STS Job Manager - {job.prefix}', 'project_id': services.bigquery.project, 'transfer_spec': { 'object_conditions': { 'include_prefixes': [ job.prefix ] }, 'transfer_options': { 'overwrite_objects_already_existing_in_sink': options.overwrite_dest_objects }, 'gcs_data_source': { 'bucket_name': options.source_bucket }, 'gcs_data_sink': { 'bucket_name': options.destination_bucket } }, 'schedule': { "schedule_start_date": { "year": utc_now.year, "month": utc_now.month, "day": utc_now.day }, "schedule_end_date": { "year": utc_now.year, "month": utc_now.month, "day": utc_now.day } }, 'status': 'ENABLED' } request = services.sts.transferJobs().create( body=transfer_job_body) response = request.execute() logger.info( f'Created new transfer job for `{job.prefix}`: ({response}).') return True def determine_stalled_jobs(jobs: Dict[str, Job], last_jobs: Dict[str, Job]) \ -> List[Job]: stalled_jobs: List[Job] = [] for prefix in jobs: if prefix not in last_jobs: continue current_job = jobs[prefix] last_job = last_jobs[prefix] if current_job.status != STATUS.RUNNING or \ last_job.status != STATUS.RUNNING: continue if not current_job.operation_results or not last_job.operation_results: continue current_counters = \ current_job.operation_results['metadata']['counters'] last_counters = last_job.operation_results['metadata']['counters'] if current_counters and last_counters: has_changed = False for key in current_counters: if key not in last_counters or \ current_counters[key] != last_counters[key]: has_changed = True break if not has_changed: stalled_jobs.append(current_job) return stalled_jobs def manage_jobs(jobs: Dict[str, Job], last_jobs: Dict[str, Job], services: Services, options: STSJobManagerOptions): """ Determines the number of new operations to spin-up, then spins them up. """ def num_new_jobs_to_run(): pending_job_count = 0 current_running_jobs = 0 for prefix in jobs: if jobs[prefix].status == STATUS.RUNNING: current_running_jobs += 1 elif jobs[prefix].status == STATUS.WAITING or \ jobs[prefix].status == STATUS.PAUSED: pending_job_count += 1 elif not options.no_retry_on_job_error and \ jobs[prefix].status == STATUS.ERROR: pending_job_count += 1 if options.allow_new_jobs_when_stalled: stalled_count = len(determine_stalled_jobs(jobs, last_jobs)) current_running_jobs = max(0, current_running_jobs - stalled_count) max_number_jobs_available_to_run = \ options.max_concurrent_jobs - current_running_jobs double_current_job_count = current_running_jobs * 2 if not pending_job_count: logger.info('No jobs available to run') return 0 elif current_running_jobs > options.max_concurrent_jobs: logger.info(f'Will not create any new jobs - too many are running \ (current = {current_running_jobs}, \ max = {options.max_concurrent_jobs})') return 0 elif current_running_jobs == 0 and \ max_number_jobs_available_to_run > 0: logger.info( 'Will prepare initial job, as no other jobs are running') return 1 else: logger.info('Ramping up job count') return min(max_number_jobs_available_to_run, double_current_job_count) logger.info('Managing jobs...') count = num_new_jobs_to_run() if not count: logger.info('...no new jobs to run.') return logger.info(f'...spinning up to {count} new job(s)...') run_jobs(count, services, options) logger.info('...done running jobs.') def publish_heartbeat(jobs: Dict[str, Job], last_jobs: Dict[str, Job], services: Services, options: STSJobManagerOptions, monitoring_types=monitoring_v3.types): """ Publishes status heartbeats """ def publish_timeseries_heartbeat(name: str, value: int, services: Services, project_name: str, monitoring_types=monitoring_v3.types): logger.info(f'Preparing heartbeat for `{name}` (value: {value})...') series = monitoring_types.TimeSeries() series.metric.type = name point = series.points.add() point.value.int64_value = value point.interval.end_time.seconds = int(time.time()) services.monitoring.create_time_series(project_name, [series]) logger.info(f'...published heartbeat `{name}`.') p = options.stackdriver_project if options.stackdriver_project \ else services.bigquery.project monitoring_project_name = services.monitoring.project_path(p) logger.info(f'Preparing heartbeats for `{monitoring_project_name}`...') status_count: Dict[str, int] = {} stalled_count = 0 # Ensure known statuses are published, even if 0 for status in KNOWN_STATUSES: status_count[status] = 0 # Gather raw status counts for prefix in jobs: job = jobs[prefix] # status could be unknown if job.status not in status_count: status_count[job.status] = 0 status_count[job.status] += 1 for status in status_count: name = f'custom.googleapis.com/sts_job_manager/status/{status}' count = status_count[status] publish_timeseries_heartbeat( name, count, services, monitoring_project_name, monitoring_types) for job in determine_stalled_jobs(jobs, last_jobs): logger.warn(f'Job `{job.job_name}` appears to be stalled.') stalled_count += 1 # Publish stalled count stalled_metric = 'custom.googleapis.com/sts_job_manager/metrics/stalled' publish_timeseries_heartbeat( stalled_metric, stalled_count, services, monitoring_project_name, monitoring_types) logger.info('...done publishing heartbeats.') def interval(services: Services, options: STSJobManagerOptions): """ The main state and job running interval. This runs the main lifecycle of this application. """ interval_count = 0 last_state_check = 0.0 last_manage_jobs = 0.0 last_jobs: Dict[str, Job] = {} jobs: Dict[str, Job] = {} while True: logger.info(f'Running main interval #{interval_count}...') start = time.time() job_timeout = start - last_manage_jobs >= options.job_interval metrics_timeout = start - last_state_check >= options.metrics_interval if job_timeout or metrics_timeout: last_jobs = jobs jobs = manage_state(services, options) last_state_check = time.time() if job_timeout: manage_jobs(jobs, last_jobs, services, options) # Regather metrics jobs = manage_state(services, options) last_manage_jobs = time.time() if options.publish_heartbeat: try: publish_heartbeat(jobs, last_jobs, services, options) except Exception as e: logger.error('Failed to publish heartbeat:') logger.exception(e) delta = time.time() - start + options.sleep_timeout logger.info(f'...done running main interval #{interval_count}.\n') if delta > 0: time.sleep(delta) interval_count += 1 def main(options: STSJobManagerOptions): """ The main function. """ logger.info('Initializing STS Job Manager.') services = Services() interval(services, options) if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) options = STSJobManagerOptions() options.setup_arg_parser(parser) args = parser.parse_args() options.assign_from_parsed_args(args) main(options)