fix(worker): record worker state at startup, for every pool

in_worker() read task_join_will_block, which eventlet and gevent leave unset,
and the task's current_worker_task, which they scope to one greenlet -- so a
greenlet a task spawned in those pools still took the web-process branch and
dispatched to its own worker, where it could queue behind its parent and time
out.

The worker's own startup now records it: worker_init fires in every worker's
main process before the pool starts (where solo, threads, eventlet and gevent
run tasks, and what prefork children fork from), and worker_process_init in
each prefork child. worker_ready would be too late -- prefork children are
forked before it fires. The two existing checks stay for anything that runs
tasks without that startup.
This commit is contained in:
Alex committed 2026-09-19 14:33:14 +01:00
1 parent a20402a6dc
commit 877609dfb6
2 files changed
+56 -11

No files matched your search

+34 -11
View File
@@ -13,6 +13,7 @@ from celery.signals import (
setup_logging,
task_postrun,
task_prerun,
worker_init,
worker_process_init,
worker_ready,
)
@@ -173,27 +174,49 @@ celery = make_celery()
celery.config_from_object("docsgpt.celeryconfig")
#: Set once this process starts as a worker; see :func:`_mark_worker_process`.
_IS_WORKER_PROCESS = False
@worker_init.connect
@worker_process_init.connect
def _mark_worker_process(*args, **kwargs):
"""Record that this process runs tasks, for :func:`in_worker`.
``worker_init`` fires in every worker's main process before its pool
starts: that is where solo, threads, eventlet and gevent run tasks, and
what prefork children fork from. ``worker_process_init`` covers prefork
children however they were started.
"""
global _IS_WORKER_PROCESS
_IS_WORKER_PROCESS = True
def in_worker() -> bool:
"""True anywhere in a Celery worker process, on any thread.
"""True anywhere in a Celery worker process, on any thread or greenlet.
``current_worker_task`` alone is not enough: Celery records the executing
task on the thread that runs it, so a thread the task starts sees none and
would take the web-process branch — dispatching to the worker it is running
in and blocking on the result. Celery refuses that ``get()`` ("Never call
result.get() within a task!"), or, where joins are allowed, it waits on a
queue only this busy process serves.
task on the thread (or greenlet) that runs it, so one the task starts sees
none and would take the web-process branch — dispatching to the worker it
is running in and blocking on the result. Celery refuses that ``get()``
("Never call result.get() within a task!"), or, where joins are allowed,
it waits on a queue that only this busy process may be able to serve.
``task_join_will_block`` is process-wide and set for every blocking pool
(prefork, solo, threads) — exactly the condition under which dispatching
and waiting goes wrong. eventlet/gevent leave it unset, so the task's own
thread still counts through ``current_worker_task``.
The worker's own startup (:func:`_mark_worker_process`) answers for every
pool. ``task_join_will_block`` — process-wide, set for every blocking pool
— and the task's own ``current_worker_task`` still count for a process
that runs tasks without having gone through that startup.
Returns:
bool: Whether this call is running inside a worker process.
"""
from celery.result import task_join_will_block
return task_join_will_block() or celery.current_worker_task is not None
return (
_IS_WORKER_PROCESS
or task_join_will_block()
or celery.current_worker_task is not None
)
#: Task-name prefix the package carried before the rename to ``docsgpt``.
+22
View File
@@ -317,6 +317,28 @@ class TestInWorker:
with denied_join_result():
assert self._ask_from_a_new_thread() is True
def test_true_on_any_thread_of_a_non_blocking_pool_worker(self, monkeypatch):
# eventlet/gevent leave the join flag unset and scope the current task
# to one greenlet, so only the worker's own startup can say this
# process is a worker. The lifecycle signal records that for every
# thread and greenlet in it.
import docsgpt.celery_init as celery_init
monkeypatch.setattr(celery_init, "_IS_WORKER_PROCESS", False)
assert self._ask_from_a_new_thread() is False
celery_init.worker_init.send(sender=None)
assert self._ask_from_a_new_thread() is True
def test_prefork_children_record_it_on_their_own_start(self, monkeypatch):
import docsgpt.celery_init as celery_init
monkeypatch.setattr(celery_init, "_IS_WORKER_PROCESS", False)
celery_init.worker_process_init.send(sender=None)
assert self._ask_from_a_new_thread() is True
def test_true_on_the_task_thread_of_a_non_blocking_pool(self):
# eventlet/gevent pools leave the process flag unset; the thread
# running the task still knows it is in one.