<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

###############################################################################
#  Copyright 2016 Kitware Inc.
#
#  Licensed under the Apache License, Version 2.0 ( the "License" );
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
###############################################################################
from __future__ import absolute_import
import logging
import importlib

from functools import wraps
import json
import threading

import requests

from girder_client import GirderClient, HttpError

from celery.signals import before_task_publish, task_prerun, task_failure, \
    task_success

from celery.canvas import maybe_signature
from celery import current_task
from celery.utils.log import get_task_logger

import cumulus.celery


logger = get_task_logger(__name__)


# We need to use thread local storage to save the "current user task", celery's
# current_task is not always what we need.
thread_local = threading.local()

# These key we used to store our taskflow data in the celery headers
TASKFLOW_HEADER = 'taskflow'
TASKFLOW_TASK_ID_HEADER = 'taskflow_task_id'
TASKFLOW_RETRY_HEADER = 'taskflow_retries'


# The states that a taskflow can be in, there are likely to be more
class TaskState:
    CREATED = 'created'
    RUNNING = 'running'
    ERROR = 'error'
    COMPLETE = 'complete'


class TaskFlowState:
    CREATED = 'created'
    RUNNING = 'running'
    ERROR = 'error'
    UNEXPECTEDERROR = 'unexpectederror'
    TERMINATING = 'terminating'
    TERMINATED = 'terminated'
    DELETING = 'deleting'
    DELETED = 'deleted'
    COMPLETE = 'complete'


def _get_task_logger(id, girder_api_url, girder_token):
    logger = logging.getLogger('task.%s' % id)
    logger.setLevel(logging.INFO)
    # Only add new new handler if we don't already have one.
    if not logger.handlers:
        url = '%s/tasks/%s/log' % (girder_api_url, id)
        logger.addHandler(TaskFlowLogHandler(girder_token, url))

    return logger


def _create_girder_client(girder_api_url, girder_token):
    client = GirderClient(apiUrl=girder_api_url)
    client.token = girder_token

    return client


def task(func):
    """
    Taskflow decortator used to inject the taskflow object as the first
    postional argument of the task. In the future probably also need to support
    providing the celery task as well.
    """
    # First apply our own wrapper
    @wraps(func)
    def wrapped(celery_task, *args, **kwargs):
        # Extract the taskflow from the headers
        taskflow = celery_task.request.headers[TASKFLOW_HEADER]
        taskflow = to_taskflow(taskflow)
        task_id = celery_task.request.headers[TASKFLOW_TASK_ID_HEADER]

        # Get the current state of the taskflow so we know if we are terminating
        # Only do this if this task is not associated with termination ...
        if 'terminate' not in taskflow or not taskflow['terminate']:
            status = taskflow.status()
            if status == TaskFlowState.TERMINATING or \
                    status == TaskFlowState.UNEXPECTEDERROR:
                return

        setattr(celery_task, 'taskflow', taskflow)
        setattr(celery_task, 'logger', _get_task_logger(
            task_id, taskflow.girder_api_url, taskflow.girder_token))

        return func(celery_task, *args, **kwargs)

    # Now apply the celery decorator
    celery_task = cumulus.celery.command.task(
        wrapped, bind=True)

    return celery_task


def load_class(class_name):
    module, cls = class_name.rsplit('.', 1)
    try:
        imported = importlib.import_module(module)
    except ImportError:
        raise

    try:
        constructor = getattr(imported, cls)
    except AttributeError:
        raise

    return constructor


def to_taskflow(taskflow):
    """
    Utility method to ensure we have a taskflow instance rather than a simple
    dictionary.
    """
    if taskflow is not None:
        if isinstance(taskflow, dict):
            constr = load_class(taskflow['_type'])
            taskflow = constr(**taskflow)

    return taskflow


class TaskFlowLogHandler(logging.Handler):

    def __init__(self, girder_token, url):
        super(TaskFlowLogHandler, self).__init__()
        self._url = url
        self._headers = {'Girder-Token':  girder_token}

    def emit(self, record):
        json_str = json.dumps(record.__dict__, default=str)
        r = requests.post(self._url, headers=self._headers, data=json_str)
        r.raise_for_status()


class TaskFlow(dict):
    """
    This is the base class users derive that taskflows from. In the future more
    utility methods can be added for example to update data/results associated
    with the taskflow.
    """
    def __init__(self, id=None, girder_token=None,
                 girder_api_url=None, **kwargs):
        """
        Constructs a new TaskFlow instance
        """
        super(TaskFlow, self).__init__(
            girder_api_url=girder_api_url,
            girder_token=girder_token,
            id=id, **kwargs)
        cls = self.__class__
        self['_type'] = '%s.%s' % (cls.__module__, cls.__name__)<commit_msg>Improve JSON encoding of log entries<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

###############################################################################
#  Copyright 2016 Kitware Inc.
#
#  Licensed under the Apache License, Version 2.0 ( the "License" );
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
###############################################################################
from __future__ import absolute_import
import logging
import importlib
import traceback
import types

