<commit_before>import six

from girder.api import access
from girder.api.describe import describeRoute, Description
from girder.api.rest import ensureTokenScopes, filtermodel, loadmodel, Resource
from girder.constants import AccessType, TokenScope
from girder.models.model_base import ValidationException
from girder.plugins.worker import utils
from . import constants


class WorkerTask(Resource):
    def __init__(self):
        super(WorkerTask, self).__init__()

        self.resourceName = 'worker_task'

        self.route('GET', (), self.listTasks)
        self.route('POST', (':id', 'execution'), self.executeTask)

    @access.public
    @describeRoute(
        Description('List all available tasks that can be executed.')
        .pagingParams(defaultSort='name')
    )
    @filtermodel(model='item')
    def listTasks(self, params):
        limit, offset, sort = self.getPagingParameters(params, 'name')

        cursor = self.model('item').find({
            'meta.workerTaskSpec': {'$exists': True}
        }, sort=sort)

        return list(self.model('item').filterResultsByPermission(
            cursor, self.getCurrentUser(), level=AccessType.READ, limit=limit, offset=offset,
            flags=constants.ACCESS_FLAG_EXECUTE_TASK))

    def _validateTask(self, item):
        """
        Some basic validation of the task spec.
        """
        if 'workerTaskSpec' not in item.get('meta'):
            raise ValidationException('Item (%s) does not contain a worker task specification.')
        spec = item['meta']['workerTaskSpec']

        if not isinstance(spec, dict):
            raise ValidationException('Task spec should be a JSON object.')

        inputs = spec.get('inputs', [])
        outputs = spec.get('outputs', [])

        if not isinstance(inputs, (list, tuple)):
            raise ValidationException('Task inputs must be a list.')

        if not isinstance(outputs, (list, tuple)):
            raise ValidationException('Task outputs must be a list.')

        if 'mode' not in spec:
            raise ValidationException('Task must contain a "mode" field.')

        # Ensure that format and type keys exist in every task IO spec,
        # the worker complains otherwise.
        for ioSpec in inputs + outputs:
            ioSpec['format'] = ioSpec.get('format', 'none')
            ioSpec['type'] = ioSpec.get('type', 'none')

        return spec

    def _transformInputs(self, inputs, token):
        """
        Validates and sanitizes the input bindings. If they are Girder inputs, adds
        the necessary token info. If the token does not allow DATA_READ, or if the user
        does not have read access to the resource, raises an AccessException.
        """
        transformed = {}
        for k, v in six.viewitems(inputs):
            if v['mode'] == 'girder':
                ensureTokenScopes(token, TokenScope.DATA_READ)
                rtype = v.get('resource_type', 'file')
                if rtype not in {'file', 'item', 'folder'}:
                    raise ValidationException('Invalid input resource_type: %s.' % rtype)

                resource = self.model(rtype).load(
                    v['id'], level=AccessType.READ, user=self.getCurrentUser(), exc=True)

                transformed[k] = utils.girderInputSpec(
                    resource, resourceType=rtype, token=token, dataFormat='none')
            elif v['mode'] == 'inline':
                transformed[k] = {
                    'mode': v['mode'],
                    'data': v['data'],
                    'type': 'none',
                    'format': 'none'
                }
            else:
                raise ValidationException('Invalid input mode: %s.' % v['mode'])

        return transformed

    def _transformOutputs(self, outputs, token):
        """
        Validates and sanitizes the input bindings. If they are Girder inputs, adds
        the necessary token info. If the token does not allow DATA_WRITE, or if the user
        does not have write access to the destination, raises an AccessException.
        """
        transformed = {}
        for k, v in six.viewitems(outputs):
            if v['mode'] == 'girder':
                ensureTokenScopes(token, TokenScope.DATA_WRITE)
                ptype = v.get('parent_type', 'folder')
                if ptype not in {'item', 'folder'}:
                    raise ValidationException('Invalid output parent type: %s.' % ptype)

                parent = self.model(ptype).load(
                    v['parent_id'], level=AccessType.WRITE, user=self.getCurrentUser(), exc=True)

                transformed[k] = utils.girderOutputSpec(
                    parent, parentType=ptype, token=token, name=v.get('name'), dataFormat='none')
            else:
                raise ValidationException('Invalid output mode: %s.' % v['mode'])

        return transformed


    @access.user(scope=constants.TOKEN_SCOPE_EXECUTE_TASK)
    @loadmodel(
        model='item', level=AccessType.READ, requiredFlags=constants.ACCESS_FLAG_EXECUTE_TASK)
    @filtermodel(model='job', plugin='jobs')
    @describeRoute(
        Description('Execute a task described by an item.')
        .param('id', 'The ID of the item representing the task specification.', paramType='path')
        .param('jobTitle', 'Title for this job execution.', required=False)
        .param('inputs', 'The input bindings for the task.', required=False)
        .param('outputs', 'The output bindings for the task.', required=False)
        .param('includeJobInfo', 'Whether to track the task using a job record.', required=False,
               dataType='boolean', default=True)
    )
    def executeTask(self, item, params):
        includeJobInfo = self.boolParam('includeJobInfo', params, default=True)
        title = params.get('jobTitle', item['name'])
        task = self._validateTask(item)

        jobModel = self.model('job', 'jobs')

        job = jobModel.createJob(
            title=title, type='worker_task', handler='worker_handler', user=self.getCurrentUser())

        # If this is a user auth token, we make an IO-enabled token
        token = self.getCurrentToken()<commit_msg>Configurable handler for worker tasks<commit_after>import six

from girder.api import access
from girder.api.describe import describeRoute, Description
from girder.api.rest import ensureTokenScopes, filtermodel, loadmodel, Resource
from girder.constants import AccessType, TokenScope
from girder.models.model_base import ValidationException
from girder.plugins.worker import utils
from . import constants


