<commit_before>import logging
import signal
import datetime, time
import os, sys
import json_parser as json
import commands

from pyres.exceptions import NoQueueError
from pyres.job import Job
from pyres import ResQ, Stat, __version__



logger = logging.getLogger(__name__)

class Worker(object):
    """Defines a worker. The ``pyres_worker`` script instantiates this Worker
    class and passes a comma-separated list of queues to listen on.::

       >>> from pyres.worker import Worker
       >>> Worker.run([queue1, queue2], server="localhost:6379")

    """
    
    job_class = Job
    
    def __init__(self, queues=(), server="localhost:6379", password=<PASSWORD>):
        self.queues = queues
        self.validate_queues()
        self._shutdown = False
        self.child = None
        self.pid = os.getpid()
        self.hostname = os.uname()[1]

        if isinstance(server, basestring):
            self.resq = ResQ(server=server, password=password)
        elif isinstance(server, ResQ):
            self.resq = server
        else:
            raise Exception("Bad server argument")

    def validate_queues(self):
        """Checks if a worker is given at least one queue to work on."""
        if not self.queues:
            raise NoQueueError("Please give each worker at least one queue.")

    def register_worker(self):
        self.resq.redis.sadd('resque:workers', str(self))
        #self.resq._redis.add("worker:#{self}:started", Time.now.to_s)
        self.started = datetime.datetime.now()

    def _set_started(self, dt):
        if dt:
            key = int(time.mktime(dt.timetuple()))
            self.resq.redis.set("resque:worker:%s:started" % self, key)
        else:
            self.resq.redis.delete("resque:worker:%s:started" % self)

    def _get_started(self):
        datestring = self.resq.redis.get("resque:worker:%s:started" % self)
        #ds = None
        #if datestring:
        #    ds = datetime.datetime.strptime(datestring, '%Y-%m-%d %H:%M:%S')
        return datestring

    started = property(_get_started, _set_started)

    def unregister_worker(self):
        self.resq.redis.srem('resque:workers', str(self))
        self.started = None
        Stat("processed:%s" % self, self.resq).clear()
        Stat("failed:%s" % self, self.resq).clear()

    def prune_dead_workers(self):
        all_workers = Worker.all(self.resq)
        known_workers = self.worker_pids()
        for worker in all_workers:
            host, pid, queues = worker.id.split(':')
            if host != self.hostname:
                continue
            if pid in known_workers:
                continue
            logger.warning("pruning dead worker: %s" % worker)
            worker.unregister_worker()

    def startup(self):
        self.register_signal_handlers()
        self.prune_dead_workers()
        self.register_worker()

    def register_signal_handlers(self):
        signal.signal(signal.SIGTERM, self.shutdown_all)
        signal.signal(signal.SIGINT, self.shutdown_all)
        signal.signal(signal.SIGQUIT, self.schedule_shutdown)
        signal.signal(signal.SIGUSR1, self.kill_child)

    def shutdown_all(self, signum, frame):
        self.schedule_shutdown(signum, frame)
        self.kill_child(signum, frame)

    def schedule_shutdown(self, signum, frame):
        self._shutdown = True

    def kill_child(self, signum, frame):
        if self.child:
            logger.info("Killing child at %s" % self.child)
            os.kill(self.child, signal.SIGKILL)

    def __str__(self):
        if getattr(self,'id', None):
            return self.id
        return '%s:%s:%s' % (self.hostname, self.pid, ','.join(self.queues))

    def work(self, interval=5):
        """Invoked by ``run`` method. ``work`` listens on a list of queues and
        sleeps for ``interval`` time.

        ``interval`` -- Number of seconds the worker will wait until processing
        the next job. Default is "5".

        Whenever a worker finds a job on the queue it first calls ``reserve`` on
        that job to make sure another worker won't run it, then *forks* itself
        to work on that job.

        Finally, the ``process`` method actually processes the job by eventually
        calling the Job instance's ``perform`` method.

        """
        setproctitle('pyres_worker-%s: Starting' % __version__)
        self.startup()

        while True:
            if self._shutdown:
                logger.info('shutdown scheduled')
                break

            job = self.reserve(interval)

            if job:
                logger.info('picked up job')
                logger.debug('job details: %s' % job)
                self.child = os.fork() 
                if self.child:
                    setproctitle("pyres_worker%s: Forked %s at %s" %
                                 (__version__,
                                  self.child,
                                  datetime.datetime.now()))
                    logger.info('Forked %s at %s' % (self.child,
                                                      datetime.datetime.now()))

                    try:
                        os.waitpid(self.child, 0)
                    except OSError as ose:
                        import errno

                        if ose.errno != errno.EINTR:
                            raise ose
                    #os.wait()
                    logger.debug('done waiting')
                else:
                    setproctitle("pyres_worker-%s: Processing %s since %s" %
                                 (__version__, job._queue,
                                  datetime.datetime.now()))
                    logger.info('Processing %s since %s' %
                                 (job._queue, datetime.datetime.now()))
                    self.process(job)
                    os._exit(0)
                self.child = None
            else:
                if interval == 0:
                    break
                #procline @paused ? "Paused" : "Waiting for #{@queues.join(',')}"
                setproctitle("pyres_worker-%s: Waiting for %s " %
                             (__version__, ','.join(self.queues)))
                #time.sleep(interval)
        self.unregister_worker()
<commit_msg>before and after fork hooks for worker<commit_after>import logging
import signal
import datetime, time
import os, sys
import json_parser as json
import commands

from pyres.exceptions import NoQueueError
from pyres.job import Job
from pyres import ResQ, Stat, __version__



logger = logging.getLogger(__name__)

