mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 16:14:16 +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
@@ -72,6 +72,32 @@ class TestInsideAWorker:
|
||||
assert vector == [1.0, 2.0]
|
||||
celery.send_task.assert_not_called()
|
||||
|
||||
def test_a_thread_started_inside_the_worker_embeds_locally(self):
|
||||
"""The task's own thread is not the only one in a worker.
|
||||
|
||||
Graph extraction and per-source retrieval both fan out to thread pools
|
||||
inside tasks. The check used to read the task off the current thread
|
||||
only, so from those threads it dispatched to the worker it was running
|
||||
in -- and Celery refuses that ``get()`` inside a worker, failing the
|
||||
call and latching the 30s dispatch cooldown for every caller after it.
|
||||
"""
|
||||
from celery.result import denied_join_result
|
||||
|
||||
from docsgpt.celery_init import celery
|
||||
|
||||
local = MagicMock()
|
||||
local.embed_documents.return_value = [[1.0, 2.0]]
|
||||
client = DelegatedEmbeddings("some/model")
|
||||
vectors = []
|
||||
with denied_join_result():
|
||||
with patch("docsgpt.vectorstore.base.build_local_embeddings", return_value=local):
|
||||
with patch.object(celery, "send_task") as send_task:
|
||||
thread = threading.Thread(target=lambda: vectors.append(client.embed_query("hi")))
|
||||
thread.start()
|
||||
thread.join()
|
||||
assert vectors == [[1.0, 2.0]]
|
||||
send_task.assert_not_called()
|
||||
|
||||
def test_the_local_model_is_built_once(self):
|
||||
local = MagicMock()
|
||||
local.embed_documents.return_value = [[1.0]]
|
||||
|
||||
Reference in new issue
Block a user