diff --git a/server/website/website/management/commands/runcelery.py b/server/website/website/management/commands/runcelery.py new file mode 100644 index 0000000..0da037b --- /dev/null +++ b/server/website/website/management/commands/runcelery.py @@ -0,0 +1,70 @@ +# +# OtterTune - runcelery.py +# +# Copyright (c) 2017-18, Carnegie Mellon University Database Group +# +import os + +from django.core.management import call_command +from django.core.management.base import BaseCommand +from django.utils import autoreload +from fabric.api import hide, local, settings + + +class Command(BaseCommand): + help = 'Start celery and celerybeat using the auto-reloader.' + + def add_arguments(self, parser): + parser.add_argument( + '--loglevel', + metavar='LOGLEVEL', + help='Logging level, choose between DEBUG, INFO, WARNING, ERROR, CRITICAL, or FATAL. ' + 'Defaults to DEBUG if settings.DEBUG is true, otherwise INFO.') + parser.add_argument( + '--pool', + metavar='POOL_CLS', + default='threads', + help='Pool implementation: prefork (default), eventlet, gevent, solo or threads. ' + 'Default: threads.') + parser.add_argument( + '--celery-pidfile', + metavar='PIDFILE', + default='celery.pid', + help='File used to store the process pid. The program will not start if this ' + 'file already exists and the pid is still alive. Default: celery.pid.') + parser.add_argument( + '--celerybeat-pidfile', + metavar='PIDFILE', + default='celerybeat.pid', + help='File used to store the process pid. The program will not start if this ' + 'file already exists and the pid is still alive. Default: celerybeat.pid.') + parser.add_argument( + '--celery-options', + metavar='OPTIONS', + help="A comma-separated list of additional options to pass to celery, " + "see 'python manage.py celery worker --help' for all available options. " + "IMPORTANT: the option's initial -/-- must be omitted. " + "Example: '--celery-options purge,include=foo.tasks,q'.") + parser.add_argument( + '--celerybeat-options', + metavar='OPTIONS', + help="A comma-separated list of additional options to pass to celerybeat, " + "see 'python manage.py celerybeat --help' for all available options. " + "IMPORTANT: the option's initial -/-- must be omitted. " + "Example: '--celerybeat-options uid=123,q'.") + + def inner_run(self, *args, **options): # pylint: disable=unused-argument + autoreload.raise_last_exception() + + for pidfile in (options['celery_pidfile'], options['celerybeat_pidfile']): + if os.path.exists(pidfile): + with open(pidfile, 'r') as f: + pid = f.read().strip() + with settings(warn_only=True), hide('commands'): # pylint: disable=not-context-manager + local('kill -9 {}'.format(pid)) + local('rm -f {}'.format(pidfile)) + call_command('startcelery', silent=True, pipe='', **options) + self.stdout.write(self.style.SUCCESS("Successfully reloaded celery and celerybeat.")) + + def handle(self, *args, **options): + autoreload.main(self.inner_run, None, options) diff --git a/server/website/website/management/commands/startcelery.py b/server/website/website/management/commands/startcelery.py index b0e79ce..04cde65 100644 --- a/server/website/website/management/commands/startcelery.py +++ b/server/website/website/management/commands/startcelery.py @@ -17,11 +17,12 @@ class Command(BaseCommand): max_wait_sec = 15 def add_arguments(self, parser): - parser.add_argument( + group = parser.add_mutually_exclusive_group() + group.add_argument( '--celery-only', action='store_true', help='Start celery only (skip celerybeat).') - parser.add_argument( + group.add_argument( '--celerybeat-only', action='store_true', help='Start celerybeat only (skip celery).') @@ -63,7 +64,8 @@ class Command(BaseCommand): "IMPORTANT: the option's initial -/-- must be omitted. " "Example: '--celerybeat-options uid=123,q'.") - def _parse_suboptions(self, suboptions): + @staticmethod + def _parse_suboptions(suboptions): suboptions = suboptions or '' parsed = [] for opt in suboptions.split(','): @@ -73,6 +75,21 @@ class Command(BaseCommand): return parsed def handle(self, *args, **options): + # Stealth option to disable stdout + if options.get('silent', False): + self.stdout = open(os.devnull, 'w') + + # Stealth option that assigns where to pipe initial output + pipe = options.get('pipe', None) + if pipe is None: + pipe = '> /dev/null 2>&1' + try: + if 'celery' in settings.LOGGING['loggers']['celery']['handlers']: + logfile = settings.LOGGING['handlers']['celery']['filename'] + pipe = '>> {} 2>&1'.format(logfile) + except KeyError: + pass + loglevel = options['loglevel'] or ('DEBUG' if settings.DEBUG else 'INFO') celery_options = [ '--loglevel={}'.format(loglevel), @@ -84,10 +101,7 @@ class Command(BaseCommand): '--pidfile={}'.format(options['celerybeat_pidfile']), ] + self._parse_suboptions(options['celerybeat_options']) - pipe = '' if 'console' in settings.LOGGING['loggers']['celery']['handlers'] \ - else '> /dev/null 2>&1' - - with lcd(settings.PROJECT_ROOT), hide('commands'): + with lcd(settings.PROJECT_ROOT), hide('commands'): # pylint: disable=not-context-manager if not options['celerybeat_only']: local(self.celery_cmd( cmd='celery worker', opts=' '.join(celery_options), pipe=pipe))