zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-16 12:47:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refact: improve task resume mechanism for graphrag (#14096)

Browse Source 
### 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
This commit is contained in:
Minal Mahala
2026-04-15 15:07:28 +05:30
committed by GitHub
parent 3364d86e6b
commit f930389311
3 changed files with 415 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

56
rag/graphrag/general/index.py
View File

@@ -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)")

46
rag/svr/task_executor.py
View File

@@ -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:

316
test/unit_test/rag/graphrag/test_checkpoint_resume.py Normal file
View File

@@ -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