From 4c7b8700623ed9c3d3018b64482157f9d640352e Mon Sep 17 00:00:00 2001 From: wangxiaolei Date: Tue, 11 Aug 2026 13:40:23 +0800 Subject: [PATCH] fix: _multiple_retrieve_thread not fail fast (#40466) --- api/core/rag/retrieval/dataset_retrieval.py | 24 ++++- .../rag/retrieval/test_dataset_retrieval.py | 90 +++++++++++++++++++ .../extensions/otel/test_retrieval_tracing.py | 44 +++++++++ 3 files changed, 157 insertions(+), 1 deletion(-) diff --git a/api/core/rag/retrieval/dataset_retrieval.py b/api/core/rag/retrieval/dataset_retrieval.py index 5c1f2b6e6d9..784ac5fb5bb 100644 --- a/api/core/rag/retrieval/dataset_retrieval.py +++ b/api/core/rag/retrieval/dataset_retrieval.py @@ -9,6 +9,7 @@ from collections.abc import Generator, Mapping from typing import Any, Union, cast from flask import Flask, current_app +from opentelemetry.trace import get_current_span from sqlalchemy import and_, func, literal, or_, select, update from sqlalchemy.orm import Session, sessionmaker @@ -1204,8 +1205,9 @@ class DatasetRetrieval: attachment_ids: list[str] | None, cancel_event: threading.Event | None, thread_exceptions: list[Exception] | None, + skip_on_error: bool = False, ) -> None: - """Collect errors only after they pass through the traced retrieval method.""" + """Collect errors after tracing, or skip dataset-level failures when requested.""" try: self._run_retriever_thread( flask_app=flask_app, @@ -1218,6 +1220,25 @@ class DatasetRetrieval: attachment_ids=attachment_ids, ) except Exception as exc: + if skip_on_error: + logger.warning( + "Skipping dataset retrieval because retriever failed, dataset_id=%s, error_type=%s, error=%s", + dataset_id, + type(exc).__name__, + str(exc), + ) + span = get_current_span() + if span and span.is_recording(): + span.add_event( + "dataset_retrieval.skipped", + attributes={ + "dataset_id": dataset_id, + "error.type": type(exc).__name__, + "error.message": str(exc), + }, + ) + return + if cancel_event: cancel_event.set() if thread_exceptions is not None: @@ -1878,6 +1899,7 @@ class DatasetRetrieval: "attachment_ids": [attachment_id] if attachment_id else None, "cancel_event": cancel_event, "thread_exceptions": retrieval_thread_exceptions, + "skip_on_error": True, }, ) threads.append(retrieval_thread) diff --git a/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval.py b/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval.py index f9922643e96..808cc9bec18 100644 --- a/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval.py +++ b/api/tests/unit_tests/core/rag/retrieval/test_dataset_retrieval.py @@ -3894,6 +3894,96 @@ class TestKnowledgeRetrievalRegression: assert cancel_event.is_set() assert thread_exceptions == [expected_error] + def test_run_retriever_thread_safely_skips_failed_dataset_when_requested(self, caplog): + dataset_retrieval = DatasetRetrieval() + all_documents: list[Document] = [] + cancel_event = threading.Event() + thread_exceptions: list[Exception] = [] + expected_error = RuntimeError("retrieval failed") + + with _patched_retriever_session(): + with patch.object(dataset_retrieval, "_retriever", side_effect=expected_error): + dataset_retrieval._run_retriever_thread_safely( + flask_app=_FakeFlaskApp(), + dataset_id="dataset-1", + query="test query", + top_k=3, + all_documents=all_documents, + document_ids_filter=None, + metadata_condition=None, + attachment_ids=None, + cancel_event=cancel_event, + thread_exceptions=thread_exceptions, + skip_on_error=True, + ) + + assert not cancel_event.is_set() + assert thread_exceptions == [] + assert "dataset_id=dataset-1" in caplog.text + assert "Skipping dataset retrieval because retriever failed" in caplog.text + + def test_multiple_retrieve_thread_skips_failed_dataset(self, mock_dataset, caplog): + dataset_retrieval = DatasetRetrieval() + flask_app = Flask(__name__) + successful_dataset = Dataset( + id=str(uuid4()), + provider="dify", + indexing_technique="high_quality", + ) + document = Document( + page_content="successful doc", + metadata={ + "doc_id": "doc1", + "score": 0.95, + "document_id": str(uuid4()), + "dataset_id": successful_dataset.id, + }, + provider="dify", + ) + + def fake_retriever( + flask_app, + session, + dataset_id, + query, + top_k, + all_documents, + document_ids_filter, + metadata_condition, + attachment_ids, + ): + if dataset_id == mock_dataset.id: + raise RuntimeError("dataset unavailable") + all_documents.append(document) + + all_documents: list[Document] = [] + + with ( + patch.object(dataset_retrieval, "_retriever", side_effect=fake_retriever), + _patched_retriever_session(), + ): + dataset_retrieval._multiple_retrieve_thread( + flask_app=flask_app, + available_datasets=[mock_dataset, successful_dataset], + metadata_condition=None, + metadata_filter_document_ids=None, + all_documents=all_documents, + tenant_id=str(uuid4()), + reranking_enable=False, + reranking_mode="reranking_model", + reranking_model=None, + weights=None, + top_k=3, + score_threshold=0.0, + query="test query", + attachment_id=None, + dataset_count=2, + ) + + assert all_documents == [document] + assert f"dataset_id={mock_dataset.id}" in caplog.text + assert "Skipping dataset retrieval because retriever failed" in caplog.text + class _FakeFlaskApp: def app_context(self): diff --git a/api/tests/unit_tests/extensions/otel/test_retrieval_tracing.py b/api/tests/unit_tests/extensions/otel/test_retrieval_tracing.py index 79273d07019..824028316fe 100644 --- a/api/tests/unit_tests/extensions/otel/test_retrieval_tracing.py +++ b/api/tests/unit_tests/extensions/otel/test_retrieval_tracing.py @@ -121,3 +121,47 @@ def test_retriever_thread_exception_sets_error_span_and_is_collected( assert retrieval_span.status.status_code == StatusCode.ERROR assert cancel_event.is_set() assert thread_exceptions == [expected_error] + + +def test_retriever_thread_exception_emits_skip_event_when_requested( + app, + memory_span_exporter, + tracer_provider_with_memory_exporter, +) -> None: + retrieval = DatasetRetrieval() + cancel_event = threading.Event() + thread_exceptions: list[Exception] = [] + expected_error = RuntimeError("retrieval failed") + dataset_id = str(uuid4()) + + with ( + patch("extensions.otel.decorators.base.dify_config.ENABLE_OTEL", True), + patch("core.rag.retrieval.dataset_retrieval.session_factory.create_session"), + patch.object(retrieval, "_retriever", side_effect=expected_error), + get_tracer(__name__).start_as_current_span("dataset-retrieval-parent") as parent_span, + ): + retrieval._run_retriever_thread_safely( + flask_app=app, + dataset_id=dataset_id, + query="test query", + top_k=4, + all_documents=[], + document_ids_filter=None, + metadata_condition=None, + attachment_ids=None, + cancel_event=cancel_event, + thread_exceptions=thread_exceptions, + skip_on_error=True, + ) + + retrieval_span = next( + span + for span in memory_span_exporter.get_finished_spans() + if span.name.endswith("DatasetRetrieval._run_retriever_thread") + ) + skip_event = next(event for event in parent_span.events if event.name == "dataset_retrieval.skipped") + assert retrieval_span.status.status_code == StatusCode.ERROR + assert skip_event.attributes["dataset_id"] == dataset_id + assert skip_event.attributes["error.message"] == "retrieval failed" + assert not cancel_event.is_set() + assert thread_exceptions == []