mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 00:16:46 +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
@@ -17,8 +17,6 @@ import signal
|
||||
import threading
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from celery import current_task
|
||||
|
||||
from docsgpt.agents.tools.artifact_ref import resolve_artifact_id
|
||||
from docsgpt.agents.tools.attachment_bridge import (
|
||||
AttachmentBridgeError,
|
||||
@@ -26,6 +24,7 @@ from docsgpt.agents.tools.attachment_bridge import (
|
||||
match_attachment,
|
||||
)
|
||||
from docsgpt.agents.tools.base import Tool
|
||||
from docsgpt.celery_init import in_worker
|
||||
from docsgpt.core.json_schema_utils import (
|
||||
JsonSchemaValidationError,
|
||||
normalize_json_schema_payload,
|
||||
@@ -229,9 +228,9 @@ class ReadDocumentTool(Tool):
|
||||
# (floored at DOCUMENT_PARSE_TIMEOUT).
|
||||
timeout = parse_timeout_for_size(self._input_size)
|
||||
|
||||
# ``current_task`` is a Celery proxy: truthy only while this runs inside a worker task,
|
||||
# falsy in the web process (the bare proxy is NOT identity-None, so test truthiness).
|
||||
if current_task:
|
||||
# Process-wide, not the thread-local ``current_task``: a thread a task starts has no
|
||||
# task of its own, and dispatching from there is the self-deadlock described above.
|
||||
if in_worker():
|
||||
from docsgpt.worker import run_parse_document
|
||||
|
||||
try:
|
||||
|
||||
@@ -172,6 +172,29 @@ def _run_version_check(*args, **kwargs):
|
||||
celery = make_celery()
|
||||
celery.config_from_object("docsgpt.celeryconfig")
|
||||
|
||||
|
||||
def in_worker() -> bool:
|
||||
"""True anywhere in a Celery worker process, on any thread.
|
||||
|
||||
``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_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``.
|
||||
|
||||
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
|
||||
|
||||
#: Task-name prefix the package carried before the rename to ``docsgpt``.
|
||||
|
||||
|
||||
|
||||
@@ -10,8 +10,9 @@ Celery and the vector comes back. The API pays a broker round trip per query
|
||||
and no resident model.
|
||||
|
||||
Inside a worker there is nothing to delegate to -- dispatching would queue work
|
||||
behind the task already running and wait on itself -- so a call made while a
|
||||
task is executing runs locally, on a model this process loads once and caches.
|
||||
behind the task already running and wait on itself -- so a call made anywhere in
|
||||
a worker process, including from a thread a task started, runs locally, on a
|
||||
model this process loads once and caches.
|
||||
``DOCUMENT_PARSE_QUEUE`` exists for the same reason on the parsing side.
|
||||
|
||||
Production deployments should point ``EMBEDDINGS_BASE_URL`` at a real embedding
|
||||
@@ -79,11 +80,11 @@ def _forget(result) -> None:
|
||||
|
||||
|
||||
def _in_worker() -> bool:
|
||||
"""True when a Celery task is executing in this process."""
|
||||
"""True anywhere in a Celery worker process -- on any thread, not only the task's."""
|
||||
try:
|
||||
from docsgpt.celery_init import celery
|
||||
from docsgpt.celery_init import in_worker
|
||||
|
||||
return celery.current_worker_task is not None
|
||||
return in_worker()
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
Reference in new issue
Block a user