Repository: pulp/pulpcore

Issue:
TASK_DIAGNOSTICS enabled should automatically profile all tasks
Currently when the TASK_DIAGNOSTICS option is enabled, Pulp creates a directory `/var/tmp/pulp/<task-id>/` and writes a data file w/ memory usage over time for the task.

This is really useful for debugging purposes.  But it would be even more useful if tasks would automatically be profiled and the profile written to the same directory.  Otherwise adding this profiling requires tedious manual code changes.

This also opens the door for easy in-situ debugging by users, who can then easily enable profiling themselves on their own tasks, and analyze them or provide them to developers (us) for analysis without necessarily requiring tedious reproduction of the setup ourselves.


Discussion:

Resolving patch:
diff --git a/pulpcore/app/settings.py b/pulpcore/app/settings.py
--- a/pulpcore/app/settings.py
+++ b/pulpcore/app/settings.py
@@ -294,6 +294,8 @@
 
 DEFAULT_AUTO_FIELD = "django.db.models.AutoField"
 
+TASK_DIAGNOSTICS = False
+
 # HERE STARTS DYNACONF EXTENSION LOAD (Keep at the very bottom of settings.py)
 # Read more at https://dynaconf.readthedocs.io/en/latest/guides/django.html
 import dynaconf  # noqa
diff --git a/pulpcore/constants.py b/pulpcore/constants.py
--- a/pulpcore/constants.py
+++ b/pulpcore/constants.py
@@ -1,5 +1,8 @@
+from pathlib import Path
 from types import SimpleNamespace
 
+VAR_TMP_PULP = Path("/var/tmp/pulp")
+
 #: All valid task states.
 TASK_STATES = SimpleNamespace(
     WAITING="waiting",
diff --git a/pulpcore/tasking/pulpcore_worker.py b/pulpcore/tasking/pulpcore_worker.py
--- a/pulpcore/tasking/pulpcore_worker.py
+++ b/pulpcore/tasking/pulpcore_worker.py
@@ -6,10 +6,13 @@
 import logging
 import os
 import random
+import resource
 import select
 import signal
 import socket
 import sys
+import threading
+import time
 import traceback
 from contextlib import suppress
 from datetime import timedelta
@@ -35,6 +38,7 @@
 from pulpcore.constants import (  # noqa: E402: module level not at top of file
     TASK_STATES,
     TASK_INCOMPLETE_STATES,
+    VAR_TMP_PULP,
 )
 
 from pulpcore.exceptions import AdvisoryLockError  # noqa: E402: module level not at top of file
@@ -341,6 +345,27 @@ def run_forever(self):
             self.shutdown()
 
 
+def task_diagnostics_dir(task_pk):
+    # create a directory in /var/tmp/pulp/<task_id>/ and return it
+    taskdata_dir = VAR_TMP_PULP / str(task_pk)
+    taskdata_dir.mkdir(parents=True, exist_ok=True)
+    return taskdata_dir
+
+
+def write_memory_usage(path):
+    _logger.info("Writing task memory data to {}".format(path))
+
+    with open(path, "w") as file:
+        file.write("# Seconds\tMemory in MB\n")
+        seconds = 0
+        while True:
+            current_mb_in_use = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024
+            file.write(f"{seconds}\t{current_mb_in_use:.2f}\n")
+            file.flush()
+            time.sleep(5)
+            seconds += 5
+
+
 def child_signal_handler(sig, frame):
     # Reset signal handlers to default
     # If you kill the process a second time it's not graceful anymore.
@@ -358,6 +383,15 @@ def _perform_task(task_pk, task_working_dir_rel_path):
     signal.signal(signal.SIGINT, child_signal_handler)
     signal.signal(signal.SIGTERM, child_signal_handler)
     signal.signal(signal.SIGUSR1, child_signal_handler)
+    if settings.TASK_DIAGNOSTICS:
+        diagnostics_dir = task_diagnostics_dir(task_pk)
+        mem_diagnostics_path = diagnostics_dir / "memory.datum"
+        # It would be better to have this recording happen in the parent process instead of here
+        # https://github.com/pulp/pulpcore/issues/2337
+        mem_diagnostics_thread = threading.Thread(
+            target=write_memory_usage, args=(mem_diagnostics_path,), daemon=True
+        )
+        mem_diagnostics_thread.start()
     # All processes need to create their own postgres connection
     connection.connection = None
     task = Task.objects.get(pk=task_pk)
@@ -367,28 +401,45 @@ def _perform_task(task_pk, task_working_dir_rel_path):
     user = get_users_with_perms(task).first()
     _set_current_user(user)
     set_guid(task.logging_cid)
-    try:
-        _logger.info(_("Starting task %s"), task.pk)
-
-        # Execute task
-        module_name, function_name = task.name.rsplit(".", 1)
-        module = importlib.import_module(module_name)
-        func = getattr(module, function_name)
-        args = json.loads(task.args) or ()
-        kwargs = json.loads(task.kwargs) or {}
-        os.chdir(task_working_dir_rel_path)
-        result = func(*args, **kwargs)
-        if asyncio.iscoroutine(result):
-            _logger.debug(_("Task is coroutine %s"), task.pk)
-            loop = asyncio.get_event_loop()
-            loop.run_until_complete(result)
-
-    except Exception:
-        exc_type, exc, tb = sys.exc_info()
-        task.set_failed(exc, tb)
-        _logger.info(_("Task %s failed (%s)"), task.pk, exc)
-        _logger.info("\n".join(traceback.format_list(traceback.extract_tb(tb))))
+
+    def execute_task(task):
+        try:
+            _logger.info(_("Starting task %s"), task.pk)
+
+            # Execute task
+            module_name, function_name = task.name.rsplit(".", 1)
+            module = importlib.import_module(module_name)
+            func = getattr(module, function_name)
+            args = json.loads(task.args) or {}
+            kwargs = json.loads(task.kwargs) or {}
+            os.chdir(task_working_dir_rel_path)
+            result = func(*args, **kwargs)
+            if asyncio.iscoroutine(result):
+                _logger.debug(_("Task is coroutine %s"), task.pk)
+                loop = asyncio.get_event_loop()
+                loop.run_until_complete(result)
+
+        except Exception:
+            exc_type, exc, tb = sys.exc_info()
+            task.set_failed(exc, tb)
+            _logger.info(_("Task %s failed (%s)"), task.pk, exc)
+            _logger.info("\n".join(traceback.format_list(traceback.extract_tb(tb))))
+        else:
+            task.set_completed()
+            _logger.info(_("Task completed %s"), task.pk)
+
+    # set up profiling
+    if settings.TASK_DIAGNOSTICS and importlib.util.find_spec("pyinstrument") is not None:
+        from pyinstrument import Profiler
+
+        with Profiler() as profiler:
+            execute_task(task)
+
+        profile_file = diagnostics_dir / "pyinstrument.html"
+        _logger.info("Writing task profile data to {}".format(profile_file))
+        with open(profile_file, "w+") as f:
+            f.write(profiler.output_html())
     else:
-        task.set_completed()
-        _logger.info(_("Task completed %s"), task.pk)
+        execute_task(task)
+
     os.environ.pop("PULP_TASK_ID")


Test patch: