From f930389311420b989b9b7419d42de32c0e03f252 Mon Sep 17 00:00:00 2001 From: Minal Mahala Date: Wed, 15 Apr 2026 15:07:28 +0530 Subject: [PATCH] Refact: improve task resume mechanism for graphrag (#14096) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What problem does this PR solve? Addresses review feedback on #14074 (Checkpoint mechanism for long-running workflow jobs, issue #12494). **Changes based on @yuzhichang's review:** 1. **Renamed `checkpoint_service.py` → `task_checkpoint.py`** as suggested. 2. **Replaced Redis with direct docEngine queries** as suggested — the subgraph already gets persisted to the doc store by `generate_subgraph()`, so we just query for it instead of maintaining a separate checkpoint in Redis. This is simpler, has no extra dependency, and uses a single source of truth. **Changes based on CodeRabbit review:** 3. **Fixed `source_id` query format mismatch** — subgraphs are stored with `source_id: [doc_id]` (list), but the original query used `source_id: doc_id` (string). Now follows the same pattern as `does_graph_contains()` in `rag/graphrag/utils.py`: filter by `knowledge_graph_kwd` only, then match `source_id` in Python. This avoids ambiguity across Elasticsearch / Infinity / OceanBase backends. ### Changes | File | Change | |---|---| | `api/db/services/task_checkpoint.py` (new) | `load_subgraph_from_store()` and `has_raptor_chunks()` — docEngine-based checkpoint queries | | `rag/graphrag/general/index.py` | `build_one()` calls `load_subgraph_from_store()` before running LLM extraction | | `rag/svr/task_executor.py` | RAPTOR per-doc loop calls `has_raptor_chunks()` before processing | | `test/unit_test/rag/graphrag/test_checkpoint_resume.py` (new) | 10 unit tests covering subgraph loading, source_id filtering, edge cases | ### How it works - **GraphRAG:** Before running expensive LLM entity/relation extraction for a doc, checks the doc store for an existing subgraph (saved by a previous interrupted run). If found, loads it directly and skips LLM calls. - **RAPTOR:** Before processing a doc, checks if RAPTOR chunks (`raptor_kwd="raptor"`) already exist for it. If yes, skips. ### Testing - 10 new unit tests — all passing - Full existing suite: 617 passed ### Type of change - [x] New Feature (non-breaking change which adds functionality) - [x] Refactoring --- rag/graphrag/general/index.py | 56 ++++ rag/svr/task_executor.py | 46 ++- .../rag/graphrag/test_checkpoint_resume.py | 316 ++++++++++++++++++ 3 files changed, 415 insertions(+), 3 deletions(-) create mode 100644 test/unit_test/rag/graphrag/test_checkpoint_resume.py diff --git a/rag/graphrag/general/index.py b/rag/graphrag/general/index.py index 1a43ebe092..2dc8bd4204 100644 --- a/rag/graphrag/general/index.py +++ b/rag/graphrag/general/index.py @@ -43,6 +43,56 @@ from common.misc_utils import thread_pool_exec from rag.nlp import rag_tokenizer, search from rag.utils.redis_conn import RedisDistributedLock from common import settings +from common.doc_store.doc_store_base import OrderByExpr + + + +async def load_subgraph_from_store(tenant_id: str, kb_id: str, doc_id: str): + """Load a previously saved subgraph from the doc store. + + Filters directly by source_id (== doc_id) and knowledge_graph_kwd in the + query so the doc store index does the heavy lifting. Expects at most one + matching chunk per doc_id (as written by generate_subgraph). + Returns a networkx Graph on hit, or None on miss. + """ + fields = ["content_with_weight", "source_id"] + condition = { + "knowledge_graph_kwd": ["subgraph"], + "removed_kwd": "N", + "source_id": [doc_id], + } + try: + res = await thread_pool_exec( + settings.docStoreConn.search, + fields, [], condition, [], OrderByExpr(), + 0, 1, search.index_name(tenant_id), [kb_id] + ) + field_map = settings.docStoreConn.get_fields(res, fields) + for cid, row in field_map.items(): + content = row.get("content_with_weight", "") + if not content: + continue + try: + data = json.loads(content) + sg = nx.node_link_graph(data, edges="edges") + sg.graph["source_id"] = [doc_id] + logging.info( + "Checkpoint hit: subgraph for doc %s (tenant=%s kb=%s) found at chunk %s", + doc_id, tenant_id, kb_id, cid, + ) + return sg + except Exception: + logging.exception( + "Failed to parse subgraph JSON for doc %s chunk %s", doc_id, cid + ) + except Exception: + logging.exception("Failed to load subgraph from store for doc %s", doc_id) + return None + logging.info( + "Checkpoint miss: no subgraph for doc %s (tenant=%s kb=%s)", + doc_id, tenant_id, kb_id, + ) + return None async def run_graphrag( @@ -242,6 +292,12 @@ async def run_graphrag_for_kb( deadline = max(120, len(chunks) * 60 * 10) if enable_timeout_assertion else 10000000000 async with semaphore: + # CHECKPOINT: bounded by semaphore so doc-store lookups respect max_parallel_docs + existing_sg = await load_subgraph_from_store(tenant_id, kb_id, doc_id) + if existing_sg: + subgraphs[doc_id] = existing_sg + callback(msg=f"[GraphRAG] doc:{doc_id} subgraph found in store, skipping LLM extraction.") + return try: msg = f"[GraphRAG] build_subgraph doc:{doc_id}" callback(msg=f"{msg} start (chunks={len(chunks)}, timeout={deadline}s)") diff --git a/rag/svr/task_executor.py b/rag/svr/task_executor.py index c9adc990bc..dde7ef00f2 100644 --- a/rag/svr/task_executor.py +++ b/rag/svr/task_executor.py @@ -768,6 +768,40 @@ async def run_dataflow(task: dict): dsl=str(pipeline)) +async def has_raptor_chunks(doc_id: str, tenant_id: str, kb_id: str) -> bool: + """Return True if RAPTOR chunks already exist for doc_id in the doc store. + + Queries directly for raptor_kwd="raptor" rows so a non-RAPTOR leading + chunk cannot produce a false-negative result. Uses thread_pool_exec so + the blocking doc-store call does not stall the event loop. + """ + from common.doc_store.doc_store_base import OrderByExpr + from rag.nlp import search as nlp_search + try: + condition = {"doc_id": doc_id, "raptor_kwd": ["raptor"]} + res = await thread_pool_exec( + settings.docStoreConn.search, + ["raptor_kwd"], [], condition, [], OrderByExpr(), + 0, 1, nlp_search.index_name(tenant_id), [kb_id] + ) + field_map = settings.docStoreConn.get_fields(res, ["raptor_kwd"]) + found = bool(field_map) + if found: + logging.info( + "Checkpoint hit: RAPTOR chunks for doc %s (tenant=%s kb=%s) already exist", + doc_id, tenant_id, kb_id, + ) + else: + logging.info( + "Checkpoint miss: no RAPTOR chunks for doc %s (tenant=%s kb=%s)", + doc_id, tenant_id, kb_id, + ) + return found + except Exception: + logging.exception("Failed to check RAPTOR chunks for doc %s", doc_id) + return False + + @timeout(3600) async def run_raptor_for_kb(row, kb_parser_config, chat_mdl, embd_mdl, vector_size, callback=None, doc_ids=[]): fake_doc_id = GRAPH_RAPTOR_FAKE_DOC_ID @@ -825,6 +859,12 @@ async def run_raptor_for_kb(row, kb_parser_config, chat_mdl, embd_mdl, vector_si if raptor_config.get("scope", "file") == "file": for x, doc_id in enumerate(doc_ids): + # CHECKPOINT: skip docs that already have RAPTOR chunks in the doc store + if await has_raptor_chunks(doc_id, row["tenant_id"], row["kb_id"]): + callback(msg=f"[RAPTOR] doc:{doc_id} already has RAPTOR chunks, skipping.") + callback(prog=(x + 1.) / len(doc_ids)) + continue + chunks = [] skipped_chunks = 0 for d in settings.retriever.chunk_list(doc_id, row["tenant_id"], [str(row["kb_id"])], @@ -836,15 +876,15 @@ async def run_raptor_for_kb(row, kb_parser_config, chat_mdl, embd_mdl, vector_si logging.warning(f"RAPTOR: Chunk missing vector field '{vctr_nm}' in doc {doc_id}, skipping") continue chunks.append((d["content_with_weight"], np.array(d[vctr_nm]))) - + if skipped_chunks > 0: callback(msg=f"[WARN] Skipped {skipped_chunks} chunks without vector field '{vctr_nm}' for doc {doc_id}. Consider re-parsing the document with the current embedding model.") - + if not chunks: logging.warning(f"RAPTOR: No valid chunks with vectors found for doc {doc_id}") callback(msg=f"[WARN] No valid chunks with vectors found for doc {doc_id}, skipping") continue - + await generate(chunks, doc_id) callback(prog=(x + 1.) / len(doc_ids)) else: diff --git a/test/unit_test/rag/graphrag/test_checkpoint_resume.py b/test/unit_test/rag/graphrag/test_checkpoint_resume.py new file mode 100644 index 0000000000..f8abde08db --- /dev/null +++ b/test/unit_test/rag/graphrag/test_checkpoint_resume.py @@ -0,0 +1,316 @@ +# +# Copyright 2025 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. +# + +"""Tests for GraphRAG/RAPTOR checkpoint/resume logic. + +Calls the real implementations: + - load_subgraph_from_store (rag/graphrag/general/index.py) + - has_raptor_chunks (rag/svr/task_executor.py) + +Both modules are loaded via importlib with their infrastructure dependencies +mocked, so the actual query logic, pagination, and error handling are exercised +without needing running services. +""" + +import importlib.util +import json +import pathlib +import sys +import warnings +from unittest.mock import MagicMock + +# Suppress deprecation warnings from third-party libraries (e.g. huggingface_hub) +# that are triggered during module import but are not related to the code under test. +warnings.filterwarnings("ignore", category=UserWarning, module="huggingface_hub") + +import networkx as nx +import pytest + +# --------------------------------------------------------------------------- +# Additional sys.modules mocks needed beyond what conftest already provides. +# +# conftest.py (same directory) mocks the heavy packages listed in +# _modules_to_mock. We need a few more to satisfy index.py and +# task_executor.py's import-time dependencies. +# --------------------------------------------------------------------------- +_EXTRA_MOCKS = [ + # for index.py + "api.db.services.document_service", + # for task_executor.py + "api.db", + "api.db.services.knowledgebase_service", + "api.db.services.pipeline_operation_log_service", + "api.db.joint_services", + "api.db.joint_services.memory_message_service", + "api.db.joint_services.tenant_model_service", + "api.db.services.doc_metadata_service", + "api.db.services.llm_service", + "api.db.services.file2document_service", + "api.db.db_models", + "common.metadata_utils", + "common.log_utils", + "common.config_utils", + "common.versions", + "common.token_utils", + "common.signal_utils", + "common.exceptions", + "common.constants", + "rag.utils.base64_image", + + "rag.prompts.generator", + "rag.raptor", + "rag.app", + "rag.graphrag.utils", +] +for _m in _EXTRA_MOCKS: + if _m not in sys.modules: + sys.modules[_m] = MagicMock() + +# --------------------------------------------------------------------------- +# Load the real implementations via importlib. +# --------------------------------------------------------------------------- +_ROOT = pathlib.Path(__file__).parents[4] + + +def _load_module(dotted_name: str, rel_path: str): + path = _ROOT / rel_path + spec = importlib.util.spec_from_file_location(dotted_name, path) + mod = importlib.util.module_from_spec(spec) + sys.modules[dotted_name] = mod + spec.loader.exec_module(mod) + return mod + + +_index_mod = _load_module("rag.graphrag.general.index", "rag/graphrag/general/index.py") +_executor_mod = _load_module("rag.svr.task_executor", "rag/svr/task_executor.py") + +load_subgraph_from_store = _index_mod.load_subgraph_from_store +has_raptor_chunks = _executor_mod.has_raptor_chunks + +# settings is a MagicMock installed by conftest; grab it to monkeypatch docStoreConn. +import common.settings as _settings # noqa: E402 + +# Ensure docStoreConn is a MagicMock so monkeypatch.setattr works in all environments. +if not isinstance(_settings.docStoreConn, MagicMock): + _settings.docStoreConn = MagicMock() + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + +def _make_subgraph(doc_id: str) -> nx.Graph: + sg = nx.Graph() + sg.add_node("ENTITY_A", description="test entity A", source_id=[doc_id]) + sg.add_node("ENTITY_B", description="test entity B", source_id=[doc_id]) + sg.add_edge("ENTITY_A", "ENTITY_B", description="related", source_id=[doc_id], weight=1.0, keywords=[]) + sg.graph["source_id"] = [doc_id] + return sg + + +def _to_store_content(sg: nx.Graph) -> str: + return json.dumps(nx.node_link_data(sg, edges="edges"), ensure_ascii=False) + + +def _single_page_mocks(field_map: dict): + """search + get_fields mocks that simulate a single-page result.""" + sentinel = object() + call_count = {"n": 0} + + def _get_fields(_res, _fields): + call_count["n"] += 1 + return field_map if call_count["n"] == 1 else {} + + return MagicMock(return_value=sentinel), MagicMock(side_effect=_get_fields) + + +# --------------------------------------------------------------------------- +# Tests for load_subgraph_from_store (rag/graphrag/general/index.py) +# --------------------------------------------------------------------------- + +class TestLoadSubgraphFromStore: + + @pytest.mark.p1 + @pytest.mark.asyncio + async def test_loads_existing_subgraph(self, monkeypatch): + """Subgraph present in the store is returned as nx.Graph.""" + doc_id = "doc_001" + sg = _make_subgraph(doc_id) + field_map = {"chunk_001": {"content_with_weight": _to_store_content(sg), "source_id": [doc_id]}} + s, gf = _single_page_mocks(field_map) + monkeypatch.setattr(_settings.docStoreConn, "search", s) + monkeypatch.setattr(_settings.docStoreConn, "get_fields", gf) + + result = await load_subgraph_from_store("t1", "kb1", doc_id) + + assert result is not None and isinstance(result, nx.Graph) + assert result.has_node("ENTITY_A") and result.has_node("ENTITY_B") + assert result.graph["source_id"] == [doc_id] + + @pytest.mark.p1 + @pytest.mark.asyncio + async def test_returns_none_when_no_subgraph(self, monkeypatch): + """Empty store returns None without raising.""" + s, gf = _single_page_mocks({}) + monkeypatch.setattr(_settings.docStoreConn, "search", s) + monkeypatch.setattr(_settings.docStoreConn, "get_fields", gf) + + assert await load_subgraph_from_store("t1", "kb1", "doc_missing") is None + + @pytest.mark.p2 + @pytest.mark.asyncio + async def test_passes_doc_id_in_search_condition(self, monkeypatch): + """source_id (== doc_id) is included in the search condition so the doc + store filters results directly rather than fetching all subgraphs.""" + captured = {} + + def _capture(fields, filters, condition, *_a, **_kw): + captured["condition"] = condition + return object() + + sg = _make_subgraph("doc_b") + monkeypatch.setattr(_settings.docStoreConn, "search", _capture) + monkeypatch.setattr(_settings.docStoreConn, "get_fields", + MagicMock(return_value={"chunk_b": {"content_with_weight": _to_store_content(sg), "source_id": ["doc_b"]}})) + + result = await load_subgraph_from_store("t1", "kb1", "doc_b") + assert result is not None and result.graph["source_id"] == ["doc_b"] + assert captured["condition"]["source_id"] == ["doc_b"] + + @pytest.mark.p2 + @pytest.mark.asyncio + async def test_skips_malformed_json_returns_none(self, monkeypatch): + """Malformed JSON is logged and skipped; None is returned (not raised).""" + field_map = {"chunk_bad": {"content_with_weight": "not valid json{{{", "source_id": ["doc_bad"]}} + s, gf = _single_page_mocks(field_map) + monkeypatch.setattr(_settings.docStoreConn, "search", s) + monkeypatch.setattr(_settings.docStoreConn, "get_fields", gf) + + assert await load_subgraph_from_store("t1", "kb1", "doc_bad") is None + + @pytest.mark.p2 + @pytest.mark.asyncio + async def test_issues_single_query_with_limit_one(self, monkeypatch): + """Exactly one search call is issued with limit=1 — the doc store index + does the filtering, so no pagination is required.""" + doc_id = "doc_single" + sg = _make_subgraph(doc_id) + search_calls: list[tuple] = [] + + def _search(fields, filters, condition, order, orderby, offset, limit, *_a, **_kw): + search_calls.append((offset, limit)) + return object() + + monkeypatch.setattr(_settings.docStoreConn, "search", _search) + monkeypatch.setattr(_settings.docStoreConn, "get_fields", + MagicMock(return_value={"chunk_t": {"content_with_weight": _to_store_content(sg), "source_id": [doc_id]}})) + + result = await load_subgraph_from_store("t1", "kb1", doc_id) + assert result is not None + assert len(search_calls) == 1, "must issue exactly one query" + assert search_calls[0] == (0, 1), "must use offset=0, limit=1" + + @pytest.mark.p2 + @pytest.mark.asyncio + async def test_doc_store_exception_returns_none(self, monkeypatch): + """A doc-store exception is caught; None is returned safely.""" + monkeypatch.setattr(_settings.docStoreConn, "search", MagicMock(side_effect=RuntimeError("db down"))) + assert await load_subgraph_from_store("t1", "kb1", "doc_001") is None + + +# --------------------------------------------------------------------------- +# Tests for has_raptor_chunks (rag/svr/task_executor.py) +# --------------------------------------------------------------------------- + +class TestHasRaptorChunks: + + @pytest.mark.p1 + @pytest.mark.asyncio + async def test_returns_true_when_raptor_chunk_exists(self, monkeypatch): + """Doc store returns a RAPTOR row -> True.""" + monkeypatch.setattr(_settings.docStoreConn, "search", MagicMock(return_value=object())) + monkeypatch.setattr(_settings.docStoreConn, "get_fields", + MagicMock(return_value={"chunk_r": {"raptor_kwd": "raptor"}})) + + assert await has_raptor_chunks("doc_001", "t1", "kb1") is True + + @pytest.mark.p1 + @pytest.mark.asyncio + async def test_returns_false_when_no_raptor_chunks(self, monkeypatch): + """Doc store returns empty -> False.""" + monkeypatch.setattr(_settings.docStoreConn, "search", MagicMock(return_value=object())) + monkeypatch.setattr(_settings.docStoreConn, "get_fields", MagicMock(return_value={})) + + assert await has_raptor_chunks("doc_001", "t1", "kb1") is False + + @pytest.mark.p1 + @pytest.mark.asyncio + async def test_queries_specifically_for_raptor_kwd(self, monkeypatch): + """raptor_kwd is in the search condition so non-RAPTOR leading chunks + cannot produce a false-negative.""" + captured = {} + + def _capture(fields, filters, condition, *_a, **_kw): + captured["condition"] = condition + return object() + + monkeypatch.setattr(_settings.docStoreConn, "search", _capture) + monkeypatch.setattr(_settings.docStoreConn, "get_fields", MagicMock(return_value={})) + + await has_raptor_chunks("doc_001", "t1", "kb1") + assert captured["condition"] == {"doc_id": "doc_001", "raptor_kwd": ["raptor"]} + + @pytest.mark.p2 + @pytest.mark.asyncio + async def test_returns_false_on_doc_store_exception(self, monkeypatch): + """Exception is caught; False is returned without crashing.""" + monkeypatch.setattr(_settings.docStoreConn, "search", MagicMock(side_effect=RuntimeError("db down"))) + assert await has_raptor_chunks("doc_001", "t1", "kb1") is False + + +# --------------------------------------------------------------------------- +# End-to-end workflow test +# --------------------------------------------------------------------------- + +class TestCheckpointResumeWorkflow: + + @pytest.mark.p1 + @pytest.mark.asyncio + async def test_resume_finds_completed_docs_skips_new_ones(self, monkeypatch): + """3 docs completed before crash; on resume each is found, new doc is not.""" + completed = ["doc_1", "doc_2", "doc_3"] + field_map = { + f"chunk_{d}": {"content_with_weight": _to_store_content(_make_subgraph(d)), "source_id": [d]} + for d in completed + } + # The doc store filters by source_id (doc_id) directly, so get_fields + # should return only the matching chunk for each call. + def _get_fields_by_doc(res, fields): + # res is the sentinel from search; extract the doc_id it was called with + return {k: v for k, v in field_map.items() if v["source_id"] == [_get_fields_by_doc.last_doc_id]} + + def _search(fields, filters, condition, *_a, **_kw): + _get_fields_by_doc.last_doc_id = (condition or {}).get("source_id", [None])[0] + return object() + + monkeypatch.setattr(_settings.docStoreConn, "search", _search) + monkeypatch.setattr(_settings.docStoreConn, "get_fields", _get_fields_by_doc) + + for doc_id in completed: + result = await load_subgraph_from_store("t1", "kb1", doc_id) + assert result is not None and result.graph["source_id"] == [doc_id] + + assert await load_subgraph_from_store("t1", "kb1", "doc_4_new") is None