Add X-TASK-DIAGNOSTICS header support

This allows individual tasks to be profiled. The value must be a comma separated list of "memory", "pyinstrument", "memray".

closes: #6725

Author: dkliban

.github/workflows/scripts/install.sh

@@ -120,7 +120,7 @@ if [ "$TEST" = "azure" ]; then
       - ./azurite:/etc/pulp\
     command: "azurite-blob --blobHost 0.0.0.0"' vars/main.yaml
   sed -i -e '$a azure_test: true\
-pulp_scenario_settings: {"api_root_rewrite_header": "X-API-Root", "content_origin": null, "domain_enabled": true, "rest_framework__default_permission_classes": ["pulpcore.plugin.access_policy.DefaultAccessPolicy"]}\
+pulp_scenario_settings: {"api_root_rewrite_header": "X-API-Root", "content_origin": null, "domain_enabled": true, "rest_framework__default_permission_classes": ["pulpcore.plugin.access_policy.DefaultAccessPolicy"], "task_diagnostics": ["memory"]}\
 pulp_scenario_env: {}\
 ' vars/main.yaml
 fi

CHANGES/6725.feature

@@ -0,0 +1,2 @@
+Added support for profiling individual tasks with X-TASK-DIAGNOSTICS header.
+The TASK_DIAGNOSTICS setting now determines which profilers are available to users of the REST API.

docs/admin/reference/settings.md

@@ -344,18 +344,15 @@ Defaults to `'REMOTE_USER'`.
 
 ### TASK\_DIAGNOSTICS
 
-When enabled, each task records various diagnostics for analysis by a developer.
-Some of these diagnostics add runtime overhead, and should probably not be enabled in production environments without careful thought and supervision.
+When enabled, users are allowed to request various diagnostics for analysis by a developer.
 
-The list of possible diagnostics that can be used are `["memory", "pyinstrument", "memray"]`.
+Administrators can enable support for `["memory", "pyinstrument", "memray"]`.
 Any subset of those can be used, or all of them, or none of them.
 
-A value of `True` enables all diagnostics.
-
 !!! note
     `memray` and `pyinstrument` diagnostics require additional packages to be installed.
 
-Defaults to `[]` (no diagnostics).
+Defaults to `[]` (no diagnostics enabled).
 
 See [task diagnostics documentation] for more details.
 

docs/user/guides/troubleshoot-tasks.md

