zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-24 17:10:12 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(task_executor): fix Langfuse flush/shutdown deadlock that freezes document parsing (#16502)

Browse Source
This commit is contained in:
Öndery
2026-07-07 14:06:30 +03:00
committed by GitHub
parent 6cd03d7a70
commit 28a41ed070
10 changed files with 519 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
api/apps/restful_apis/document_api.py
View File

@@ -14,6 +14,7 @@
# limitations under the License.
#
from io import BytesIO
from datetime import datetime
import logging
import json
import os
@@ -1473,6 +1474,11 @@ def _run_sync(user_id: str, req):
has_unfinished_task = any((task.progress or 0) < 1 for task in tasks)
if str(doc.run) in [TaskStatus.RUNNING.value, TaskStatus.CANCEL.value] or has_unfinished_task:
cancel_all_task_of(doc_id)
# Append a "stopped by user" marker so the history is preserved and
# the document no longer looks like it is still waiting in the queue.
cancel_doc_msg = f"\n{datetime.now().strftime('%H:%M:%S')} Task stopped by user."
info["progress_msg"] = (doc.progress_msg or "") + cancel_doc_msg
logging.debug("Appended cancellation marker to progress_msg on cancel for doc %s", doc_id)
else:
return RetCode.DATA_ERROR, "Cannot cancel a task that is not in RUNNING status"
if all([rerun_with_delete, str(doc.run) == TaskStatus.DONE.value]):
@@ -1698,14 +1704,17 @@ async def stop_parse_documents(tenant_id, dataset_id):
continue
cancel_all_task_of(doc_id)
cancel_doc_msg = f"\n{datetime.now().strftime('%H:%M:%S')} Task stopped by user."
DocumentService.update_by_id(
doc_id,
{
"run": str(TaskStatus.CANCEL.value),
"progress": 0,
"chunk_num": 0,
"progress_msg": (doc.progress_msg or "") + cancel_doc_msg,
},
)
logging.debug("Appended cancellation marker to progress_msg on stop-parse for doc %s", doc_id)
index_name = search.index_name(tenant_id)
if settings.docStoreConn.index_exist(index_name, doc.kb_id):
settings.docStoreConn.delete({"doc_id": doc.id}, index_name, doc.kb_id)

7
api/apps/restful_apis/task_api.py
View File

@@ -89,7 +89,12 @@ async def _cancel_task(task_id):
if doc_id and doc_id not in (CANVAS_DEBUG_DOC_ID, GRAPH_RAPTOR_FAKE_DOC_ID):
_, doc = DocumentService.get_by_id(doc_id)
if doc and str(doc.run) in (TaskStatus.RUNNING.value, TaskStatus.SCHEDULE.value):
DocumentService.update_by_id(doc_id, {"run": TaskStatus.CANCEL.value, "progress": 0})
cancel_doc_msg = f"\n{datetime.now().strftime('%H:%M:%S')} Task stopped by user."
DocumentService.update_by_id(
doc_id,
{"run": TaskStatus.CANCEL.value, "progress": 0, "progress_msg": (doc.progress_msg or "") + cancel_doc_msg},
)
logging.debug("Appended cancellation marker to progress_msg on task cancel: task_id=%s doc_id=%s", task_id, doc_id)
except Exception as e:
logging.warning("Failed to update document run status for task %s: %s", task_id, str(e))

2
api/db/services/canvas_service.py
View File

@@ -337,7 +337,7 @@ async def completion(tenant_id, agent_id, session_id=None, **kwargs):
"files": files,
"user_id": user_id,
"inputs": inputs,
# Used by Canvas.run to correlate RAGFlow's Langfuse generations by session.
# Forwarded to upstream LLM providers as the `user` field for session correlation.
"session_id": session_id,
}
if chat_template_kwargs is not None:

91
api/db/services/document_service.py
View File

@@ -16,6 +16,7 @@
import logging
import random
from datetime import datetime
from time import monotonic
import xxhash
from peewee import fn, Case, JOIN
@@ -954,6 +955,10 @@ class DocumentService(CommonService):
doc_progress = doc.progress if doc and doc.progress else 0.0
special_task_running = False
priority = 0
# Count this document's own not-yet-started tasks per priority so
# they can be excluded from the "tasks ahead in the queue" figure
# for the matching priority queue.
own_queued_by_priority = {}
for t in tsks:
task_type = (t.task_type or "").lower()
if task_type in PIPELINE_SPECIAL_PROGRESS_FREEZE_TASK_TYPES:
@@ -962,6 +967,8 @@ class DocumentService(CommonService):
finished = False
if t.progress == -1:
bad += 1
if (t.progress or 0) == 0:
own_queued_by_priority[t.priority] = own_queued_by_priority.get(t.priority, 0) + 1
prg += t.progress if t.progress >= 0 else 0
if (t.progress_msg or "").strip():
msg.append(t.progress_msg)
@@ -991,9 +998,14 @@ class DocumentService(CommonService):
if msg:
info["progress_msg"] = msg
if msg.endswith("created task graphrag") or msg.endswith("created task raptor") or msg.endswith("created task mindmap"):
info["progress_msg"] += "\n%d tasks are ahead in the queue..." % get_queue_length(priority)
# Exclude this document's own queued tasks in the same
# priority queue: they are not "ahead" of itself, they
# ARE the work being waited on.
queue_ahead = max(0, get_queue_length(priority) - own_queued_by_priority.get(priority, 0))
info["progress_msg"] += "\n%d tasks are ahead in the queue..." % queue_ahead
else:
info["progress_msg"] = "%d tasks are ahead in the queue..." % get_queue_length(priority)
queue_ahead = max(0, get_queue_length(priority) - own_queued_by_priority.get(priority, 0))
info["progress_msg"] = "%d tasks are ahead in the queue..." % queue_ahead
info["update_time"] = current_timestamp()
info["update_date"] = get_format_time()
(cls.model.update(info).where((cls.model.id == d["id"]) & ((cls.model.run.is_null(True)) | (cls.model.run != TaskStatus.CANCEL.value))).execute())
@@ -1165,8 +1177,79 @@ def queue_per_doc_raptor_task(doc, priority):
return task["id"]
# Short-lived per-priority cache for the genuine queued-task backlog so the
# per-document progress sync does not issue a COUNT query for every document
# each cycle. Keyed by priority (None means "all priorities").
_PENDING_TASK_COUNT_CACHE = {}
_PENDING_TASK_COUNT_TTL_SECONDS = 3.0
def get_pending_task_count(priority=None):
"""Count tasks that are genuinely still waiting to be processed.
A task counts as "waiting" when it has not started yet (progress == 0) and
its document is neither cancelled nor failed. We deliberately do NOT require
the document to be RUNNING with progress in [0, 1): special tasks (graphrag/
raptor/mindmap) are queued via ``begin2parse(keep_progress=True)`` while the
document's own progress may already be 1, so requiring RUNNING/progress<1
would undercount them and wrongly drop the cap to 0 while Redis lag is still
non-zero. Only cancelled documents (run == CANCEL) and failed ones
(progress < 0) are excluded, plus soft-deleted (invalid) documents.
When ``priority`` is given, only tasks queued at that priority are counted,
so the figure stays consistent with the per-priority Redis queue it caps.
Returns None when the count cannot be determined, so callers can fall back
to the raw Redis stream lag.
"""
now = monotonic()
cached = _PENDING_TASK_COUNT_CACHE.get(priority)
if cached and cached.get("expire_at", 0.0) > now:
return cached["value"]
try:
query = (
Task.select(fn.COUNT(Task.id))
.join(Document, on=(Task.doc_id == Document.id))
.where(
(Task.progress == 0)
& ((Document.run.is_null(True)) | (Document.run != TaskStatus.CANCEL.value))
& (Document.progress >= 0)
& (Document.status == StatusEnum.VALID.value)
)
)
if priority is not None:
query = query.where(Task.priority == priority)
count = int(query.scalar() or 0)
except Exception:
logging.exception("get_pending_task_count failed")
return None
_PENDING_TASK_COUNT_CACHE[priority] = {"value": count, "expire_at": now + _PENDING_TASK_COUNT_TTL_SECONDS}
return count
def get_queue_length(priority, suffix="common"):
"""Return how many tasks are ahead in the processing queue.
The Redis stream consumer-group ``lag`` counts every message that has not
yet been delivered to a task executor, including messages whose tasks were
already cancelled/stopped. Those messages only stop counting once an
executor happens to read them, so after a user stops parsing the lag can
stay inflated indefinitely and parsing appears to hang forever
("N tasks are ahead in the queue...").
To keep the figure honest, the raw lag is capped by the number of tasks
that are genuinely still waiting in the database, which self-heals the
moment work is cancelled or completes.
"""
group_info = REDIS_CONN.queue_info(settings.get_svr_queue_name(priority, suffix), SVR_CONSUMER_GROUP_NAME)
if not group_info:
lag = int(group_info.get("lag", 0) or 0) if group_info else 0
# Nothing queued in Redis: the answer is 0 regardless of the DB backlog, so
# short-circuit to avoid a COUNT/JOIN on every progress-sync cycle.
if lag <= 0:
return 0
return int(group_info.get("lag", 0) or 0)
pending = get_pending_task_count(priority)
if pending is None:
return lag
return min(lag, pending)

