<commit_before># pytask
# File: pytask/tasks.py
# Desc: monitoring and cleanup tasks designed to keep a pytask cluster in shape

import logging
from time import time

import gevent

from .pytask import Task
from .helpers import run_loop


class Monitor(Task):
    '''
    A pytask which monitors pytasks! Checks all task state in Redis at a configured
    interval, will re-queue tasks which timeout.
    '''

    NAME = 'pytask/monitor'

    def __init__(self, **task_data):
        self.loop_interval = task_data.get('loop_interval', 10)
        self.task_timeout = task_data.get('task_timeout', 60)

        self.logger = logging.getLogger('pytask-monitoring')

    def start(self):
        self.loop = gevent.spawn(run_loop, self.check_tasks, self.loop_interval)
        self.loop.get()

    def stop(self):
        self.loop.kill()

    def check_tasks(self):
        self.logger.info('Checking tasks...')

        task_ids = self.helpers.get_task_ids()

        for task_id in task_ids:
            task = self.helpers.get_task(task_id)

            if not task:
                self.logger.critical('Task ID in set but no hash!: {0}'.format(task_id))
                continue

            # Ensure running tasks are actually running
            if task.get('state') == 'RUNNING':
                update_time = task.get('last_update', 0)
                diff = time() - float(update_time)

                if diff > self.task_timeout:
                    # Local task that is still RUNNING, just remove/cleanup
                    if task.get('local') == 'true':
                        self.logger.info(
                            'Removing timed out local task: {0}'.format(task_id)
                        )
                        self.helpers.remove_task(task_id)

                    # Normal task, reset the state and restart it
                    else:
                        self.logger.warning(
                            'Restarting timed out task: {0}'.format(task_id)
                        )
                        self.helpers.set_task(task_id, 'state', 'WAIT')
                        self.helpers.restart_task(task_id)


class Cleanup(Task):
    '''
    A pytask which cleans up other pytasks! Checks all task state in Redis at configured
    interval, removing where required.
    '''

    NAME = 'pytask/cleanup'

    def __init__(self, **task_data):
        self.loop_interval = task_data.get('loop_interval', 60)

        self.logger = logging.getLogger('pytask-cleaning')

    def start(self):
        self.loop = gevent.spawn(run_loop, self.cleanup_tasks, self.loop_interval)
        self.loop.get()

    def stop(self):
        self.loop.kill()

    def cleanup_tasks(self):
        self.logger.info('Cleaning tasks...')

        task_ids = self.helpers.get_task_ids()

        for task_id in task_ids:
            task = self.helpers.get_task(task_id)

            if not task:
                self.logger.critical(
                    'Removing task ID in set but no hash!: {0}'.format(task_id)
                )
                self.helpers.remove_task(task_id)
                continue

            state = task.get('state')
            if state != 'RUNNING':
                self.logger.warning(
                    'Task {0} is {1} but in the task set, removing ID from set'.format(
                        task_id, state
                    )
                )
                self.redis.srem(self.helpers.TASK_SET, task_id)<commit_msg>Make `Cleanup` task just do cleanup, and optionally take a module:attr string to handle task cleanup.<commit_after># pytask
# File: pytask/tasks.py
# Desc: monitoring and cleanup tasks designed to keep a pytask cluster in shape

import logging
from time import time
from importlib import import_module

import gevent

from .pytask import Task
from .helpers import run_loop


class Monitor(Task):
    '''
    A pytask which monitors pytasks! Checks all task state in Redis at a configured
    interval, will re-queue tasks which timeout.
    '''

    NAME = 'pytask/monitor'

    def __init__(self, loop_interval=10, task_timeout=60):
        self.loop_interval = loop_interval
        self.task_timeout = task_timeout

        self.logger = logging.getLogger('pytask-monitoring')

    def start(self):
        self.loop = gevent.spawn(run_loop, self.check_tasks, self.loop_interval)
        self.loop.get()

    def stop(self):
        self.loop.kill()

    def check_tasks(self):
        self.logger.info('Checking tasks...')

        task_ids = self.helpers.get_task_ids()

        for task_id in task_ids:
            task = self.helpers.get_task(task_id)

            if not task:
                self.logger.critical('Task ID in set but no hash!: {0}'.format(task_id))
                continue

            # Ensure running tasks are actually running
            if task.get('state') == 'RUNNING':
                update_time = task.get('last_update', 0)
                diff = time() - float(update_time)

                if diff > self.task_timeout:
                    # Local task that is still RUNNING, just remove/cleanup
                    if task.get('local') == 'true':
                        self.logger.info(
                            'Removing timed out local task: {0}'.format(task_id)
                        )
                        self.helpers.remove_task(task_id)

                    # Normal task, reset the state and restart it
                    else:
                        self.logger.warning(
                            'Restarting timed out task: {0}'.format(task_id)
                        )
                        self.helpers.set_task(task_id, 'state', 'WAIT')
                        self.helpers.restart_task(task_id)


class Cleanup(Task):
    '''
    A pytask which cleans up other pytasks! Checks all task state in Redis at configured
    interval, removing where required.
    '''

    NAME = 'pytask/cleanup'

    def __init__(self, task_handler=None):
        if task_handler:
            module, attr = task_handler.split(':')
            self.task_handler = getattr(import_module(module), attr)

        self.logger = logging.getLogger('pytask-cleaning')

    def start(self):
        while True:
            task_id = self.redis.brpop(self.helpers.END_QUEUE)
            task = self.helpers.get_task(task_id)

            self.logger.info('Cleaning {0} task: {1}'.format(task.get('state'), task_id))
            self.helpers.remove_task(task_id)

            if self.task_handler:
                self.task_handler(task_id, task)
{'fraction_non_alphanumeric': 0.0956879475849459, 'fraction_numerical': 0.003504494895627, 'mean_word_length': 1.8123393316195373, 'pattern_counts': {'":': 0, '<': 3, '<?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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1303753', 'n_tokens_mistral': 1915, 'n_tokens_neox': 1831, 'n_words': 478}