@@ -47,7 +47,33 @@ Cancel a running task:
 pulp task cancel --href "$TASK_HREF"
 ```
 
+## Profiling Tasks
+
+Task profiling is enabled by an administrator using the [TASK_DIAGNOSTICS] setting.
+
+Users can submit a `X-TASK-DIAGNOSTICS` header with their API requests that generate tasks.
+The header value should be a comma separated list of : "memory", "pyinstrument",or "memray".
+For example:
+
+```bash
+`pulp --header X-TASK-DIAGNOSTICS:pyinstrument,memory rpm publication create --repository foo
+````
+
+Once the task finishes, the files generated by the profiler can be found using the following command:
+
+```bash
+pulp task profile-artifact-urls --href "$TASK_HREF"
+{
+  "memory_profile": "/pulp/content/default/115341ffbc5c32b379142936bd85ab658a83209a0ab03f495d0448bf1f9ffee0/0197af04-272a-71cd-b90d-174be8d8ddd8?expires=1750992034&validate_token=e56ce8188a4c16bcb36ddc916d314e623d959b3ae1ebd84a79ebafac47bc49ea:46fcedc3ca8e826fd578016cb79fa9842e472572122c312e59a6f3682d50c4aa",
+  "pyinstrument_profile": "/pulp/content/default/115341ffbc5c32b379142936bd85ab658a83209a0ab03f495d0448bf1f9ffee0/0197af04-2769-78a1-9b1d-96854d8573a0?expires=1750992034&validate_token=b58aba0031f662de5666ef2bc1682093ba07d3d9cefcfbfe254ad762b34e9d0e:377c1b92d327a1432a46b64d747ecb50e830aadd79892c770e8249d087eaecdd"
+}
+```
+
+The files can then be downloaded from Pulp using the paths returned by the above command.
+
 ## Tracing workloads
 
-To help users better trace workloads in Pulp, Pulp provides [support for correlation
-ids](site:pulpcore/docs/user/guides/correlation-id/).
+To help users better trace workloads in Pulp, Pulp provides support for [correlation ids].
+
+[correlation ids]: site:pulpcore/docs/user/guides/correlation-id/
+[TASK_DIAGNOSTICS]: site:pulpcore/docs/admin/reference/settings/#task_diagnostics

pulpcore/app/migrations/0132_task_profile_options.py

@@ -0,0 +1,19 @@
+# Generated by Django 4.2.10 on 2025-06-26 19:44
+
+import django.contrib.postgres.fields
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('core', '0131_distribution_checkpoint_publication_checkpoint'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='task',
+            name='profile_options',
+            field=django.contrib.postgres.fields.ArrayField(base_field=models.TextField(), null=True, size=None),
+        ),
+    ]

pulpcore/app/models/task.py

@@ -147,6 +147,7 @@ class Task(BaseModel, AutoAddObjPermsMixin):
     pulp_domain = models.ForeignKey("Domain", default=get_domain_pk, on_delete=models.CASCADE)
     versions = HStoreField(default=dict)
 
+    profile_options = ArrayField(models.TextField(), null=True)
     profile_artifacts = models.ManyToManyField("Artifact", through=ProfileArtifact)
 
     immediate = models.BooleanField(default=False, null=True)

pulpcore/app/settings.py

@@ -173,6 +173,7 @@
     "django.middleware.clickjacking.XFrameOptionsMiddleware",
     "pulpcore.middleware.DomainMiddleware",
     "pulpcore.middleware.APIRootRewriteMiddleware",
+    "pulpcore.middleware.TaskProfilerMiddleware",
 ]
 
 AUTHENTICATION_BACKENDS = [

pulpcore/middleware.py

@@ -1,4 +1,5 @@
 import time
+from contextvars import ContextVar
 
 from django.http.response import Http404
 from django.conf import settings
@@ -13,6 +14,8 @@
     set_domain,
 )
 
+x_task_diagnostics_var = ContextVar("x_profile_task")
+
 
 class DomainMiddleware:
     """
@@ -143,3 +146,25 @@ def _process_path(request, response):
         match = getattr(request, "resolver_match", "")
         route = getattr(match, "route", "")
         return route
+
+
+class TaskProfilerMiddleware:
+    """
+    Middleware that looks for the presence of X-TASK-DIAGNOSTICS header and provides its value
+    as a ContextVar to the dispatch() method in the tasking system. It enables generating
+    profiling data of tasks dispatched via API.
+    """
+
+    def __init__(self, get_response):
+        self.get_response = get_response
+
+    def __call__(self, request):
+        if "HTTP_X_TASK_DIAGNOSTICS" in request.META:
+            task_diagnostics = request.META["HTTP_X_TASK_DIAGNOSTICS"]
+            ctx_token = x_task_diagnostics_var.set(task_diagnostics.split(","))
+            try:
+                return self.get_response(request)
+            finally:
+                x_task_diagnostics_var.reset(ctx_token)
+        else:
+            return self.get_response(request)

pulpcore/openapi/__init__.py

@@ -201,6 +201,19 @@ def _get_serializer_name(self, serializer, direction, bypass_extensions=False):
             name = name + "Response"
         return name
 
+    def get_override_parameters(self):
+        parameters = super().get_override_parameters()
+        parameters.append(
+            OpenApiParameter(
+                name="X-Task-Diagnostics",
+                location=OpenApiParameter.HEADER,
+                required=False,
+                type={"type": "array", "items": {"type": "string"}},
+                description="List of profilers to use on tasks.",
+            )
+        )
+        return parameters
+
     def map_parsers(self):
         """
         Get request parsers.

pulpcore/tasking/_util.py

@@ -127,21 +127,26 @@ def perform_task(task_pk, task_working_dir_rel_path):
     set_domain(task.pulp_domain)
     os.chdir(task_working_dir_rel_path)
 
-    if settings.TASK_DIAGNOSTICS:
-        _execute_task_and_profile(task)
+    if task.profile_options:
+        profilers = set(task.profile_options) & set(settings.TASK_DIAGNOSTICS)
+        if unavailable_profilers := set(task.profile_options) - set(settings.TASK_DIAGNOSTICS):
+            _logger.warning(
+                "Requested task diagnostic profilers are not available: %s", unavailable_profilers
+            )
+        _execute_task_and_profile(task, profilers)
     else:
         execute_task(task)
 
 
-def _execute_task_and_profile(task):
+def _execute_task_and_profile(task, profile_options):
     with tempfile.TemporaryDirectory(dir=settings.WORKING_DIRECTORY) as temp_dir:
         _execute_task = execute_task
 
-        if settings.TASK_DIAGNOSTICS is True or "memory" in settings.TASK_DIAGNOSTICS:
+        if "memory" in profile_options:
             _execute_task = _memory_diagnostic_decorator(temp_dir, _execute_task)
-        if settings.TASK_DIAGNOSTICS is True or "pyinstrument" in settings.TASK_DIAGNOSTICS:
+        if "pyinstrument" in profile_options:
             _execute_task = _pyinstrument_diagnostic_decorator(temp_dir, _execute_task)
-        if settings.TASK_DIAGNOSTICS is True or "memray" in settings.TASK_DIAGNOSTICS:
+        if "memray" in profile_options:
             _execute_task = _memray_diagnostic_decorator(temp_dir, _execute_task)
 
         _execute_task(task)