<commit_before>from pipelines.models import OperationRun
from polyaxon.celery_api import CeleryTask


class OperationTask(CeleryTask):
    """Base operation celery task with basic logging."""
    _operation = None

    def __call__(self, *args, **kwargs):
        self._operation_run = OperationRun.objects.get(id=kwargs['operation_run_id'])
        self._operation_run.on_run()
        self.max_retries = self._operation_run.operation.max_retries
        self.countdown = self._operation_run.operation.get_countdown(self.request.retries)

        super(OperationTask, self).__call__(*args, **kwargs)

    def on_failure(self, exc, task_id, args, kwargs, exc_info):
        """Update query status and send email notification to a user"""
        super(OperationTask, self).on_failure(exc, task_id, args, kwargs, exc_info)
        self._operation_run.on_failure()

    def on_retry(self, exc, task_id, args, kwargs, exc_info):
        super(OperationTask, self).on_retry(exc, task_id, args, kwargs, exc_info)
        self._operation_run.on_retry()

    def on_success(self, retval, task_id, args, kwargs):
        """Send email notification and a file, if requested to do so by a user"""
        super(OperationTask, self).on_success(retval, task_id, args, kwargs)
        self._operation_run.on_success()<commit_msg>Add class based task<commit_after>from billiard.exceptions import SoftTimeLimitExceeded
from polyaxon_schemas.utils import to_list

from pipelines.models import OperationRun
from polyaxon.celery_api import CeleryTask


class OperationRunTaskException(Exception):
    pass


class OperationTask(CeleryTask):
    """Base operation celery task with basic logging."""
    _operation_run = None

    def __call__(self, *args, **kwargs):
        try:
            self._operation_run = OperationRun.objects.get(id=kwargs['operation_run_id'])
        except (KeyError, OperationRun.DoesNotExist):
            raise OperationRunTaskException
        self._operation_run.on_run()
        self.max_retries = self._operation_run.operation.max_retries
        self.countdown = self._operation_run.operation.get_countdown(self.request.retries)

        super(OperationTask, self).__call__(*args, **kwargs)

    def on_failure(self, exc, task_id, args, kwargs, exc_info):
        """Update query status and send email notification to a user"""
        super(OperationTask, self).on_failure(exc, task_id, args, kwargs, exc_info)
        if isinstance(exc, OperationRunTaskException):
            return
        self._operation_run.on_failure()

    def on_retry(self, exc, task_id, args, kwargs, exc_info):
        super(OperationTask, self).on_retry(exc, task_id, args, kwargs, exc_info)
        self._operation_run.on_retry()

    def on_success(self, retval, task_id, args, kwargs):
        """Send email notification and a file, if requested to do so by a user"""
        super(OperationTask, self).on_success(retval, task_id, args, kwargs)
        self._operation_run.on_success()


class ClassBasedTask(object):
    """A class based task to use for operation tasks."""
    retry_for = None

    @classmethod
    def run(cls, task_bind, *args, **kwargs):
        retry_for = cls.retry_for or []
        retry_for = to_list(retry_for)
        if SoftTimeLimitExceeded not in retry_for:
            retry_for.append(SoftTimeLimitExceeded)
        try:
            return cls._run(task_bind, *args, **kwargs)
        except tuple(retry_for) as exc:
            if task_bind.request.retries < task_bind.max_retries:
                raise task_bind.retry(countdown=task_bind.countdown)
            else:
                raise exc

    @staticmethod
    def _run(task_bind, *args, **kwargs):
        raise NotImplemented
{'fraction_non_alphanumeric': 0.09653132562516806, 'fraction_numerical': 0.0, 'mean_word_length': 3.4444444444444446, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 3, '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': '3290371', 'n_tokens_mistral': 1098, 'n_tokens_neox': 1098, 'n_words': 271}