refactor: overhaul task executor with layered architecture and comprehensive test suite (#15471)

## Summary

Decomposes the monolithic `task_executor.py` (1945 lines) into a 6-layer
architecture with clear separation of concerns. The refactored code is
functionally equivalent to the original, verified through 400 passing
tests and a production-vs-dry-run comparison framework.

## Architecture

```
entry (task_manager)
  └─ orchestration (task_handler)
       ├─ services (chunk_service, embedding_service, dataflow_service, raptor_service, post_processor)
       │    └─ utilities (chunk_builder, chunk_post_processor, embedding_utils)
       └─ infrastructure (task_context, recording_context, interceptor)
```

Key design decisions:
- **TaskContext** — typed facade over raw task dict, injects rate
limiters + callbacks via composition
- **RecordingContext + Comparator** — enables side-by-side production vs
dry-run execution for safe migration
- **NullRecordingContext** — zero-allocation no-op for production, uses
`__slots__`
- **WriteOperationInterceptor** — FIFO replay of previous runs function
returns for comparison mode

## Migration Strategy

The original `handle_task()` in `task_executor.py` uses a 3-way switch
via `TE_RUN_MODE`:
- `TE_RUN_MODE=0` (default) → runs refactored code
- `TE_RUN_MODE=1` → runs both original + refactored, compares all
intermediate results
- `TE_RUN_MODE=2` → runs original code (fallback)

The comparison mode (`TE_RUN_MODE=1`) records ~40 intermediate values
(chunks, vectors, token counts, func return values) from the production
run and replays them during dry-run, then uses `ContextComparator` to
report mismatches.

## Functional Equivalence Fixes

All divergences between original and refactored code were identified and
fixed:
- Timeout decorators (handle/build_chunks/raptor/embedding)
- NullRecordingContext leak in finally block causing RuntimeError
- MinIO None-binary check with proper FileNotFoundError
- Dataflow dispatch after embedding binding + init_kb
- Memory task missing return after processing
- RAPTOR checkpoint progress reporting
- Tag cache (get_tags_from_cache/set_tags_to_cache) restoration
- dataflow_id correction in _load_dsl
- Language default Chinese, dead code guard removal
- embed_chunks made async with proper thread_pool_exec
- Full GraphRAG default configuration (10 parameters)
- Hardcoded q_768_vec fallback removal in RAPTOR

## Test Changes

- 20 new tests covering table parser manual mode, tag cache, embedding
edge cases, RAPTOR checkpoint, dataflow_id correction, storage binary
None, cancel cleanup, metadata=None boundary
- Unified `make_task_context`/`make_task_dict` factories eliminated 10+
duplicated helpers
- DataflowService tests migrated from internal method mocks to IO
boundary mocks (real orchestration code executes)
- Parametrized duplicate build_chunks post-processor tests
- 7 raptor tests modernized to @pytest.mark.asyncio
- Mock count per test reduced through boundary-level mocking strategy

**Test count: 400 passing, 0 warnings, 0 skips**

## Files Changed

| File | Change |
|------|--------|
| `rag/svr/task_executor.py` | +1 line (NullRecordingContext fix) |
| `rag/svr/task_executor_refactor/task_handler.py` | Orchestration
layer, 8 logic fixes |
| `rag/svr/task_executor_refactor/chunk_service.py` | +timeout +
None-check |
| `rag/svr/task_executor_refactor/embedding_service.py` | sync→async
rewrite |
| `rag/svr/task_executor_refactor/dataflow_service.py` | dataflow_id fix
+ timeout |
| `rag/svr/task_executor_refactor/raptor_service.py` | checkpoint fix +
assert |
| `rag/svr/task_executor_refactor/chunk_post_processor.py` | tag cache
restore |
| `rag/svr/task_executor_refactor/task_context.py` | language default
fix |
| `test/.../conftest.py` | +294 lines shared helpers |
| `test/.../*.py` | 15 test files refactored, 20 new tests |

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

rag/svr/task_executor.py

@@ -1728,6 +1728,7 @@ async def handle_task():
         elif run_mode == "0": # use refactor-ed version
             # switch to refactor-ed version
             logging.info(f"-----run refactor-ed task executor:{task_id}, {task.get('name', '')}, doc id:{task.get('doc_id', '')}")
+            set_recording_context(NullRecordingContext())
             await TaskManager.run_refactored_task(task, chat_limiter, minio_limiter, chunk_limiter,
                                embed_limiter,kg_limiter, set_progress, has_canceled)
         else: # original version

rag/svr/task_executor_refactor/chunk_post_processor.py

@@ -41,7 +41,7 @@ from api.db.services.doc_metadata_service import DocMetadataService
 from api.db.services.llm_service import LLMBundle
 from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance
 from rag.prompts.generator import gen_metadata, keyword_extraction, question_proposal, content_tagging
-from rag.graphrag.utils import get_llm_cache, set_llm_cache
+from rag.graphrag.utils import get_llm_cache, set_llm_cache, get_tags_from_cache, set_tags_to_cache
 
 
 async def extract_keywords(docs: List[Dict], ctx: TaskContext) -> None:
@@ -243,10 +243,14 @@ async def apply_tags(docs: List[Dict], ctx: TaskContext) -> None:
     S = 1000
     st = timer()
     examples = []
-    all_tags = settings.retriever.all_tags_in_portion(tenant_id, kb_ids, S)
+    all_tags = get_tags_from_cache(kb_ids)
+    if not all_tags:
+        all_tags = settings.retriever.all_tags_in_portion(tenant_id, kb_ids, S)
+        set_tags_to_cache(kb_ids, all_tags)
+    else:
+        all_tags = json.loads(all_tags)
     chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, ctx.llm_id)
     with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_model:
