# Copyright 2022-2023 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. import uuid, hashlib, datetime, json, configparser, math import constants from google.cloud import firestore from google.cloud import tasks_v2 from google.protobuf import duration_pb2 from google.api_core.client_info import ClientInfo USER_AGENT = 'cloud-solutions/datacatalog-tag-engine-v2' class TaskManager: """Class for creating and managing work requests in the form of cloud tasks cloud_run_sa = Cloud Run service account project = Cloud Run project id (e.g. tag-engine-project) region = Cloud Run region (e.g. us-central1) queue_name = Cloud Task queue (e.g. tag-engine-queue) task_handler_uri = task handler uri in the Flask app hosted by Cloud Run """ def __init__(self, tag_engine_sa, tag_engine_project, tag_engine_region, tag_engine_queue, task_handler_uri, db_project, db_name): self.tag_engine_sa = tag_engine_sa self.tag_engine_project = tag_engine_project self.tag_engine_region = tag_engine_region self.tag_engine_queue = tag_engine_queue self.task_handler_uri = task_handler_uri self.db = firestore.Client(project=db_project, database=db_name, client_info=ClientInfo(user_agent=USER_AGENT)) self.tasks_per_shard = 1000 self.task_deadline = duration_pb2.Duration().FromSeconds(1800) # 30 minutes per task (max supported duration) ##################### API METHODS ################# def create_config_uuid_tasks(self, tag_creator_account, tag_invoker_account, job_uuid, config_uuid, config_type, uris): # create shards of 1000 tasks if len(uris) > self.tasks_per_shard: shards = math.ceil(len(uris) / self.tasks_per_shard) else: shards = 1 task_running_total = 0 task_counter = 0 for shard_index in range(0, shards): shard_id_raw = job_uuid + str(shard_index) shard_uuid = hashlib.md5(shard_id_raw.encode()).hexdigest() self._create_shard(job_uuid, shard_uuid) for uri_index, uri_val in enumerate(uris[task_running_total:], task_running_total): # create the task if isinstance(uri_val, str): task_id_raw = job_uuid + uri_val + str(datetime.datetime.utcnow()) if isinstance(uri_val, tuple): task_id_raw = job_uuid + ''.join(uri_val) + str(datetime.datetime.utcnow()) # uri_val is a tuple when it contains a gcs path task_id = hashlib.md5(task_id_raw.encode()).hexdigest() task_uuid = self._record_config_uuid_task(job_uuid, shard_uuid, task_id, config_uuid, config_type, uri_val) self._create_config_uuid_task(tag_creator_account, tag_invoker_account, job_uuid, shard_uuid, task_uuid, task_id, config_uuid, config_type, uri_val) task_counter += 1 task_running_total += 1 if task_counter == self.tasks_per_shard: self._update_shard_tasks(job_uuid, shard_uuid, task_counter) task_counter = 0 break # update shard with last task_counter if task_counter > 0: self._update_shard_tasks(job_uuid, shard_uuid, task_counter) def create_tag_extract_tasks(self, tag_creator_account, tag_invoker_account, job_uuid, config_uuid, config_type, tag_extract_list): # create shards of 5000 records if len(tag_extract_list) > self.tasks_per_shard: shards = math.ceil(len(tag_extract_list) / self.tasks_per_shard) else: shards = 1 task_running_total = 0 task_counter = 0 for shard_index in range(0, shards): shard_id_raw = job_uuid + str(shard_index) shard_uuid = hashlib.md5(shard_id_raw.encode()).hexdigest() self._create_shard(job_uuid, shard_uuid) for extract_index, extract_val in enumerate(tag_extract_list[task_running_total:], task_running_total): print('task_running_total: ', task_running_total) print('extract_index: ', extract_index) print('extract_val: ', extract_val) # create the task task_id_raw = job_uuid + ''.join(str(extract_val)) + str(datetime.datetime.utcnow()) task_id = hashlib.md5(task_id_raw.encode()).hexdigest() #print('task_id: ', task_id) task_uuid = self._record_tag_extract_task(job_uuid, shard_uuid, task_id, config_uuid, config_type, extract_val) self._create_tag_extract_task(tag_creator_account, tag_invoker_account, job_uuid, shard_uuid, task_uuid, task_id, config_uuid, config_type, extract_val) task_counter += 1 task_running_total += 1 if task_counter == self.tasks_per_shard: self._update_shard_tasks(job_uuid, shard_uuid, task_counter) task_counter = 0 break # update shard with last task_counter if task_counter > 0: self._update_shard_tasks(job_uuid, shard_uuid, task_counter) def update_task_status(self, shard_uuid, task_uuid, status): if status == 'RUNNING': self._set_task_running(shard_uuid, task_uuid) self._set_rollup_tasks_running(shard_uuid) if status == 'SUCCESS': self._set_task_success(shard_uuid, task_uuid) self._set_rollup_tasks_success(shard_uuid) if status == 'ERROR': self._set_task_failed(shard_uuid, task_uuid) self._set_rollup_tasks_failed(shard_uuid) ################ INTERNAL PROCESSING METHODS ################# def _create_shard(self, job_uuid, shard_uuid): print('*** _create_shard ***') print('job_uuid: ' + job_uuid + ', shard_uuid: ' + shard_uuid) shard_ref = self.db.collection('shards').document(shard_uuid) shard_ref.set({ 'shard_uuid': shard_uuid, 'job_uuid': job_uuid, 'tasks_ran': 0, 'tasks_success': 0, 'tasks_failed': 0, 'creation_time': datetime.datetime.utcnow() }) def _update_shard_tasks(self, job_uuid, shard_uuid, task_counter): #print('*** _update_shard ***') self.db.collection('shards').document(shard_uuid).update({'task_count': task_counter}); def _record_config_uuid_task(self, job_uuid, shard_uuid, task_id, config_uuid, config_type, uri): #print('*** _record_config_uuid_task ***') task_uuid = uuid.uuid1().hex task_ref = self.db.collection('shards').document(shard_uuid).collection('tasks').document(task_uuid) task_ref.set({ 'task_uuid': task_uuid, # task identifier in Firestore 'task_id': task_id, # cloud task identifier, based on uri 'shard_uuid': shard_uuid, # shard which this task belongs to 'job_uuid': job_uuid, 'config_uuid': config_uuid, 'config_type': config_type, 'uri': uri, 'status': 'PENDING', 'creation_time': datetime.datetime.utcnow() }) #print('created task record ' + task_uuid + ' in shard ' + shard_uuid) return task_uuid def _record_tag_extract_task(self, job_uuid, shard_uuid, task_id, config_uuid, config_type, extract): print('*** _record_task ***') task_uuid = uuid.uuid1().hex task_ref = self.db.collection('shards').document(shard_uuid).collection('tasks').document(task_uuid) task_ref.set({ 'task_uuid': task_uuid, # task identifier in Firestore 'task_id': task_id, # cloud task identifier, based on uri 'shard_uuid': shard_uuid, # shard which this task belongs to 'job_uuid': job_uuid, 'config_uuid': config_uuid, 'config_type': config_type, 'tag_extract': extract, 'status': 'PENDING', 'creation_time': datetime.datetime.utcnow() }) #print('created task record ' + task_uuid + ' in shard ' + shard_uuid) return task_uuid def _create_config_uuid_task(self, tag_creator_account, tag_invoker_account, job_uuid, shard_uuid, task_uuid, task_id, \ config_uuid, config_type, uri): success = True payload = {'job_uuid': job_uuid, 'shard_uuid': shard_uuid, 'task_uuid': task_uuid, 'config_uuid': config_uuid, \ 'config_type': config_type, 'uri': uri, 'tag_creator_account': tag_creator_account, \ 'tag_invoker_account': tag_invoker_account} client = tasks_v2.CloudTasksClient() parent = client.queue_path(self.tag_engine_project, self.tag_engine_region, self.tag_engine_queue) task = { 'name': parent + '/tasks/' + task_id, 'dispatch_deadline': self.task_deadline, 'http_request': { 'http_method': 'POST', 'url': self.task_handler_uri, 'headers': {'content-type': 'application/json'}, 'body': json.dumps(payload).encode(), 'oidc_token': {'service_account_email': self.tag_engine_sa, 'audience': self.task_handler_uri} } } print('task request:', task) try: task = client.create_task(parent=parent, task=task) #print('task response:', task) except Exception as e: print('Error: could not create task for uri', self.task_handler_uri, '. Error: ', e) self._set_task_failed(shard_uuid, task_uuid) success = False return success def _create_tag_extract_task(self, tag_creator_account, tag_invoker_account, job_uuid, shard_uuid, task_uuid, task_id, config_uuid, config_type, extract): success = True payload = {'job_uuid': job_uuid, 'shard_uuid': shard_uuid, 'task_uuid': task_uuid, 'config_uuid': config_uuid, \ 'config_type': config_type, 'tag_extract': extract, 'tag_creator_account': tag_creator_account, \ 'tag_invoker_account': tag_invoker_account} client = tasks_v2.CloudTasksClient() parent = client.queue_path(self.tag_engine_project, self.tag_engine_region, self.tag_engine_queue) task = { 'name': parent + '/tasks/' + task_id, 'dispatch_deadline': self.task_deadline, 'http_request': { 'http_method': 'POST', 'url': self.task_handler_uri, 'headers': {'content-type': 'application/json'}, 'body': json.dumps(payload).encode(), 'oidc_token': {'service_account_email': self.tag_engine_sa, 'audience': self.task_handler_uri} } } print('task request:', task) try: task = client.create_task(parent=parent, task=task) print('task response:', task) except Exception as e: print('Error: could not create task for uri', self.task_handler_uri, '. Error: ', e) self._set_task_failed(shard_uuid, task_uuid) success = False return success def _set_task_running(self, shard_uuid, task_uuid): #print('*** _set_task_running ***') task_ref = self.db.collection('shards').document(shard_uuid).collection('tasks').document(task_uuid) task_ref.set({ 'status': 'RUNNING', 'start_time': datetime.datetime.utcnow() }, merge=True) print('set task running.') def _set_rollup_tasks_running(self, shard_uuid): shard_ref = self.db.collection('shards').document(shard_uuid) shard_ref.update({'tasks_running': firestore.Increment(1)}) def _set_task_success(self, shard_uuid, task_uuid): print('*** _set_task_success ***') task_ref = self.db.collection('shards').document(shard_uuid).collection('tasks').document(task_uuid) task_ref.set({ 'status': 'SUCCESS', 'end_time': datetime.datetime.utcnow() }, merge=True) print('set task success.') def _set_rollup_tasks_success(self, shard_uuid): shard_ref = self.db.collection('shards').document(shard_uuid) shard_ref.update({'tasks_success': firestore.Increment(1), 'tasks_running': firestore.Increment(-1)}) def _set_task_failed(self, shard_uuid, task_uuid): print('*** _set_task_failed ***') task_ref = self.db.collection('shards').document(shard_uuid).collection('tasks').document(task_uuid) task_ref.set({ 'status': 'ERROR', 'end_time': datetime.datetime.utcnow() }, merge=True) print('set task failed.') def _set_rollup_tasks_failed(self, shard_uuid): shard_ref = self.db.collection('shards').document(shard_uuid) shard_ref.update({'tasks_failed': firestore.Increment(1), 'tasks_running': firestore.Increment(-1)})