root: celery refactor (#6095)

* root: celery refactor

cleanup deprecation messages by configuring celery with a single object

run celery as django management command

Signed-off-by: Jens Langhammer <jens@goauthentik.io>

* improve debug experience

Signed-off-by: Jens Langhammer <jens@goauthentik.io>

* fix lint

Signed-off-by: Jens Langhammer <jens@goauthentik.io>

* add debugpy to dev dependencies

Signed-off-by: Jens Langhammer <jens@goauthentik.io>

* fix task_always_eager

Signed-off-by: Jens Langhammer <jens@goauthentik.io>

---------

Signed-off-by: Jens Langhammer <jens@goauthentik.io>

.vscode/launch.json

@@ -0,0 +1,27 @@
+{
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python: PDB attach Server",
+            "type": "python",
+            "request": "attach",
+            "connect": {
+                "host": "localhost",
+                "port": 6800
+            },
+            "justMyCode": true,
+            "django": true
+        },
+        {
+            "name": "Python: PDB attach Worker",
+            "type": "python",
+            "request": "attach",
+            "connect": {
+                "host": "localhost",
+                "port": 6900
+            },
+            "justMyCode": true,
+            "django": true
+        },
+    ]
+}

authentik/admin/api/workers.py

@@ -19,7 +19,7 @@ class WorkerView(APIView):
     def get(self, request: Request) -> Response:
         """Get currently connected worker count."""
         count = len(CELERY_APP.control.ping(timeout=0.5))
-        # In debug we run with `CELERY_TASK_ALWAYS_EAGER`, so tasks are ran on the main process
+        # In debug we run with `task_always_eager`, so tasks are ran on the main process
         if settings.DEBUG:  # pragma: no cover
             count += 1
         return Response({"count": count})

authentik/core/management/commands/worker.py

@@ -0,0 +1,40 @@
+"""Run worker"""
+from sys import exit as sysexit
+from tempfile import tempdir
+
+from celery.apps.worker import Worker
+from django.core.management.base import BaseCommand
+from django.db import close_old_connections
+from structlog.stdlib import get_logger
+
+from authentik.lib.config import CONFIG
+from authentik.root.celery import CELERY_APP
+
+LOGGER = get_logger()
+
+
+class Command(BaseCommand):
+    """Run worker"""
+
+    def handle(self, **options):
+        close_old_connections()
+        if CONFIG.y_bool("remote_debug"):
+            import debugpy
+
+            debugpy.listen(("0.0.0.0", 6900))  # nosec
+        worker: Worker = CELERY_APP.Worker(
+            no_color=False,
+            quiet=True,
+            optimization="fair",
+            max_tasks_per_child=1,
+            autoscale=(3, 1),
+            task_events=True,
+            beat=True,
+            schedule_filename=f"{tempdir}/celerybeat-schedule",
+            queues=["authentik", "authentik_scheduled", "authentik_events"],
+        )
+        for task in CELERY_APP.tasks:
+            LOGGER.debug("Registered task", task=task)
+
+        worker.start()
+        sysexit(worker.exitcode)

authentik/events/monitored_tasks.py

@@ -41,6 +41,7 @@ class TaskResult:
 
     def with_error(self, exc: Exception) -> "TaskResult":
         """Since errors might not always be pickle-able, set the traceback"""
+        # TODO: Mark exception somehow so that is rendered as <pre> in frontend
         self.messages.append(exception_to_string(exc))
         return self
 

authentik/lib/default.yml

@@ -26,6 +26,7 @@ redis:
   cache_timeout_reputation: 300
 
 debug: false
+remote_debug: false
 
 log_level: info
 

authentik/root/celery.py

@@ -130,11 +130,7 @@ class LivenessProbe(bootsteps.StartStopStep):
         HEARTBEAT_FILE.touch()
 
 