-
         docs_to_tag = []
         for doc in docs:
             if ctx.has_canceled_func(ctx.id):

rag/svr/task_executor_refactor/chunk_service.py

@@ -34,6 +34,7 @@ from typing import Any, Dict, List
 
 import xxhash
 from common import settings
+from common.connection_utils import timeout
 from common.constants import PAGERANK_FLD, TAG_FLD
 from common.misc_utils import thread_pool_exec
 from common.float_utils import normalize_overlapped_percent
@@ -85,6 +86,7 @@ class ChunkService:
         """
         self._task_context = ctx
 
+    @timeout(60 * 80, 1)
     async def build_chunks(
         self,
         storage_binary: bytes,
@@ -190,6 +192,7 @@ class ChunkService:
 
         st = timer()
 
+        @timeout(60)
         async def upload_to_minio(document, chunk):
             try:
                 d = copy.deepcopy(document)

rag/svr/task_executor_refactor/dataflow_service.py

@@ -39,6 +39,7 @@ from api.db.services.document_service import DocumentService
 from api.db.services.doc_metadata_service import DocMetadataService
 from api.db.services.pipeline_operation_log_service import PipelineOperationLogService
 from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance
+from common.connection_utils import timeout
 from common.constants import LLMType, PipelineTaskType
 from common.metadata_utils import update_metadata_to
 from common.misc_utils import thread_pool_exec
@@ -102,9 +103,10 @@ class DataflowService:
             task_dataset_id = ctx.kb_id
 
             # Load DSL
-            dsl = await self._load_dsl(dataflow_id)
+            dsl, corrected_id = await self._load_dsl(dataflow_id)
             if dsl is None:
                 return
+            dataflow_id = corrected_id
 
             # Run pipeline
             pipeline = Pipeline(
@@ -193,17 +195,23 @@ class DataflowService:
                 await self._billing_hook.on_pipeline_error()
             raise
 
-    async def _load_dsl(self, dataflow_id: str) -> Optional[str]:
-        """Load dataflow DSL from service."""
+    async def _load_dsl(self, dataflow_id: str) -> tuple:
+        """Load dataflow DSL from service.
+
+        Returns:
+            Tuple of (dsl, corrected_dataflow_id).
+            When task_type is not 'dataflow', the dataflow_id is corrected
+            from the pipeline log's pipeline_id.
+        """
         ctx = self._task_context
         if ctx.task_type == "dataflow":
             e, cvs = UserCanvasService.get_by_id(dataflow_id)
             assert e, "User pipeline not found."
-            return cvs.dsl
+            return cvs.dsl, dataflow_id
         else:
             e, pipeline_log = PipelineOperationLogService.get_by_id(dataflow_id)
             assert e, "Pipeline log not found."
-            return pipeline_log.dsl
+            return pipeline_log.dsl, pipeline_log.pipeline_id
 
     @staticmethod
     def _get_output_type(chunks: Dict) -> str:
@@ -235,6 +243,7 @@ class DataflowService:
             return [{"text": [chunks["html"]]}] if chunks["html"] else []
         return []
 
+    @timeout(60)
     async def _embed_chunks(
         self, chunks: List[Dict], token_consumption: int
     ) -> Tuple[Optional[List[Dict]], int]:

rag/svr/task_executor_refactor/embedding_service.py

@@ -19,11 +19,12 @@ Embedding Service Module.
 Provides [`EmbeddingService`](rag/svr/task_executor_refactor/embedding_service.py:42) for vector embedding operations.
 """
 
