mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-16 00:14:31 +00:00
fix(worker): know you are in a worker from any thread, not only the task's
Celery records the executing task on the thread that runs it, so a thread that
task starts sees none. The embeddings client and read_document both decided
"am I in a worker?" from that alone, and from any other thread took the
web-process branch: dispatch to the worker they were running in and block on
the result. Celery refuses that get() ("Never call result.get() within a
task!"), so the embed failed and latched the 30s dispatch cooldown for every
caller after it; with joins allowed, read_document would instead wait on a
parsing queue only its own busy process serves.
Threads inside tasks are not hypothetical: per-source retrieval fans out to a
pool, so a scheduled or webhook agent searching several sources embedded from
pool threads. Graph extraction did too, which failed every chunk of a build.
in_worker() in celery_init answers for the whole process. Celery's
task_join_will_block is process-wide and set for every blocking pool (prefork,
solo, threads) -- exactly the condition under which dispatch-and-wait goes
wrong; eventlet/gevent leave it unset, so the task's own thread still counts
through current_worker_task. Verified with real workers on each blocking pool:
from a thread a task started, the old check dispatched and hit the error, the
new one embedded locally.
This commit is contained in:
1 parent
5e67f8c927
commit
15bdda8554
6 files changed
+117
-16
No files matched your search
@@ -276,3 +276,55 @@ class TestReclaimIsSkippedForEmbeds:
|
||||
from docsgpt.vectorstore.embeddings_delegated import EMBED_TASK
|
||||
|
||||
assert EMBED_TASK in _NO_RECLAIM_TASKS
|
||||
|
||||
|
||||
@pytest.mark.unit
|
||||
class TestInWorker:
|
||||
"""Whether code runs inside a worker must not depend on which thread asks.
|
||||
|
||||
Celery records the executing task on the thread that runs it, so a thread
|
||||
that task starts sees no task at all. Code deciding "am I in the worker?"
|
||||
from that alone takes the web-process branch there: it dispatches to the
|
||||
worker it is running in and blocks on the result, which Celery refuses
|
||||
("Never call result.get() within a task!") or, where joins are allowed,
|
||||
waits on a queue only this busy process serves.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _ask_from_a_new_thread():
|
||||
import threading
|
||||
|
||||
from docsgpt.celery_init import in_worker
|
||||
|
||||
seen = []
|
||||
thread = threading.Thread(target=lambda: seen.append(in_worker()))
|
||||
thread.start()
|
||||
thread.join()
|
||||
return seen[0]
|
||||
|
||||
def test_false_outside_a_worker(self):
|
||||
from docsgpt.celery_init import in_worker
|
||||
|
||||
assert in_worker() is False
|
||||
assert self._ask_from_a_new_thread() is False
|
||||
|
||||
def test_true_on_a_thread_started_inside_a_worker(self):
|
||||
# Blocking pools (prefork, solo, threads) mark the whole process as one
|
||||
# where joining a task would block; ``denied_join_result`` sets exactly
|
||||
# that flag.
|
||||
from celery.result import denied_join_result
|
||||
|
||||
with denied_join_result():
|
||||
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.
|
||||
from unittest.mock import PropertyMock
|
||||
|
||||
from docsgpt.celery_init import celery, in_worker
|
||||
|
||||
with patch.object(
|
||||
type(celery), "current_worker_task", new_callable=PropertyMock, return_value=object()
|
||||
):
|
||||
assert in_worker() is True
|
||||
Reference in new issue
Block a user