43
api/db/services/tenant_llm_service.py
View File

@@ -534,27 +534,34 @@ class LLM4Tenant:
def close(self):
"""Release resources held by this LLM4Tenant instance.
This method should be called when the instance is no longer needed
to properly release resources such as:
- Langfuse tracing client (flush and shutdown)
- Underlying model instance resources (HTTP sessions, etc.)
IMPORTANT: do NOT call ``langfuse.flush()`` or ``langfuse.shutdown()``
here. ``close()`` runs once per task, synchronously, on the asyncio
event-loop thread of the task executor. Two problems follow:
- ``flush()`` blocks on an unbounded ``queue.join()`` in the underlying
OpenTelemetry span processor. If the exporter cannot drain (slow or
unreachable Langfuse, or an already-shutdown processor) it never
returns.
- ``shutdown()`` permanently tears down the process-wide Langfuse /
OpenTelemetry tracer provider that every ``LLMBundle`` shares. After
the first task shuts it down, every subsequent ``flush()`` blocks
forever.
Because this runs on the event loop, a single stuck ``flush()`` freezes
the entire task executor: all in-flight parse tasks stop making
progress and no new tasks are ever picked up (observed as document
parsing being stuck with every executor thread parked on a lock).
Langfuse already exports spans from its own background processor and
flushes at process exit, so releasing the reference is sufficient here.
"""
# Flush and shutdown Langfuse client if it was initialized
if self.langfuse:
try:
self.langfuse.flush()
if hasattr(self.langfuse, "shutdown"):
self.langfuse.shutdown()
except Exception:
# Ignore errors during cleanup
pass
finally:
self.langfuse = None
# Drop the Langfuse reference WITHOUT flushing/shutting down the shared
# client (see the docstring above for why this would deadlock).
self.langfuse = None
# Release underlying model instance if it has a close method
if self.mdl and hasattr(self.mdl, "close") and callable(getattr(self.mdl, "close")):
if self.mdl and callable(getattr(self.mdl, "close", None)):
try:
self.mdl.close()
except Exception:
# Ignore errors during cleanup
pass
logging.warning("LLM4Tenant.close: error while closing model instance", exc_info=True)