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:
Alex committed 2026-09-19 14:07:42 +01:00
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]]