from functools import wraps
import json
import threading

import requests

from girder_client import GirderClient, HttpError

from celery.signals import before_task_publish, task_prerun, task_failure, \
    task_success

from celery.canvas import maybe_signature
from celery import current_task
from celery.utils.log import get_task_logger

import cumulus.celery


logger = get_task_logger(__name__)


# We need to use thread local storage to save the "current user task", celery's
# current_task is not always what we need.
thread_local = threading.local()

# These key we used to store our taskflow data in the celery headers
TASKFLOW_HEADER = 'taskflow'
TASKFLOW_TASK_ID_HEADER = 'taskflow_task_id'
TASKFLOW_RETRY_HEADER = 'taskflow_retries'


# The states that a taskflow can be in, there are likely to be more
class TaskState:
    CREATED = 'created'
    RUNNING = 'running'
    ERROR = 'error'
    COMPLETE = 'complete'


class TaskFlowState:
    CREATED = 'created'
    RUNNING = 'running'
    ERROR = 'error'
    UNEXPECTEDERROR = 'unexpectederror'
    TERMINATING = 'terminating'
    TERMINATED = 'terminated'
    DELETING = 'deleting'
    DELETED = 'deleted'
    COMPLETE = 'complete'


def _get_task_logger(id, girder_api_url, girder_token):
    logger = logging.getLogger('task.%s' % id)
    logger.setLevel(logging.INFO)
    # Only add new new handler if we don't already have one.
    if not logger.handlers:
        url = '%s/tasks/%s/log' % (girder_api_url, id)
        logger.addHandler(TaskFlowLogHandler(girder_token, url))

    return logger


def _create_girder_client(girder_api_url, girder_token):
    client = GirderClient(apiUrl=girder_api_url)
    client.token = girder_token

    return client


def task(func):
    """
    Taskflow decortator used to inject the taskflow object as the first
    postional argument of the task. In the future probably also need to support
    providing the celery task as well.
    """
    # First apply our own wrapper
    @wraps(func)
    def wrapped(celery_task, *args, **kwargs):
        # Extract the taskflow from the headers
        taskflow = celery_task.request.headers[TASKFLOW_HEADER]
        taskflow = to_taskflow(taskflow)
        task_id = celery_task.request.headers[TASKFLOW_TASK_ID_HEADER]

        # Get the current state of the taskflow so we know if we are terminating
        # Only do this if this task is not associated with termination ...
        if 'terminate' not in taskflow or not taskflow['terminate']:
            status = taskflow.status()
            if status == TaskFlowState.TERMINATING or \
                    status == TaskFlowState.UNEXPECTEDERROR:
                return

        setattr(celery_task, 'taskflow', taskflow)
        setattr(celery_task, 'logger', _get_task_logger(
            task_id, taskflow.girder_api_url, taskflow.girder_token))

        return func(celery_task, *args, **kwargs)

    # Now apply the celery decorator
    celery_task = cumulus.celery.command.task(
        wrapped, bind=True)

    return celery_task


def load_class(class_name):
    module, cls = class_name.rsplit('.', 1)
    try:
        imported = importlib.import_module(module)
    except ImportError:
        raise

    try:
        constructor = getattr(imported, cls)
    except AttributeError:
        raise

    return constructor


def to_taskflow(taskflow):
    """
    Utility method to ensure we have a taskflow instance rather than a simple
    dictionary.
    """
    if taskflow is not None:
        if isinstance(taskflow, dict):
            constr = load_class(taskflow['_type'])
            taskflow = constr(**taskflow)

    return taskflow

class LogRecordEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, type):
            return obj.__name__
        elif isinstance(obj, types.TracebackType):
            return traceback.format_tb(obj)[0]
        else:
            return str(obj)

class TaskFlowLogHandler(logging.Handler):

    def __init__(self, girder_token, url):
        super(TaskFlowLogHandler, self).__init__()
        self._url = url
        self._headers = {'Girder-Token':  girder_token}

    def emit(self, record):
        json_str = json.dumps(record.__dict__, cls=LogRecordEncoder)
        r = requests.post(self._url, headers=self._headers, data=json_str)
        r.raise_for_status()


class TaskFlow(dict):
    """
    This is the base class users derive that taskflows from. In the future more
    utility methods can be added for example to update data/results associated
    with the taskflow.
    """
    def __init__(self, id=None, girder_token=None,
                 girder_api_url=None, **kwargs):
        """
        Constructs a new TaskFlow instance
        """
        super(TaskFlow, self).__init__(
            girder_api_url=girder_api_url,
            girder_token=girder_token,
            id=id, **kwargs)
        cls = self.__class__
        self['_type'] = '%s.%s' % (cls.__module__, cls.__name__)
{'fraction_non_alphanumeric': 0.10607638888888889, 'fraction_numerical': 0.0018229166666666667, 'mean_word_length': 3.1955571740713764, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '5188916', 'n_tokens_mistral': 3429, 'n_tokens_neox': 3347, 'n_words': 1036}