class Worker(object):
    """Defines a worker. The ``pyres_worker`` script instantiates this Worker
    class and passes a comma-separated list of queues to listen on.::

       >>> from pyres.worker import Worker
       >>> Worker.run([queue1, queue2], server="localhost:6379")

    """
    
    job_class = Job
    
    def __init__(self, queues=(), server="localhost:6379", password=<PASSWORD>):
        self.queues = queues
        self.validate_queues()
        self._shutdown = False
        self.child = None
        self.pid = os.getpid()
        self.hostname = os.uname()[1]

        if isinstance(server, basestring):
            self.resq = ResQ(server=server, password=password)
        elif isinstance(server, ResQ):
            self.resq = server
        else:
            raise Exception("Bad server argument")

    def validate_queues(self):
        """Checks if a worker is given at least one queue to work on."""
        if not self.queues:
            raise NoQueueError("Please give each worker at least one queue.")

    def register_worker(self):
        self.resq.redis.sadd('resque:workers', str(self))
        #self.resq._redis.add("worker:#{self}:started", Time.now.to_s)
        self.started = datetime.datetime.now()

    def _set_started(self, dt):
        if dt:
            key = int(time.mktime(dt.timetuple()))
            self.resq.redis.set("resque:worker:%s:started" % self, key)
        else:
            self.resq.redis.delete("resque:worker:%s:started" % self)

    def _get_started(self):
        datestring = self.resq.redis.get("resque:worker:%s:started" % self)
        #ds = None
        #if datestring:
        #    ds = datetime.datetime.strptime(datestring, '%Y-%m-%d %H:%M:%S')
        return datestring

    started = property(_get_started, _set_started)

    def unregister_worker(self):
        self.resq.redis.srem('resque:workers', str(self))
        self.started = None
        Stat("processed:%s" % self, self.resq).clear()
        Stat("failed:%s" % self, self.resq).clear()

    def prune_dead_workers(self):
        all_workers = Worker.all(self.resq)
        known_workers = self.worker_pids()
        for worker in all_workers:
            host, pid, queues = worker.id.split(':')
            if host != self.hostname:
                continue
            if pid in known_workers:
                continue
            logger.warning("pruning dead worker: %s" % worker)
            worker.unregister_worker()

    def startup(self):
        self.register_signal_handlers()
        self.prune_dead_workers()
        self.register_worker()

    def register_signal_handlers(self):
        signal.signal(signal.SIGTERM, self.shutdown_all)
        signal.signal(signal.SIGINT, self.shutdown_all)
        signal.signal(signal.SIGQUIT, self.schedule_shutdown)
        signal.signal(signal.SIGUSR1, self.kill_child)

    def shutdown_all(self, signum, frame):
        self.schedule_shutdown(signum, frame)
        self.kill_child(signum, frame)

    def schedule_shutdown(self, signum, frame):
        self._shutdown = True

    def kill_child(self, signum, frame):
        if self.child:
            logger.info("Killing child at %s" % self.child)
            os.kill(self.child, signal.SIGKILL)

    def __str__(self):
        if getattr(self,'id', None):
            return self.id
        return '%s:%s:%s' % (self.hostname, self.pid, ','.join(self.queues))

    def work(self, interval=5):
        """Invoked by ``run`` method. ``work`` listens on a list of queues and
        sleeps for ``interval`` time.

        ``interval`` -- Number of seconds the worker will wait until processing
        the next job. Default is "5".

        Whenever a worker finds a job on the queue it first calls ``reserve`` on
        that job to make sure another worker won't run it, then *forks* itself
        to work on that job.

        Finally, the ``process`` method actually processes the job by eventually
        calling the Job instance's ``perform`` method.

        """
        setproctitle('pyres_worker-%s: Starting' % __version__)
        self.startup()

        while True:
            if self._shutdown:
                logger.info('shutdown scheduled')
                break

            job = self.reserve(interval)

            if job:
                logger.debug('picked up job')
                logger.debug('job details: %s' % job)
                self.before_fork(job)
                self.child = os.fork() 
                if self.child:
                    setproctitle("pyres_worker%s: Forked %s at %s" %
                                 (__version__,
                                  self.child,
                                  datetime.datetime.now()))
                    logger.info('Forked %s at %s' % (self.child,
                                                      datetime.datetime.now()))

                    try:
                        os.waitpid(self.child, 0)
                    except OSError as ose:
                        import errno

                        if ose.errno != errno.EINTR:
                            raise ose
                    #os.wait()
                    logger.debug('done waiting')
                else:
                    setproctitle("pyres_worker-%s: Processing %s since %s" %
                                 (__version__, job._queue,
                                  datetime.datetime.now()))
                    logger.info('Processing %s since %s' %
                                 (job._queue, datetime.datetime.now()))
                    self.after_fork(job)
                    self.process(job)
                    os._exit(0)
                self.child = None
            else:
                if interval == 0:
                    break
                #procline @paused ? "Paused" : "Waiting for #{@queues.join(',')}"
                setproctitle("pyres_worker-%s: Waiting for %s " %
                             (__version__, ','.join(self.queues)))
                #time.sleep(interval)
        self.unregister_worker()

    def before_fork(self, job):
        """
        hook for making changes immediately before forking to process
        a job
        """
        pass

    def after_fork(self, job):
        """
        hook for making changes immediately after forking to process a
        job
        """
        pass

{'fraction_non_alphanumeric': 0.10736077855873948, 'fraction_numerical': 0.0026260909863288794, 'mean_word_length': 1.7660756248664815, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4775028', 'n_tokens_mistral': 3715, 'n_tokens_neox': 3534, 'n_words': 906}