mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-03 11:11:58 +00:00
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:
1 parent
a20402a6dc
commit
877609dfb6
2 files changed
+56
-11
No files matched your search
+34
-11
@@ -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``.
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
Reference in new issue
Block a user