UN-3266 [FEAT] Async Executor Backend for Prompt Studio #1849
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,105 @@ | ||
| """Lightweight Celery app for dispatching tasks to worker-v2 workers. | ||
|
|
||
| The Django backend already has a Celery app for internal tasks (beat, | ||
| periodic tasks, etc.) whose broker URL is set via CELERY_BROKER_URL. | ||
| Workers use the same broker. This module provides a second Celery app | ||
| instance that reuses the same broker URL (from Django settings) but | ||
| bypasses Celery's env-var-takes-priority behaviour so it can coexist | ||
| with the main Django Celery app in the same process. | ||
|
|
||
| Problem: Celery reads the ``CELERY_BROKER_URL`` environment variable | ||
| with highest priority — overriding constructor args, ``conf.update()``, | ||
| and ``config_from_object()``. | ||
|
|
||
| Solution: Subclass Celery and override ``connection_for_write`` / | ||
| ``connection_for_read`` so they always use our explicit broker URL, | ||
| bypassing the config resolution chain entirely. | ||
| """ | ||
|
|
||
| import logging | ||
| from urllib.parse import quote_plus | ||
|
|
||
| from celery import Celery | ||
| from django.conf import settings | ||
| from kombu import Queue | ||
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
| _worker_app: Celery | None = None | ||
|
|
||
|
|
||
| class _WorkerDispatchCelery(Celery): | ||
|
Contributor
There was a problem hiding this comment. Why celery here ? We already moved it from backend . What this methods do here?
Contributor
Author
There was a problem hiding this comment. @muhammad-ali-e The backend Celery worker handles fire-and-forget callback tasks that run after the executor worker finishes. Here's the flow: Backend dispatches task → Executor Worker (does the heavy lifting) Why these run on the backend (not the executor worker):
|
||
| """Celery subclass that forces an explicit broker URL. | ||
|
|
||
| Works around Celery's env-var-takes-priority behaviour where | ||
| ``CELERY_BROKER_URL`` always overrides per-app configuration. | ||
| The connection methods are the actual points where Celery opens | ||
| AMQP/Redis connections, so overriding them is both sufficient | ||
| and safe. | ||
| """ | ||
|
|
||
| _explicit_broker: str | None = None | ||
|
|
||
| def connection_for_write(self, url=None, *args, **kwargs): | ||
| return super().connection_for_write(url or self._explicit_broker, *args, **kwargs) | ||
|
|
||
| def connection_for_read(self, url=None, *args, **kwargs): | ||
| return super().connection_for_read(url or self._explicit_broker, *args, **kwargs) | ||
|
|
||
|
|
||
| def get_worker_celery_app() -> Celery: | ||
| """Get or create a Celery app for dispatching to worker-v2 workers. | ||
|
|
||
| The app uses: | ||
| - Same broker as the workers (built from CELERY_BROKER_BASE_URL, | ||
| CELERY_BROKER_USER, CELERY_BROKER_PASS via Django settings) | ||
| - Same PostgreSQL result backend as the Django Celery app | ||
|
|
||
| Returns: | ||
| Celery app configured for worker-v2 dispatch. | ||
| """ | ||
| global _worker_app | ||
| if _worker_app is not None: | ||
| return _worker_app | ||
|
Comment on lines
+61
to
+63
Contributor
There was a problem hiding this comment.
if _worker_app is not None:
return _worker_appUnder gunicorn with threaded workers (or any multi-threaded Django deployment), two threads can simultaneously see The idiomatic Python fix is to use a module-level lock: import threading
_worker_app: Celery | None = None
_worker_app_lock = threading.Lock()
def get_worker_celery_app() -> Celery:
global _worker_app
if _worker_app is not None:
return _worker_app
with _worker_app_lock:
if _worker_app is None: # re-check inside lock
...
_worker_app = app
return _worker_appPrompt To Fix With AIThis is a comment left during a code review.
Path: backend/backend/worker_celery.py
Line: 61-63
Comment:
**Unsynchronised singleton initialisation — race condition under concurrent requests**
`get_worker_celery_app()` uses the classic double-check-without-lock pattern:
```python
if _worker_app is not None:
return _worker_app
```
Under gunicorn with threaded workers (or any multi-threaded Django deployment), two threads can simultaneously see `_worker_app is None` and both proceed to create a new `_WorkerDispatchCelery` instance. The second assignment overwrites the first (last-writer-wins), so each thread may end up holding a reference to a *different* object than what ends up in the module global. This is benign in practice because both instances are configured identically, but it is wasteful and could cause subtle issues if Celery connection pools are per-instance.
The idiomatic Python fix is to use a module-level lock:
```python
import threading
_worker_app: Celery | None = None
_worker_app_lock = threading.Lock()
def get_worker_celery_app() -> Celery:
global _worker_app
if _worker_app is not None:
return _worker_app
with _worker_app_lock:
if _worker_app is None: # re-check inside lock
...
_worker_app = app
return _worker_app
```
How can I resolve this? If you propose a fix, please make it concise. |
||
|
|
||
| # Reuse the broker URL already built by Django settings (base.py) | ||
| # from CELERY_BROKER_BASE_URL + CELERY_BROKER_USER + CELERY_BROKER_PASS | ||
| broker_url = settings.CELERY_BROKER_URL | ||
|
|
||
| # Reuse the same PostgreSQL result backend as Django's Celery app | ||
| result_backend = ( | ||
| f"db+postgresql://{settings.DB_USER}:" | ||
| f"{quote_plus(settings.DB_PASSWORD)}" | ||
| f"@{settings.DB_HOST}:{settings.DB_PORT}/" | ||
| f"{settings.CELERY_BACKEND_DB_NAME}" | ||
|
Comment on lines
+70
to
+74
Contributor
There was a problem hiding this comment.
Prompt To Fix With AIThis is a comment left during a code review.
Path: backend/backend/worker_celery.py
Line: 70-74
Comment:
**`settings.DB_USER` is not URL-encoded in the result backend URL**
`quote_plus` is correctly applied to `DB_PASSWORD`, but `DB_USER` is interpolated raw. If the database username contains any URL-special characters (e.g. `@`, `:`, `/`), the resulting connection string would be malformed and the Celery result backend would fail to connect. Apply the same `quote_plus` encoding to `settings.DB_USER` for consistency and correctness, just as is done for `settings.DB_PASSWORD`.
How can I resolve this? If you propose a fix, please make it concise. |
||
| ) | ||
|
|
||
| app = _WorkerDispatchCelery( | ||
| "worker-dispatch", | ||
| set_as_current=False, | ||
| fixups=[], | ||
| ) | ||
| # Store the explicit broker URL for use in connection overrides | ||
| app._explicit_broker = broker_url | ||
|
|
||
| app.conf.update( | ||
| result_backend=result_backend, | ||
| task_queues=[Queue("executor")], | ||
| task_serializer="json", | ||
| accept_content=["json"], | ||
| result_serializer="json", | ||
| result_extended=True, | ||
| ) | ||
|
Comment on lines
+85
to
+92
Contributor
There was a problem hiding this comment. Configured queue name
The queue declared on the app ( Either align the queue name to Prompt To Fix With AIThis is a comment left during a code review.
Path: backend/backend/worker_celery.py
Line: 85-92
Comment:
**Configured queue name `"executor"` doesn't match the actual dispatch queue**
`get_worker_celery_app()` registers `task_queues=[Queue("executor")]`, but `ExecutionDispatcher._get_queue()` (in `sdk1/execution/dispatcher.py`) constructs the actual queue name as `celery_executor_{executor_name}` — for the legacy executor this becomes `"celery_executor_legacy"`.
The queue declared on the app (`"executor"`) never matches the queue used by `send_task`, so this `task_queues` setting has no practical effect. While `send_task` with an explicit `queue` parameter bypasses queue routing and the task is delivered correctly, the misconfigured `task_queues` setting means any queue-routing policies (e.g. prefetch limits, fair scheduling) configured on `"executor"` will not apply.
Either align the queue name to `"celery_executor_legacy"` (or the appropriate prefix), or remove the stale `task_queues` declaration from this app's config if it is intentionally unused.
How can I resolve this? If you propose a fix, please make it concise. |
||
|
|
||
| _worker_app = app | ||
| # Log broker host only (mask credentials) | ||
| safe_broker = broker_url.split("@")[-1] if "@" in broker_url else broker_url | ||
| safe_backend = ( | ||
| result_backend.split("@")[-1] if "@" in result_backend else result_backend | ||
| ) | ||
| logger.info( | ||
| "Created worker dispatch Celery app (broker=%s, result_backend=%s)", | ||
| safe_broker, | ||
| safe_backend, | ||
| ) | ||
| return _worker_app | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I don’t think this should be allowed when the organization is missing. Also how it works with an empty org_id?
cc: @vishnuszipstack