Source code for bioconda_utils.bot.worker

"""
Celery Worker Setup
"""

import abc
import asyncio
import logging
import os
import re
from functools import wraps
import subprocess
from importlib import import_module

import aiohttp

from celery import Celery, Task
from celery.signals import celeryd_init

from kombu import serialization
import simplejson

from ..githubhandler import GitHubAppHandler, GitHubHandler
from ..githandler import install_gpg_key
from ..utils import RepoData, setup_logger
from .config import (
    APP_ID, APP_KEY, CODE_SIGNING_KEY, BOT_NAME, REPODATA_TIMEOUT,
    APP_CLIENT_ID, APP_CLIENT_SECRET
)

logger = logging.getLogger(__name__)  # pylint: disable=invalid-name


[docs]class AsyncTask(Task): """Task class with support for async tasks We override celery.Task with our own version, with some extra features and defaults: - Since we already use a lot of async stuff elsewhere, it's useful to allow the ``run`` method of tasks be ``async``. This Task class detects if the method provided is a coroutine and runs it inside the asyncio event loop. >>> @app.task() >>> async def mytask(self, bind=True): >>> await self.async_init() >>> ... - Provide access to a GitHubAppHandler instance shared at least within the worker process. This is a little tedious. Since the task may be spawned some time after the webook that created it was triggered, the tokens we got inside the webserver may have timed out. In an attempt to avoid wasting API calls to create those tokens continuously, the Task class maintains a copy. - Default to ``acks_late = True``. The reason we use Celery at all is so that spawned tasks can survive a shutdown of the app. """ #: Our tasks should be re-run if they don't finish acks_late = True #: Access the Github API ghapi: "GitHubHandler" = None #: Access Github App API ghappapi: "GitHubAppHandler" = None #: Stores the async run method when the sync run wrapper is installed _async_run = None
[docs] def bind(self, app=None): """Intercept binding of task to (celery) app Here we take the half-finished generated Task class and replace the async run method with a sync run method that executes the original method inside the asyncio loop. """ if asyncio.iscoroutinefunction(self.run): # only for async funcs @wraps(self.run) def sync_run(*args, **kwargs): largs = list(args) # need list so that pre-run can modify self.loop.run_until_complete(self.async_pre_run(largs, kwargs)) return self.loop.run_until_complete(self._async_run(*largs, **kwargs)) # swap run method with wrapper defined above self._async_run, self.run = self.run, sync_run if not self.loop.is_running(): self.loop.run_until_complete(self.async_init()) super().bind(app)
[docs] async def async_init(self): """Init things that need to be run inside the loop This happens during binding -> on load. """ if not self.ghappapi: self.ghappapi = GitHubAppHandler(aiohttp.ClientSession(), BOT_NAME, APP_KEY, APP_ID, APP_CLIENT_ID, APP_CLIENT_SECRET)
[docs] async def async_pre_run(self, args, _kwargs): """Per-call async initialization Prepares the `ghapi` property for tasks. FIXME: doesn't replace kwargs """ for num, arg in enumerate(args): if isinstance(arg, GitHubHandler): args[num] = await self.ghappapi.get_github_api( False, arg.user, arg.repo, arg.installation)
[docs] @abc.abstractmethod def run(self, *_args, **_kwargs): """The tasks actual run method. Will be replaced during bind"""
@property def loop(self): """Get the async loop - creating a new one if necessary""" try: return asyncio.get_event_loop() except RuntimeError: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) return loop
[docs]def custom_dumps(string): """Serialize **s** to JSON accepting **for_json** serializer method""" return simplejson.dumps(string, for_json=True)
[docs]def custom_loads(string): """Deserialize **s** recreating objects JSON objects (dicts) containing a __type__ and a __module__ field are turned into objects by loading and instantiating the type, passing the result dict from obj.for_json() to __init__(). """ def decode(obj): if isinstance(obj, dict): try: typ = obj.pop('__type__') mod = import_module(obj.pop('__module__')) klass = getattr(mod, typ) return klass(**obj) except KeyError: pass return obj return simplejson.loads(string, object_hook=decode)
# Register a custom serializer. We do this so we can conveniently # transfer objects without resorting to pickling. serialization.register('custom_json', custom_dumps, custom_loads, content_type='application/x-bioconda-json', content_encoding='utf8') # Instantiate Celery app, setting our AsyncTask as default # task class and loading the tasks from tasks.py capp = Celery( # pylint: disable=invalid-name task_cls=AsyncTask, include=['bioconda_utils.bot.tasks'] ) # Celery must be configured at module level to catch worker as well # Settings are suggestions from CloudAMPQ capp.conf.update( # Set the URL to the AMQP broker using environment variable broker_url=os.environ.get('CLOUDAMQP_URL'), # Limit the number of connections to the pool. This should # be 2 when running on Heroku to avoid running out of free # connections on CloudAMPQ. # # broker_pool_limit=2, # need two so we can inspect broker_heartbeat=None, broker_connection_timeout=30, # We don't feed back our tasks results result_backend='rpc://', event_queue_expires=60, worker_prefetch_multiplier=1, worker_concurrency=1, task_serializer='custom_json', accept_content=['custom_json', 'json'] #task_acks_late=true )
[docs]@celeryd_init.connect def setup_new_celery_process(sender=None, conf=None, **_kwargs): """This hook is called when a celery worker is initialized Here we make sure that the GPG signing key is installed """ install_gpg_key(CODE_SIGNING_KEY) RepoData().set_timeout(REPODATA_TIMEOUT)