import collections import importlib import logging import threading import types import sys from ..models import BatchErrorException, TaskAddCollectionResult, TaskAddStatus from ..custom.custom_errors import CreateTasksErrorException from ..operations.task_operations import TaskOperations MAX_TASKS_PER_REQUEST = 100 _LOGGER = logging.getLogger(__name__) class _TaskWorkflowManager(object): """Worker class for one add_collection request :param ~TaskOperations task_operations: Parent object which instantiated this :param str job_id: The ID of the job to which the task collection is to be added. :param tasks_to_add: The collection of tasks to add. :type tasks_to_add: list of :class:`TaskAddParameter ` :param task_add_collection_options: Additional parameters for the operation :type task_add_collection_options: :class:`TaskAddCollectionOptions ` :param dict custom_headers: headers that will be added to the request :param bool raw: returns the direct response alongside the deserialized response """ def __init__( self, client, original_add_collection, job_id, tasks_to_add, task_add_collection_options=None, custom_headers=None, raw=False, **kwargs): # Append operations thread safe - Only read once all threads have completed # List of tasks which failed to add due to a returned client error self._failure_tasks = collections.deque() # List of unknown exceptions which occurred during requests. self._errors = collections.deque() # synchronized through lock variables self.error = None # Only written once all threads have completed self._max_tasks_per_request = MAX_TASKS_PER_REQUEST self._tasks_to_add = collections.deque(tasks_to_add) self._error_lock = threading.Lock() self._max_tasks_lock = threading.Lock() self._pending_queue_lock = threading.Lock() # Variables to be used for task add_collection requests self._client = client self._original_add_collection = original_add_collection self._job_id = job_id self._task_add_collection_options = task_add_collection_options self._custom_headers = custom_headers self._raw = raw self._kwargs = dict(**kwargs) def _bulk_add_tasks(self, results_queue, chunk_tasks_to_add): """Adds a chunk of tasks to the job Retry chunk if body exceeds the maximum request size and retry tasks if failed due to server errors. :param results_queue: Queue to place the return value of the request :type results_queue: collections.deque :param chunk_tasks_to_add: Chunk of at most 100 tasks with retry details :type chunk_tasks_to_add: list[~TrackedCloudTask] """ try: add_collection_response = self._original_add_collection( self._client, self._job_id, chunk_tasks_to_add, self._task_add_collection_options, self._custom_headers, self._raw) except BatchErrorException as e: # In case of a chunk exceeding the MaxMessageSize split chunk in half # and resubmit smaller chunk requests # TODO: Replace string with constant variable once available in SDK if e.error.code == "RequestBodyTooLarge": # pylint: disable=no-member # In this case the task is misbehaved and will not be able to be added due to: # 1) The task exceeding the max message size # 2) A single cell of the task exceeds the per-cell limit, or # 3) Sum of all cells exceeds max row limit if len(chunk_tasks_to_add) == 1: failed_task = chunk_tasks_to_add.pop() self._errors.appendleft(e) _LOGGER.error("Failed to add task with ID %s due to the body" " exceeding the maximum request size", failed_task.id) else: # Assumption: Tasks are relatively close in size therefore if one batch exceeds size limit # we should decrease the initial task collection size to avoid repeating the error # Midpoint is lower bounded by 1 due to above base case midpoint = int(len(chunk_tasks_to_add) / 2) # Restrict one thread at a time to do this compare and set, # therefore forcing max_tasks_per_request to be strictly decreasing with self._max_tasks_lock: if midpoint < self._max_tasks_per_request: self._max_tasks_per_request = midpoint _LOGGER.info("Amount of tasks per request reduced from %s to %s due to the" " request body being too large", str(self._max_tasks_per_request), str(midpoint)) # Not the most efficient solution for all cases, but the goal of this is to handle this # exception and have it work in all cases where tasks are well behaved # Behavior retries as a smaller chunk and # appends extra tasks to queue to be picked up by another thread . self._tasks_to_add.extendleft(chunk_tasks_to_add[midpoint:]) self._bulk_add_tasks(results_queue, chunk_tasks_to_add[:midpoint]) # Retry server side errors elif 500 <= e.response.status_code <= 599: self._tasks_to_add.extendleft(chunk_tasks_to_add) else: # Re-add to pending queue as unknown status / don't have result self._tasks_to_add.extendleft(chunk_tasks_to_add) # Unknown State - don't know if tasks failed to add or were successful self._errors.appendleft(e) except Exception as e: # pylint: disable=broad-except # Re-add to pending queue as unknown status / don't have result self._tasks_to_add.extendleft(chunk_tasks_to_add) # Unknown State - don't know if tasks failed to add or were successful self._errors.appendleft(e) else: try: add_collection_response = add_collection_response.output except AttributeError: pass for task_result in add_collection_response.value: # pylint: disable=no-member if task_result.status == TaskAddStatus.server_error: # Server error will be retried with self._pending_queue_lock: for task in chunk_tasks_to_add: if task.id == task_result.task_id: self._tasks_to_add.appendleft(task) elif (task_result.status == TaskAddStatus.client_error and not task_result.error.code == "TaskExists"): # Client error will be recorded unless Task already exists self._failure_tasks.appendleft(task_result) else: results_queue.appendleft(task_result) def task_collection_thread_handler(self, results_queue): """Main method for worker to run Pops a chunk of tasks off the collection of pending tasks to be added and submits them to be added. :param collections.deque results_queue: Queue for worker to output results to """ # Add tasks until either we run out or we run into an unexpected error while self._tasks_to_add and not self._errors: max_tasks = self._max_tasks_per_request # local copy chunk_tasks_to_add = [] with self._pending_queue_lock: while len(chunk_tasks_to_add) < max_tasks and self._tasks_to_add: chunk_tasks_to_add.append(self._tasks_to_add.pop()) if chunk_tasks_to_add: self._bulk_add_tasks(results_queue, chunk_tasks_to_add) # Only define error if all threads have finished and there were failures with self._error_lock: if threading.active_count() == 1 and (self._failure_tasks or self._errors): self.error = CreateTasksErrorException( "One or more tasks failed to be added", self._failure_tasks, self._tasks_to_add, self._errors) def _handle_output(results_queue): """Scan output for exceptions If there is an output from an add task collection call add it to the results. :param results_queue: Queue containing results of attempted add_collection's :type results_queue: collections.deque :return: list of TaskAddResults :rtype: list[~TaskAddResult] """ results = [] while results_queue: queue_item = results_queue.pop() results.append(queue_item) return results def build_new_add_collection(original_add_collection): def bulk_add_collection( self, job_id, value, task_add_collection_options=None, custom_headers=None, raw=False, threads=0, **operation_config): """Adds a collection of tasks to the specified job. Note that each task must have a unique ID. The Batch service may not return the results for each task in the same order the tasks were submitted in this request. If the server times out or the connection is closed during the request, the request may have been partially or fully processed, or not at all. In such cases, the user should re-issue the request. Note that it is up to the user to correctly handle failures when re-issuing a request. For example, you should use the same task IDs during a retry so that if the prior operation succeeded, the retry will not create extra tasks unexpectedly. If the response contains any tasks which failed to add, a client can retry the request. In a retry, it is most efficient to resubmit only tasks that failed to add, and to omit tasks that were successfully added on the first attempt. :param job_id: The ID of the job to which the task collection is to be added. :type job_id: str :param value: The collection of tasks to add. The total serialized size of this collection must be less than 4MB. If it is greater than 4MB (for example if each task has 100's of resource files or environment variables), the request will fail with code 'RequestBodyTooLarge' and should be retried again with fewer tasks. :type value: list of :class:`TaskAddParameter ` :param task_add_collection_options: Additional parameters for the operation :type task_add_collection_options: :class:`TaskAddCollectionOptions ` :param dict custom_headers: headers that will be added to the request :param bool raw: returns the direct response alongside the deserialized response :param int threads: number of threads to use in parallel when adding tasks. If specified and greater than 0, will start additional threads to submit requests and wait for them to finish. Otherwise will submit add_collection requests sequentially on main thread :return: :class:`TaskAddCollectionResult ` or :class:`ClientRawResponse` if raw=true :rtype: :class:`TaskAddCollectionResult ` or :class:`ClientRawResponse` :raises: :class:`CreateTasksErrorException` """ results_queue = collections.deque() # deque operations(append/pop) are thread-safe task_workflow_manager = _TaskWorkflowManager( self, original_add_collection, job_id, value, task_add_collection_options, custom_headers, raw, **operation_config) # multi-threaded behavior if threads: if threads < 0: raise ValueError("Threads must be positive or 0") active_threads = [] for i in range(threads): active_threads.append(threading.Thread( target=task_workflow_manager.task_collection_thread_handler, args=(results_queue,))) active_threads[-1].start() for thread in active_threads: thread.join() # single-threaded behavior else: task_workflow_manager.task_collection_thread_handler(results_queue) if task_workflow_manager.error: raise task_workflow_manager.error # pylint: disable=raising-bad-type else: submitted_tasks = _handle_output(results_queue) return TaskAddCollectionResult(value=submitted_tasks) bulk_add_collection.metadata = {'url': '/jobs/{jobId}/addtaskcollection'} return bulk_add_collection def patch_client(): try: models = sys.modules['azure.batch.models'] except KeyError: models = importlib.import_module('azure.batch.models') setattr(models, 'CreateTasksErrorException', CreateTasksErrorException) sys.modules['azure.batch.models'] = models operations_modules = importlib.import_module('azure.batch.operations') operations_modules.TaskOperations.add_collection = build_new_add_collection(operations_modules.TaskOperations.add_collection)