<commit_before>import ConfigParser
import atexit
import daemon
import errno
import json
import logging
import os
import pika
import signal
import smtplib
import socket
import sys
import time
import traceback
from argparse import ArgumentParser, FileType
from email.mime.text import MIMEText

__version__ = '0.1'


class AMQPWorker(object):
    PRIORITY_QUEUE_POLL_DELAY = 0.5  # In seconds
    MAX_SLEEP_TIME = 64

    def __init__(self, server, receive_queue, worker_func, error_queue=None,
                 complete_queue=None, working_dir=None, is_daemon=False,
                 log_file=None, pid_file=None, email_from=None,
                 email_host=None, email_subject=None, email_to=None):
        """Initialize an AMQPWorker.

        :param worker_func: is called with message from receive_queue where the
            dictionary keywords become the keyword arguments. If this function
            returns a dictionary, that dictionary is added as a message to the
            complete_queue. Alternatively an iterable of dictionaries can be
            returned, each of which will be added to the complete_queue.
        :param receive_queue: is the queue from which to receive jobs. If this
            argument is of a type list, then each item should be a queue, with
            the highest priority queues listed first.
        :param complete_queue: is the queue to put the job results. If this
            argument is of a type list, then each item should be a queue, with
            the highest priority queues listed first.

        """
        def fullpath(path):
            return os.path.abspath(os.path.expanduser(path)) if path else None

        self.server = server
        self.receive_queue = receive_queue
        self.worker_func = worker_func
        self.complete_queue = complete_queue
        self.is_daemon = is_daemon
        self.log_file = fullpath(log_file)
        self.pid_file = fullpath(pid_file)
        self.working_dir = fullpath(working_dir)
        self.error_queue = error_queue
        self.connection = self.channel = None
        self.use_priority = isinstance(receive_queue, list)
        if email_from and email_to:
            self.email = {'host': email_host or 'localhost',
                          'subject': email_subject or 'AMQPWorker Exception',
                          'from': email_from, 'to': email_to}
        else:
            self.email = None
        # Configure logger
        log_format = '%(asctime)s %(levelname)s %(name)s %(message)s'
        logging.basicConfig(stream=sys.stdout, level=logging.INFO,
                            format=log_format)
        # Only log pika ERRORS
        logging.getLogger('pika').setLevel(logging.ERROR)

    def _start(self):
        sleep_time = 1
        iterations = 0
        running = True
        self.last_exc = None
        while running:
            if iterations > 0:
                logging.info('Retrying in {0} seconds'.format(sleep_time))
                time.sleep(sleep_time)
                sleep_time = min(self.MAX_SLEEP_TIME, sleep_time * 2)
            iterations += 1

            try:
                self.initialize_connection()
                sleep_time = 1
                self.receive_jobs()
            except socket.error as error:
                logging.warning('Error connecting to rabbitmq: {0!s}'
                                .format(error))
            except pika.exceptions.AMQPConnectionError:
                logging.warning('Lost connection to rabbitmq')
            except Exception as exc:
                if self.email and (type(exc) is not type(self.last_exc)
                                   or vars(exc) != vars(self.last_exc)):
                    self.email_message(traceback.format_exc())
                self.last_exc = exc
                traceback.print_exc()
            except KeyboardInterrupt:
                logging.info('Goodbye!')
                running = False
            finally:
                # Force disconnect to release the jobs
                if self.connection and not self.connection.is_closed:
                    # This call complains when daemon mode dies, oh well
                    self.connection.close()
                self.connection = None

    def consume_callback(self, channel, method, _, message):
        return_messages = None
        try:
            kwargs = json.loads(message)
            if '_priority' in kwargs:
                try:
                    priority = int(kwargs['_priority'])
                except ValueError:
                    priority = None
                del kwargs['_priority']
            else:
                priority = None
            return_messages = self.worker_func(**kwargs)
        except Exception as exc:
            if self.error_queue:
                # Save the original in the error queue
                self.channel.basic_publish(
                    exchange='', body=message, routing_key=self.error_queue,
                    properties=pika.BasicProperties(delivery_mode=2))
                logging.error('Message moved to error_queue: {0}'.format(exc))
            else:
                logging.error('Failed on {0}. Reason: {1}'
                              .format(message, exc))
            traceback.print_exc()
            if self.email:
                self.email_message(traceback.format_exc())

        if return_messages is not None:
            # Determine the return queue
            if isinstance(self.complete_queue, list):
                max_priority = len(self.complete_queue)
                if priority is not None:
                    priority = min(max(0, priority), max_priority - 1)
                    queue = self.complete_queue[priority]
                else:
                    queue = self.complete_queue[max_priority / 2]<commit_msg>Add support to retry failed jobs. Release version 0.2.<commit_after>import ConfigParser
import atexit
import daemon
import errno
import json
import logging
import os
import pika
import signal
import smtplib
import socket
import sys
import time
import traceback
from argparse import ArgumentParser, FileType
from email.mime.text import MIMEText

__version__ = '0.2'


