mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
fix(persistence): avoid duplicate index dumps
This commit is contained in:
jinli.yl
parent
626c850ccb
commit
5848c500e6
4 changed files with 69 additions and 0 deletions
|
|
@ -19,6 +19,7 @@ class LocalFileGraph(BaseFileGraph):
|
|||
self._inverse: dict[str, set[str]] = {} # real target → sources
|
||||
self._pending: dict[str, set[str]] = {} # virtual target → sources
|
||||
self._graph_file: Path = self.component_metadata_path / f"{self.name}.jsonl.zst"
|
||||
self._needs_persist = True
|
||||
|
||||
# -- Lifecycle ---------------------------------------------------------
|
||||
|
||||
|
|
@ -35,13 +36,17 @@ class LocalFileGraph(BaseFileGraph):
|
|||
if line.strip():
|
||||
node = FileNode.model_validate_json(line)
|
||||
self._nodes[node.path] = node
|
||||
self._needs_persist = False
|
||||
self.logger.debug(f"Loaded {len(self._nodes)} nodes from {self._graph_file}")
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to load {self._graph_file}: {e}")
|
||||
|
||||
async def dump(self) -> None:
|
||||
if not self._needs_persist:
|
||||
return
|
||||
try:
|
||||
write_jsonl_zst(self._graph_file, (n.model_dump_json() for n in self._nodes.values()))
|
||||
self._needs_persist = False
|
||||
self.logger.info(f"Saved {len(self._nodes)} nodes to {self._graph_file}")
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to write {self._graph_file}: {e}")
|
||||
|
|
@ -78,6 +83,8 @@ class LocalFileGraph(BaseFileGraph):
|
|||
# -- Node CRUD ---------------------------------------------------------
|
||||
|
||||
async def upsert_nodes(self, nodes: list[FileNode]) -> None:
|
||||
if nodes:
|
||||
self._needs_persist = True
|
||||
for node in nodes:
|
||||
path = node.path
|
||||
old = self._nodes.get(path)
|
||||
|
|
@ -96,6 +103,7 @@ class LocalFileGraph(BaseFileGraph):
|
|||
node = self._nodes.pop(path, None)
|
||||
if node is None:
|
||||
continue
|
||||
self._needs_persist = True
|
||||
for target in self._targets(node):
|
||||
self._remove_edge(path, target)
|
||||
demoted = self._inverse.pop(path, None)
|
||||
|
|
@ -119,6 +127,7 @@ class LocalFileGraph(BaseFileGraph):
|
|||
self._inverse.clear()
|
||||
self._pending.clear()
|
||||
self._graph_file.unlink(missing_ok=True)
|
||||
self._needs_persist = False
|
||||
|
||||
# -- Link access -------------------------------------------------------
|
||||
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ class NxFileGraph(BaseFileGraph):
|
|||
self._graph = nx.MultiDiGraph()
|
||||
self.component_metadata_path.mkdir(parents=True, exist_ok=True)
|
||||
self._graph_file: Path = self.component_metadata_path / f"{self.name}.pkl"
|
||||
self._needs_persist = True
|
||||
|
||||
# -- Lifecycle ---------------------------------------------------------
|
||||
|
||||
|
|
@ -34,16 +35,20 @@ class NxFileGraph(BaseFileGraph):
|
|||
try:
|
||||
with open(self._graph_file, "rb") as f:
|
||||
self._graph = pickle.load(f)
|
||||
self._needs_persist = False
|
||||
self.logger.info(f"Loaded {self._real_count()} nodes from {self._graph_file}")
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to load {self._graph_file}: {e}")
|
||||
|
||||
async def dump(self) -> None:
|
||||
if not self._needs_persist:
|
||||
return
|
||||
try:
|
||||
tmp = self._graph_file.with_suffix(".tmp")
|
||||
with open(tmp, "wb") as f:
|
||||
pickle.dump(self._graph, f, protocol=pickle.HIGHEST_PROTOCOL)
|
||||
tmp.replace(self._graph_file)
|
||||
self._needs_persist = False
|
||||
self.logger.info(f"Saved {self._real_count()} nodes to {self._graph_file}")
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to write {self._graph_file}: {e}")
|
||||
|
|
@ -69,6 +74,8 @@ class NxFileGraph(BaseFileGraph):
|
|||
# -- Node CRUD ---------------------------------------------------------
|
||||
|
||||
async def upsert_nodes(self, nodes: list[FileNode]) -> None:
|
||||
if nodes:
|
||||
self._needs_persist = True
|
||||
for node in nodes:
|
||||
path = node.path
|
||||
if self._graph.has_node(path):
|
||||
|
|
@ -80,6 +87,7 @@ class NxFileGraph(BaseFileGraph):
|
|||
for path in paths:
|
||||
if not self._graph.has_node(path):
|
||||
continue
|
||||
self._needs_persist = True
|
||||
self._graph.remove_edges_from(list(self._graph.out_edges(path, keys=True)))
|
||||
self._graph.nodes[path].pop("node", None) # demote to virtual
|
||||
if self._graph.in_degree(path) == 0:
|
||||
|
|
@ -101,6 +109,7 @@ class NxFileGraph(BaseFileGraph):
|
|||
async def clear(self):
|
||||
self._graph.clear()
|
||||
self._graph_file.unlink(missing_ok=True)
|
||||
self._needs_persist = False
|
||||
|
||||
# -- Link access -------------------------------------------------------
|
||||
|
||||
|
|
|
|||
|
|
@ -52,6 +52,10 @@ class BM25Index(BaseKeywordIndex):
|
|||
|
||||
# IDF cache; invalidated whenever live-doc count or postings change.
|
||||
self._idf_cache: dict[int, float] = {}
|
||||
# A fresh component still needs one persistence pass so an empty stale
|
||||
# snapshot can be removed. Successful dump/load operations clear this
|
||||
# flag; mutations set it again.
|
||||
self._needs_persist = True
|
||||
|
||||
# -- Properties -----------------------------------------------------------
|
||||
|
||||
|
|
@ -245,6 +249,8 @@ class BM25Index(BaseKeywordIndex):
|
|||
if not docs_dict:
|
||||
return
|
||||
|
||||
self._needs_persist = True
|
||||
|
||||
new_doc_ids: list[str] = []
|
||||
new_doc_lens: list[int] = []
|
||||
new_doc_token_ids: list[np.ndarray] = []
|
||||
|
|
@ -272,6 +278,8 @@ class BM25Index(BaseKeywordIndex):
|
|||
|
||||
async def delete_docs(self, doc_ids: list[str]) -> None:
|
||||
"""Lazy-delete a batch of doc_ids; physical reclaim happens in optimize_index."""
|
||||
if doc_ids:
|
||||
self._needs_persist = True
|
||||
for doc_id in doc_ids:
|
||||
self._remove_doc(doc_id)
|
||||
self._idf_cache = {}
|
||||
|
|
@ -354,8 +362,11 @@ class BM25Index(BaseKeywordIndex):
|
|||
|
||||
async def dump(self) -> None:
|
||||
"""Persist the index via temp file + atomic rename to avoid torn writes."""
|
||||
if not self._needs_persist:
|
||||
return
|
||||
if self.n_docs == 0 and not self.vocab:
|
||||
self.index_file.unlink(missing_ok=True)
|
||||
self._needs_persist = False
|
||||
return
|
||||
try:
|
||||
self.index_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
|
@ -363,6 +374,7 @@ class BM25Index(BaseKeywordIndex):
|
|||
with open(tmp, "wb") as f:
|
||||
pickle.dump(self._snapshot(), f)
|
||||
tmp.replace(self.index_file)
|
||||
self._needs_persist = False
|
||||
self.logger.info(f"Saved {self.n_docs} docs to {self.index_file}")
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to write {self.index_file}: {e}")
|
||||
|
|
@ -376,6 +388,7 @@ class BM25Index(BaseKeywordIndex):
|
|||
with open(self.index_file, "rb") as f:
|
||||
data = pickle.load(f)
|
||||
self._restore(data)
|
||||
self._needs_persist = False
|
||||
self.logger.info(f"Loaded {self.n_docs} docs from {self.index_file}")
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to load index: {e}")
|
||||
|
|
@ -394,6 +407,7 @@ class BM25Index(BaseKeywordIndex):
|
|||
self._posting_tfs = {}
|
||||
self._idf_cache = {}
|
||||
self.index_file.unlink(missing_ok=True)
|
||||
self._needs_persist = False
|
||||
|
||||
# -- Compaction -----------------------------------------------------------
|
||||
|
||||
|
|
@ -479,3 +493,4 @@ class BM25Index(BaseKeywordIndex):
|
|||
self._posting_doc_idxs = new_posting_idxs
|
||||
self._posting_tfs = new_posting_tfs
|
||||
self._idf_cache = {}
|
||||
self._needs_persist = True
|
||||
|
|
|
|||
|
|
@ -16,6 +16,8 @@ import pytest
|
|||
|
||||
from reme.components.file_store import FaissLocalFileStore, LocalFileStore, ZvecLocalFileStore
|
||||
from reme.components.file_store import local_file_store as local_file_store_module
|
||||
from reme.components.file_graph import local_file_graph as local_file_graph_module
|
||||
from reme.components.keyword_index import bm25_index as bm25_index_module
|
||||
from reme.components.embedding_store import LocalEmbeddingStore
|
||||
from reme.schema import FileChunk, FileNode
|
||||
from reme.utils.jsonl_zst import read_jsonl_zst, write_jsonl_zst
|
||||
|
|
@ -283,6 +285,40 @@ def test_keyword_only_upsert_removes_old_chunks_and_docs():
|
|||
run(go())
|
||||
|
||||
|
||||
def test_close_does_not_rewrite_keyword_index_and_graph_after_store_dump(monkeypatch):
|
||||
"""A store cascade leaves child components clean for their later close()."""
|
||||
bm25_writes = 0
|
||||
graph_writes = 0
|
||||
real_pickle_dump = bm25_index_module.pickle.dump
|
||||
real_graph_write = local_file_graph_module.write_jsonl_zst
|
||||
|
||||
def count_bm25_write(*args, **kwargs):
|
||||
nonlocal bm25_writes
|
||||
bm25_writes += 1
|
||||
return real_pickle_dump(*args, **kwargs)
|
||||
|
||||
def count_graph_write(*args, **kwargs):
|
||||
nonlocal graph_writes
|
||||
graph_writes += 1
|
||||
return real_graph_write(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(bm25_index_module.pickle, "dump", count_bm25_write)
|
||||
monkeypatch.setattr(local_file_graph_module, "write_jsonl_zst", count_graph_write)
|
||||
|
||||
async def go():
|
||||
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
|
||||
store = _new_local_store("t_single_close_dump")
|
||||
await store.start()
|
||||
await store.upsert([(node("note.md"), [chunk("note", "note.md", "persist once")])])
|
||||
|
||||
await store.close()
|
||||
|
||||
assert bm25_writes == 1
|
||||
assert graph_writes == 1
|
||||
|
||||
run(go())
|
||||
|
||||
|
||||
def test_start_does_not_health_check_embedding_without_backfill():
|
||||
"""Hot startup keeps local vector retrieval independent of provider health."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue