From 1d141aff1880aeb30239807ea8e71f5ebfdcdf8b Mon Sep 17 00:00:00 2001 From: buua436 Date: Mon, 3 Aug 2026 19:15:44 +0800 Subject: [PATCH] fix: optimize dataflow indexing and logs (#17737) --- .../pipeline_operation_log_service.py | 19 ++++++- common/settings.py | 4 +- rag/flow/pipeline.py | 5 ++ .../task_executor_refactor/chunk_service.py | 34 ++++++++----- rag/utils/es_conn.py | 4 +- rag/utils/infinity_conn.py | 2 +- rag/utils/ob_conn.py | 2 +- rag/utils/opensearch_conn.py | 4 +- .../test_pipeline_operation_log_service.py | 49 +++++++++++++++++++ 9 files changed, 103 insertions(+), 20 deletions(-) create mode 100644 test/unit_test/api/db/services/test_pipeline_operation_log_service.py diff --git a/api/db/services/pipeline_operation_log_service.py b/api/db/services/pipeline_operation_log_service.py index 4b700827b9..00806df5ce 100644 --- a/api/db/services/pipeline_operation_log_service.py +++ b/api/db/services/pipeline_operation_log_service.py @@ -16,6 +16,7 @@ import json import logging import os +import re from datetime import datetime, timedelta from peewee import fn @@ -50,6 +51,22 @@ _PIPELINE_TASK_TYPE_TO_FINISH_FIELD = { PipelineTaskType.STRUCTURE: "structure_task_finish_at", } +_EMBEDDING_VECTOR_FIELD = re.compile(r"^q_\d+_vec$") + + +def _remove_embedding_vectors(value): + """Remove index-only embedding vectors from a runtime pipeline snapshot.""" + if isinstance(value, dict): + for key in list(value): + if _EMBEDDING_VECTOR_FIELD.fullmatch(str(key)): + del value[key] + else: + _remove_embedding_vectors(value[key]) + elif isinstance(value, list): + for item in value: + _remove_embedding_vectors(item) + return value + class PipelineOperationLogService(CommonService): model = PipelineOperationLog @@ -205,7 +222,7 @@ class PipelineOperationLogService(CommonService): progress_msg=progress_msg, process_begin_at=process_begin_at, process_duration=process_duration, - dsl=json.loads(dsl), + dsl=_remove_embedding_vectors(json.loads(dsl)), task_type=task_type, operation_status=operation_status, avatar=avatar, diff --git a/common/settings.py b/common/settings.py index bf540dbc91..0bccddfceb 100644 --- a/common/settings.py +++ b/common/settings.py @@ -125,7 +125,7 @@ OS = {} GCS = {} DOC_MAXIMUM_SIZE: int = 128 * 1024 * 1024 -DOC_BULK_SIZE: int = 4 +DOC_BULK_SIZE: int = 32 EMBEDDING_BATCH_SIZE: int = 16 PARALLEL_DEVICES: int = 0 @@ -395,7 +395,7 @@ def init_settings(): global DOC_MAXIMUM_SIZE, DOC_BULK_SIZE, EMBEDDING_BATCH_SIZE DOC_MAXIMUM_SIZE = int(os.environ.get("MAX_CONTENT_LENGTH", 128 * 1024 * 1024)) - DOC_BULK_SIZE = int(os.environ.get("DOC_BULK_SIZE", 4)) + DOC_BULK_SIZE = int(os.environ.get("DOC_BULK_SIZE", 32)) EMBEDDING_BATCH_SIZE = int(os.environ.get("EMBEDDING_BATCH_SIZE", 16)) os.environ["DOTNET_SYSTEM_GLOBALIZATION_INVARIANT"] = "1" diff --git a/rag/flow/pipeline.py b/rag/flow/pipeline.py index d5efa7b427..3829a760ed 100644 --- a/rag/flow/pipeline.py +++ b/rag/flow/pipeline.py @@ -49,6 +49,11 @@ class Pipeline(Graph): if has_canceled(self.task_id): progress = -1 message += "[CANCEL]" + # Progress-only callbacks are used for fine-grained updates during + # tokenization and embedding. Do not persist an empty message as a log + # entry; otherwise the task log is filled with timestamp-only lines. + if not str(message or "").strip(): + return try: bin = REDIS_CONN.get(log_key) obj = json.loads(bin.encode("utf-8")) if bin else [] diff --git a/rag/svr/task_executor_refactor/chunk_service.py b/rag/svr/task_executor_refactor/chunk_service.py index 449b846c05..67c1679bec 100644 --- a/rag/svr/task_executor_refactor/chunk_service.py +++ b/rag/svr/task_executor_refactor/chunk_service.py @@ -314,7 +314,7 @@ class ChunkService: ) -> bool: """Insert mother chunks in batches.""" for b in range(0, len(mothers), doc_bulk_size): - await self._intercept_doc_store_insert(mothers[b : b + doc_bulk_size], search.index_name(task_tenant_id), task_dataset_id) + await self._intercept_doc_store_insert(mothers[b : b + doc_bulk_size], search.index_name(task_tenant_id), task_dataset_id, refresh=False) if self._task_context.has_canceled_func(task_id): self._task_context.progress_cb(-1, msg="Task has been canceled.") @@ -328,13 +328,13 @@ class ChunkService: else: return await thread_pool_exec(settings.docStoreConn.delete, condition, index_name, task_dataset_id) - async def _intercept_doc_store_insert(self, chunks: list, index_name: str, task_dataset_id: str) -> Any: + async def _intercept_doc_store_insert(self, chunks: list, index_name: str, task_dataset_id: str, refresh: str | bool = "wait_for") -> Any: if self._task_context.write_interceptor: if self._task_context.doc_id == GRAPH_RAPTOR_FAKE_DOC_ID: # raptor - non-determinisic return self._task_context.write_interceptor.intercept("docStoreConn.insert", []) return self._task_context.write_interceptor.intercept("docStoreConn.insert") else: - return await thread_pool_exec(settings.docStoreConn.insert, chunks, index_name, task_dataset_id) + return await thread_pool_exec(settings.docStoreConn.insert, chunks, index_name, task_dataset_id, refresh) async def _insert_main_chunks( self, @@ -345,8 +345,13 @@ class ChunkService: doc_bulk_size: int, ) -> bool: """Insert main chunks in batches with cancellation handling.""" + # Persist task chunk IDs periodically instead of once per bulk request. + # This keeps the task resumable while avoiding one MySQL transaction for + # every small document-store batch. + checkpoint_batches = max(1, 256 // doc_bulk_size) + last_checkpoint = 0 for b in range(0, len(chunks), doc_bulk_size): - doc_store_result = await self._intercept_doc_store_insert(chunks[b : b + doc_bulk_size], search.index_name(task_tenant_id), task_dataset_id) + doc_store_result = await self._intercept_doc_store_insert(chunks[b : b + doc_bulk_size], search.index_name(task_tenant_id), task_dataset_id, refresh=False) if self._task_context.has_canceled_func(task_id): # Roll back partial RAPTOR summary inserts @@ -362,13 +367,20 @@ class ChunkService: self._task_context.progress_cb(-1, msg=error_message) raise Exception(error_message) - # Update chunk IDs in task - chunk_ids = [chunk["id"] for chunk in chunks[: b + doc_bulk_size]] - if not await self._update_task_chunk_ids(task_id, chunk_ids): - # Roll back on failure - await self._rollback_insertion(task_tenant_id, task_dataset_id, chunk_ids) - self._task_context.progress_cb(-1, msg=f"Chunk updates failed since task {task_id} is unknown.") - return False + batch_end = min(b + doc_bulk_size, len(chunks)) + is_last_batch = batch_end == len(chunks) + if is_last_batch or batch_end - last_checkpoint >= checkpoint_batches * doc_bulk_size: + chunk_ids = [chunk["id"] for chunk in chunks[:batch_end]] + if not await self._update_task_chunk_ids(task_id, chunk_ids): + # Roll back on failure + await self._rollback_insertion(task_tenant_id, task_dataset_id, chunk_ids) + self._task_context.progress_cb(-1, msg=f"Chunk updates failed since task {task_id} is unknown.") + return False + last_checkpoint = batch_end + + refresh_idx = getattr(settings.docStoreConn, "refresh_idx", None) + if callable(refresh_idx): + await thread_pool_exec(refresh_idx, search.index_name(task_tenant_id)) return True diff --git a/rag/utils/es_conn.py b/rag/utils/es_conn.py index 7dd6366880..4c8bfc24a6 100644 --- a/rag/utils/es_conn.py +++ b/rag/utils/es_conn.py @@ -334,7 +334,7 @@ class ESConnection(ESConnectionBase): self.logger.error(f"ESConnection.search timeout for {ATTEMPT_TIME} times!") raise Exception("ESConnection.search timeout.") - def insert(self, documents: list[dict], index_name: str, knowledgebase_id: str = None) -> list[str]: + def insert(self, documents: list[dict], index_name: str, knowledgebase_id: str = None, refresh: str | bool = "wait_for") -> list[str]: # Refers to https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html operations = [] for d in documents: @@ -351,7 +351,7 @@ class ESConnection(ESConnectionBase): for _ in range(ATTEMPT_TIME): try: res = [] - r = self.es.bulk(index=index_name, operations=operations, refresh="wait_for", timeout="60s") + r = self.es.bulk(index=index_name, operations=operations, refresh=refresh, timeout="60s") if re.search(r"False", str(r["errors"]), re.IGNORECASE): return res diff --git a/rag/utils/infinity_conn.py b/rag/utils/infinity_conn.py index e2980f19e2..33b8cb8250 100644 --- a/rag/utils/infinity_conn.py +++ b/rag/utils/infinity_conn.py @@ -379,7 +379,7 @@ class InfinityConnection(InfinityConnectionBase): chunk["id"] = chunk_id return chunk - def insert(self, documents: list[dict], index_name: str, knowledgebase_id: str = None) -> list[str]: + def insert(self, documents: list[dict], index_name: str, knowledgebase_id: str = None, refresh: str | bool = "wait_for") -> list[str]: """ # Save input to file to test inserting from file in GO import datetime diff --git a/rag/utils/ob_conn.py b/rag/utils/ob_conn.py index 3e1b480412..3ad0d07e6f 100644 --- a/rag/utils/ob_conn.py +++ b/rag/utils/ob_conn.py @@ -1014,7 +1014,7 @@ class OBConnection(OBConnectionBase): logger.exception(f"OBConnection.get({chunk_id}) got exception") raise e - def insert(self, documents: list[dict], index_name: str, knowledgebase_id: str = None) -> list[str]: + def insert(self, documents: list[dict], index_name: str, knowledgebase_id: str = None, refresh: str | bool = "wait_for") -> list[str]: if not documents: return [] diff --git a/rag/utils/opensearch_conn.py b/rag/utils/opensearch_conn.py index 97d510308f..30d195b3e5 100644 --- a/rag/utils/opensearch_conn.py +++ b/rag/utils/opensearch_conn.py @@ -507,7 +507,7 @@ class OSConnection(DocStoreConnection): logger.error(f"OSConnection.get timeout for {ATTEMPT_TIME} times!") raise Exception("OSConnection.get timeout.") - def insert(self, documents: list[dict], indexName: str, knowledgebaseId: str = None) -> list[str]: + def insert(self, documents: list[dict], indexName: str, knowledgebaseId: str = None, refresh: str | bool = "wait_for") -> list[str]: # Refers to https://opensearch.org/docs/latest/api-reference/document-apis/bulk/ operations = [] for d in documents: @@ -525,7 +525,7 @@ class OSConnection(DocStoreConnection): for _ in range(ATTEMPT_TIME): try: res = [] - r = self.os.bulk(index=(indexName), body=operations, refresh="wait_for", timeout=60) + r = self.os.bulk(index=(indexName), body=operations, refresh=refresh, timeout=60) if re.search(r"False", str(r["errors"]), re.IGNORECASE): return res diff --git a/test/unit_test/api/db/services/test_pipeline_operation_log_service.py b/test/unit_test/api/db/services/test_pipeline_operation_log_service.py new file mode 100644 index 0000000000..f8a5523eb7 --- /dev/null +++ b/test/unit_test/api/db/services/test_pipeline_operation_log_service.py @@ -0,0 +1,49 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from api.db.services.pipeline_operation_log_service import _remove_embedding_vectors + + +def test_remove_embedding_vectors_preserves_pipeline_runtime_outputs(): + dsl = { + "components": { + "Tokenizer:0": { + "obj": { + "params": { + "outputs": { + "chunks": { + "value": [ + { + "text": "content", + "q_1024_vec": [0.1, 0.2], + "metadata": {"q_3_vec": [1, 2, 3], "source": "test"}, + } + ] + }, + "embedding_token_consumption": {"value": 12}, + } + } + } + } + }, + "path": ["Tokenizer:0"], + } + + result = _remove_embedding_vectors(dsl) + + chunk = result["components"]["Tokenizer:0"]["obj"]["params"]["outputs"]["chunks"]["value"][0] + assert chunk == {"text": "content", "metadata": {"source": "test"}} + assert result["components"]["Tokenizer:0"]["obj"]["params"]["outputs"]["embedding_token_consumption"]["value"] == 12 + assert result["path"] == ["Tokenizer:0"]