class AMQPWorker(object):
    PRIORITY_QUEUE_POLL_DELAY = 0.5  # In seconds
    MAX_SLEEP_TIME = 64

    def __init__(self, server, receive_queue, worker_func, error_queue=None,
                 complete_queue=None, working_dir=None, is_daemon=False,
                 log_file=None, pid_file=None, email_from=None, max_retries=3,
                 email_host=None, email_subject=None, email_to=None):
        """Initialize an AMQPWorker.

        :param worker_func: is called with message from receive_queue where the
            dictionary keywords become the keyword arguments. If this function
            returns a dictionary, that dictionary is added as a message to the
            complete_queue. Alternatively an iterable of dictionaries can be
            returned, each of which will be added to the complete_queue.
        :param receive_queue: is the queue from which to receive jobs. If this
            argument is of a type list, then each item should be a queue, with
            the highest priority queues listed first.
        :param complete_queue: is the queue to put the job results. If this
            argument is of a type list, then each item should be a queue, with
            the highest priority queues listed first.
        :param max_retries: the number of times to requeue a failed job before
            treating it as an error.

        """
        def fullpath(path):
            return os.path.abspath(os.path.expanduser(path)) if path else None

        self.server = server
        self.receive_queue = receive_queue
        self.worker_func = worker_func
        self.complete_queue = complete_queue
        self.is_daemon = is_daemon
        self.log_file = fullpath(log_file)
        self.pid_file = fullpath(pid_file)
        self.working_dir = fullpath(working_dir)
        self.error_queue = error_queue
        self.connection = self.channel = None
        self.use_priority = isinstance(receive_queue, list)
        self.retries = int(max_retries)
        if email_from and email_to:
            self.email = {'host': email_host or 'localhost',
                          'subject': email_subject or 'AMQPWorker Exception',
                          'from': email_from, 'to': email_to}
        else:
            self.email = None
        # Configure logger
        log_format = '%(asctime)s %(levelname)s %(name)s %(message)s'
        logging.basicConfig(stream=sys.stdout, level=logging.INFO,
                            format=log_format)
        # Only log pika ERRORS
        logging.getLogger('pika').setLevel(logging.ERROR)

    def _start(self):
        sleep_time = 1
        iterations = 0
        running = True
        self.last_exc = None
        while running:
            if iterations > 0:
                logging.info('Retrying in {0} seconds'.format(sleep_time))
                time.sleep(sleep_time)
                sleep_time = min(self.MAX_SLEEP_TIME, sleep_time * 2)
            iterations += 1

            try:
                self.initialize_connection()
                sleep_time = 1
                self.receive_jobs()
            except socket.error as error:
                logging.warning('Error connecting to rabbitmq: {0!s}'
                                .format(error))
            except pika.exceptions.AMQPConnectionError:
                logging.warning('Lost connection to rabbitmq')
            except Exception as exc:
                traceback.print_exc()
                if self.email and (type(exc) is not type(self.last_exc)
                                   or vars(exc) != vars(self.last_exc)):
                    self.email_message(traceback.format_exc())
                self.last_exc = exc
            except KeyboardInterrupt:
                logging.info('Goodbye!')
                running = False
            finally:
                # Force disconnect to release the jobs
                if self.connection and not self.connection.is_closed:
                    # This call complains when daemon mode dies, oh well
                    self.connection.close()
                self.connection = None

    def consume_callback(self, channel, method, _, message):
        def get_int_value(value, items, default=None):
            """Remove and return integer value from dictionary items."""
            if value in items:
                try:
                    retval = int(items[value])
                except ValueError:
                    retval = default
                del items[value]
            else:
                retval = default
            return retval

        return_messages = None
        try:
            kwargs = json.loads(message)
            _orig = kwargs.copy()
            attempt = get_int_value('_attempt', kwargs, default=0)
            priority = get_int_value('_priority', kwargs)
            try:
                return_messages = self.worker_func(**kwargs)
            except Exception as exc:
                if attempt < self.retries:  # Increase attempt and retry
                    _orig['_attempt'] = attempt + 1
                    self.channel.basic_publish(
                        exchange='', body=json.dumps(_orig),
                        routing_key=method.routing_key,
                        properties=pika.BasicProperties(delivery_mode=2))
                    logging.warning(
                        'Requeue (attempt {0}): {1}'.format(attempt + 1, exc))
                else:
                    raise
        except Exception as exc:
            if self.error_queue:
                # Save the original in the error queue
                self.channel.basic_publish(
                    exchange='', body=message, routing_key=self.error_queue,
                    properties=pika.BasicProperties(delivery_mode=2))
                logging.error('Message moved to error_queue: {0}'.format(exc))
            else:
                logging.error('Failed on {0}. Reason: {1}'
                              .format(message, exc))
            traceback.print_exc()
            if self.email:
                self.email_message(traceback.format_exc())

        if return_messages is not None:
            # Determine the return queue
            if isinstance(self.complete_queue, list):
                max_priority = len(self.complete_queue)
                if priority is not None:
                    priority = min(max(0, priority), max_priority - 1)
                    queue = self.complete_queue[priority]
                else:
                    queue = self.complete_queue[max_priority / 2]
{'fraction_non_alphanumeric': 0.0721625225507883, 'fraction_numerical': 0.004000313750098047, 'mean_word_length': 1.7615334632878492, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1622913', 'n_tokens_mistral': 3326, 'n_tokens_neox': 3163, 'n_words': 945}