class WorkerTask(Resource):
    def __init__(self):
        super(WorkerTask, self).__init__()

        self.resourceName = 'worker_task'

        self.route('GET', (), self.listTasks)
        self.route('POST', (':id', 'execution'), self.executeTask)

    @access.public
    @describeRoute(
        Description('List all available tasks that can be executed.')
        .pagingParams(defaultSort='name')
    )
    @filtermodel(model='item')
    def listTasks(self, params):
        limit, offset, sort = self.getPagingParameters(params, 'name')

        cursor = self.model('item').find({
            'meta.workerTaskSpec': {'$exists': True}
        }, sort=sort)

        return list(self.model('item').filterResultsByPermission(
            cursor, self.getCurrentUser(), level=AccessType.READ, limit=limit, offset=offset,
            flags=constants.ACCESS_FLAG_EXECUTE_TASK))

    def _validateTask(self, item):
        """
        Some basic validation of the task spec.
        """
        if 'workerTaskSpec' not in item.get('meta'):
            raise ValidationException('Item (%s) does not contain a worker task specification.')
        spec = item['meta']['workerTaskSpec']

        if not isinstance(spec, dict):
            raise ValidationException('Task spec should be a JSON object.')

        inputs = spec.get('inputs', [])
        outputs = spec.get('outputs', [])

        if not isinstance(inputs, (list, tuple)):
            raise ValidationException('Task inputs must be a list.')

        if not isinstance(outputs, (list, tuple)):
            raise ValidationException('Task outputs must be a list.')

        if 'mode' not in spec:
            raise ValidationException('Task must contain a "mode" field.')

        # Ensure that format and type keys exist in every task IO spec,
        # the worker complains otherwise.
        for ioSpec in inputs + outputs:
            ioSpec['format'] = ioSpec.get('format', 'none')
            ioSpec['type'] = ioSpec.get('type', 'none')

        handler = item.get('meta', {}).get('workerTaskHandler') or 'worker_handler'

        return spec, handler

    def _transformInputs(self, inputs, token):
        """
        Validates and sanitizes the input bindings. If they are Girder inputs, adds
        the necessary token info. If the token does not allow DATA_READ, or if the user
        does not have read access to the resource, raises an AccessException.
        """
        transformed = {}
        for k, v in six.viewitems(inputs):
            if v['mode'] == 'girder':
                ensureTokenScopes(token, TokenScope.DATA_READ)
                rtype = v.get('resource_type', 'file')
                if rtype not in {'file', 'item', 'folder'}:
                    raise ValidationException('Invalid input resource_type: %s.' % rtype)

                resource = self.model(rtype).load(
                    v['id'], level=AccessType.READ, user=self.getCurrentUser(), exc=True)

                transformed[k] = utils.girderInputSpec(
                    resource, resourceType=rtype, token=token, dataFormat='none')
            elif v['mode'] == 'inline':
                transformed[k] = {
                    'mode': v['mode'],
                    'data': v['data'],
                    'type': 'none',
                    'format': 'none'
                }
            else:
                raise ValidationException('Invalid input mode: %s.' % v['mode'])

        return transformed

    def _transformOutputs(self, outputs, token):
        """
        Validates and sanitizes the input bindings. If they are Girder inputs, adds
        the necessary token info. If the token does not allow DATA_WRITE, or if the user
        does not have write access to the destination, raises an AccessException.
        """
        transformed = {}
        for k, v in six.viewitems(outputs):
            if v['mode'] == 'girder':
                ensureTokenScopes(token, TokenScope.DATA_WRITE)
                ptype = v.get('parent_type', 'folder')
                if ptype not in {'item', 'folder'}:
                    raise ValidationException('Invalid output parent type: %s.' % ptype)

                parent = self.model(ptype).load(
                    v['parent_id'], level=AccessType.WRITE, user=self.getCurrentUser(), exc=True)

                transformed[k] = utils.girderOutputSpec(
                    parent, parentType=ptype, token=token, name=v.get('name'), dataFormat='none')
            else:
                raise ValidationException('Invalid output mode: %s.' % v['mode'])

        return transformed


    @access.user(scope=constants.TOKEN_SCOPE_EXECUTE_TASK)
    @loadmodel(
        model='item', level=AccessType.READ, requiredFlags=constants.ACCESS_FLAG_EXECUTE_TASK)
    @filtermodel(model='job', plugin='jobs')
    @describeRoute(
        Description('Execute a task described by an item.')
        .param('id', 'The ID of the item representing the task specification.', paramType='path')
        .param('jobTitle', 'Title for this job execution.', required=False)
        .param('inputs', 'The input bindings for the task.', required=False)
        .param('outputs', 'The output bindings for the task.', required=False)
        .param('includeJobInfo', 'Whether to track the task using a job record.', required=False,
               dataType='boolean', default=True)
    )
    def executeTask(self, item, params):
        includeJobInfo = self.boolParam('includeJobInfo', params, default=True)
        title = params.get('jobTitle', item['name'])
        task, handler = self._validateTask(item)

        jobModel = self.model('job', 'jobs')

        job = jobModel.createJob(
            title=title, type='worker_task', handler=handler, user=self.getCurrentUser())

        # If this is a user auth token, we make an IO-enabled token
        token = self.getCurrentToken()
{'fraction_non_alphanumeric': 0.11489635224168407, 'fraction_numerical': 0.0, 'mean_word_length': 2.469063545150502, 'pattern_counts': {'":': 0, '<': 3, '<?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': '5377576', 'n_tokens_mistral': 3323, 'n_tokens_neox': 3193, 'n_words': 972}