RAG: Replay Celery Tasks

Posted on

A detour to look at some plumbing to support replaying/reinvoking Celery tasks.

This follows the previous entry: RAG: Source Ingesting and Model – Snippets

Background: Retrying Celery Tasks

Celery has some mechanism to retry automatically some failed/aborted tasks with settings such as task_acks_on_failure_or_timeout and task_acks_late. When used in conjunction with a broker’s configuration of its messages’ longevity and ACK policy, I can tune an async system such that aborted/killed tasks (e.g. when the Celery process is killed by a deployment) can be retried. I have some notes on Celery with RabbitMQ here: https://www.pn.therealvan.com/2024/09/01/celery-rabbitmq-notes-on-task-timings/

However, if a task failed due to a bad application or data state, retrying the task will just fail again. Because of this, most setups won’t retry a task indefinitely lest we end up with a backlog of failed Celery tasks failing over and over again.

From my past experience, when a Celery task failed due to bad application/data state, assuming we don’t retry forever, the sequence of events typically is as follows:

  • Identify the failed task somewhere. Sometimes this is not an easy job; the error logs just suddenly spit out a bunch of ERROR entries. The right ones need to be filtered out, often times done after looking at a bigger context logged in adjacent log entries. Sometimes this involved correlating data from multiple logs, systems, DBs. Quickly erected “microservice architectures” tend to make this investigation trickier.
  • Investigation of the failed tasks begins. Depending on the severity of the failures, the Celery workers may be temporarily deactivated (and let the messages queue up temporarily). Otherwise, just let them fail and collect information about the failed tasks so that we can retry later.
  • Once the underlying application/data state is “fixed,” somehow, we need to retry the failed Celery tasks. This CAN be quite a chore depending on how easy (if even possible) it is to get the application/data state to a point where we can retrigger the tasks. IDs are collected. Data models are analyzed. Logs are collected, parsed, sorted, etc. In the worst cases, new code is written to specifically call the Celery tasks with the correct parameters.

Not much can be done about most of these other than investing in making it easier to identify and log these errors when they happen AND design the flows to allow replays possible (e.g. idempotent service APIs) and easier. However, once the surrounding concerns are addressed, one part that can be improved is to make it easier to reinvoke failed Celery tasks with the same exact parameters that they were originally called with.

Celery Task Invocation with Parameters

Recent versions of Celery defaults to using JSON serialization of parameters. Therefore, if the Celery task follows the default configuration, it’s very easy to record task invocations with their exact args and kwargs. For example, these Celery tasks can be called:

@shared_task
def index_url(url: str, requested_by: int):
    ...

@shared_task
def index_urls(urls: list[str], requested_by: int, time_to_wait: int = 600):
    ...


...

# the following will call: index_url(url, 7)
index_url.apply_async(args=[url, 7])



# the following call: index_urls([url1, url2, url3], 7, time_to_wait=1200)
index_urls.apply_async(args=[[url1, url2, url3], 7], kwargs={'time_to_wait': 1200})

As long as we record these properties somewhere (e.g. a DB table, a file):

  • task name
  • the value of args
  • the value of kwargs
  • any additional keyword arguments passed to apply_async (e.g. countdown)

then we can call the task as many times as we want:

# Assume we retrieved these properties from somewhere
task_name = 'index_urls'
task_args = [[url1, url2, url3], 7]
task_kwargs = {'time_to_wait': 1200}
task_additional_params = {'countdown': 30}

from celery import Celery

# assume this is from a celery.py in our Django project.
from project.celery import app

assert isinstance(app, Celery)


# This is the magic to get an instance of the Celery task from its name
task = app.tasks.get(task_name)
if task:
    task.apply_async(args=task_args, kwargs=task_kwargs, **task_additional_params)

    # The above has the same effect as:

    # index_urls.apply_async(
    #     args=[[url1, url2, url3], 7], 
    #     kwargs={'time_to_wait': 1200}, 
    #     countdown=30
    # )

Storage of Task State and Parameters

With the necessary pieces identified, the only step left is to record these properties for each invocation of the task somewhere. When the need to replay comes, we look up the task invocation’s properties and simply follow the steps above to reinvoke the task.

Of course, some additional information are needed in order to help us filter for the records we want. For example:

  • invocation time — when was the task invoked?
  • fail time — when did it fail?
  • error details — why did it fail?

Implementation in Our RAG System

The implementation I chose for our RAG system is that the ingestion tasks invocations are stored into the PostgreSQL DB in a dedicated model/table (vasync_request.models.ProcessTaskRequest). The task ID (a UUID) is a unique indexed field. In addition, the model has some of the properties identified above are fields.

I chose to use a TextField to store the stringified JSON value representing the task parameters. Since the parameters themselves should be opaque and treated as one atomic value, a single field/column should suffice. I won’t index the field because I don’t expect to search or order based on the field’s contents.

Additional Reading

Monitoring For Failed Tasks

See https://docs.celeryq.dev/en/stable/userguide/monitoring.html#real-time-processing for hooks into Celery’s notification of failed tasks.

Result Backends

See https://docs.celeryq.dev/en/stable/getting-started/first-steps-with-celery.html#keeping-results for information about Celery’s support for task results. My experience is that the documentation is sparse, and trial-and-error experiments are the best way to figure it out.