-# Using a string here means the worker doesn't have to serialize
+CELERY_APP.config_from_object(settings.CELERY)
-# the configuration object to child processes.
-# - namespace='CELERY' means all celery-related configuration keys
-#   should have a `CELERY_` prefix.
-CELERY_APP.config_from_object(settings, namespace="CELERY")
 
 # Load task modules from all registered Django app configs.
 CELERY_APP.autodiscover_tasks()

authentik/root/settings.py

@@ -182,13 +182,13 @@ REST_FRAMEWORK = {
     },
 }
 
-REDIS_PROTOCOL_PREFIX = "redis://"
+_redis_protocol_prefix = "redis://"
-REDIS_CELERY_TLS_REQUIREMENTS = ""
+_redis_celery_tls_requirements = ""
 if CONFIG.y_bool("redis.tls", False):
-    REDIS_PROTOCOL_PREFIX = "rediss://"
+    _redis_protocol_prefix = "rediss://"
-    REDIS_CELERY_TLS_REQUIREMENTS = f"?ssl_cert_reqs={CONFIG.y('redis.tls_reqs')}"
+    _redis_celery_tls_requirements = f"?ssl_cert_reqs={CONFIG.y('redis.tls_reqs')}"
 _redis_url = (
-    f"{REDIS_PROTOCOL_PREFIX}:"
+    f"{_redis_protocol_prefix}:"
     f"{quote_plus(CONFIG.y('redis.password'))}@{quote_plus(CONFIG.y('redis.host'))}:"
     f"{int(CONFIG.y('redis.port'))}"
 )
@@ -326,12 +326,11 @@ USE_TZ = True
 
 LOCALE_PATHS = ["./locale"]
 
-# Celery settings
+CELERY = {
-# Add a 10 minute timeout to all Celery tasks.
+    "task_soft_time_limit": 600,
-CELERY_TASK_SOFT_TIME_LIMIT = 600
+    "worker_max_tasks_per_child": 50,
-CELERY_WORKER_MAX_TASKS_PER_CHILD = 50
+    "worker_concurrency": 2,
-CELERY_WORKER_CONCURRENCY = 2
+    "beat_schedule": {
-CELERY_BEAT_SCHEDULE = {
         "clean_expired_models": {
             "task": "authentik.core.tasks.clean_expired_models",
             "schedule": crontab(minute="2-59/5"),
@@ -342,11 +341,12 @@ CELERY_BEAT_SCHEDULE = {
             "schedule": crontab(minute="9-59/5"),
             "options": {"queue": "authentik_scheduled"},
         },
+    },
+    "task_create_missing_queues": True,
+    "task_default_queue": "authentik",
+    "broker_url": f"{_redis_url}/{CONFIG.y('redis.db')}{_redis_celery_tls_requirements}",
+    "result_backend": f"{_redis_url}/{CONFIG.y('redis.db')}{_redis_celery_tls_requirements}",
 }
-CELERY_TASK_CREATE_MISSING_QUEUES = True
-CELERY_TASK_DEFAULT_QUEUE = "authentik"
-CELERY_BROKER_URL = f"{_redis_url}/{CONFIG.y('redis.db')}{REDIS_CELERY_TLS_REQUIREMENTS}"
-CELERY_RESULT_BACKEND = f"{_redis_url}/{CONFIG.y('redis.db')}{REDIS_CELERY_TLS_REQUIREMENTS}"
 
 # Sentry integration
 env = get_env()
@@ -455,7 +455,7 @@ _DISALLOWED_ITEMS = [
     "INSTALLED_APPS",
     "MIDDLEWARE",
     "AUTHENTICATION_BACKENDS",
-    "CELERY_BEAT_SCHEDULE",
+    "CELERY",
 ]
 
 
@@ -466,7 +466,7 @@ def _update_settings(app_path: str):
         INSTALLED_APPS.extend(getattr(settings_module, "INSTALLED_APPS", []))
         MIDDLEWARE.extend(getattr(settings_module, "MIDDLEWARE", []))
         AUTHENTICATION_BACKENDS.extend(getattr(settings_module, "AUTHENTICATION_BACKENDS", []))
-        CELERY_BEAT_SCHEDULE.update(getattr(settings_module, "CELERY_BEAT_SCHEDULE", {}))
+        CELERY["beat_schedule"].update(getattr(settings_module, "CELERY_BEAT_SCHEDULE", {}))
         for _attr in dir(settings_module):
             if not _attr.startswith("__") and _attr not in _DISALLOWED_ITEMS:
                 globals()[_attr] = getattr(settings_module, _attr)
@@ -482,7 +482,7 @@ for _app in INSTALLED_APPS:
 _update_settings("data.user_settings")
 
 if DEBUG:
-    CELERY_TASK_ALWAYS_EAGER = True
+    CELERY["task_always_eager"] = True
     os.environ[ENV_GIT_HASH_KEY] = "dev"
     INSTALLED_APPS.append("silk")
     SILKY_PYTHON_PROFILER = True

authentik/root/test_runner.py

@@ -30,7 +30,7 @@ class PytestTestRunner:  # pragma: no cover
             self.args.append(f"--randomly-seed={kwargs['randomly_seed']}")
 
         settings.TEST = True
-        settings.CELERY_TASK_ALWAYS_EAGER = True
+        settings.CELERY["task_always_eager"] = True
         CONFIG.y_set("avatars", "none")
         CONFIG.y_set("geoip", "tests/GeoLite2-City-Test.mmdb")
         CONFIG.y_set("blueprints_dir", "./blueprints")

docker-compose.override.yml

@@ -0,0 +1,36 @@
+# This file is used for development and debugging, and should not be used for production instances
+
+version: '3.5'
+
+services:
+  flower:
+    image: ${AUTHENTIK_IMAGE:-ghcr.io/goauthentik/server}:${AUTHENTIK_TAG:-2023.5.4}
+    restart: unless-stopped
+    command: worker-status
+    environment:
+      AUTHENTIK_REDIS__HOST: redis
+      AUTHENTIK_POSTGRESQL__HOST: postgresql
+      AUTHENTIK_POSTGRESQL__USER: ${PG_USER:-authentik}
+      AUTHENTIK_POSTGRESQL__NAME: ${PG_DB:-authentik}
+      AUTHENTIK_POSTGRESQL__PASSWORD: ${PG_PASS}
+    env_file:
+      - .env
+    ports:
+      - "9001:9000"
+    depends_on:
+      - postgresql
+      - redis
+  server:
+    environment:
+      AUTHENTIK_REMOTE_DEBUG: "true"
+      PYDEVD_THREAD_DUMP_ON_WARN_EVALUATION_TIMEOUT: "true"
+    ports:
+      - 6800:6800
+  worker:
+    environment:
+      CELERY_RDB_HOST: "0.0.0.0"
+      CELERY_RDBSIG: "1"
+      AUTHENTIK_REMOTE_DEBUG: "true"
+      PYDEVD_THREAD_DUMP_ON_WARN_EVALUATION_TIMEOUT: "true"
+    ports:
+      - 6900:6900

lifecycle/ak

@@ -54,6 +54,16 @@ function cleanup {
     rm -f ${MODE_FILE}
 }
 
+function prepare_debug {
+    pip install --no-cache-dir -r /requirements-dev.txt
+    touch /unittest.xml
+    chown authentik:authentik /unittest.xml
+}
+
+if [[ "${AUTHENTIK_REMOTE_DEBUG}" == "true" ]]; then
+    prepare_debug
+fi
+
 if [[ "$1" == "server" ]]; then
     wait_for_db
     set_mode "server"
@@ -67,7 +77,7 @@ if [[ "$1" == "server" ]]; then
 elif [[ "$1" == "worker" ]]; then
     wait_for_db
     set_mode "worker"
-    check_if_root "celery -A authentik.root.celery worker -Ofair --max-tasks-per-child=1 --autoscale 3,1 -E -B -s /tmp/celerybeat-schedule -Q authentik,authentik_scheduled,authentik_events"
+    check_if_root "python -m manage worker"
 elif [[ "$1" == "worker-status" ]]; then
     wait_for_db
     celery -A authentik.root.celery flower \
@@ -75,9 +85,7 @@ elif [[ "$1" == "worker-status" ]]; then
 elif [[ "$1" == "bash" ]]; then
     /bin/bash
 elif [[ "$1" == "test-all" ]]; then
-    pip install --no-cache-dir -r /requirements-dev.txt
+    prepare_debug
-    touch /unittest.xml
-    chown authentik:authentik /unittest.xml
     check_if_root "python -m manage test authentik"
 elif [[ "$1" == "healthcheck" ]]; then
     run_authentik healthcheck $(cat $MODE_FILE)

lifecycle/gunicorn.conf.py

@@ -157,3 +157,8 @@ if not CONFIG.y_bool("disable_startup_analytics", False):
         # pylint: disable=broad-exception-caught
         except Exception:  # nosec
             pass
+
+if CONFIG.y_bool("remote_debug"):
+    import debugpy
+
+    debugpy.listen(("0.0.0.0", 6800))  # nosec

poetry.lock

@@ -1066,6 +1066,33 @@ twisted = {version = ">=22.4", extras = ["tls"]}
 [package.extras]
 tests = ["django", "hypothesis", "pytest", "pytest-asyncio"]
 
+[[package]]
+name = "debugpy"
+version = "1.6.7"
+description = "An implementation of the Debug Adapter Protocol for Python"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "debugpy-1.6.7-cp310-cp310-macosx_11_0_x86_64.whl", hash = "sha256:b3e7ac809b991006ad7f857f016fa92014445085711ef111fdc3f74f66144096"},
+    {file = "debugpy-1.6.7-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e3876611d114a18aafef6383695dfc3f1217c98a9168c1aaf1a02b01ec7d8d1e"},
+    {file = "debugpy-1.6.7-cp310-cp310-win32.whl", hash = "sha256:33edb4afa85c098c24cc361d72ba7c21bb92f501104514d4ffec1fb36e09c01a"},
+    {file = "debugpy-1.6.7-cp310-cp310-win_amd64.whl", hash = "sha256:ed6d5413474e209ba50b1a75b2d9eecf64d41e6e4501977991cdc755dc83ab0f"},
+    {file = "debugpy-1.6.7-cp37-cp37m-macosx_10_15_x86_64.whl", hash = "sha256:38ed626353e7c63f4b11efad659be04c23de2b0d15efff77b60e4740ea685d07"},
+    {file = "debugpy-1.6.7-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:279d64c408c60431c8ee832dfd9ace7c396984fd7341fa3116aee414e7dcd88d"},
+    {file = "debugpy-1.6.7-cp37-cp37m-win32.whl", hash = "sha256:dbe04e7568aa69361a5b4c47b4493d5680bfa3a911d1e105fbea1b1f23f3eb45"},
+    {file = "debugpy-1.6.7-cp37-cp37m-win_amd64.whl", hash = "sha256:f90a2d4ad9a035cee7331c06a4cf2245e38bd7c89554fe3b616d90ab8aab89cc"},
+    {file = "debugpy-1.6.7-cp38-cp38-macosx_10_15_x86_64.whl", hash = "sha256:5224eabbbeddcf1943d4e2821876f3e5d7d383f27390b82da5d9558fd4eb30a9"},
+    {file = "debugpy-1.6.7-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bae1123dff5bfe548ba1683eb972329ba6d646c3a80e6b4c06cd1b1dd0205e9b"},
+    {file = "debugpy-1.6.7-cp38-cp38-win32.whl", hash = "sha256:9cd10cf338e0907fdcf9eac9087faa30f150ef5445af5a545d307055141dd7a4"},
+    {file = "debugpy-1.6.7-cp38-cp38-win_amd64.whl", hash = "sha256:aaf6da50377ff4056c8ed470da24632b42e4087bc826845daad7af211e00faad"},
+    {file = "debugpy-1.6.7-cp39-cp39-macosx_11_0_x86_64.whl", hash = "sha256:0679b7e1e3523bd7d7869447ec67b59728675aadfc038550a63a362b63029d2c"},
+    {file = "debugpy-1.6.7-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:de86029696e1b3b4d0d49076b9eba606c226e33ae312a57a46dca14ff370894d"},
+    {file = "debugpy-1.6.7-cp39-cp39-win32.whl", hash = "sha256:d71b31117779d9a90b745720c0eab54ae1da76d5b38c8026c654f4a066b0130a"},
+    {file = "debugpy-1.6.7-cp39-cp39-win_amd64.whl", hash = "sha256:c0ff93ae90a03b06d85b2c529eca51ab15457868a377c4cc40a23ab0e4e552a3"},
+    {file = "debugpy-1.6.7-py2.py3-none-any.whl", hash = "sha256:53f7a456bc50706a0eaabecf2d3ce44c4d5010e46dfc65b6b81a518b42866267"},
+    {file = "debugpy-1.6.7.zip", hash = "sha256:c4c2f0810fa25323abfdfa36cbbbb24e5c3b1a42cb762782de64439c575d67f2"},
+]
+
 [[package]]
 name = "deepmerge"
 version = "1.1.0"
@@ -1686,13 +1713,13 @@ test = ["Cython (>=0.29.24,<0.30.0)"]
 
 [[package]]
 name = "humanize"
-version = "4.6.0"
+version = "4.7.0"
 description = "Python humanize utilities"
 optional = false
-python-versions = ">=3.7"
+python-versions = ">=3.8"
 files = [
-    {file = "humanize-4.6.0-py3-none-any.whl", hash = "sha256:401201aca462749773f02920139f302450cb548b70489b9b4b92be39fe3c3c50"},
+    {file = "humanize-4.7.0-py3-none-any.whl", hash = "sha256:df7c429c2d27372b249d3f26eb53b07b166b661326e0325793e0a988082e3889"},
-    {file = "humanize-4.6.0.tar.gz", hash = "sha256:5f1f22bc65911eb1a6ffe7659bd6598e33dcfeeb904eb16ee1e705a09bf75916"},
+    {file = "humanize-4.7.0.tar.gz", hash = "sha256:7ca0e43e870981fa684acb5b062deb307218193bca1a01f2b2676479df849b3a"},
 ]
 
 [package.extras]
@@ -3419,13 +3446,13 @@ wsproto = ">=0.14"
 
 [[package]]
 name = "twilio"
-version = "8.3.0"
+version = "8.4.0"
 description = "Twilio API client and TwiML generator"
 optional = false
 python-versions = ">=3.7.0"
 files = [
-    {file = "twilio-8.3.0-py2.py3-none-any.whl", hash = "sha256:f8f4a26e7491e015777c2c12abcc068321f12302d081fc355df486601434c311"},
+    {file = "twilio-8.4.0-py2.py3-none-any.whl", hash = "sha256:56b812b4d77dabcfdf7aa02aac966065e064beabd083621940856a6ee0d060ee"},
-    {file = "twilio-8.3.0.tar.gz", hash = "sha256:e76543b054f09304557d9bd0f9e3c21d09ca935d88f833788d43cab1f1fb67d1"},
+    {file = "twilio-8.4.0.tar.gz", hash = "sha256:23fa599223d336a19d674394535d42bd1e260f7ca350a51d02b9d902370d76ef"},
 ]
 
 [package.dependencies]
@@ -4159,4 +4186,4 @@ files = [
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "7c78d6909ba8cc5b8fb41233e2506f0b919b71e263213068e479af706a9670ce"
+content-hash = "60a0e729895ebd44235e88e0414cc64e50c41736903ea61e6fb94a542dd2bb3c"

pyproject.toml

@@ -178,6 +178,7 @@ black = "*"
 bump2version = "*"
 colorama = "*"
 coverage = { extras = ["toml"], version = "*" }
+debugpy = "*"
 django-silk = "*"
 drf-jsonschema-serializer = "*"
 importlib-metadata = "*"