fix: _multiple_retrieve_thread not fail fast (#40466)

2026-10-28 03:35:04 +08:00 · 2026-08-11 05:40:23 +00:00
parent 48451f5adf
commit 4c7b870062
3 changed files with 157 additions and 1 deletions
@@ -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):
@@ -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 == []