-import asyncio
 from typing import Any, Dict, List, Tuple
 
 import numpy as np
 from common import settings
+from common.misc_utils import thread_pool_exec
+from common.token_utils import truncate
 from rag.svr.task_executor_refactor.embedding_utils import EmbeddingUtils
 from rag.svr.task_executor_refactor.task_context import TaskContext
 
@@ -54,7 +55,7 @@ class EmbeddingService:
 
         self._embedding_batch_size = embedding_batch_size or settings.EMBEDDING_BATCH_SIZE
 
-    def embed_chunks(
+    async def embed_chunks(
         self,
         docs: List[Dict[str, Any]],
         embedding_model,
@@ -79,7 +80,8 @@ class EmbeddingService:
         # Encode titles using EmbeddingUtils for truncation
         tk_count = 0
         if len(titles) > 0 and len(titles) == len(contents):
-            vts, c = self._encode_single([titles[0]], embedding_model)
+            async with self._task_context.embed_limiter:
+                vts, c = await thread_pool_exec(embedding_model.encode, titles[0:1])
             tts = np.tile(vts[0], (len(contents), 1))
             tk_count += c
         else:
@@ -89,7 +91,12 @@ class EmbeddingService:
         vects_batches = []
         for i in range(0, len(contents), self._embedding_batch_size):
             batch = contents[i: i + self._embedding_batch_size]
-            vts, c = self._encode_batch(batch, embedding_model)
+            async with self._task_context.embed_limiter:
+                vts, c = await thread_pool_exec(
+                    self._batch_encode_wrapper,
+                    [truncate(t, embedding_model.max_length - 10) for t in batch],
+                    embedding_model,
+                )
             vects_batches.append(vts)
             tk_count += c
             if self._task_context.progress_cb:
@@ -109,19 +116,7 @@ class EmbeddingService:
 
         return tk_count, vector_size
 
-    def _encode_single(self, texts: List[str], model) -> Tuple[np.ndarray, int]:
-        """Encode a single batch of texts."""
-        return self._run_encode(texts, model)
-
-    def _encode_batch(self, texts: List[str], model) -> Tuple[np.ndarray, int]:
-        """Encode a batch of texts with rate limiting and truncation."""
-        # Use EmbeddingUtils for truncation
-        truncated = EmbeddingUtils.truncate_texts(texts, model.max_length)
-        return self._run_encode(truncated, model)
-
-    def _run_encode(self, texts: List[str], model) -> Tuple[np.ndarray, int]:
-        """Run encoding with rate limiting."""
-        async def _encode():
-            async with self._task_context.embed_limiter:
-                return model.encode(texts)
-        return asyncio.get_event_loop().run_until_complete(_encode())
+    @staticmethod
+    def _batch_encode_wrapper(txts: List[str], embedding_model) -> Tuple[np.ndarray, int]:
+        """Synchronous wrapper for batch encoding — used with thread_pool_exec."""
+        return embedding_model.encode(txts)

rag/svr/task_executor_refactor/raptor_service.py

@@ -31,6 +31,7 @@ import numpy as np
 from api.db.services.document_service import DocumentService
 from api.db.services.task_service import GRAPH_RAPTOR_FAKE_DOC_ID
 from common import settings
+from common.connection_utils import timeout
 from common.constants import PAGERANK_FLD
 from common.misc_utils import thread_pool_exec
 from common.token_utils import num_tokens_from_string
@@ -68,6 +69,7 @@ class RaptorService:
         """
         self._task_context = ctx
 
+    @timeout(3600)
     async def run_raptor_for_kb(
         self,
         kb_parser_config: Dict,
@@ -166,8 +168,8 @@ class RaptorService:
                         doc_id, tree_builder, cleanup_raptor_chunks
                     )
                     self._task_context.progress_cb(msg=f"[RAPTOR] doc:{doc_id} will remove old RAPTOR summaries after insert.")
-                    self._task_context.progress_cb(msg=f"[RAPTOR] doc:{doc_id} already has {tree_builder} RAPTOR chunks, skipping.")
-                    self._task_context.progress_cb(prog=(x + 1.) / len(doc_ids))
+                self._task_context.progress_cb(msg=f"[RAPTOR] doc:{doc_id} already has {tree_builder} RAPTOR chunks, skipping.")
+                self._task_context.progress_cb(prog=(x + 1.) / len(doc_ids))
                 continue
 
             if existing_methods:
@@ -370,7 +372,8 @@ class RaptorService:
         from rag.raptor import RecursiveAbstractiveProcessing4TreeOrganizedRetrieval as Raptor
 
         raptor_ext_config = raptor_config.get("ext") or {}
-        vctr_nm = "q_%d_vec" % len(chunks[0][1]) if chunks else "q_768_vec"
+        assert chunks, "_generate_raptor must not be called with empty chunks"
+        vctr_nm = "q_%d_vec" % len(chunks[0][1])
 
         raptor = Raptor(
             raptor_config.get("max_cluster", 64),

rag/svr/task_executor_refactor/task_context.py

@@ -224,7 +224,7 @@ class TaskContext:
         "parser_id": "",
         "parser_config": {},
         "kb_parser_config": {},
-        "language": "en",
+        "language": "Chinese",
         "llm_id": "",
         "embd_id": "",
         "from_page": 0,

rag/svr/task_executor_refactor/task_handler.py

@@ -39,6 +39,7 @@ from api.db.services.llm_service import LLMBundle
 from api.db.services.task_service import GRAPH_RAPTOR_FAKE_DOC_ID
 from common.constants import LLMType
 from common.exceptions import TaskCanceledException
+from common.connection_utils import timeout
 from common.misc_utils import thread_pool_exec
 from rag.nlp import search
 from rag.svr.task_executor_refactor.constants import CANVAS_DEBUG_DOC_ID
@@ -111,6 +112,7 @@ class TaskHandler:
                     logging.exception(
                         f"Remove doc({task_doc_id}) from docStore failed when task({task_id}) canceled, exception: {e}")
 
+    @timeout(60 * 60 * 3, 1)
     async def handle(self) -> None:
         """Handle a document processing task."""
         ctx = self._task_context
@@ -125,14 +127,6 @@ class TaskHandler:
             else:
                 # actual run - not dry run
                 await handle_save_to_memory_task(ctx.raw_task)
-
-        # Handle dataflow debug mode
-        if task_type == "dataflow" and ctx.doc_id == CANVAS_DEBUG_DOC_ID:
-            await self._run_dataflow()
-            return
-
-        if task_type.startswith("dataflow"):
-            await self._run_dataflow()
             return
 
         # Check if task is canceled
@@ -140,15 +134,25 @@ class TaskHandler:
             ctx.progress_cb(-1, msg="Task has been canceled.")
             return
 
-        # Bind embedding model
-        embedding_model = await self._bind_embedding_model()
-        if embedding_model is None:
+        # Language defaults to "Chinese" via TaskContext._DEFAULTS — safe to bind model directly.
+        # Bind embedding model (matching original do_handle_task order: bind + init_kb before routing)
+        result = await self._bind_embedding_model()
+        if result is None:
             return
+        embedding_model, vector_size = result
 
         with embedding_model:
-            vector_size = self._get_vector_size(embedding_model)
             self._init_kb(vector_size)
 
+            # Handle dataflow tasks (after init_kb, matching original behavior)
+            if task_type == "dataflow" and ctx.doc_id == CANVAS_DEBUG_DOC_ID:
+                await self._run_dataflow()
+                return
+
+            if task_type.startswith("dataflow"):
+                await self._run_dataflow()
+                return
+
             # Route to appropriate handler
             if task_type == "raptor":
                 await self._run_raptor(embedding_model, vector_size)
@@ -166,12 +170,6 @@ class TaskHandler:
                 await self._run_standard_chunking(embedding_model)
 
 
-    @classmethod
-    def _get_vector_size(cls, embedding_model: LLMBundle) -> int:
-        """Get vector size from embedding model."""
-        vts, _ = embedding_model.encode(["ok"])
-        return len(vts[0])
-
     def _init_kb(self, vector_size: int) -> None:
         """Initialize knowledge base index."""
         ctx = self._task_context
@@ -203,8 +201,12 @@ class TaskHandler:
         ctx = self._task_context
         ctx.progress_cb(1, "Clone task placeholder")
 
-    async def _bind_embedding_model(self) -> Optional[LLMBundle]:
-        """Bind embedding model to task."""
+    async def _bind_embedding_model(self) -> Optional[tuple]:
+        """Bind embedding model to task.
+
+        Returns:
+            Tuple of (embedding_model, vector_size) on success, or None on failure.
+        """
         ctx = self._task_context
         task_tenant_id = ctx.tenant_id
         task_embedding_id = ctx.embd_id
@@ -221,7 +223,7 @@ class TaskHandler:
                 )
             embedding_model = LLMBundle(task_tenant_id, embd_model_config, lang=task_language)
             vts, _ = embedding_model.encode(["ok"])
-            return embedding_model
+            return embedding_model, len(vts[0])
         except Exception as e:
             error_message = f'Fail to bind embedding model: {str(e)}'
             ctx.progress_cb(-1, msg=error_message)
@@ -340,8 +342,24 @@ class TaskHandler:
             kb_parser_config.update({
                 "graphrag": {
                     "use_graphrag": True,
-                    "entity_types": ["organization", "person", "geo", "event", "category"],
+                    "entity_types": [
+                        "organization",
+                        "person",
+                        "geo",
+                        "event",
+                        "category",
+                    ],
                     "method": "light",
+                    "batch_chunk_token_size": 4096,
+                    "retry_attempts": 2,
+                    "retry_backoff_seconds": 2.0,
+                    "retry_backoff_max_seconds": 60.0,
+                    "build_subgraph_timeout_per_chunk_seconds": 300,
+                    "build_subgraph_min_timeout_seconds": 600,
+                    "merge_timeout_seconds": 180,
+                    "resolution_timeout_seconds": 1800,
+                    "community_timeout_seconds": 1800,
+                    "lock_acquire_timeout_seconds": 600,
                 }
             })
             if ctx.write_interceptor:
@@ -400,6 +418,10 @@ class TaskHandler:
         # Get storage binary
         bucket, name = File2DocumentService.get_storage_address(doc_id=ctx.doc_id)
         binary = await self._get_storage_binary(bucket, name)
+        if binary is None:
+            raise FileNotFoundError(
+                f"Can not find file <{ctx.name}> from minio. Could you try it again."
+            )
 
         chunks = await chunk_service.build_chunks(binary)
         ctx.recording_context.record("chunks", chunks)
@@ -418,7 +440,7 @@ class TaskHandler:
         start_ts = timer()
         embedding_service = EmbeddingService(ctx=ctx)
         try:
-            token_count, vector_size = embedding_service.embed_chunks(
+            token_count, vector_size = await embedding_service.embed_chunks(
                 chunks, embedding_model, ctx.parser_config
             )
         except TaskCanceledException: