Add celery profiling (#913)

engine/engine/celery.py

@@ -1,13 +1,20 @@
+import logging
 import os
+import time
 
 import celery
 from celery.app.log import TaskFormatter
+from celery.utils.debug import memdump, sample_mem
+from celery.utils.log import get_task_logger
+from django.conf import settings
 
-# set the default Django settings module for the 'celery' program.
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "settings.prod")
 
 from django.db import connection  # noqa: E402
 
+logger = get_task_logger(__name__)
+logger.setLevel(logging.DEBUG)
+
 connection.cursor()
 from celery import Celery  # noqa: E402
 
@@ -38,6 +45,20 @@ def on_after_setup_logger(logger, **kwargs):
     for handler in logger.handlers:
         handler.setFormatter(
             TaskFormatter(
-                "%(asctime)s source=engine:celery task_id=%(task_id)s task_name=%(task_name)s name=%(name)s level=%(levelname)s %(message)s"
+                "%(asctime)s source=engine:celery worker=%(processName)s task_id=%(task_id)s task_name=%(task_name)s name=%(name)s level=%(levelname)s %(message)s"
             )
         )
+
+
+if settings.DEBUG_CELERY_TASKS_PROFILING:
+
+    @celery.signals.task_prerun.connect
+    def start_task_timer(task_id=None, task=None, *a, **kw):
+        logger.info("started: {} of {} with cpu={} at {}".format(task_id, task.name, time.perf_counter(), time.time()))
+        sample_mem()
+
+    @celery.signals.task_postrun.connect
+    def finish_task_timer(task_id=None, task=None, *a, **kw):
+        logger.info("ended: {} of {} with cpu={} at {}".format(task_id, task.name, time.perf_counter(), time.time()))
+        sample_mem()
+        memdump()

engine/requirements.txt

@@ -41,3 +41,4 @@ psycopg2-binary==2.9.3
 emoji==1.7.0
 apns2==0.7.2
 regex==2021.11.2
+psutil==5.9.4

engine/settings/base.py

@@ -39,6 +39,8 @@ MIRAGE_CIPHER_MODE = "CBC"
 # SECURITY WARNING: don't run with debug turned on in production!
 DEBUG = False
 
+DEBUG_CELERY_TASKS_PROFILING = getenv_boolean("DEBUG_CELERY_TASKS_PROFILING", False)
+
 ALLOWED_HOSTS = [item.strip() for item in os.environ.get("ALLOWED_HOSTS", "*").split(",")]
 
 # TODO: update link to up-to-date docs